What's the goal:

Port the oqd_trical.mechanical module of TrICal (Trapped-Ion Calculator, OQD's atomic-level emulator) from autograd to JAX, and replace its SciPy optimization routines with optax.

TrICal solves for the mechanical structure of a Coulomb crystal of trapped ions given the trap parameters, and simulates the resulting time-dependent dynamics. The mechanical module currently relies on autograd to compute gradients of the ion-chain potential energy, and on scipy.optimize to find the equilibrium ion positions that minimize that energy.

autograd is in maintenance mode and is no longer actively developed; JAX is its modern successor with a near-identical NumPy API, plus jit, vmap, and GPU/TPU support. Moving to JAX + optax brings TrICal in line with the rest of the OQD scientific stack and unlocks faster, hardware-accelerated structure solves.

Some details:

Documentation for TrICal is here. The relevant code lives in src/oqd_trical/mechanical, where you'll find:

• Potential-energy expressions for an ion chain in a trap (Coulomb interaction + trap pseudopotential), currently differentiated using autograd.grad / autograd.hessian.
• Equilibrium-position solvers built on top of scipy.optimize (e.g., scipy.optimize.minimize with BFGS/Newton-CG variants).
• Normal-mode calculations that diagonalize the Hessian of the potential at equilibrium.

The successful PR should:

• Replace autograd.numpy with jax.numpy, and autograd.grad / autograd.hessian with jax.grad / jax.hessian (or jax.jacfwd / jax.jacrev where appropriate). jit-compile hot paths where it's a clean win.
• Replace the scipy.optimize minimization with an optax-based loop (e.g., optax.lbfgs, optax.adam, or optax.scale_by_zoom_linesearch for line-search variants) that converges to the same equilibrium positions to within numerical tolerance.
• Preserve the public API of the mechanical module so downstream callers keep working without changes.
• Match existing test coverage; add tests where behavior is newly defined (e.g., optimizer convergence criteria).

Requirements

• All autograd uses ported to JAX.
• autograd removed from dependencies; jax and optax added.
• scipy.optimize routines in the mechanical module replaced with optax equivalents.
• Existing tests pass; equilibrium positions and normal-mode frequencies agree with the current implementation to within a documented tolerance.
• Public API of oqd_trical.mechanical is preserved (or breaking changes are clearly documented).

Possible extensions

• jit-compile the energy, gradient, and Hessian functions and benchmark against the current implementation.
• Add a GPU benchmark for solving large ion chains (e.g., N ≥ 50).
• Expose the optimizer choice (L-BFGS, Adam, etc.) and its hyperparameters as configurable arguments on the structure solver.