Physics Simulations and the JAX Ecosystem

The JAX scientific computing ecosystem has grown rapidly. Here are the key libraries and what they do:

JAX-MD: Molecular Dynamics

# JAX-MD provides differentiable molecular dynamics simulations
# Example concept (simplified):
import jax
import jax.numpy as jnp

def lennard_jones(r, epsilon=1.0, sigma=1.0):
    """Lennard-Jones potential between two particles at distance r."""
    s6 = (sigma / r) ** 6
    return 4.0 * epsilon * (s6**2 - s6)

# The potential is differentiable — get forces automatically
force = -jax.grad(lennard_jones)
print(f"Force at r=1.5: {force(1.5):.4f}")

# Vectorize over all pairs
def total_energy(positions):
    """Compute total LJ energy for a system of particles."""
    n = positions.shape[0]
    energy = 0.0
    for i in range(n):
        for j in range(i + 1, n):
            r = jnp.linalg.norm(positions[i] - positions[j])
            energy += lennard_jones(r)
    return energy

# Get forces on all particles: F = -∇E
forces = -jax.grad(total_energy)(positions)

Brax: Differentiable Rigid-Body Physics

Brax is Google's differentiable physics engine for robotics and reinforcement learning. It simulates rigid-body dynamics entirely in JAX, making the physics differentiable and GPU-accelerated.

# Brax enables training robot controllers through physics
# The entire simulation is differentiable:
# environment_step → reward → grad(policy_params)
# This is vastly faster than finite-difference methods used in
# traditional RL for continuous control.

JAXopt: Optimization Beyond SGD

# JAXopt provides scipy-style optimization with JAX
import jaxopt

# Minimize a function with L-BFGS
def rosenbrock(x):
    return (1 - x[0])**2 + 100 * (x[1] - x[0]**2)**2

solver = jaxopt.LBFGS(fun=rosenbrock, maxiter=100)
result = solver.run(jnp.array([0.0, 0.0]))
print(f"Minimum at: {result.params}")  # ≈ [1.0, 1.0]

# Projected gradient descent (constrained optimization)
solver = jaxopt.ProjectedGradient(
    fun=objective,
    projection=jaxopt.projection.projection_box,
    maxiter=200,
)

Taylor-Mode AD: High-Order Derivatives

JAX has a unique capability for computing higher-order derivatives efficiently using Taylor-mode automatic differentiation. This was highlighted in a NeurIPS 2024 best paper for solving PDEs:

import jax
import jax.numpy as jnp

# Compute Taylor coefficients of a function
# This is useful for solving PDEs with high-order derivatives
def f(x):
    return jnp.sin(x) * jnp.exp(-x)

# Get the 4th derivative efficiently
from jax import jet

# jet computes Taylor-mode derivatives
primals = jnp.array(1.0)
series = (jnp.array(1.0), jnp.array(0.0), jnp.array(0.0), jnp.array(0.0))
primal_out, series_out = jax.jet(f, (primals,), (series,))
# series_out contains higher-order derivatives
# This scales linearly with derivative order, not exponentially