vmec_jax.vmec_tomnsp

VMEC Fourier transform conventions (fixaray + tomnsps) for parity work.

VMEC does not use a plain unweighted DFT for its force/residual transforms. Instead it uses:

  • symmetry-aware theta grids (ntheta1/2/3) and endpoint weights,

  • mode normalization scalings (mscale, nscale),

  • precomputed trig tables (cosmu/sinmu, cosnv/sinnv) and their derivative companions (cosmum/sinmum, cosnvn/sinnvn).

This module implements the core pieces needed for parity diagnostics:

  • vmec_trig_tables(…) : build VMEC-style trig and weight tables

  • tomnsps_rzl(…) : real-space -> Fourier-space force transform

Scope

This is intended for diagnostics/regressions against VMEC2000 wout outputs. It currently implements the lasym=False and lasym=True tables, but the transform itself is primarily exercised for the parity kernel work.

Functions

tomnspa_rzl(*, armn_even, armn_odd, ...[, ...])

VMEC tomnspa asymmetric force transform (real-space -> Fourier-space).

tomnsps_fft_policy_override(enabled)

Temporarily override auto TOMNSPS FFT policy when the solver device is explicit.

tomnsps_masks(*, ns, mpol, include_edge[, ...])

Precompute parity/evolution masks for tomnsps/tomnspa.

tomnsps_rzl(*, armn_even, armn_odd, ...[, ...])

VMEC real-space -> Fourier-space force transform (core of tomnsps).

vmec_angle_grid(*, ntheta, nzeta, nfp, lasym)

Build the VMEC internal (theta,zeta) grid implied by read_indata.f.

vmec_theta_sizes(ntheta, *, lasym)

Reproduce VMEC read_indata.f theta sizes.

vmec_trig_tables(*, ntheta, nzeta, nfp, ...)

Build VMEC-style trig and weight tables.

Classes

TomnspsMasks(ns, mpol, include_edge, ...[, ...])

Precomputed parity/evolution masks for tomnsps/tomnspa.