Let's build a complete, runnable example: a differentiable projectile simulation where we find the optimal launch angle to hit a target.
import jax
import jax.numpy as jnp
def simulate_projectile(angle, speed, dt=0.01, gravity=9.81,
drag_coeff=0.01, num_steps=1000):
"""Simulate a projectile with air drag.
Returns the trajectory as (x, y) positions.
The simulation is fully differentiable w.r.t. angle and speed.
"""
vx0 = speed * jnp.cos(angle)
vy0 = speed * jnp.sin(angle)
def step(state, _):
x, y, vx, vy = state
# Air drag proportional to velocity squared
v = jnp.sqrt(vx**2 + vy**2) + 1e-8
ax = -drag_coeff * v * vx
ay = -gravity - drag_coeff * v * vy
vx = vx + ax * dt
vy = vy + ay * dt
x = x + vx * dt
y = y + vy * dt
# Clamp y >= 0 (ground)
y = jnp.maximum(y, 0.0)
return (x, y, vx, vy), (x, y)
init = (0.0, 0.0, vx0, vy0)
_, trajectory = jax.lax.scan(step, init, None, length=num_steps)
return trajectory # (xs, ys), each shape (num_steps,)
def landing_distance(angle, speed=30.0):
"""Where does the projectile land?"""
xs, ys = simulate_projectile(angle, speed)
# Find approximate landing point: last index where y > 0.01
above_ground = ys > 0.01
# Use weighted sum as a differentiable approximation
weights = above_ground.astype(jnp.float32)
landing_x = jnp.sum(xs * weights) / (jnp.sum(weights) + 1e-8)
return landing_x
# Gradient: how does landing distance change with launch angle?
d_distance_d_angle = jax.grad(landing_distance)
# Find the angle that hits a target at x=50
target_x = 50.0
def loss(angle):
return (landing_distance(angle) - target_x) ** 2
# Gradient descent to find optimal angle
angle = jnp.float32(0.7) # start at ~40 degrees
lr = 0.001
for i in range(200):
loss_val, grad = jax.value_and_grad(loss)(angle)
angle = angle - lr * grad
if i % 50 == 0:
dist = landing_distance(angle)
print(f"Step {i}: angle={jnp.degrees(angle):.1f}°, "
f"distance={dist:.1f}m, loss={loss_val:.4f}")
# Result
optimal_angle = angle
print(f"\\nOptimal angle: {jnp.degrees(optimal_angle):.2f}°")
print(f"Landing distance: {landing_distance(optimal_angle):.2f}m")
print(f"Target: {target_x}m")
# Bonus: sweep many angles with vmap
angles = jnp.linspace(0.1, 1.4, 100) # 5° to 80°
distances = jax.vmap(landing_distance)(angles)
best_idx = jnp.argmin(jnp.abs(distances - target_x))
print(f"\\nGrid search best: {jnp.degrees(angles[best_idx]):.1f}°")
print(f"Grid search distance: {distances[best_idx]:.1f}m")💡 Why This Matters
This example demonstrates the core superpower of JAX for science: we wrote a physics simulation with 1000 time steps, and JAX automatically differentiates through all of them. We then used gradient descent to solve an inverse problem (find the launch angle to hit a target). This same pattern applies to: optimizing rocket trajectories, fitting physical models to data, designing metamaterials, inverse rendering, and countless other scientific applications. The simulation runs on GPU, the gradients are exact (not finite-difference approximations), and vmap lets us sweep parameters trivially.