vmec_jax.geom

Geometry and metric/Jacobian utilities.

This extends the coordinate kernel (R,Z,lambda on an (s,theta,zeta) grid) with:

  • radial derivatives (via finite differences on coefficient arrays)

  • a cylindrical -> Cartesian embedding using the physical toroidal angle

  • covariant basis vectors and metric tensor g_ij

  • Jacobian sqrt(g) = e_s · (e_theta × e_phi)

This module is intentionally minimal: it contains only what we need to start validating the geometric pieces against VMEC2000 and to support a future force/residual kernel.

Notes on angles

We follow the conventions already used in vmec_jax.fourier:

  • zeta spans a single field period, i.e. zeta ∈ [0, 2π)

  • the physical toroidal angle is phi = zeta / NFP

  • the Fourier phase is m*theta - n*zeta, equivalent to m*theta - (n*NFP)*phi

Accordingly: * eval_fourier_dzeta_phys returns ∂/∂phi.

Functions

eval_geom(state, static)

Compute geometry/metric/Jacobian on the full 3D grid.

register_pytree_node_class(cls)