vmec-jax documentation¶
vmec-jax is an incremental, JAX-based rewrite of VMEC2000, targeting:
bundled-reference validation and optional executable-backed VMEC2000 checks for representative fixed-boundary and free-boundary solves, with strict field parity promoted case-by-case,
axisymmetric and non-axisymmetric
lasym=False/Truecoverage, including convergence/physics gates where strict parity is not yet promoted,end-to-end differentiability (JAX autodiff),
performance profiling and tuned default paths, with CPU/GPU benchmark results documented per case rather than implied globally,
required fast coverage enforced at the 95% gate after the latest CI-equivalent coverage ratchet,
stepwise validation against VMEC2000 output (
wout_*.nc).
User guide
- Overview
- Installation
- Quickstart
- Run the bundled CLI test
- Run the minimal showcase (recommended)
- CLI (VMEC2000-style executable)
- Boozer-coordinate CLI workflow
- Load, save, and inspect a
woutfile - Spline pressure, iota, and current profiles
- Free-boundary CLI smoke test
- Kernel parity on reference states (solver-free)
- Scalar residual parity (
fsqr/fsqz/fsql) on reference states
- Optimisation with vmec_jax
- Current reproducible workflow
- Motivation: differentiability without finite differences
- Comparison with SIMSOPT
- SIMSOPT-style objective tuples
- Quasi-helical symmetry example
- Recommended standalone scripts
- Full QA/QH/QP/QI policy sweep
- Finite-beta stage-one examples
- QI objective tuning
- QI diagnostics and validation plan
- Algorithms in detail
- Public API
- Source files
- Running the QH example
- Running the QI objective comparison
- GPU acceleration
- Further reading
- Optimization Sweep Results
- Publication structure
- Current README NFP Snapshot
- Individual Examples
- Sweep Reproduction
- Historical Compact Best Rows
- QI_optimization Input Coverage
- Minimal-Seed Showcase Snapshot
- Objective Histories
- Constrained QI Matrix
- Initial/Final State Atlases
- Summary Tables
- Publication Panel
- Finite-beta Stage-One Examples
Physics and algorithms
- Theory and conventions
- Equations and derivations
- Notation and conventions
- VMEC weighted DFT tables (
fixaray) - Two-stage DFT for
tomnsps - Ideal MHD equilibrium
- Energy principle (VMEC formulation)
- Flux coordinates and straight-field-line angle
- Internal scaling and regularity (
scalxc) - m=1 internal constraint (
lconm1) - Magnetic field representation
- Current density
- Redl Bootstrap-Current Mismatch
- Force balance in VMEC (residual form)
- Time-step control (Garabedian update)
- Multigrid interpolation (
interp.f) - Pressure and beta
- Geometric embedding and Jacobian
- Implementation mapping (
vmec_jax) - References (local)
- VMEC wiki primer (for vmec-jax parity)
- Algorithms
- Discretization summary
- Geometry pipeline
- Iteration Loop: Non-Scan vs Scan
- Free-Boundary Path (Current WP2)
- Profiles and volume integrals
- Field and energy
- Lambda-only solve (inner solve)
- Implicit differentiation
- Experimental VMEC-residual solvers (not yet VMEC2000-parity)
- Roadmap to VMEC-quality parity
- Discrete-adjoint differentiation
- Piecewise-Omnigenous Optimization Plan
- Comparison with SIMSOPT
- JXBFORCE / Mercier Diagnostics (
jdotb,DMerc,D_R)- Glasser Resistive-Interchange Criterion
- Key VMEC2000 Convention: Parity Channels For
bsubu/bsubv - JXBFORCE Low-Pass Filter (Conceptual Summary)
- From Filtered Fields To
jdotb(VMEC2000 Discretization) - Numerical Sensitivity: Why
jdotbCan Look Like Noise - Profiles: QI_nfp2 + QA_lowres (
jdotbshould be small) - Implementation Pointers (Source Code)
- Performance Notes (vmec_jax vs VMEC2000)
- Reproducing The Figures
- Bootstrap-Current Fixed-Point Plan
Validation and release
- Validation and parity with VMEC2000
- Testing, Coverage, and Repository Size Strategy
- Release checklist
- Optional validation plan
- Free-Boundary Plan
- Implementation status snapshot (March 2026)
- Program goals and parity matrix
- Current tests and benchmark coverage
- WP2 dump-to-dump alignment harness
- Scope and acceptance target
- VMEC2000 source deep dive (source-of-truth)
- Core equations and numerics to replicate
- VMEC++ cross-checks that inform implementation
- Free-boundary case inventory (local)
- Implementation work packages
- Documentation plan (required)
- Testing, validation, and benchmark plan
- Definition of done for free-boundary milestone
- Free-Boundary Coil Optimization
- Architecture
- Pedagogic forward examples
- Adjoint Validation Roadmap
- Current Status
- Low-Resolution Beta Scan
- Reproduction
- High-Resolution DIII-D Finite-Beta Benchmark
- High-Resolution LP-QA Stellarator Gate
- Single-Stage Coil-Only Optimization Validation
- Benchmarks
- Optional VMEC2000 Diagnostics
- Validation Status
- Next Implementation Steps
- Literature Anchors
- Performance notes
- Overview: cold vs warm runtime
- 2026-05-25 rerun snapshot
- Key performance decisions
- May 2026 policy validation snapshot
- CPU/GPU profiling playbook
- Exact optimizer profiling
- Current performance (representative benchmarks)
- Enable float64
- JIT boundaries and compile latency
- Scan-mode iteration (fast path)
- Experimental accelerated mode
- Optimization and gradient benchmarking
- Scan chunking (fixed NSTEP blocks)
- Live NSTEP printing (debug callback)
- DFT tomnsps (GEMM path)
- Preconditioner weight caching
- Free-boundary WP1 micro-benchmark
- WP2 free-boundary runtime controls
- Direct-coil CPU/GPU micro-benchmark snapshot
- Historical bundled example runtime/memory matrix (March 2026)
- Why the GPU can still be slower than the CPU
- Fused tridiagonal solver (scan only)
- Boundary decomposition cache + JAX-friendly initial guess
- Implementation map (performance-critical paths)
- Recent profiling snapshot (QA, 3 iterations on CPU)
- VMEC++ bad-progress restarts (optional)
- Static precomputation
- VMEC phase-stack cache
- Compilation cache
- CLI profiling (pre-iteration overhead)
- Persistent compilation cache tuning
- Batched radial smoothing
- Batched VMEC real-space synthesis
- Vectorized multigrid conversion
- Multigrid interpolation caches
- Precomputed (m,n)→signed maps
- Batched sin conversions
- Scatter-free boundary/axis enforcement
- Lambda gauge masking
- Concatenation-based mode updates
- Skip asymmetric updates when
lasym=False - Vectorized axis blending
- Cached mode scaling
- Historical free-boundary cache note
- Avoid Python objects in jitted functions
- Memory considerations
Development notes
API reference