[codex] Add ESSOS VMEC-extender workflows

[rogeriojorge]

What changed

Adds the ESSOS side of the Phase 1 VMEC-extender workflow:

• VmecExtendedField adapter around virtual_casing_jax.VirtualCasingExteriorField.
• build_vmec_extended_field(...) builder for in-memory vmec_jax state/static and file-based input/wout workflows.
• essos-vmec-extender CLI with validate, grid, and trace subcommands.
• Trace sample export via --samples-out, with comparator-ready poincare_rphiz, line_id, section_phi, and approximate arc-length connection_lengths for FIELDLINES/TORLINES comparisons.
• Bundled examples for validation, grid export, tracing, and trace-sample export.
• Small CLI benchmark: benchmarks/vmec_extender_cli_benchmark.py.
• Tests for adapter behavior, CLI parsing/output, field-line sampling order/interpolation/validation, coil-coupled validation metrics, and real bundled VMEC validation.
• CI/ReadTheDocs hardening, docs pages, README examples, and AGENTS.md guardrails.
• JAX compatibility fix for coil_perturbation.matrix_sqrt_via_spectral, replacing deprecated jnp.clip(..., a_min=...) with min=....

Physics model and sign conventions

ESSOS consumes the virtual-casing field with the same convention as the core PR:

B_total_out(x) = B_coils(x) + B_internal^VC(x)

The CLI validate command reports VMEC boundary B . n, internal/external branch identity, coil plus internal-branch normal-field diagnostics, and external-branch versus coil-normal diagnostics. Coil-coupled metrics are reported as diagnostics unless the wout/coils pair is a validated matched free-boundary benchmark.

The trace sample export converts Cartesian ESSOS trajectories to (R, physical unwrapped phi, Z), samples either saved-time strides or common unwrapped-phi sections, and flattens section-major/line-minor to match the virtual_casing_jax FIELDLINES HDF5 loader contract.

Tests run

• pytest -q tests/test_vmec_extender_cli.py
  • 20 passed
• JAX_ENABLE_X64=1 JAX_PLATFORMS=cpu pytest --cov-branch --cov-report=term-missing --cov-report=xml --cov=essos
  • 116 passed, 1 skipped, repository coverage 53%, essos/vmec_extender_cli.py coverage 99%
• flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics --exclude=examples_old
  • 0
• python -m sphinx -b html docs docs/_build/html
  • passed with 20 existing AutoAPI/docstring warnings outside the new VMEC-extender pages
• PYTHONPATH=.:/Users/rogerio/local/virtual_casing_jax:/Users/rogerio/local/vmec_jax JAX_ENABLE_X64=1 JAX_PLATFORMS=cpu python examples/trace_fieldlines_vmec_extender.py
  • wrote raw trace and sample NPZ; trace shape [2, 4, 3], sample shape [8, 3]
• benchmarks/external/run_fieldline_compare.sh --reference /Users/rogerio/local/ESSOS-vmec-extender/examples/output/vmec_extender_trace_samples.npz --candidate /Users/rogerio/local/ESSOS-vmec-extender/examples/output/vmec_extender_trace_samples.npz --max-point-relative-l2 1e-14 --max-point-rms-distance 1e-14 --max-point-max-distance 1e-14 --max-connection-relative-l2 1e-14 --max-connection-max-abs 1e-14 --out /tmp/essos_vmec_extender_trace_sample_self_compare.json
  • passed; zero Poincare and connection-length error, 8 ordered samples and 2 connection lengths compared
• git diff --check
  • passed

GitHub Actions after the latest push:

• Docs: passed
• Python 3.10 Test: passed
• Python 3.11 Test: passed
• Python 3.12 Test: passed
• codecov/patch: passed

Physics validation results

Bundled low-resolution QA VMEC/coils validation:

• normalized VMEC B . n RMS: 4.515567565145951e-17
• branch identity relative L2: 7.2795609678179e-17
• branch identity max abs: 8.881784197001252e-16
• coil plus internal B . n RMS normalized: 0.3254277443426786
• external branch versus coil B . n relative L2: 174036.41201094288

The coil-coupled numbers are intentionally reported but not gated because the bundled coil/wout pair is not yet established as a matched free-boundary benchmark with validated unit/sign conventions.

The new trace-sample export was validated against the companion virtual_casing_jax FIELDLINES comparator using an ESSOS-generated candidate NPZ. The self-comparison passed all ordered point and connection-length thresholds at zero error.

Differentiability checks

ESSOS wraps the differentiable virtual_casing_jax field without adding discontinuous objectives. The CLI trace, Poincare paths, connection lengths, and wall-hit workflows remain diagnostics. Smooth differentiability checks live in the companion virtual_casing_jax PR.

Runtime/memory benchmarks

Earlier benchmark run: PYTHONPATH=/Users/rogerio/local/virtual_casing_jax:/Users/rogerio/local/vmec_jax:/Users/rogerio/local/ESSOS-vmec-extender python benchmarks/vmec_extender_cli_benchmark.py --output-dir /tmp/essos_vmec_extender_benchmark_clean:

• validate: 8.45 s
• grid export, 4x4x4: 1.59 s
• trace, 2 seeds x 4 samples: 1.97 s
• trace shape: [2, 4, 3]
• validation metrics match the values above

Latest bundled trace example with sample export completed in about 1.99 s locally.

Documentation/plots added

• docs/vmec_extender_cli.rst
• docs/vmec_extender_trace.rst
• docs/connection_length.rst
• README CLI/API examples
• Example scripts under examples/
• No publication gallery plots yet.

Remaining limitations

• This PR depends on the companion virtual_casing_jax VMEC-extender PR.
• ESSOS docs still have pre-existing AutoAPI/docstring warnings, so docs are buildable but not warning-clean repository-wide.
• Matched free-boundary coil validation and external-code benchmark gates are still pending.
• The trace export currently reports approximate Cartesian arc length along saved ESSOS trajectories; wall-limited production comparisons should replace or augment it with boundary-classifier wall-hit connection lengths.
• Connection-length maps and publication-gallery plots are documented as intended workflows but not fully implemented as acceptance-grade examples in this PR.
• BMW/vector-potential work remains Phase 2.

[codecov]

Codecov Report

❌ Patch coverage is 83.95415% with 56 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
essos/vmec_extender.py	47.52%	46 Missing and 7 partials ⚠️
essos/fields_vmec_extender.py	0.00%	2 Missing ⚠️
essos/vmec_extender_cli.py	99.59%	1 Missing ⚠️

Files with missing lines	Coverage Δ
essos/coil_perturbation.py	95.38% <100.00%> (ø)
essos/vmec_extender_cli.py	99.59% <99.59%> (ø)
essos/fields_vmec_extender.py	0.00% <0.00%> (ø)
essos/vmec_extender.py	47.52% <47.52%> (ø)

[rogeriojorge]

Coordinated local VMEC-extender validation after commit f485774 passed against local feature branches for ESSOS, virtual_casing_jax, and vmec_jax.\n\nCommands matched the feature-branch workflow intent:\n- pytest subset: tests/test_vmec_extended_field_adapter.py tests/test_vmec_extender_cli.py tests/test_vmec_extender_real_validation.py -> 25 passed\n- examples/validate_vmec_extender.py -> status ok, normalized B.n RMS 4.515567565145951e-17, branch identity relative L2 7.2795609678179e-17\n- examples/export_vmec_extender_grid.py -> grid shape [4, 4, 4]\n- examples/trace_fieldlines_vmec_extender.py -> trace shape [2, 4, 3], samples shape [8, 3]\n- benchmarks/vmec_extender_cli_benchmark.py -> validate 2.19s, grid 0.30s, trace 1.21s\n\nHosted workflow_dispatch for .github/workflows/vmec_extender.yml cannot be triggered yet because GitHub only dispatches workflows present on the default branch. Once this workflow lands on main, the same check can be run hosted with virtual_casing_jax_ref=feature/jax-vmec-extender and vmec_jax_ref=feature/jax-vmec-extender.