vmec_jax.coords¶
Geometry evaluation on the (s,theta,zeta) grid.
This corresponds to VMEC’s totzsp* stages (Fourier synthesis) and provides the raw arrays needed for downstream metric/Jacobian and force-balance calculations.
Design notes¶
This function is intentionally a thin wrapper around the Fourier synthesis kernels. It is the first building block toward a fast, end-to-end differentiable VMEC port:
Keep everything vectorized over the radial index s.
Keep the basis precomputed and passed in as a PyTree so jax.jit works.
Return plain JAX arrays (or numpy arrays if JAX isn’t available).
Functions
|