vmec_jax.qi_optimization

Support routines for staged bounded QI optimization.

The public example script keeps the user-facing workflow visible: configure inputs, construct objective tuples, call least_squares_solve, then save/plot outputs. This module holds the reusable seed-preconditioning, stage-promotion, and checkpoint helpers used by that script and sweep drivers.

Functions

apply_qi_example_cli_overrides(namespace[, argv])	Apply optional command-line overrides to a QI example namespace.
basin_prefilter_score(metrics, targets, config)	Rank prefilter candidates by QI/iota first, engineering second.
boundary_reference_preconditioner_score(...)	Rank reference-family candidates by gates first, then exact metrics.
boundary_reference_record_is_qi_safe(record, ...)	Return whether a preconditioner summary record satisfies safe gates.
configure(context)	Install script-level constants used by the staged helper routines.
diagnostic_float(record, key)	Return a scalar diagnostic value, or nan when it is unavailable.
engineering_promotion_score(record)	Rank already-gated candidates by QI first and mirror second.
jsonable(value)	Convert NumPy/JAX-like values into JSON-serializable containers.
make_basin_prefilter_options(config, *[, ctx])
make_qi_optimization_context([context, strict])	Build a typed staged-QI helper context from script constants.
materialize_qi_stage_inputs(...)	Write root-level stage input files from an optimization result.
promotion_score(record)	Lower score means a better exact-diagnostic QI candidate.
qi_diagnostics_for_result(stage_result, *, ...)
qi_diagnostics_for_run(run, *, ...[, ...])	Independent QI diagnostics for a raw fixed-boundary VMEC run.
qi_stage_modes(*, max_mode, ...[, repeats, ...])	Return the stage mode sequence for the QI example workflow.
run_basin_prefilter(input_file, output_dir, ...)	Run a bounded large-step prefilter and return the selected input deck.
run_boundary_reference_preconditioner(...[, ctx])	Scan same-NFP reference-family boundary jumps and return the selected input.
run_qi_stage_policy(active_input_file, ...)	Run the guarded staged QI policy and return (result, promotion_log).
run_target_helicity_seed_preconditioner(...)	Insert deterministic 1e-5 target-helicity modes before local QI solves.
save_raw_seed_initial_artifacts(input_file, ...)	Save the unpreconditioned VMEC input deck and its solved WOUT.
stage_modes_for(stage, *[, ctx])
stage_promotes_candidate(stage, promotion, ...)	Apply the script's staged promotion rule to exact diagnostics.
target_helicity_seed_terms(*, max_mode[, ...])	Return deterministic low-order perturbations for circular/minimal seeds.
write_qi_stage_checkpoint(stage_output_dir, ...)	Persist QI stage metrics before root finalization can be reached.

Classes

QIOptimizationContext(alpha, ...)

Explicit staged-QI controls shared by helper routines.