Differentiable simulations#

The whole time loop is a JAX function of its inputs, so derivatives of any scalar computed from the output with respect to the physical inputs are available by automatic differentiation. This is what distinguishes JAX-in-Cell from a conventional particle-in-cell code and it is the basis of the optimisation and inference examples.

Which inputs are differentiable#

Differentiable inputs are the floating-point parameters that enter the compiled program as array arguments rather than as static values. They are marked in the tables of the user guide; the complete list is:

Domain

length, length_y, length_z, timestep_over_spatialstep_times_c

Solver

filter_alpha

Every population

grid_points_per_Debye_length, weight, charge_over_elementary_charge, perturbation_amplitude_{x,y,z}, perturbation_wavenumber_{x,y,z}, vth_over_c_{x,y,z}, drift_speed_{x,y,z}, initial_positions, initial_velocities

Ion populations additionally

mass_over_proton_mass, ion_temperature_over_electron_temperature_{x,y,z}

Integer parameters (counts, mode numbers used as integers, algorithm switches, boundary codes, seeds) are static. perturbation_wavenumber_x is differentiable because it enters as a real multiplier of \(2\pi/L\); non-integer values are allowed and simply produce a displacement that is not periodic in the box.

Taking a gradient#

Build a Simulation once, then write the quantity of interest as a function of a dictionary of runtime inputs and differentiate that function:

import jax.numpy as jnp
from jax import grad, jit
from jaxincell import Simulation, load_parameters

sim = Simulation(load_parameters("examples/input.toml"))
steps = sim.domain_parameters["total_steps"]

def mean_field(drift_speed):
    output = sim.run({"electrons": {"electrons0": {"drift_speed_x": drift_speed}}})
    E = jnp.mean(output["electric_field"][:, :, 0], axis=1)
    return jnp.mean(E[steps // 2:])

value = mean_field(6e7)
derivative = grad(mean_field)(6e7)

grad here is reverse-mode differentiation: JAX records the forward run and replays it backwards. The gradient of a scalar with respect to any number of inputs costs a small multiple of one forward run, independent of how many inputs there are, which is what makes optimisation over many parameters feasible. The first evaluation compiles a second program (the backward pass) and is slower; later evaluations are not.

The runtime inputs can be nested dictionaries with several parameters and species, and the function can take a JAX array or a dictionary of arrays as its argument, for example an array of drift speeds for several populations. jax.grad and jax.jacfwd follow the usual JAX rules for pytrees.

JAX gradient compared with finite differences

(a) Time-averaged electric field as a function of the electron drift speed for the two-stream configuration of examples/input.toml (400 steps, 3000 particles per species), with the tangent from the JAX gradient. (b) One-sided finite differences \([f(v_d + \epsilon) - f(v_d)]/\epsilon\) against the step \(\epsilon\), compared with the gradient: they agree for steps below about \(10^2\) m/s and depart for larger steps, where the difference no longer samples the local slope of this noisy objective. Generated by docs/scripts/fig_autodiff.py.#

Forward mode and the implicit scheme#

jax.jvp and jax.jacfwd (forward mode) work with both integrators. Reverse mode (jax.grad, jax.vjp) works with the explicit scheme only: the implicit scheme uses a lax.while_loop for the Picard iteration, and JAX cannot differentiate a while loop in reverse mode. For a handful of parameters forward mode is just as good, and it does not need to store the forward trajectory.

Memory#

Reverse mode through a lax.scan of total_steps iterations stores the residuals of every step. For the phase-space arrays this is of the same order as the output itself, so if the forward run fits comfortably in memory the gradient usually does too; if it does not, reduce total_steps, the number of particles, or use forward mode.

Why the gradients are meaningful#

A particle-in-cell step contains operations that are not smooth: the cell index of a particle is an integer, absorbing walls zero out charges, the velocity is clipped at \(0.99c\), and the Picard iteration count depends on the data. JAX differentiates through all of them by treating the piecewise-constant parts as having zero derivative and propagating derivatives through the continuous parts: the spline weights, the field updates, the Boris rotation and the interpolation are all smooth in the positions and velocities, so the derivative of a smooth diagnostic with respect to a smooth input is recovered. The test suite checks that gradients with respect to every differentiable parameter are finite.

Two effects limit what a gradient can tell you:

  • The output of a simulation with a finite number of particles is a noisy function of its inputs. The derivative of a noisy function is noisier still, so differentiate quantities averaged over time or over many particles (energies, mode amplitudes, growth rates fitted over a window), not instantaneous point values.

  • A derivative is local. Growth rates change smoothly with the drift speed, so the gradient in the example above is informative; the saturated state after a strongly nonlinear phase can depend on the inputs in a way that no finite-difference step captures either.

Reusing compiled programs#

The gradient program is compiled for the static configuration of the Simulation object. Changing a differentiable input, including through the argument of the differentiated function, does not recompile; changing anything else does. When optimising, keep one Simulation object and pass the design variables through the runtime inputs.

Examples#

  • examples/auto-differentiability.py: the gradient check shown above.

  • examples/optimize_two_stream_saturation.py: minimise the saturated field energy over the ion temperature with scipy.optimize.least_squares (or Optax, commented out).

  • examples/inference_two_stream.py: recover the drift speed of a two-stream configuration from the growth rate of its electric-field energy with a damped Newton iteration driven by forward-mode derivatives.

See Examples.