vmec_jax.optimization_workflow

Reusable teaching workflow helpers for fixed-boundary optimizations.

The functions in this module are intentionally small building blocks for the standalone examples. Users should still construct objective lists explicitly in the scripts, as in SIMSOPT’s LeastSquaresProblem.from_tuples workflow, but the mechanical VMEC/JAX stage setup, mode continuation, saving, and plotting live here instead of being repeated in every example.

Functions

abs_mean_iota_ceiling_objective(maximum[, ...])	Smooth upper-bound penalty enforcing abs(mean_iota) <= maximum.
abs_mean_iota_floor_objective(target[, ...])	Smooth lower-bound penalty enforcing abs(mean_iota) >= target.
annotate_qs_final_history(final_result, *, label)	Attach final optimization metadata without writing artifacts.
aspect_objective(target[, weight])	Aspect-ratio least-squares objective.
boozer_b_target_from_wout(wout_path, *, ...)	Return Boozer |B| target spectra from a VMEC wout file.
build_fixed_boundary_objective_stage(cfg, ...)	Build one VMEC/JAX optimization stage from an objective list.
build_quasi_isodynamic_objective_stage(cfg, ...)	Build one QI stage while sharing one Boozer transform across QI terms.
combine_qs_stage_histories(*, label, ...)	Merge per-stage histories into one optimization history.
describe_boundary_mode_limits(stage_mode)	Return a compact label for a boundary-mode stage descriptor.
interpolate_indata_boundary(seed_indata, ...)	Interpolate selected VMEC boundary Fourier coefficients.
least_squares_solve(vmec, problem, *, ...[, ...])	Solve a SIMSOPT-style fixed-boundary least-squares problem.
lgradb_objective(*, threshold[, weight, ...])	Differentiable minimum-L_grad_B penalty objective.
mean_iota(ctx, state)	Mean rotational transform on full-mesh surfaces, excluding the axis.
mean_iota_objective(target[, weight])	Mean full-mesh rotational-transform objective.
normalize_boundary_mode_limits(stage_mode)	Normalize an int/tuple/dict stage descriptor into mode limits.
objectives_track_iota(objectives[, target_iota])	Return true when optimization history should record mean iota.
optimization_output_paths(output_dir)	Return the canonical final-artifact paths for an optimization run.
prepare_simple_omnigenity_seed_input(...[, ...])	Write and return a simple omnigenity seed input path when enabled.
print_qs_final_summary(result, *[, ...])	Print the final scalar diagnostics from an optimization result.
print_qs_problem_summary(*, method, ...)	Print the problem summary used by the standalone examples.
qi_boozer_b_target_objective(*, target_bmnc)	Boozer |B| spectrum target evaluated on the shared QI field.
qi_lgradb_objective(*, threshold[, weight, ...])	QI field-quality term penalizing small local L_grad_B.
qi_max_elongation_constraint(*, threshold[, ...])	Signed LCFS elongation constraint max_elongation - threshold <= 0.
qi_max_elongation_objective(*, threshold[, ...])	Boundary elongation upper-bound objective.
qi_mirror_ratio_constraint(*, threshold[, ...])	Signed mirror-ratio constraint mirror_ratio - threshold <= 0.
qi_mirror_ratio_objective(*, threshold[, ...])	Mirror-ratio upper-bound objective evaluated from Boozer |B| modes.
qi_residual_ceiling_objective(*, maximum[, ...])	Soft-wall objective that penalizes QI residuals above maximum.
qs_stage_budget(*, stage_mode, max_mode, ...)	Outer residual/Jacobian budget for one fixed-boundary stage.
qs_stage_modes(*, max_mode, ...)	Repeated mode-continuation sequence used by the example scripts.
quasi_isodynamic_field_objective([weight, ...])	Smooth QI residual term from quasi_isodynamic_residual_from_state.
quasisymmetry_objective(*, helicity_m, ...)	Quasisymmetry residual objective for QA, QH, or QP.
rebuild_for_optimization_resolution(indata, ...)	Set VMEC spectral resolution to at least max(min_vmec_mode, max_mode+2).
repeated_stage_modes(*, max_mode, ...[, repeats])	Same-mode repeated continuation used by the QI example.
residuals_from_objectives(objectives, ctx)	Create the state residual callback consumed by FixedBoundaryExactOptimizer.
run_fixed_boundary_objective_optimization(*, ...)	Run a fixed-boundary objective list through one or more mode stages.
run_quasi_isodynamic_objective_optimization(*, ...)	Run a QI objective list through repeated or direct mode stages.
save_optimization_result(result, *[, ...])	Save initial/final inputs, wouts, and history from a solve result.
save_qs_final_outputs(*, output_dir, ...[, ...])	Save initial/final inputs, wouts, and history.
save_qs_stage_artifacts(*, stage_dir, ...[, ...])	Save stage input files and optionally wout files.
simple_omnigenity_seed_indata(indata, *, ...)	Return an input deck with a simple deterministic omnigenity seed boundary.
write_qi_workflow_stage_checkpoint(*, ...)	Write a resumable checkpoint after one QI continuation stage.

Classes

AbsMeanIotaCeiling(maximum, *[, softness])	Smooth upper-bound objective for abs(mean_iota).
AbsMeanIotaFloor(target, *[, softness])	Smooth lower-bound objective for abs(mean_iota).
AspectRatio()	Aspect-ratio objective object.
AugmentedLagrangianConstraint(objective[, ...])	Wrap a non-negative violation objective as an augmented-Lagrangian term.
BDotB(*[, surfaces, normalize, mmax_force, ...])	VMEC bdotb profile objective from the differentiable JXBFORCE path.
BDotGradV(*[, surfaces, normalize, ...])	VMEC bdotgradv profile objective from the differentiable JXBFORCE path.
BVector(*[, s_index, normalize])	Cartesian magnetic-field vector objective on one radial surface.
BetaTotal()	Total-beta objective for finite-beta studies.
BoozerBTarget(*, target_bmnc[, target_bmns, ...])	Boozer |B| spectrum-matching objective for QI steering.
BoundaryModeLimits(mode[, max_m, max_n, label])	Boundary-parameter mode limits for one optimization stage.
DMerc(*[, minimum, softness, mmax_force, ...])	Smooth lower-bound objective for VMEC Mercier stability.
FixedBoundaryObjectiveStage(mode, ctx, ...)	Prepared optimizer and metadata for one active boundary-mode stage.
FixedBoundaryOptimizationResult(...)	Result returned by run_fixed_boundary_objective_optimization().
FixedBoundaryVMEC(input_file, cfg, indata, ...)	Small fixed-boundary optimization object used by the examples.
GlasserResistiveInterchange(*[, maximum, ...])	Smooth upper-bound objective for the Glasser resistive criterion.
JDotB(*[, surfaces, normalize, mmax_force, ...])	VMEC jdotb profile objective from the differentiable JXBFORCE path.
JVector(*[, surfaces, normalize, ...])	Flux-coordinate current-density vector objective from JXBFORCE channels.
LeastSquaresProblem(objective_terms, ...] =, ...)	Least-squares objective assembled from (function, target, weight) tuples.
LgradB(*, threshold[, s_index, ntheta, ...])	Minimum-L_grad_B penalty object usable in QS or QI examples.
MagneticWell(*[, minimum, softness])	Smooth lower-bound objective for the vacuum magnetic-well proxy.
MaxElongation(*, threshold[, ntheta, nphi, ...])	Maximum LCFS elongation penalty object for solved VMEC states.
MeanIota()	Mean rotational-transform objective object.
MirrorRatio(*, threshold[, surfaces, mboz, ...])	Maximum mirror-ratio penalty object for solved VMEC states.
ObjectiveTerm(name, evaluate, object], ...)	One weighted least-squares objective block.
OptimizationOutputPaths(initial_input, ...)	Canonical files written by fixed-boundary optimization examples.
QIObjectiveTerm(name, evaluate[, qi_options])	One field-quality objective that shares a Boozer/QI field evaluation.
QuasiIsodynamicOptions(surfaces[, mboz, ...])	Boozer/QI sampling options shared by QI objective terms.
QuasiIsodynamicResidual(options)	Smooth QI residual object using a shared Boozer field evaluation.
QuasiIsodynamicResidualCeiling(*, maximum[, ...])	Soft upper-bound objective for preserving a low-QI basin during cleanup.
QuasisymmetryRatioResidual(*, helicity_m, ...)	QS residual object for QA/QH/QP objectives.
RedlBootstrapMismatch(*, helicity_n, ...[, ...])	Redl bootstrap-current mismatch objective for finite-beta studies.
StageContext(static, indata, boundary_input, ...)	Objects needed by objective callbacks for one mode-continuation stage.
ToroidalCurrent(*[, surfaces, normalize, ...])	Integrated toroidal-current profile from VMEC's Mercier path.
ToroidalCurrentGradient(*[, surfaces, ...])	Radial derivative of ToroidalCurrent used by VMEC Mercier terms.
VMECMirrorRatio(*, threshold[, surfaces, ...])	Fast mirror-ratio penalty evaluated directly from VMEC |B|.
VolavgB()	Volume-averaged magnetic-field objective for finite-beta studies.