vmec_jax.quasi_isodynamic

Differentiable quasi-isodynamic residuals.

The reference QI scripts in omnigenity_optimization measure whether the bounce width and field-line profile of magnetic wells are independent of field-line label. Their implementation uses branch finding and splines, which are useful diagnostics but not a smooth objective for JAX autodiff. This module keeps the same target and replaces hard branch extraction with smooth level-set widths:

width(alpha, B_j) = integral H(B_j - B(alpha, phi)) dphi

where H is a logistic approximation to the step function. A QI surface has widths that are independent of alpha for every level B_j. The residual can also include branch-width and profile-consistency terms. The default weights are calibrated so the smooth metric ranks the seed, the published omnigenity_optimization QI result, and current vmec_jax candidates in the same order as the branch-squash/stretch/shuffle diagnostic from the reference Goodman et al. omnigenity workflow, while remaining differentiable.

Functions

boozer_output_from_state(*, state, static, ...)	Evaluate differentiable Boozer |B| output directly from a VMEC state.
boundary_max_elongation_from_rz(R, Z, *[, ...])	Estimate maximum LCFS cross-section elongation from R(theta, phi), Z.
jax_sigmoid(x)	Stable logistic used by QI objectives and their adjoints.
lgradb_from_state(*, state, static, indata, ...)	Evaluate the magnetic-gradient scale length on a VMEC surface.
lgradb_penalty_from_state(*, state, static, ...)	Penalize short magnetic-gradient scale length on a VMEC surface.
max_elongation_penalty_from_state(*, state, ...)	Penalize excessive LCFS elongation from a solved VMEC state.
mirror_ratio_penalty_from_boozer_modes(*, ...)	Penalize the maximum mirror ratio from Boozer |B| modes.
mirror_ratio_penalty_from_boozer_output(booz, *)	Evaluate mirror_ratio_penalty_from_boozer_modes() from Boozer output.
mirror_ratio_penalty_from_state(*, state, ...)	Penalize mirror ratio from a solved VMEC state without any QI options.
quasi_isodynamic_residual_from_boozer_modes(*, ...)	Evaluate a smooth QI residual from Boozer |B| Fourier modes.
quasi_isodynamic_residual_from_boozer_output(booz, *)	Evaluate the smooth QI residual from a booz_xform_jax output dict.
quasi_isodynamic_residual_from_state(*, ...)	Evaluate a differentiable QI residual directly from a solved VMEC state.