vmec_jax.solve

Fixed-boundary solvers.

The first solver milestone is a robust “inner solve” for the VMEC lambda field with R/Z held fixed. This is useful for:

  • validating the magnetic energy objective against VMEC2000 wout files,

  • building toward a full fixed-boundary equilibrium solve.

Notes

This module intentionally avoids optional dependencies (e.g. jaxopt). The current implementation uses gradient descent with a simple backtracking line search.

Functions

first_step_diagnostics(state0, static, *, ...)

Return a first-step diagnostic bundle (single force/precondition/update eval).

solve_fixed_boundary_gd(state0, static, *, ...)

Minimize a VMEC-style energy objective over (R,Z,lambda) coefficients.

solve_fixed_boundary_gn_vmec_residual(...[, ...])

Fixed-boundary solve using a Gauss-Newton (normal-equations) step on VMEC residuals.

solve_fixed_boundary_lbfgs(state0, static, ...)

Fixed-boundary solve using L-BFGS (no external deps).

solve_fixed_boundary_lbfgs_vmec_residual(...)

Fixed-boundary solve by minimizing a VMEC-style force-residual objective.

solve_fixed_boundary_residual_iter(state0, ...)

VMEC-style fixed-point update loop using preconditioned force residuals.

solve_lambda_gd(state0, static, *, phipf, ...)

Solve for VMEC lambda (scaled coefficients) with fixed R/Z.

Classes

SolveFixedBoundaryResult(state, n_iter, ...)

SolveLambdaResult(state, n_iter, wb_history, ...)

SolveVmecResidualResult(state, n_iter, ...)