Public API¶

vmec_jax.api is the recommended import surface for most users. It keeps the common solve, I/O, plotting, diagnostic, and example-level optimization entrypoints stable while lower-level numerical kernels continue to evolve.

Typical usage:

import vmec_jax.api as vj

fixed = vj.run_fixed_boundary("examples/data/input.circular_tokamak")
freeb = vj.run_free_boundary("examples/data/input.cth_like_free_bdy_lasym_small")

The primary optimization examples intentionally use the top-level import, which mirrors this public surface plus the broader scientific namespace:

import vmec_jax as vj

Both vmec_jax and vmec_jax.api expose the documented workflow objects used in the scripts, including example_paths/load_example and ExampleData for bundled cases, FixedBoundaryVMEC, objective wrappers, LeastSquaresProblem, least_squares_solve, QI diagnostics, BoozerBTarget/boozer_b_target_from_wout homotopy helpers, and plotting helpers. Lower-level solver kernels, force assembly routines, and replay internals remain submodule-level APIs.

Glasser resistive-interchange support is available through the public import surface as vj.GlasserResistiveInterchange. Add it to a least-squares problem as an upper-bound penalty with a zero tuple target:

glasser = vj.GlasserResistiveInterchange(maximum=0.0, softness=1.0e-3)
objective_tuples = [
    # ...
    (glasser.J, 0.0, GLASSER_WEIGHT),
]

The sign convention is the Glasser-Greene-Johnson necessary condition D_R <= 0; maximum=0.0 penalizes positive D_R. Because the criterion divides by magnetic shear squared, use it only on nonzero-shear surfaces when interpreting physics diagnostics. See JXBFORCE / Mercier Diagnostics (jdotb, DMerc, D_R) for the normalization and references.

The module exports the following user-facing helpers:

Public, user-facing API for vmec_jax.

This module intentionally re-exports a small set of functions that cover the common workflows:

I/O: load input, read/write wout
Solve: fixed- and free-boundary drivers
Plotting: VMEC-style surfaces and magnetic-field magnitude

Advanced users can import lower-level kernels directly from submodules.

class vmec_jax.api.AbsMeanIotaCeiling(maximum: float, *, softness: float = 0.001)

Bases: object

Smooth upper-bound objective for abs(mean_iota).

J(ctx: StageContext, state)

name = 'abs_iota_ceiling'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.AbsMeanIotaFloor(target: float, *, softness: float = 0.001)

Bases: object

Smooth lower-bound objective for abs(mean_iota).

J(ctx: StageContext, state)

name = 'abs_iota_floor'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.AspectRatio

Bases: object

Aspect-ratio objective object.

J(ctx: StageContext, state)

name = 'aspect'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.AugmentedLagrangianConstraint(objective: object, multiplier: float = 0.0, penalty: float = 1.0, softness: float = 0.0, name: str | None = None)

Bases: object

Wrap a non-negative violation objective as an augmented-Lagrangian term.

The wrapped objective should expose a signed constraint residual g(x) with feasibility g(x) <= 0. MirrorRatio and MaxElongation provide this signed form automatically through their constraint hooks while preserving their usual non-negative penalty behavior when used as ordinary least-squares terms.

For an inequality constraint g(x) <= 0 this wrapper adds the projected Powell-Hestenes-Rockafellar residual

sqrt(mu) * max(g(x) + lambda / mu, 0)

where lambda is the current multiplier and mu is the penalty. The constant term in the augmented Lagrangian is omitted because it does not affect minimizers. Update multipliers only from exact accepted diagnostics using updated().

J(_ctx: StageContext, _state)

multiplier: float = 0.0

name: str | None = None

objective: object

penalty: float = 1.0

property requires_qi_field: bool

softness: float = 0.0

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

to_qi_term(residual_weight: float) → QIObjectiveTerm

updated(violation: float, *, penalty_growth: float = 1.0, max_penalty: float | None = None) → AugmentedLagrangianConstraint

Return a copy with the inequality multiplier updated.

violation should be max(g(x), 0) measured from the exact accepted final state, not a trial-point residual. The multiplier is projected to be non-negative.

class vmec_jax.api.BDotB(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

VMEC bdotb profile objective from the differentiable JXBFORCE path.

name = 'bdotb'

profile_key = 'bdotb'

class vmec_jax.api.BDotGradV(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

VMEC bdotgradv profile objective from the differentiable JXBFORCE path.

name = 'bdotgradv'

profile_key = 'bdotgradv'

class vmec_jax.api.BVector(*, s_index: int = -1, normalize: float = 1.0)

Bases: object

Cartesian magnetic-field vector objective on one radial surface.

The residual vector is (Bx, By, Bz) flattened over (theta, zeta) on ctx.static.grid. s_index=-1 targets the boundary surface.

J(ctx: StageContext, state)

name = 'B_vector'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.BetaTotal

Bases: object

Total-beta objective for finite-beta studies.

J(ctx: StageContext, state)

name = 'betatotal'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.BootstrapCurrentIteration(iteration: int, mismatch_norm: float, current_update_norm: float, curtor: float, ac_aux_s: tuple[float, ...], ac_aux_f: tuple[float, ...], beta_total: float | None = None, aspect: float | None = None, mean_iota: float | None = None, fsq_total: float | None = None, effective_damping: float | None = None, current_update_limited: bool = False, unlimited_current_update_norm: float | None = None, max_current_update_norm: float | None = None)

Bases: object

Serializable per-iteration diagnostics for bootstrap fixed points.

ac_aux_f: tuple[float, ...]

ac_aux_s: tuple[float, ...]

aspect: float | None = None

beta_total: float | None = None

current_update_limited: bool = False

current_update_norm: float

curtor: float

effective_damping: float | None = None

fsq_total: float | None = None

iteration: int

max_current_update_norm: float | None = None

mean_iota: float | None = None

mismatch_norm: float

unlimited_current_update_norm: float | None = None

class vmec_jax.api.BootstrapCurrentOptions(helicity_n: int, surfaces: tuple[float, ...] | None = None, n_current: int = 50, policy: Literal['low_beta', 'lagged_pressure', 'integrating_factor'] = 'integrating_factor', damping: float = 0.5, max_current_update_norm: float | None = None, return_best_evaluated_on_max_iter: bool = False, anderson_depth: int = 0, mismatch_tol: float = 0.001, current_tol: float = 0.001, max_fixed_point_iter: int = 8, pcurr_type: str = 'cubic_spline_ip')

Bases: object

Configuration for Redl-to-VMEC current-profile fixed-point updates.

anderson_depth: int = 0

current_tol: float = 0.001

damping: float = 0.5

helicity_n: int

max_current_update_norm: float | None = None

max_fixed_point_iter: int = 8

mismatch_tol: float = 0.001

n_current: int = 50

pcurr_type: str = 'cubic_spline_ip'

policy: Literal['low_beta', 'lagged_pressure', 'integrating_factor'] = 'integrating_factor'

return_best_evaluated_on_max_iter: bool = False

surfaces: tuple[float, ...] | None = None

class vmec_jax.api.BootstrapCurrentResult(indata: InData, history: tuple[BootstrapCurrentIteration, ...], converged: bool, reason: str, last_run: Any | None = None, last_diagnostics: dict[str, Any] | None = None, returned_best_evaluated: bool = False, best_evaluated_iteration: int | None = None, best_evaluated_mismatch_norm: float | None = None)

Bases: object

Result returned by bootstrap_current_fixed_point().

best_evaluated_iteration: int | None = None

best_evaluated_mismatch_norm: float | None = None

converged: bool

history: tuple[BootstrapCurrentIteration, ...]

indata: InData

last_diagnostics: dict[str, Any] | None = None

last_run: Any | None = None

reason: str

returned_best_evaluated: bool = False

class vmec_jax.api.BoozConfig(enabled: bool = False, mbooz: int = 32, nbooz: int = 32, surfaces: tuple[float, ...] | None = None, jit: bool = False)

Bases: object

Configuration for a booz_xform_jax run.

enabled: bool = False

jit: bool = False

mbooz: int = 32

nbooz: int = 32

surfaces: tuple[float, ...] | None = None

class vmec_jax.api.BoozerBTarget(*, target_bmnc, target_bmns=None, normalize: bool = True, include_b00: bool = False, qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

Boozer |B| spectrum-matching objective for QI steering.

This term is intended as a differentiable homotopy/steering objective, not as a final QI diagnostic. It compares the current Boozer |B| spectrum against a reference spectrum on the same Boozer mode grid. By default each surface is normalized by its (m,n)=(0,0) coefficient so the term matches field shape rather than absolute field strength.

J(_ctx: StageContext, _state)

name = 'boozer_b_target'

requires_qi_field = True

to_qi_term(residual_weight: float) → QIObjectiveTerm

class vmec_jax.api.BoundaryModeLimits(mode: int, max_m: int | None = None, max_n: int | None = None, label: str | None = None)

Bases: object

Boundary-parameter mode limits for one optimization stage.

mode controls VMEC spectral-resolution extension and continuation bookkeeping. max_m and max_n optionally restrict the free boundary coefficients independently, enabling schedules such as toroidal-first stages with max_m=1 and max_n=4 before a full max_m=max_n=4 cleanup.

label: str | None = None

max_m: int | None = None

max_n: int | None = None

mode: int

class vmec_jax.api.DMerc(*, minimum: float = 0.0, softness: float = 0.001, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: object

Smooth lower-bound objective for VMEC Mercier stability.

The residual is a per-surface smooth penalty for DMerc < minimum on interior radial surfaces. It uses the differentiable state-level Mercier path for both stellarator-symmetric and LASYM equilibria.

J(ctx: StageContext, state)

name = 'DMerc'

terms(ctx: StageContext, state)

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.ExampleData(input_path: 'Path', wout_path: 'Optional[Path]', cfg: 'VMECConfig', indata: 'any', static: 'VMECStatic', wout: 'Optional[WoutData]', state: 'Optional[any]')

Bases: object

cfg: VMECConfig

indata: any

input_path: Path

state: any | None

static: VMECStatic

wout: WoutData | None

wout_path: Path | None

class vmec_jax.api.FiniteBetaTargets(aspect_ratio: float, min_iota: float, min_average_iota: float, max_iota: float, volavgB: float, beta_total: float, aspect_weight: float = 1.0, iota_weight: float = 1.0, max_iota_weight: float = 1.0, volavgB_weight: float = 1.0, beta_weight: float = 1.0)

Bases: object

Targets and weights for stage-one finite-beta fixed-boundary objectives.

aspect_ratio: float

aspect_weight: float = 1.0

beta_total: float

beta_weight: float = 1.0

iota_weight: float = 1.0

max_iota: float

max_iota_weight: float = 1.0

min_average_iota: float

min_iota: float

volavgB: float

volavgB_weight: float = 1.0

class vmec_jax.api.FixedBoundaryOptimizationResult(stage_records: list[tuple[int, FixedBoundaryExactOptimizer, ndarray, dict]], final_optimizer: FixedBoundaryExactOptimizer, final_result: dict, stage_modes: list[int])

Bases: object

Result returned by run_fixed_boundary_objective_optimization().

final_optimizer: FixedBoundaryExactOptimizer

property final_params: ndarray: Optimized boundary parameter vector for the final stage.

final_result: dict

property final_stage: tuple[int, FixedBoundaryExactOptimizer, ndarray, dict]: Last mode-continuation stage record.

property final_state: Final VMEC state if the optimizer stored one.

property history: dict: Final optimizer history dictionary written by save_history.

property history_entries: tuple[dict, ...]: Per-callback objective samples from the full solve.

property initial_optimizer: FixedBoundaryExactOptimizer: Optimizer object for the first stage.

property initial_params: ndarray: Initial boundary parameter vector for the first stage.

property initial_result: dict: Raw optimizer result dictionary for the first stage.

property initial_stage: tuple[int, FixedBoundaryExactOptimizer, ndarray, dict]

First mode-continuation stage record.

The tuple is (mode, optimizer, params0, result). Examples keep this explicit so users can choose which stage to save or inspect.

property initial_state: Initial VMEC state if the optimizer stored one.

property objective_history: ndarray: Objective values over full-solve callbacks as a NumPy array.

property stage_histories: tuple[dict, ...]: Per-stage history dictionaries in mode-continuation order.

property stage_initial_params: tuple[ndarray, ...]: Initial boundary-parameter vectors for each stage.

stage_modes: list[int]

property stage_optimizers: tuple[FixedBoundaryExactOptimizer, ...]: Optimizer objects for each mode-continuation stage.

stage_records: list[tuple[int, FixedBoundaryExactOptimizer, ndarray, dict]]

property stage_results: tuple[dict, ...]: Raw optimizer result dictionaries in mode-continuation order.

property stage_timing_summaries: tuple[dict[str, object], ...]: Small timing/iteration summaries for each stage.

property summary: dict[str, object]: Compact result summary for example reports and notebooks.

property timing_summary: dict[str, object]: Small timing/iteration summary for reports and examples.

class vmec_jax.api.FixedBoundaryRun(cfg: VMECConfig, indata: any, static: VMECStatic, state: any, result: any | None, flux: any, profiles: dict, signgs: int)

Bases: object

Container returned by run_fixed_boundary.

cfg: VMECConfig

flux: any

indata: any

profiles: dict

result: any | None

signgs: int

state: any

static: VMECStatic

class vmec_jax.api.FixedBoundaryVMEC(input_file: Path, cfg: object, indata: object, max_mode: int, min_vmec_mode: int = 5, output_dir: Path = PosixPath('results/optimization'), project_input_boundary_to_max_mode: bool = False, include: tuple[str, ...] = ('rc', 'zs'), fix: tuple[str, ...] = ('rc00',))

Bases: object

Small fixed-boundary optimization object used by the examples.

This object is intentionally lighter than SIMSOPT’s full Vmec graph, but it plays the same role in the example workflow: it owns the VMEC input deck, resolution policy, active boundary parameterization, and output path. Objective objects are then assembled into a LeastSquaresProblem and solved by least_squares_solve().

cfg: object

fix: tuple[str, ...] = ('rc00',)

classmethod from_input(input_file, *, max_mode: int, min_vmec_mode: int = 5, output_dir: Path | str = PosixPath('results/optimization'), project_input_boundary_to_max_mode: bool = False, simple_seed: bool = False, simple_seed_perturbation: float = 1e-05, include: Sequence[str] = ('rc', 'zs'), fix: Sequence[str] = ('rc00',)) → FixedBoundaryVMEC

Load a VMEC input file and apply the optimization resolution policy.

simple_seed=True replaces the boundary by the standard near-circular three-coefficient seed plus deterministic tiny active higher-mode perturbations. This is useful for stress-testing whether QA/QH/QP/QI examples can leave the zero-transform branch without changing the raw input deck on disk.

include: tuple[str, ...] = ('rc', 'zs')

indata: object

input_file: Path

max_mode: int

min_vmec_mode: int = 5

output_dir: Path = PosixPath('results/optimization')

project_input_boundary_to_max_mode: bool = False

class vmec_jax.api.GlasserResistiveInterchange(*, maximum: float = 0.0, softness: float = 0.001, shear_epsilon: float = 0.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: object

Smooth upper-bound objective for the Glasser resistive criterion.

The Glasser-Greene-Johnson necessary condition for resistive interchange stability is D_R <= 0. This objective returns one smooth penalty per interior radial surface for D_R > maximum. The optional shear_epsilon regularizes the 1 / shear**2 factor for optimization; keep it small and inspect glasser_shear_valid diagnostics for near-zero-shear surfaces.

J(ctx: StageContext, state)

name = 'D_R'

terms(ctx: StageContext, state)

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.JDotB(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

VMEC jdotb profile objective from the differentiable JXBFORCE path.

name = 'jdotb'

profile_key = 'jdotb'

class vmec_jax.api.JVector(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

Flux-coordinate current-density vector objective from JXBFORCE channels.

The returned vector contains (J^theta, J^zeta) = (itheta/sqrtg, izeta/sqrtg) flattened over the selected full-mesh surfaces and angular grid. It is a VMEC-coordinate current-density diagnostic, not a Cartesian vector.

J(ctx: StageContext, state)

name = 'J_vector'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.LeastSquaresProblem(objective_terms: tuple[~vmec_jax.optimization_workflow.ObjectiveTerm, ...]=(), qi_objective_terms: tuple[~vmec_jax.optimization_workflow.QIObjectiveTerm, ...]=(), metadata: dict[str, object]=<factory>, qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

Least-squares objective assembled from (function, target, weight) tuples.

As in SIMSOPT, tuple weight is an objective weight. Internally the residual is sqrt(weight) * (function - target). The examples pass objective .J methods here, which keeps physics targets and weights in the tuple list rather than in the driver/solver call.

classmethod from_tuples(tuples: Sequence[tuple[Callable, float | ndarray, float]])

Create a problem from SIMSOPT-style (J, target, weight) tuples.

J is usually an objective object’s .J method, such as (AspectRatio().J, 6.0, 1.0). Plain callables are also accepted when they use the workflow signature J(ctx, state). weight is the least-squares objective weight, not the residual multiplier.

property is_qi: bool: Whether the problem contains Boozer-space QI field objectives.

metadata: dict[str, object]

property objective_count: int: Total number of objective terms assembled from user tuples.

property objective_names: tuple[str, ...]: All objective names in residual assembly order.

objective_terms: tuple[ObjectiveTerm, ...] = ()

property qi_objective_names: tuple[str, ...]: Names of Boozer/QI objective terms assembled from tuples.

qi_objective_terms: tuple[QIObjectiveTerm, ...] = ()

qi_options: QuasiIsodynamicOptions | None = None

property scalar_objective_names: tuple[str, ...]: Names of state-space objective terms assembled from tuples.

property summary: dict[str, object]: Compact description of the assembled least-squares problem.

class vmec_jax.api.LgradB(*, threshold: float, s_index: int = -1, ntheta: int = 9, nphi: int = 7, smooth_penalty: float = 0.0)

Bases: object

Minimum-L_grad_B penalty object usable in QS or QI examples.

J(ctx: StageContext, state)

name = 'LgradB'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

to_qi_term(residual_weight: float) → QIObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.MagneticWell(*, minimum: float = 0.0, softness: float = 0.001)

Bases: object

Smooth lower-bound objective for the vacuum magnetic-well proxy.

The well follows the SIMSOPT/VMEC convention (dV/ds(0) - dV/ds(1)) / dV/ds(0) using the differentiable half-mesh volume derivative reconstructed from the VMEC state. Positive values are favorable; this objective returns a smooth penalty when the well falls below minimum.

J(ctx: StageContext, state)

name = 'magnetic_well'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

well(ctx: StageContext, state)

class vmec_jax.api.MaxElongation(*, threshold: float, ntheta: int = 48, nphi: int = 16, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0, qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

Maximum LCFS elongation penalty object for solved VMEC states.

J(ctx: StageContext, state)

name = 'max_elongation'

property requires_qi_field: bool

to_constraint_qi_term() → QIObjectiveTerm

to_constraint_term() → ObjectiveTerm

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

to_qi_term(residual_weight: float) → QIObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.MeanIota

Bases: object

Mean rotational-transform objective object.

J(ctx: StageContext, state)

name = 'iota'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

class vmec_jax.api.MirrorRatio(*, threshold: float, surfaces=(1.0,), mboz: int = 18, nboz: int = 18, ntheta: int = 96, nphi: int = 96, surface_index: int | None = None, phimin: float = 0.0, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0, normalize_surfaces: bool = True, jit_booz: bool = True, qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

Maximum mirror-ratio penalty object for solved VMEC states.

J(ctx: StageContext, state)

name = 'mirror_ratio'

property requires_qi_field: bool: Reuse a shared QI Boozer field only when explicitly requested.

to_constraint_qi_term() → QIObjectiveTerm

to_constraint_term() → ObjectiveTerm

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

to_qi_term(residual_weight: float) → QIObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.ObjectiveTerm(name: str, evaluate: ~typing.Callable[[~vmec_jax.optimization_workflow.StageContext, object], object], target: float | ~numpy.ndarray = 0.0, weight: float = 1.0, total: ~typing.Callable[[~vmec_jax.optimization_workflow.StageContext, object], object] | None = None, track_iota: bool = False, metadata: dict[str, object] = <factory>, prepare: ~typing.Callable[[~vmec_jax.optimization_workflow.StageContext], ~vmec_jax.optimization_workflow.ObjectiveTerm] | None = None)

Bases: object

One weighted least-squares objective block.

The callback receives (ctx, state) and returns a scalar or vector. The residual minimized by the optimizer is weight * (value - target). weight is the internal residual multiplier; public objective tuples use SIMSOPT semantics and are converted to sqrt(tuple_weight) by LeastSquaresProblem.from_tuples().

bind(ctx: StageContext) → ObjectiveTerm: Return a stage-specialized term when the objective has static setup.

evaluate: Callable[[StageContext, object], object]

metadata: dict[str, object]

name: str

prepare: Callable[[StageContext], ObjectiveTerm] | None = None

residual(ctx: StageContext, state) → object

target: float | ndarray = 0.0

total: Callable[[StageContext, object], object] | None = None

track_iota: bool = False

weight: float = 1.0

class vmec_jax.api.OptimizationOutputPaths(initial_input: Path, final_input: Path, initial_wout: Path, final_wout: Path, history: Path)

Bases: object

Canonical files written by fixed-boundary optimization examples.

as_dict() → dict[str, Path]: Return path names in the same order used by the example reports.

final_input: Path

final_wout: Path

history: Path

initial_input: Path

initial_wout: Path

class vmec_jax.api.QIDiagnosticOptions(surfaces: object | None = None, mboz: int | None = None, nboz: int | None = None, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.02, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, shuffle_profile_nphi_out: int | None = None, weighted_shuffle_profile_weight: float = 0.0, weighted_shuffle_profile_softness: float = 0.02, aligned_profile_weight: float = 0.0, aligned_profile_softness: float = 0.02, aligned_profile_trap_level: float = 0.65, aligned_profile_trap_softness: float = 0.05, phimin: float = 0.0, jit_booz: bool = True, include_legacy: bool = True, legacy_nphi_out: int | None = None, mirror_threshold: float = 0.21, mirror_ntheta: int = 128, mirror_nphi: int = 128, mirror_surface_index: int | None = None, elongation_threshold: float = 8.0, elongation_ntheta: int = 64, elongation_nphi: int = 24, include_lgradb: bool = False, lgradb_threshold: float = 0.3, lgradb_surface_index: int = -1, lgradb_ntheta: int = 9, lgradb_nphi: int = 7, lgradb_smooth_penalty: float = 0.0, fail_on_error: bool = False)

Bases: object

Resolution and threshold controls for QI diagnostic records.

aligned_profile_softness: float = 0.02

aligned_profile_trap_level: float = 0.65

aligned_profile_trap_softness: float = 0.05

aligned_profile_weight: float = 0.0

branch_width_softness: float = 0.02

branch_width_weight: float = 0.5

elongation_nphi: int = 24

elongation_ntheta: int = 64

elongation_threshold: float = 8.0

fail_on_error: bool = False

include_bounce_endpoints: bool = False

include_legacy: bool = True

include_lgradb: bool = False

jit_booz: bool = True

legacy_nphi_out: int | None = None

lgradb_nphi: int = 7

lgradb_ntheta: int = 9

lgradb_smooth_penalty: float = 0.0

lgradb_surface_index: int = -1

lgradb_threshold: float = 0.3

mboz: int | None = None

mirror_nphi: int = 128

mirror_ntheta: int = 128

mirror_surface_index: int | None = None

mirror_threshold: float = 0.21

n_bounce: int = 51

nalpha: int = 31

nboz: int | None = None

nphi: int = 151

phimin: float = 0.0

profile_weight: float = 0.1

shuffle_profile_nphi_out: int | None = None

shuffle_profile_softness: float = 0.02

shuffle_profile_weight: float = 1.0

softness: float = 0.02

surfaces: object | None = None

weighted_shuffle_profile_softness: float = 0.02

weighted_shuffle_profile_weight: float = 0.0

width_weight: float = 1.0

class vmec_jax.api.QIObjectiveTerm(name: str, evaluate: Callable[[StageContext, object, dict], tuple[object, object]], qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

One field-quality objective that shares a Boozer/QI field evaluation.

evaluate: Callable[[StageContext, object, dict], tuple[object, object]]

name: str

qi_options: QuasiIsodynamicOptions | None = None

residual_and_total(ctx: StageContext, state, field: dict) → tuple[object, object]

class vmec_jax.api.QIOptimizationContext(alpha: float, continuation_nfev: int, inner_max_iter: int, jit_booz: bool, max_elongation: float, max_mirror_ratio: float, max_mode: int, max_nfev: int, method: str, min_vmec_mode: int, mirror_surface_index: object, mirror_weight: float, opt_qi_resolution: dict, output_dir: Path, qi_gate_legacy_max: float, qi_gate_smooth_max: float, qi_options: object, qi_weight: float, scalar_cost_only_trials: bool | None, scipy_lsmr_maxiter: int | None, solver_device: str | None, stage_modes: tuple, stage_repeats: int, surfaces: object, target_abs_iota_min: float, target_aspect: float, trial_ftol: float, use_ess: bool, use_mode_continuation: bool)

Bases: object

Explicit staged-QI controls shared by helper routines.

The example scripts still show the optimization workflow directly. This context only replaces the older configure(globals()) helper plumbing so source users can reason about which controls are passed into staged helpers.

alpha: float

continuation_nfev: int

inner_max_iter: int

jit_booz: bool

max_elongation: float

max_mirror_ratio: float

max_mode: int

max_nfev: int

method: str

min_vmec_mode: int

mirror_surface_index: object

mirror_weight: float

opt_qi_resolution: dict

output_dir: Path

qi_gate_legacy_max: float

qi_gate_smooth_max: float

qi_options: object

qi_weight: float

scalar_cost_only_trials: bool | None

scipy_lsmr_maxiter: int | None

solver_device: str | None

stage_modes: tuple

stage_repeats: int

surfaces: object

target_abs_iota_min: float

target_aspect: float

trial_ftol: float

use_ess: bool

use_mode_continuation: bool

class vmec_jax.api.QISeedSuitabilityTargets(smooth_qi_max: float | None = 0.002, legacy_qi_max: float | None = 0.002, target_aspect: float | None = 5.0, aspect_relative_tolerance: float = 0.35, abs_iota_min: float | None = 0.41, mirror_ratio_max: float | None = 0.21, max_elongation: float | None = 8.0)

Bases: object

Promotion gates used to compare solved QI seed candidates.

None disables a gate. The defaults match the lightweight QI audit and optimization examples in this repository: smooth/legacy QI first, then nonzero transform, aspect, mirror, and elongation cleanup gates.

abs_iota_min: float | None = 0.41

aspect_relative_tolerance: float = 0.35

legacy_qi_max: float | None = 0.002

max_elongation: float | None = 8.0

mirror_ratio_max: float | None = 0.21

smooth_qi_max: float | None = 0.002

target_aspect: float | None = 5.0

class vmec_jax.api.QuasiIsodynamicOptions(surfaces: object, mboz: int = 18, nboz: int = 18, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.02, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, shuffle_profile_nphi_out: int | None = None, weighted_shuffle_profile_weight: float = 0.0, weighted_shuffle_profile_softness: float = 0.02, aligned_profile_weight: float = 0.0, aligned_profile_softness: float = 0.02, aligned_profile_trap_level: float = 0.65, aligned_profile_trap_softness: float = 0.05, phimin: float = 0.0, jit_booz: bool = True)

Bases: object

Boozer/QI sampling options shared by QI objective terms.

aligned_profile_softness: float = 0.02

aligned_profile_trap_level: float = 0.65

aligned_profile_trap_softness: float = 0.05

aligned_profile_weight: float = 0.0

branch_width_softness: float = 0.02

branch_width_weight: float = 0.5

include_bounce_endpoints: bool = False

jit_booz: bool = True

mboz: int = 18

n_bounce: int = 51

nalpha: int = 31

nboz: int = 18

nphi: int = 151

phimin: float = 0.0

profile_weight: float = 0.1

shuffle_profile_nphi_out: int | None = None

shuffle_profile_softness: float = 0.02

shuffle_profile_weight: float = 1.0

softness: float = 0.02

surfaces: object

weighted_shuffle_profile_softness: float = 0.02

weighted_shuffle_profile_weight: float = 0.0

width_weight: float = 1.0

class vmec_jax.api.QuasiIsodynamicResidual(options: QuasiIsodynamicOptions)

Bases: object

Smooth QI residual object using a shared Boozer field evaluation.

J(_ctx: StageContext, _state)

name = 'qi'

requires_qi_field = True

to_qi_term(residual_weight: float) → QIObjectiveTerm

class vmec_jax.api.QuasiIsodynamicResidualCeiling(*, maximum: float, smooth_penalty: float = 0.0, qi_options: QuasiIsodynamicOptions | None = None)

Bases: object

Soft upper-bound objective for preserving a low-QI basin during cleanup.

J(_ctx: StageContext, _state)

name = 'qi_ceiling'

requires_qi_field = True

to_qi_term(residual_weight: float) → QIObjectiveTerm

class vmec_jax.api.QuasisymmetryRatioResidual(*, helicity_m: int, helicity_n: int, surfaces)

Bases: object

QS residual object for QA/QH/QP objectives.

J(ctx: StageContext, state)

name = 'qs'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.RedlBootstrapMismatch(*, helicity_n: int, ne_coeffs, Te_coeffs, Ti_coeffs=None, Zeff_coeffs=1.0, surfaces: Sequence[float] | None = None, n_lambda: int = 32, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

Redl bootstrap-current mismatch objective for finite-beta studies.

Polynomial profile coefficients follow SIMSOPT ProfilePolynomial ordering. ne_coeffs are in m^-3 and Te_coeffs/Ti_coeffs in eV. The residual block is normalized as in SIMSOPT’s VmecRedlBootstrapMismatch objective.

J(ctx: StageContext, state)

name = 'redl_bootstrap_mismatch'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.ToroidalCurrent(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

Integrated toroidal-current profile from VMEC’s Mercier path.

The profile key is torcur and follows VMEC’s Mercier normalization: signgs * 2*pi * <B_u> on the full radial mesh. This is a state-derived current profile, not just the prescribed input ICURV.

name = 'torcur'

profile_key = 'torcur'

class vmec_jax.api.ToroidalCurrentGradient(*, surfaces: Sequence[float] | None = None, normalize: float = 1.0, mmax_force: int | None = None, nmax_force: int | None = None)

Bases: _MercierProfileObjective

Radial derivative of ToroidalCurrent used by VMEC Mercier terms.

name = 'torcur_prime'

profile_key = 'ip'

class vmec_jax.api.VMECMirrorRatio(*, threshold: float, surfaces=(1.0,), surface_index: int | None = None, ntheta: int | None = None, nphi: int | None = None, nzeta: int | None = None, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0, normalize_surfaces: bool = True, bmag_floor: float = 1e-300)

Bases: object

Fast mirror-ratio penalty evaluated directly from VMEC |B|.

The scalar mirror ratio (Bmax - Bmin) / (Bmax + Bmin) does not require Boozer coordinates. This objective samples the VMEC/JAX real-space field on the solver grid and avoids the booz_xform_jax transform used by MirrorRatio. It is intended as a lightweight optimization term; final quasisymmetry/omnigenity review should still use Boozer-coordinate contour and spectral diagnostics.

J(ctx: StageContext, state)

name = 'mirror_ratio'

property requires_qi_field: bool

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

total(ctx: StageContext, state)

class vmec_jax.api.VolavgB

Bases: object

Volume-averaged magnetic-field objective for finite-beta studies.

J(ctx: StageContext, state)

name = 'volavgB'

to_objective_term(*, target, residual_weight: float) → ObjectiveTerm

vmec_jax.api.annotate_qi_seed_suitability(record: dict[str, Any], *, targets: QISeedSuitabilityTargets | None = None) → dict[str, Any]

Return record with deterministic QI seed-ranking and gate fields.

The ranking score intentionally combines the differentiable smooth QI residual with the legacy Goodman-style branch-shuffle diagnostic. Missing core QI metrics are ranked last and are also reported in qi_gate_failures/qi_failure_reasons.

vmec_jax.api.apply_current_profile_to_indata(indata: InData, *, ac_aux_s: Any, ac_aux_f: Any, curtor: Any, pcurr_type: str = 'cubic_spline_ip') → InData: Return a copy of indata with a VMEC current profile applied.

vmec_jax.api.b_cartesian_from_state(state, static, indata=None, *, wout=None, signgs: int | None = None, s_index: int = -1, use_wout_bsup: bool = False, use_vmec_synthesis: bool = False)

Compute Cartesian magnetic field from a solved VMEC state.

The returned array is evaluated on static.grid and has shape (ntheta, nzeta, 3) for the selected radial index. The helper is pure JAX-compatible when indata is supplied, so callers can use it inside jax.jvp or as a residual function for FixedBoundaryExactOptimizer.

Parameters:

state – Solved VMECState.
static – Static object whose grid defines the angular points.
indata – VMEC input data. Required unless wout is supplied.
wout – Optional WoutData or wout-like object providing flux/current fields.
signgs – VMEC orientation sign. If omitted, uses wout.signgs when available and otherwise defaults to 1.
s_index – Radial index to return. -1 is the boundary surface.
use_wout_bsup – If True and wout is supplied, use stored bsup* Fourier data. This is a reference-parity mode rather than an optimization path.
use_vmec_synthesis – Forwarded to vmec_bcovar_half_mesh_from_wout.

vmec_jax.api.bmag_from_state_physical(state, static, indata=None, *, theta: ndarray, phi: ndarray, s_index: int, signgs: int | None = None, phipf: ndarray | None = None, chipf: ndarray | None = None, lamscale: float | None = None, flux_is_internal: bool | None = None, sqrtg_floor: float | None = None, bmag_floor: float | None = None, eps: float = 1e-14) → ndarray

Compute B magnitude on a surface using physical toroidal angle phi.

Notes

If phipf/chipf are provided, they override indata-derived profiles.
If indata is None and no flux profiles are provided, this raises.
bmag_floor adds a small positive value inside the sqrt for smoother gradients.

vmec_jax.api.bmag_from_wout(wout, *, theta: ndarray, zeta: ndarray, s_index: int) → ndarray: Return B magnitude on a surface from wout Nyquist Fourier coefficients.

vmec_jax.api.bootstrap_current_fixed_point(indata: InData, *, options: BootstrapCurrentOptions, solve_fn=None, diagnostics_fn=None, ne_coeffs=None, Te_coeffs=None, Ti_coeffs=None, Zeff_coeffs=1.0, run_kwargs: dict[str, Any] | None = None) → BootstrapCurrentResult

Run a VMEC/Redl fixed-point iteration for a self-consistent current.

The loop is deliberately callback-friendly: production use can rely on the default vmec_jax solve and Redl diagnostic callbacks, while tests and workflow studies can inject cheap deterministic callbacks. The optimized quantity is the VMEC current profile only; plasma boundary coefficients are not touched by this helper.

vmec_jax.api.bootstrap_current_update_to_indata(indata: InData, *, s: Any, current_derivative: Any | None = None, current: Any | None = None, signgs: int, pcurr_type: str = 'cubic_spline_ip') → tuple[InData, dict[str, Any]]: Convert a bootstrap-current update and apply it to an InData copy.

vmec_jax.api.booz_xform_inputs_from_state(*, state, static, indata, signgs: int, use_nyq_from_grid: bool = True, trig: VmecTrigTables | None = None, flux: FluxProfiles | None = None, profiles_half: dict | None = None) → BoozXformInputs: Construct booz_xform_jax inputs from a VMEC state using JAX kernels.

vmec_jax.api.boozer_b_target_from_wout(wout_path: str | Path, *, surfaces, mboz: int, nboz: int) → dict[str, ndarray | int]

Return Boozer |B| target spectra from a VMEC wout file.

The returned bmnc_b/bmns_b arrays use the same surface-major shape as booz_xform_jax’s differentiable API, so they can be passed directly to BoozerBTarget.

vmec_jax.api.boundary_from_input_convention(boundary: BoundaryCoeffs, modes: ModeTable, *, lasym: bool, apply_m1_constraint: bool = False) → BoundaryCoeffs: Convert raw input-convention boundary coefficients to solver convention.

vmec_jax.api.boundary_input_from_indata(indata: InData, modes: ModeTable) → BoundaryCoeffs

Build dense boundary coefficients in the raw input convention.

This matches the Fourier coefficient convention used by VMEC namelists and by SurfaceRZFourier-facing wrappers before VMEC’s internal sign handling and optional theta flip are applied.

vmec_jax.api.closed_theta_grid(ntheta: int) → ndarray: Theta grid including the 2π endpoint (good for closed cross-sections).

vmec_jax.api.damp_current_profile(old: Any, new: Any, damping: float) → Any: Return a damped current-profile update.

vmec_jax.api.diagnostic_float(record, key): Return a scalar diagnostic value, or nan when it is unavailable.

vmec_jax.api.dpsi_ds_from_vmec_phiedge(phiedge: Any, *, signgs: int) → Any

Return the flux derivative convention used in the bootstrap update.

Landreman’s VMEC-current note gives

\[d\psi/ds = \mathrm{signgs}\,\Phi_\mathrm{edge}/(2\pi).\]

PHIEDGE in VMEC input is the toroidal flux at the LCFS, not divided by 2*pi.

vmec_jax.api.example_paths(case: str, *, root: str | Path | None = None) → tuple[Path, Path | None]: Return (input_path, wout_path) for a bundled example case.

vmec_jax.api.export_vtk_surface_and_fieldline(*, input_path: str | Path, wout_path: str | Path | None = None, outdir: str | Path = 'vtk_out', s_index: int = -1, hi_res: bool = False, export_volume: bool = False) → dict[str, Path]

Export one surface + a fieldline trace to VTK for ParaView.

This helper reads a VMEC wout_*.nc and writes: - surface_b.vts: structured grid with Bx/By/Bz/Bmag on a surface - fieldline.vtp: a single fieldline polyline on that surface

If export_volume=True, a coarse volume grid volume.vts is also written (with Bmag only).

vmec_jax.api.finite_beta_global_residuals_from_state(*, state, static, indata, signgs: int, targets: FiniteBetaTargets) → jax.numpy.ndarray: Build global finite-beta residuals for stage-one surface optimization.

vmec_jax.api.finite_beta_scalars_from_state(*, state, static, indata, signgs: int) → dict[str, Any]: Return JAX-differentiable finite-beta scalar diagnostics from a VMEC state.

vmec_jax.api.fix_matplotlib_3d(ax): Fix 3D matplotlib aspect so structures do not look distorted.

vmec_jax.api.flux_profiles_from_indata(indata: InData, s, *, signgs: int) → FluxProfiles

Construct simple flux profiles (phipf/chipf) from &INDATA.

This is a deliberately minimal port:

toroidal flux uses PHIEDGE and polynomial APHI (default aphi=[1]), following magnetic_fluxes.f:torflux_deriv and profil1d.f.
poloidal flux derivative follows magnetic_fluxes.f:polflux_deriv, i.e. piota(tf) * torflux_deriv(s) for non-RFP (RFP uses polflux_deriv = 1).

vmec_jax.api.glasser_resistive_interchange_from_mercier_terms(*, DMerc, shear, H=None, Dcurr=None, tjb=None, tbb=None, jdotb=None, bdotb=None, shear_epsilon: float = 0.0) → dict[str, Any]

Return the Glasser resistive-interchange criterion from Mercier terms.

Landreman & Jorge rewrite the Glasser-Greene-Johnson necessary condition for resistive interchange stability as

D_R = -DMerc + 4*pi^2/iota_prime^2 * (H - iota_prime^2/(8*pi^2))^2.

The Mercier implementation in this module uses the VMEC/Ichiguchi normalization in which shear = d iota / d Phi and Dshear = shear**2 / 4. Since Phi = 2*pi*psi, the equivalent normalized expression is

D_R = -DMerc + (H - shear**2/2)**2 / shear**2.

The strict resistive-interchange necessary condition is D_R <= 0 on surfaces with nonzero magnetic shear. glasser_shear_valid marks the surfaces on which this division is physically meaningful. A positive shear_epsilon regularizes the denominator for smooth optimization penalties while preserving the validity mask.

If H is not supplied, it is reconstructed from VMEC profile data as shear * (tjb - (jdotb / bdotb) * tbb) when those profiles are available. If only the VMEC Mercier current term is available, the helper falls back to H = -Dcurr; this is exact when the surface-averaged parallel-current ratio equals the toroidal-current derivative.

vmec_jax.api.integrate_current_derivative(s: Any, current_derivative: Any) → Any: Integrate a current-derivative profile with I(0)=0.

vmec_jax.api.interpolate_indata_boundary(seed_indata, reference_indata, lam: float, *, keys: Sequence[str] = ('RBC', 'ZBS', 'RBS', 'ZBC'), max_mode: int | None = None, require_same_nfp: bool = True, preserve_seed_scalars: Sequence[str] = ('NFP', 'LASYM'))

Interpolate selected VMEC boundary Fourier coefficients.

This helper implements a deterministic global-to-local preconditioner for far-seed QI optimization. lam=0 preserves the seed boundary for the selected keys, while lam=1 uses the reference boundary. Scalar VMEC metadata remains seed-owned for entries in preserve_seed_scalars so a reference family can be used without accidentally changing the user’s field-period count or symmetry flag.

If max_mode is given, selected boundary coefficient dictionaries are projected to modes with abs(m) <= max_mode and abs(n) <= max_mode.

vmec_jax.api.jsonable(value): Convert NumPy/JAX-like values into JSON-serializable containers.

vmec_jax.api.least_squares_solve(vmec: FixedBoundaryVMEC, problem: LeastSquaresProblem, *, stage_modes: Sequence[int], max_nfev: int, continuation_nfev: int, method: str = 'scipy', ftol: float = 0.0001, gtol: float = 0.0001, xtol: float = 0.0001, use_ess: bool = False, ess_alpha: float = 1.2, label: str = 'Fixed-boundary optimization', use_mode_continuation: bool = True, inner_max_iter: int = 120, inner_ftol: float = 1e-09, trial_max_iter: int = 120, trial_ftol: float = 1e-09, solver_device: str | None = None, exact_path: str | None = None, scipy_tr_solver: str | None = 'lsmr', scipy_lsmr_maxiter: int | None = None, lbfgs_step_bound: float | None = None, scalar_step_bound: float | None = None, scalar_cost_only_trials: bool | None = None, save_stage_inputs: bool = True, save_stage_wouts: bool = False, save_rerun_wouts: bool = False, save_final_outputs: bool = True) → FixedBoundaryOptimizationResult

Solve a SIMSOPT-style fixed-boundary least-squares problem.

The examples use this as the common public workflow:

create a FixedBoundaryVMEC,
assemble a LeastSquaresProblem from (J, target, weight) tuples,
choose stage modes and optimizer settings,
call this function.

vmec_jax.api.lgradb_from_state(*, state, static, indata, signgs: int, s_index: int = -1, ntheta: int = 9, nphi: int = 7, flux_local=None)

Evaluate the magnetic-gradient scale length on a VMEC surface.

This is the JAX-native analogue of the L_grad_B diagnostic used in the reference SIMSOPT/omnigenity_optimization examples. It computes

L_grad_B = |B| sqrt(2 / (nabla B : nabla B))

from the Cartesian magnetic field vector on a small VMEC grid. Angular derivatives use periodic centered differences and the radial derivative uses the same differentiable finite-difference operator as the geometry kernel. The returned arrays are differentiable with respect to the VMEC state and boundary parameters.

vmec_jax.api.lgradb_penalty_from_state(*, state, static, indata, signgs: int, threshold: float = 0.3, s_index: int = -1, ntheta: int = 9, nphi: int = 7, smooth_penalty: float = 0.0, flux_local=None)

Penalize short magnetic-gradient scale length on a VMEC surface.

The residual follows the reference omnigenity scripts:

max(1/L_grad_B - 1/threshold, 0) / sqrt(ntheta*nphi).

Use this as an independent least-squares block, e.g. with residual weight sqrt(0.01) to match the QI reference script.

vmec_jax.api.load_example(case: str, *, root: str | Path | None = None, with_wout: bool = True, grid=None) → ExampleData: Load a bundled example case (config + static + optional wout/state).

vmec_jax.api.load_input(path: str | Path): Convenience wrapper around load_config.

vmec_jax.api.load_wout(path: str | Path) → WoutData: Convenience wrapper around read_wout.

vmec_jax.api.magnetic_well_from_state(*, state, static, indata, signgs: int) → Any: Return the differentiable VMEC magnetic-well proxy for an equilibrium.

vmec_jax.api.make_qi_optimization_context(context: dict | None = None, /, *, strict: bool = False, **overrides) → QIOptimizationContext

Build a typed staged-QI helper context from script constants.

strict=True is recommended for standalone examples: every required control must be present in context or overrides. The default strict=False preserves compatibility with older scripts that installed constants with configure() and relied on module globals.

vmec_jax.api.max_elongation_penalty_from_state(*, state, static, threshold: float = 8.0, ntheta: int = 64, nphi: int = 24, s_index: int = -1, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0): Penalize excessive LCFS elongation from a solved VMEC state.

vmec_jax.api.minimal_fixed_boundary_indata(*, nfp: int, r0: float = 1.0, rbc01: float = 0.2, zbs01: float = 0.2, mpol: int = 5, ntor: int = 5, ns_array: int | list[int] = 35, niter_array: int | list[int] = 1500, ftol_array: float | list[float] = 1e-13, phiedge: float = 0.083) → InData

Return a minimal fixed-boundary VMEC seed used by optimization examples.

The boundary has only the circular/elliptic seed coefficients RBC(0,0), RBC(0,1), and ZBS(0,1). Optimization examples can then activate higher Fourier coefficients through their selected max_mode and continuation policy, so the same simple template can be used to demonstrate QA, QH, QP, and QI optimization from a seed far from the target magnetic-field structure.

vmec_jax.api.mirror_ratio_penalty_from_boozer_modes(*, bmnc_b, xm_b, xn_b, nfp: int, bmns_b=None, threshold: float = 0.21, weights: Iterable[float] | None = None, ntheta: int = 128, nphi: int = 128, phimin: float = 0.0, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0)

Penalize the maximum mirror ratio from Boozer |B| modes.

This is the JAX-native analogue of the MirrorRatioPen diagnostic used in the reference omnigenity_optimization QI script. For every supplied Boozer surface it evaluates |B|(theta_B, phi_B) on a uniform grid and computes

M = (Bmax - Bmin) / (Bmax + Bmin).

The least-squares residual is max(0, M - threshold) per surface. Set smooth_extrema and/or smooth_penalty to positive values when a fully smooth softmax/softplus surrogate is preferred.

vmec_jax.api.mirror_ratio_penalty_from_boozer_output(booz, *, nfp: int | None = None, threshold: float = 0.21, weights: Iterable[float] | None = None, ntheta: int = 128, nphi: int = 128, phimin: float = 0.0, smooth_extrema: float = 0.0, smooth_penalty: float = 0.0): Evaluate mirror_ratio_penalty_from_boozer_modes() from Boozer output.

vmec_jax.api.optimization_output_paths(output_dir: str | Path) → OptimizationOutputPaths: Return the canonical final-artifact paths for an optimization run.

vmec_jax.api.parse_booz_surfaces(value: object) → tuple[float, ...] | None: Parse all or a comma/space-separated list of surface values.

vmec_jax.api.plot_3d_boundary_comparison(wout_initial, wout_final, *, outdir=None) → Path

Plot initial/final LCFS 3-D surfaces colored by |B|.

wout_initial and wout_final can be loaded WoutData objects or paths to wout_*.nc files. The returned path points to boundary_comparison.png in outdir.

vmec_jax.api.plot_bmag_contours(wout_initial, wout_final, *, outdir=None) → Path

Plot line contours of |B| on the initial/final LCFS.

wout_initial and wout_final can be loaded WoutData objects or paths to wout_*.nc files. The returned path points to bmag_surface.png in outdir.

vmec_jax.api.plot_boozer_bmag_contours_from_state(state, *, static, indata, signgs: int, outdir: str | Path, filename: str = 'boozer_bmag_surface.png', surfaces=(1.0,), surface_index: int = -1, mboz: int = 18, nboz: int = 18, ntheta: int = 128, nphi: int = 256, phimin: float = 0.0, title: str | None = None): Write a Boozer-coordinate line-contour plot of |B| for QI review.

vmec_jax.api.plot_boozer_lcfs_bmag_comparison(wout_initial, wout_final, *, outdir: str | Path | None = None, name: str = 'boozer_lcfs_bmag_comparison', mbooz: int = 32, nbooz: int = 32, ntheta: int = 128, nphi: int = 256, jit: bool = False) → Path

Plot initial/final LCFS |B| contours in Boozer coordinates.

This helper is intended for optimization examples. It runs booz_xform_jax on each WOUT file using the outermost Boozer surface, then plots line contours in (phi_B, theta_B). VMEC-angle contour plots are useful for parity checks; quasisymmetry and omnigenity should be judged in Boozer coordinates.

vmec_jax.api.plot_boozmn(boozmn_path: str | Path, outdir: str | Path | None = None, *, name: str | None = None, show: bool = False) → dict[str, Path]: Generate polished Boozer-coordinate diagnostic plots from boozmn.

vmec_jax.api.plot_boozmn_bmag_contours(boozmn, *, outdir: str | Path | None = None, name: str | None = None, ntheta: int = 128, nphi: int = 256) → Path: Plot mid-radius and outermost Boozer |B| line contours from boozmn.

vmec_jax.api.plot_boozmn_mode_families(boozmn, *, outdir: str | Path | None = None, name: str | None = None, max_modes: int = 80) → Path: Plot radial Boozer |B| mode amplitudes grouped by symmetry family.

vmec_jax.api.plot_boozmn_spectrum(boozmn, *, outdir: str | Path | None = None, name: str | None = None, surface_index: int = -1, nmodes: int = 40) → Path: Plot the largest Boozer |B| Fourier amplitudes on one surface.

vmec_jax.api.plot_objective_history(history_path, *, outdir=None) → Path: Plot objective, aspect-ratio, and optional iota history from JSON.

vmec_jax.api.plot_wout(wout_path: str | Path, outdir: str | Path | None = None, *, name: str | None = None, s_plot_ignore: float = 0.2, show: bool = False) → dict

Generate diagnostic plots from a VMEC wout_*.nc file.

Replicates the output of the standalone vmecPlot2.py script in a vectorised, vmec_jax-native form. Four figures are written:

<name>_VMECparams.pdf — 9-panel profile + |B| overview
<name>_poloidal_plot.png — LCFS cross-sections at multiple toroidal angles
<name>_VMECsurfaces.pdf — nested flux-surface cross-sections (8 panels)
<name>_VMEC_3Dplot.png — 3-D LCFS surface coloured by |B|

Parameters:

wout_path – Path to the wout_*.nc file.
outdir – Directory to save figures. Defaults to the directory containing wout_path.
name – Base name for output files. Defaults to the wout stem with the leading wout_ stripped (e.g. wout_nfp4_QH.nc → nfp4_QH).
s_plot_ignore – Fraction of flux surfaces near the axis to ignore when plotting DMerc.
show – If True, call plt.show() after saving all figures.

Returns:

{"vmec_params", "poloidal_plot", "vmec_surfaces", "3d_plot"} mapping to saved Path objects.

Return type:

dict

vmec_jax.api.prepare_simple_omnigenity_seed_input(input_file, output_dir, *, max_mode: int, min_vmec_mode: int = 5, enabled: bool = True, include: Sequence[str] = ('rc', 'zs'), fix: Sequence[str] = ('rc00',), perturbation: float = 1e-05, filename: str = 'input.simple_seed'): Write and return a simple omnigenity seed input path when enabled.

vmec_jax.api.profiles_from_wout(wout) → dict[str, ndarray]: Return common radial profiles from wout.

vmec_jax.api.qi_boozer_b_target_objective(*, target_bmnc, target_bmns=None, weight: float = 1.0, normalize: bool = True, include_b00: bool = False, qi_options: QuasiIsodynamicOptions | None = None) → QIObjectiveTerm: Boozer |B| spectrum target evaluated on the shared QI field.

vmec_jax.api.qi_cleanup_candidate_promotable(candidate: dict[str, Any], *, reference: dict[str, Any] | None = None, targets: QISeedSuitabilityTargets | None = None, require_seed_gate: bool = True, require_engineering_gate: bool = False, require_mirror_improvement: bool = True, mirror_improvement_min: float = 0.0) → dict[str, Any]

Annotate whether a QI cleanup candidate should replace a reference.

Mirror-ratio and elongation cleanup terms are engineering constraints, not definitions of QI. This helper encodes the promotion rule used by the example optimizations: do not promote a candidate that improves mirror by destroying the smooth/legacy QI, aspect, or transform gates, and do not advance a mirror-ramp stage unless the mirror ratio actually decreases relative to the previously accepted state.

vmec_jax.api.qi_diagnostics_for_result(stage_result, *, mirror_threshold, mirror_surface_index, smooth_qi_max=None, legacy_qi_max=None, ctx: QIOptimizationContext | None = None)

vmec_jax.api.qi_diagnostics_from_boozer_output(booz: dict[str, Any], *, options: QIDiagnosticOptions | None = None, nfp: int | None = None, weights: Iterable[float] | None = None) → dict[str, Any]

Evaluate QI diagnostics from an existing Boozer output dictionary.

The returned record is flat and JSON/CSV-friendly. State-only metrics (elongation and LgradB) are left as None by this Boozer-only entry point.

vmec_jax.api.qi_diagnostics_from_state(*, state: Any, static: Any, indata: Any, signgs: int, surfaces: Any | None = None, options: QIDiagnosticOptions | None = None, weights: Iterable[float] | None = None, flux_local: Any = None, prof_local: Any = None, pressure_local: Any = None, jit_booz: bool | None = None, booz_constants: Any = None, booz_grids: Any = None, surface_indices: Any = None) → dict[str, Any]: Evaluate a complete QI diagnostic record from a solved VMEC state.

vmec_jax.api.qi_max_elongation_constraint(*, threshold: float, ntheta: int = 48, nphi: int = 16, smooth_extrema: float = 0.0, qi_options: QuasiIsodynamicOptions | None = None) → QIObjectiveTerm: Signed LCFS elongation constraint max_elongation - threshold <= 0.

vmec_jax.api.qi_mirror_ratio_constraint(*, threshold: float, ntheta: int = 96, nphi: int = 96, surface_index: int | None = None, phimin: float = 0.0, smooth_extrema: float = 0.0, normalize_surfaces: bool = True, qi_options: QuasiIsodynamicOptions | None = None) → QIObjectiveTerm: Signed mirror-ratio constraint mirror_ratio - threshold <= 0.

vmec_jax.api.qi_promotion_score(record: dict[str, Any], *, targets: QISeedSuitabilityTargets | None = None, require_legacy_source: bool = False, objective_key: str = 'objective_final', wall_time_key: str = 'total_wall_time_s') → tuple[object, ...]

Lexicographic score for promoting final QI optimization candidates.

Seed audits intentionally rank QI-like starts before engineering cleanup so potentially useful seeds are not hidden. Final promotion is stricter: a candidate that preserves QI while satisfying aspect/iota/mirror/elongation gates should beat a lower scalar objective that destroys an engineering gate. Lower tuples are better.

vmec_jax.api.qi_stage_modes(*, max_mode: int, use_mode_continuation: bool, continuation_nfev: int, repeats: int = 3, policy: str = 'lower') → list[int]

Return the stage mode sequence for the QI example workflow.

policy="lower" uses the same lower-mode continuation semantics as the QA/QH/QP examples. policy="lower-repeat" repeats each lower-mode rung, which is useful for far circular seeds because each active spectral shell gets cleanup passes before adding more degrees of freedom. policy="repeat" preserves the older QI behavior of repeating only the final mode, which can still be useful when the input is already in the right basin.

vmec_jax.api.qs_stage_modes(*, max_mode: int, use_mode_continuation: bool, continuation_nfev: int) → list[int]: Repeated mode-continuation sequence used by the example scripts.

vmec_jax.api.quasi_isodynamic_residual_from_boozer_modes(*, bmnc_b, xm_b, xn_b, iota_b, nfp: int, weights: Iterable[float] | None = None, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.01, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, shuffle_profile_nphi_out: int | None = None, weighted_shuffle_profile_weight: float = 0.0, weighted_shuffle_profile_softness: float = 0.02, aligned_profile_weight: float = 0.0, aligned_profile_softness: float = 0.02, aligned_profile_trap_level: float = 0.65, aligned_profile_trap_softness: float = 0.05, phimin: float = 0.0)

Evaluate a smooth QI residual from Boozer |B| Fourier modes.

Parameters:

bmnc_b – Cosine Boozer |B| coefficients with shape (nsurf, nmodes).
xm_b – Boozer mode numbers. xn_b should use the physical toroidal mode convention, matching booz_xform_jax and BOOZ_XFORM.
xn_b – Boozer mode numbers. xn_b should use the physical toroidal mode convention, matching booz_xform_jax and BOOZ_XFORM.
iota_b – Rotational transform on the same Boozer surfaces.
nfp – Number of field periods.
weights – Optional surface weights. They are applied as square roots to the least-squares residual vector.
nphi – Sampling resolution along toroidal angle, field-line label, and normalized bounce level.
nalpha – Sampling resolution along toroidal angle, field-line label, and normalized bounce level.
n_bounce – Sampling resolution along toroidal angle, field-line label, and normalized bounce level.
include_bounce_endpoints – If true, include normalized bounce levels 0 and 1 in the smooth level-set residuals, matching the legacy Goodman-style branch-shuffle diagnostic. The default keeps the historical smooth objective behavior and samples only interior bounce levels.
softness – Logistic smoothing width in normalized |B| units. Smaller values approach hard branch widths but increase stiffness.
width_weight – Relative weight for the smooth level-set occupancy width residual.
branch_width_weight – Relative weight for a branch-based trapped-well width residual. This follows the reference omnigenity objective more closely than the occupancy width: each field line is split at its well minimum, each side is made monotone with a cumulative maximum, and level crossings are computed with a smooth inverse.
branch_width_softness – Normalized |B| smoothing width for branch level crossings.
profile_weight – Small relative weight for field-line profile consistency. Width-only and branch-width-only surrogates can rank some QH-like candidates too favorably; keeping a small profile term restores the legacy branch-shuffle ranking without making this term dominate the objective.
shuffle_profile_weight – Relative weight for a differentiable branch-shuffle profile residual. This term follows the reference Goodman/omnigenity_optimization diagnostic more directly than the occupancy-width residual: it builds left/right trapped-branch crossings, shifts them so every field-line label has the mean bounce width, and compares that shuffled well to the original profile.
shuffle_profile_softness – Logistic smoothing width used to estimate branch crossing locations for shuffle_profile_weight.
shuffle_profile_nphi_out – Optional dense output grid for the branch-shuffle profile residual. If set, the shuffled and original wells are compared on this many toroidal samples, matching the legacy arr_out=True Goodman objective more closely than the default base nphi grid.
weighted_shuffle_profile_weight – Relative weight for a branch-shuffle profile residual whose mean bounce widths are weighted by a differentiable proxy for the legacy squash/stretch quality weights. This is useful when the unweighted smooth residual ranks high-mirror false positives ahead of the Goodman-style branch diagnostic.
weighted_shuffle_profile_softness – Logistic smoothing width used by weighted_shuffle_profile_weight.
aligned_profile_weight – Relative weight for a differentiable trapped-well profile residual. Each field-line profile is circularly shifted by its smooth minimum before comparing against the mean over field-line label. This is a smooth surrogate for the branch/shuffle profile comparison in the reference QI scripts.
aligned_profile_softness – Temperature used for the smooth circular argmin that locates each well minimum in normalized |B| units.
aligned_profile_trap_level – Logistic trapped-region window applied to the aligned profiles. Values below aligned_profile_trap_level receive the most weight.
aligned_profile_trap_softness – Logistic trapped-region window applied to the aligned profiles. Values below aligned_profile_trap_level receive the most weight.
phimin – Start of the toroidal interval. The interval length is one field period.

Returns:

residuals1d is suitable for least-squares optimization. total is its squared norm.

Return type:

dict

vmec_jax.api.quasi_isodynamic_residual_from_boozer_output(booz, *, nfp: int | None = None, weights: Iterable[float] | None = None, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.01, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, shuffle_profile_nphi_out: int | None = None, weighted_shuffle_profile_weight: float = 0.0, weighted_shuffle_profile_softness: float = 0.02, aligned_profile_weight: float = 0.0, aligned_profile_softness: float = 0.02, aligned_profile_trap_level: float = 0.65, aligned_profile_trap_softness: float = 0.05, phimin: float = 0.0): Evaluate the smooth QI residual from a booz_xform_jax output dict.

vmec_jax.api.quasi_isodynamic_residual_from_state(*, state, static, indata, signgs: int, surfaces, weights: Iterable[float] | None = None, mboz: int = 12, nboz: int = 12, nphi: int = 151, nalpha: int = 31, n_bounce: int = 51, include_bounce_endpoints: bool = False, softness: float = 0.02, width_weight: float = 1.0, branch_width_weight: float = 0.5, branch_width_softness: float = 0.01, profile_weight: float = 0.1, shuffle_profile_weight: float = 1.0, shuffle_profile_softness: float = 0.02, shuffle_profile_nphi_out: int | None = None, weighted_shuffle_profile_weight: float = 0.0, weighted_shuffle_profile_softness: float = 0.02, aligned_profile_weight: float = 0.0, aligned_profile_softness: float = 0.02, aligned_profile_trap_level: float = 0.65, aligned_profile_trap_softness: float = 0.05, phimin: float = 0.0, flux_local=None, prof_local=None, pressure_local=None, jit_booz: bool = False, booz_constants=None, booz_grids=None, surface_indices=None)

Evaluate a differentiable QI residual directly from a solved VMEC state.

This uses vmec_jax.booz_xform_inputs_from_state followed by the functional booz_xform_jax API. booz_xform_jax is an optional runtime dependency; install it to use this state-level objective.

vmec_jax.api.rank_qi_seed_records(records: Iterable[dict[str, Any]], *, targets: QISeedSuitabilityTargets | None = None) → list[dict[str, Any]]: Annotate and rank QI seed records by smooth+legacy QI quality.

vmec_jax.api.read_booz_config(path: str | Path) → BoozConfig

Read optional &BOOZ_XFORM_JAX settings from a VMEC input file.

The group is intentionally separate from &INDATA so conventional VMEC inputs remain readable by VMEC2000 while vmec_jax can carry Boozer defaults in the same text file.

vmec_jax.api.read_indata(path: str | Path) → InData: Read &INDATA from a VMEC input file.

vmec_jax.api.read_wout(path: str | Path) → WoutData: Read a subset of wout_*.nc needed for regression comparisons.

vmec_jax.api.redl_current_derivative_update(*, s: Any, jdotB_redl: Any, bdotb: Any, dpsi_ds: Any, dpds: Any | None = None, previous_current: Any | None = None, policy: Literal['low_beta', 'lagged_pressure'] = 'low_beta') → Any

Return a VMEC current-derivative update I'(s).

low_beta applies the common approximation in which the pressure-gradient correction is neglected. lagged_pressure evaluates the pressure-gradient term using the previous equilibrium/current profile:

\[I'_{k+1} = R_k - \mu_0 I_k p'_k/\langle B^2\rangle_k.\]

vmec_jax.api.redl_current_integrating_factor_update(*, s: Any, jdotB_redl: Any, bdotb: Any, dpsi_ds: Any, dpds: Any) → dict[str, Any]

Solve the Redl/VMEC current ODE with an integrating factor.

The equation is

\[I'(s) + a(s) I(s) = R(s),\qquad a(s)=\mu_0 p'(s)/\langle B^2\rangle.\]

With I(0)=0,

\[I(s) = \exp[-A(s)] \int_0^s \exp[A(t)] R(t)\,dt, \quad A(s)=\int_0^s a(t)\,dt.\]

The derivative returned is computed from the ODE, I' = R - a I.

vmec_jax.api.redl_current_rhs(*, jdotB_redl: Any, bdotb: Any, dpsi_ds: Any) → Any: Return the Redl source term for the VMEC current ODE.

\[R(s) = 2\pi \psi'(s)\,\langle J.B\rangle_\mathrm{Redl}/\langle B^2\rangle.\]

vmec_jax.api.repeated_stage_modes(*, max_mode: int, use_mode_continuation: bool, continuation_nfev: int, repeats: int = 5) → list[int]

Same-mode repeated continuation used by the QI example.

Unlike qs_stage_modes(), repeated same-mode continuation has no lower-mode stages, so a zero continuation_nfev should not disable the repeated max-mode sequence.

vmec_jax.api.residual_scalars_from_state(*, state, static, indata, signgs: int, wout=None, use_vmec_synthesis: bool = True)

Compute VMEC-style invariant residual scalars (fsqr/fsqz/fsql) from a state.

This uses the residual pipeline:: bcovar -> forces -> tomnsps -> getfsq

and is intentionally input-only: flux profiles and pressure are derived from indata rather than a reference wout.

vmec_jax.api.resolve_boozmn_path(*, source_path: Path, outdir: Path | None = None, output: Path | None = None) → Path: Return the default boozmn_*.nc path for a VMEC input or WOUT file.

vmec_jax.api.run_booz_xform(wout_path: str | Path, *, output_path: str | Path | None = None, outdir: str | Path | None = None, mbooz: int = 32, nbooz: int = 32, surfaces: tuple[float, ...] | None = None, jit: bool = False, verbose: bool = True) → Path: Run booz_xform_jax from a WOUT file and write boozmn_*.nc.

vmec_jax.api.run_boundary_reference_preconditioner(input_file, output_dir, config, *, ctx: QIOptimizationContext | None = None): Scan same-NFP reference-family boundary jumps and return the selected input.

vmec_jax.api.run_fixed_boundary(input_path: str | Path, *, solver: str = 'vmec2000_iter', solver_mode: str | None = None, max_iter: int | object = <object object>, step_size: float | object = <object object>, history_size: int = 10, gn_damping: float | None = None, gn_cg_tol: float | None = None, gn_cg_maxiter: int = 80, use_initial_guess: bool = False, vmec_project: bool = True, use_restart_triggers: bool | None = None, vmecpp_restart: bool = False, use_direct_fallback: bool | None = None, multigrid: bool | None = None, multigrid_use_input_niter: bool = True, verbose: bool = True, jit_forces: bool | str = True, jit_precompile: bool | None = None, use_scan: bool | None = None, performance_mode: bool = True, scan_wout_corrector: bool | None = None, stage_transition_heuristic: bool | None = None, stage_transition_factor: float = 50.0, stage_transition_scale: float = 0.5, grid=None, ns_override: int | None = None, restart_state: any | None = None, restart_wout_path: str | Path | None = None, restart_solver_state: dict | None = None, cli_fixed_boundary_mode: bool = False, solver_device: str | None = None, external_field_provider_kind: str | None = None, external_field_provider_static: Any = None, external_field_provider_params: Any = None, free_boundary_activate_fsq: float | None = None, limit_update_rms: bool | None = None, _auto_cli_fixed_boundary_mode: bool = True, _solver_device_context_active: bool = False)

vmec_jax.api.run_free_boundary(input_path: str | Path, **kwargs)

Run a free-boundary vmec_jax solve.

Parameters:

input_path – Path to a VMEC-style input.* file with LFREEB = T and a valid MGRID_FILE entry.
**kwargs – Forwarded directly to run_fixed_boundary(). Common options include max_iter, verbose, use_initial_guess, vmec_project, solver_mode, jit_forces, and limit_update_rms.

Returns:

Run container with the parsed input, static data, final state, and solver diagnostics.

Return type:

FixedBoundaryRun

Raises:

ValueError – If the input deck is not a free-boundary case.

Notes

This wrapper intentionally shares the internal implementation with run_fixed_boundary(). The only behavioral difference is that run_free_boundary validates the mode up front, which makes scripts and examples clearer and avoids silently running the wrong branch.

vmec_jax.api.run_qi_stage_policy(active_input_file, output_dir, *, solve_qi_stage, make_qi_problem, boundary_reference_preconditioner, mirror_ramp_stages, ctx: QIOptimizationContext | None = None)

Run the guarded staged QI policy and return (result, promotion_log).

The public script still defines the objectives and the solve function. This helper only handles repeated staged promotion, exact diagnostics, and checkpoint bookkeeping.

vmec_jax.api.run_target_helicity_seed_preconditioner(input_file, output_dir, config, *, ctx: QIOptimizationContext | None = None)

Insert deterministic 1e-5 target-helicity modes before local QI solves.

The source VMEC input is left untouched. Existing nonzero coefficients are preserved by default, so reviewed QI reference inputs are not perturbed while circular/minimal seeds get a reproducible non-axisymmetric derivative seed.

vmec_jax.api.save_optimization_result(result: FixedBoundaryOptimizationResult, *, output_dir: str | Path | None = None, paths: OptimizationOutputPaths | None = None) → OptimizationOutputPaths

Save initial/final inputs, wouts, and history from a solve result.

The examples use this for the mechanical file writes only. Diagnostics, plotting, and any extra exports should remain explicit in the user script.

vmec_jax.api.save_raw_seed_initial_artifacts(input_file, input_out, wout_out, *, ctx: QIOptimizationContext | None = None): Save the unpreconditioned VMEC input deck and its solved WOUT.

vmec_jax.api.signgs_from_sqrtg(sqrtg, *, axis_index: int = 1) → int

Infer VMEC’s signgs (+1 or -1) from a signed Jacobian array.

We intentionally compute this outside of any jitted objective; signgs should be treated as a fixed convention, not something to differentiate through.

vmec_jax.api.simple_omnigenity_seed_indata(indata, *, max_mode: int, include: Sequence[str] = ('rc', 'zs'), fix: Sequence[str] = ('rc00',), perturbation: float = 1e-05, r0: float | None = None, rbc01: float | None = None, zbs01: float | None = None)

Return an input deck with a simple deterministic omnigenity seed boundary.

The seed keeps only the near-circular base shape RBC(0,0), RBC(0,1), and ZBS(0,1) from the source deck unless explicit values are supplied. Every other active optimizable boundary coefficient with max(abs(m), abs(n)) <= max_mode is set to a deterministic ±perturbation value. This avoids exactly-zero Jacobian columns when examples start far from a QA/QH/QP/QI warm-start boundary.

vmec_jax.api.state_from_wout(wout: WoutData) → jax.tree_util.register_pytree_node_class

Build a VMECState from wout Fourier coefficients.

Notes

VMEC’s wout files do not store the internal lambda coefficients in the same units VMEC uses in bcovar / totzsps.

In wrout.f VMEC writes (schematically, for each radial surface js):

lmns_wout(:,js) = (lmns_internal(:,js) / phipf(js)) * lamscale

to preserve an older output convention.

For parity-style kernels that re-use VMEC’s bcovar formulas, we therefore invert this scaling when constructing the state:

lmns_internal = lmns_wout * phipf / lamscale

vmec_jax.api.surface_rz_from_wout_physical(wout, *, theta: ndarray, phi: ndarray, s_index: int, nyq: bool = False) → tuple[ndarray, ndarray]

Return R,Z on a surface using physical toroidal angle phi.

This matches vmecPlot2’s convention: phase = m*theta - xn*phi, where xn already includes the nfp factor.

vmec_jax.api.surface_stack(wout, *, theta: ndarray, zeta_list: Iterable[float], s_index: int) → tuple[ndarray, ndarray]: Stack R,Z slices for multiple zeta values.

vmec_jax.api.target_helicity_seed_terms(*, max_mode, amplitude=1e-05): Return deterministic low-order perturbations for circular/minimal seeds.

vmec_jax.api.vmec_current_profile_from_bootstrap_update(*, s: Any, current_derivative: Any | None = None, current: Any | None = None, signgs: int, pcurr_type: str = 'cubic_spline_ip') → dict[str, Any]: Return VMEC current-profile arrays and CURTOR from bootstrap data.

vmec_jax.api.vmecplot2_bmag_grid(wout, *, s_index: int, ntheta: int = 30, nzeta: int = 65, zeta_max: float | None = None) → tuple[ndarray, ndarray, ndarray]

Return (theta, zeta, B) on a grid matching vmecPlot2.py defaults.

Parameters:: zeta_max – Upper bound of the toroidal angle range. Defaults to 2π (full toroidal circle). Pass 2π/nfp to restrict to one field period.

vmec_jax.api.vmecplot2_cross_section_indices(nzeta: int) → ndarray: Indices used by vmecPlot2 for cross sections (0,2,4,6).

vmec_jax.api.vmecplot2_lcfs_3d_grid(wout, *, s_index: int, ntheta: int = 80, nzeta: int | None = None) → tuple[ndarray, ndarray, ndarray, ndarray, ndarray]: Return (theta, phi, R, Z, B) grids matching vmecPlot2.py 3D defaults.

vmec_jax.api.vmecplot2_surface_grid(wout, *, s_index: int, ntheta: int = 200, nzeta: int = 8) → tuple[ndarray, ndarray, ndarray, ndarray]: Return (theta, zeta, R, Z) grids matching vmecPlot2.py surface defaults.

vmec_jax.api.wout_from_fixed_boundary_run(run: FixedBoundaryRun, *, include_fsq: bool = True, path: str | Path | None = None, fast_bcovar: bool | None = None) → WoutData

Build a minimal VMEC-style WoutData from a fixed-boundary run.

This is the in-memory counterpart to write_wout_from_fixed_boundary_run(). The fast bcovar path is the default; set fast_bcovar=False to force the legacy force-kernel output path for debugging.

vmec_jax.api.write_axisym_overview(case: str, *, outdir: str | Path | None = None) → Path: Write a quick axisymmetric overview plot from bundled reference wout.

vmec_jax.api.write_bmag_parity_figures(*, input_path: str | Path, wout_path: str | Path, outdir: str | Path | None = None, s_index: int | None = None) → Path: Write magnetic-field-magnitude parity figures for wout vs vmec_jax.

vmec_jax.api.write_bsub_parity_figures(*, input_path: str | Path, wout_path: str | Path, outdir: str | Path | None = None, s_index: int | None = None) → Path: Write (bsubu, bsubv) parity figures comparing wout vs vmec_jax geometry.

vmec_jax.api.write_bsup_parity_figures(*, input_path: str | Path, wout_path: str | Path, outdir: str | Path | None = None, s_index: int | None = None) → Path: Write (bsupu, bsupv) parity figures comparing wout vs vmec_jax geometry.

vmec_jax.api.write_indata(path: str | Path, indata: InData) → None

Write a VMEC &INDATA namelist block.

The output is intended for reproducible round-tripping through read_indata(), not for preserving the exact original whitespace or comments from the source file.

vmec_jax.api.write_wout_from_fixed_boundary_run(path: str | Path, run: FixedBoundaryRun, *, include_fsq: bool = True, fast_bcovar: bool | None = None): Write a minimal VMEC-style wout_*.nc from a fixed-boundary run.

vmec_jax.api.zeta_grid(nzeta: int, *, endpoint: bool = False) → ndarray: Uniform zeta grid over one field period.

vmec_jax.api.zeta_grid_field_period(nzeta: int, *, nfp: int) → ndarray: Uniform zeta grid over one field period (0..2π/nfp).