VMEC JAX workflow

This page is the concrete workflow contract for optional vmex -> booz_xform_jax -> dkx coupling. It does not make either geometry package a hard dependency. Default CI must still pass when both packages are absent.

The current public lane is a Boozer-spectrum proxy transport-objective gradient gate. It is not a claim of full VMEC-boundary-to-SFINCS kinetic transport gradients.

Preflight

Run the lightweight status scaffold first:

python examples/optimization/vmex_workflow_status.py --json

This reports shallow importability for vmex and booz_xform_jax, the no-overclaim gate, a no-optional-dependency Boozer-spectrum autodiff readiness gate, and the exact command for the optional proxy-gradient gate. It does not import either optional backend.

The status JSON also contains no_solve_provenance_gate. This gate is a machine-readable assertion that the workflow is still a proxy-gradient lane: kinetic_solve_executed is false, the differentiated object is the Boozer-spectrum transport-like scalar, and full VMEC-boundary-to-SFINCS kinetic gradients remain deferred. Both workflow examples use the same shared dkx.workflows.geometry_adapters gate, so the skip-safe status path and the file-backed proxy-gradient summary enforce the same scalar-contract boundary.

The same payload carries kinetic_transport_scalar_contract and a copy of its gate in no_solve_provenance_gate.kinetic_transport_scalar_contract_gate. This is the forward contract for a future VMEC/Boozer-to-kinetic-transport scalar. It lists required_kinetic_transport_scalar_stages in machine-readable form:

vmec_source,
vmec_equilibrium_or_wout,
boozer_transform,
sfincs_geometry_adapter,
kinetic_operator_assembly,
linear_kinetic_solve,
transport_scalar_reduction,
gradient_validation.

Each stage records its current public role, differentiability boundary, current status, and evidence required before the stage can support a full kinetic transport scalar. The current public scalar is explicitly boozer_spectrum_proxy_not_kinetic; kinetic_transport_scalar_claimed and kinetic_solve_executed must both remain false in this lane.

The existing end-to-end example has the same skip-safe backend contract:

python examples/autodiff/vmex_to_boozer_sfincs_pipeline.py --check-backends --json

In --check-backends mode, the JSON payload includes backend_readiness_gate. That gate is intentionally synthetic: it evaluates a small Boozer spectrum through a dkx proxy transport objective, checks the full spectral JAX gradient against centered finite differences, and checks a JVP against the gradient dot product. It is a local backend-readiness/sensitivity check for the downstream differentiable objective, not evidence that vmex or booz_xform_jax executed.

To persist a provenance record without running optional geometry code:

python examples/autodiff/vmex_to_boozer_sfincs_pipeline.py \
  --check-backends \
  --summary-json workflow-summary.json

Optional install pattern

Use editable installs for local research checkouts, or install equivalent packages into the active environment:

python -m pip install -e /path/to/vmex
python -m pip install -e /path/to/booz_xform_jax

The status scaffold should then report both optional backends as available.

Proxy-gradient gate

Run the documented file-backed workflow with an explicit wout file:

python examples/autodiff/vmex_to_boozer_sfincs_pipeline.py \
  --wout /path/to/wout_circular_tokamak.nc \
  --mboz 3 \
  --nboz 3 \
  --surface 0.5 \
  --steps 0 \
  --summary-json workflow-summary.json

The same input can be supplied as:

export DKX_VMEX_WOUT=/path/to/wout_circular_tokamak.nc
python examples/autodiff/vmex_to_boozer_sfincs_pipeline.py \
  --mboz 3 \
  --nboz 3 \
  --surface 0.5 \
  --steps 0

If vmex example decks are available, the example can build the VMEC-like object before the Boozer transform:

python examples/autodiff/vmex_to_boozer_sfincs_pipeline.py \
  --vmec-case circular_tokamak \
  --vmec-max-iter 1 \
  --steps 0

Differentiability contract

Differentiated in this lane:

scaled VMEC-like magnetic spectral arrays,
the booz_xform_jax transform,
dkx.workflows.geometry_adapters.boozer_spectrum_proxy_transport_objective.

Setup or provenance only, not differentiated:

VMEC file I/O,
vmex fixed-boundary setup used to produce a wout-like object,
dkx VMEC file adapters and scheme-5 parity readers.

Explicit non-claims:

no full VMEC-boundary-to-SFINCS kinetic transport gradients,
no gradient through the SFINCS kinetic transport solve in this lane,
no production solver dependency on vmex or booz_xform_jax.

Future kinetic scalar contract:

kinetic_transport_scalar_contract.required_stages is the authoritative list of stages that must be present before a VMEC/Boozer-to-SFINCS kinetic scalar can be claimed.
kinetic_transport_scalar_contract.current_public_scalar separates differentiated proxy stages from setup-only and not-covered stages.
kinetic_transport_scalar_contract.no_overclaim_gate.status must be "pass" in default CI. It fails if the proxy lane claims a kinetic solve, requires optional geometry packages in default CI, drops a required stage, or promotes the full kinetic scalar while deferred stages remain.

Gates

Run the VMEC/Boozer workflow and adapter gates:

python -m pytest tests/test_vmex_workflow.py tests/test_jax_geometry_adapters.py -q

These gates also include a no-solve invariant check that the normalized Boozer proxy transport objective is unchanged by global \(|B|\) spectrum scaling and is exactly zero, with zero gradient, for a constant-\(B\) spectrum.

Optional ecosystem benchmark CLIs are not part of the stable VMEC workflow. The stable examples use the in-tree JAX differentiable geometry and solver paths; external solver-library adoption studies are handled on research branches until they satisfy the documented accuracy, runtime, memory, and dependency gates.

The optional VMEC/Boozer integration tests use pytest.importorskip or a skip-status payload. Missing optional packages therefore record a skipped lane, not a failed default installation.

In a file-backed optional run, the written workflow-summary.json must also show no_solve_provenance_gate.status == "pass". For that path the gate requires explicit provenance fields for the source wout or in-memory VMEC object, selected surface, Boozer resolution, objective grid shape, and spectral scale. This lets downstream users audit what was differentiated without mistaking the proxy scalar for a SFINCS kinetic transport solve. The same JSON must show no_solve_provenance_gate.required_kinetic_transport_scalar_stages and kinetic_transport_scalar_contract.no_overclaim_gate.status == "pass" so the future kinetic-scalar lane cannot silently lose required stages or boundaries.

Promotion rule

This lane is complete enough for documented research use when:

--check-backends --json returns a valid workflow contract,
backend_readiness_gate.status == "pass" in the no-optional-dependency preflight/backend-contract payload,
no_solve_provenance_gate.status == "pass" and no_solve_provenance_gate.kinetic_solve_executed == false in both preflight and file-backed summary payloads,
the proxy-gradient gate writes a summary JSON with numerical_gradient_gate.status == "pass" on at least one explicit wout fixture when optional packages are installed,
default tests pass without requiring optional geometry packages,
docs continue to state the differentiability boundary and non-claims exactly.

The next blocker for a production optimization claim is a pure-JAX scheme-5 geometry path and a transport-objective gradient gate. Until that exists, keep this page scoped to the Boozer-spectrum proxy workflow.