[rogeriojorge]

Summary

• expose Diffrax's adjoint choice through simulation with its existing checkpointed discrete adjoint as the default
• add one 71-line example that first verifies a simple derivative through a simulation, then solves a physically interpretable inverse problem
• keep all physics/model source unchanged

This is stacked on #11 because differentiated and optimized runs should not execute progress-meter host callbacks.

Example

Examples/1D_autodiff.py differentiates the final electric-field energy with respect to the counter-streaming drift. It compares:

• reverse-mode discretize-then-differentiate with RecursiveCheckpointAdjoint
• forward mode with ForwardMode
• a centered finite difference

It then creates a synthetic electric-energy time trace at drift 1.0 and uses reverse-mode gradients to infer that physical drift from a starting value of 1.1.

Numerical rationale

Diffrax recommends checkpointed reverse mode for scalar objectives with many inputs and ForwardMode for forward-mode derivatives. Its continuous BacksolveAdjoint is approximate and explicitly not recommended. I tested it as requested, but it is also not hashable under SPECTRAX's current outer static JIT, so this PR does not complicate the simulation API to support a discouraged path.

Skene and Burns (2025) automate efficient discrete adjoints for sparse spectral PDE solvers by reversing high-level solver control flow and supplying efficient VJPs for structured operators. JAX already supplies FFT VJPs, while Diffrax supplies checkpointed discrete time adjoints; this example follows the same discretize-then-differentiate principle without claiming to reproduce Dedalus's symbolic sparse-solver machinery.

• Skene and Burns: https://arxiv.org/abs/2506.14792
• Diffrax adjoints: https://docs.kidger.site/diffrax/api/adjoints/
• JAX checkpointing: https://docs.jax.dev/en/latest/notebooks/autodiff_remat.html

Validation

• pytest -q: 3 passed
• example runtime: 8.7 s on local CPU
• final log-energy derivative:
  • reverse: -0.732348217136557
  • forward: -0.7323482171365575
  • centered finite difference: -0.7323482851528951
• inverse problem: inferred drift 1.00000636 for target 1.0
• loss fell monotonically from 3.30e-1 to 2.59e-8 in 12 iterations

[rogeriojorge]

Superseded by #44, which adds the adjoint pass-through and tests that reverse mode, forward mode and finite differences agree, with Orszag–Tang optimisation, sensitivity and gradient-benchmark examples. The 1D drift-inference example here was not carried over. Closing without merging; the branch is kept.