JAX - Differentiable Physics & PDEs

JAX is uniquely suited for physics because it can differentiate through numerical solvers. This guide covers how to implement traditional PDE solvers that are "optimization-friendly" and how to build neural-hybrid physical models.

When to Use

• Solving Navier-Stokes, Wave, or Heat equations on GPU.
• Implementing Physics-Informed Neural Networks (PINNs).
• Performing Inverse Design (finding material properties from observations).
• Creating differentiable simulations for robotics or climate modeling.
• Sensitivity analysis of physical systems.

Core Principles

1. Differentiation through the Solver

In JAX, if you write an Euler or Runge-Kutta integrator using jax.numpy, you can automatically calculate ∂Result/∂InitialCondition or ∂Result/∂Viscosity.

2. Staggered Grids & Vmap