Differentiable alpha loss fraction and in-memory Vmec construction

[rogeriojorge]

Makes the alpha-particle loss fraction usable as a differentiable objective driven by a VMEC equilibrium whose spectral coefficients are live JAX arrays. This unblocks vmex --trace wout_XXX.nc and boundary optimization against alpha losses in vmex, which holds the wout arrays in memory as tracers.

Gap 1: Vmec could only be built from a netCDF file

Vmec.__init__ opened the wout with Dataset(...). A caller already holding bmnc, rmnc, bsup*, ... as traced JAX arrays had to write a file first, which severs the gradient path.

API

Vmec.from_arrays(nfp, ns, bmnc, xm, xn, rmnc, zmns, bsubsmns, bsubumnc, bsubvmnc,
                 bsupumnc, bsupvmnc, gmnc, xm_nyq, xn_nyq, Aminor_p,
                 ntheta=50, nphi=50, close=True, range_torus='full torus')

Arguments carry the wout variable names and are stored as passed — no np.asarray round trip — so tracers reach B, AbsB and the traced trajectories. nfp, ns and the mode numbers set array shapes and must be concrete. VMEC_WOUT_ARRAYS names the array arguments for callers that want to forward them in bulk. Vmec(wout_filename) reads the same values from the file and hands them to the same code path, so the file route is unchanged.

Verified against examples/input_files/wout_LandremanPaul2021_QA_reactorScale_lowres.nc, comparing a file-built Vmec with one rebuilt from its own arrays:

quantity	result
B at 3 points	bit-identical
AbsB at 3 points	bit-identical
surface.gamma	bit-identical
(nfp, ns, mpol, ntor)	(2, 50, 4, 5) both
traced trajectories, 4 alphas, tmax=4e-6	bit-identical

Gap 2: the loss fraction had an identically zero gradient

Tracing.loss_fraction builds trajectories_r >= r_max, then argmax, bincount, cumsum. All piecewise-constant in the trajectories, so jax.grad returns exactly zero and the objective is only reachable by gradient-free optimizers.

Measured with 8 alphas launched at s=0.85, tmax=5e-6, 30 saves, atol=rtol=1e-5, differentiating with respect to a scale on every B-related coefficient (bmnc, bsub*, bsup*) passed through Vmec.from_arrays, at r_max=0.88:

	value	jax.grad
loss_fraction (exact)	0.2500	0.0
soft_loss_fraction, width=0.05	0.371216	-5.671810e-02
soft_loss_fraction, width=0.02	0.306422	-1.341416e-01
soft_loss_fraction, width=0.005	0.256971	-8.715592e-02

The surrogate

Tracing.soft_loss_fraction(r_max=0.99, width=0.02)

For each particle i with radial coordinate r_i(t):

r_soft_i = sum_t r_i(t) * softmax_t(r_i(t) / width)
surrogate = mean_i sigmoid((r_soft_i - r_max) / width)

r_soft is the softmax-weighted mean over time: a soft maximum that spreads the gradient over every sample near the radial excursion peak. jnp.max is subdifferentiable but routes the gradient through the single peak sample, which is far too sparse to optimize against. It is also preferred over a plain width * logsumexp(r/width), which carries a width * log(n_times) offset fixed by how densely the trajectories are saved rather than by the orbit; the softmax-weighted mean is exact for a flat profile and its bias on a peaked one is O(width) with no n_times dependence. Non-finite saves past a terminating event are treated as having reached the boundary.

The single width keyword sets both the soft-maximum temperature and the crossing sharpness, both in units of s.

Convergence in value

Same trace as above, r_max=0.88, exact loss_fraction = 0.2500:

width	surrogate	exact	|diff|
0.1	0.421020	0.2500	0.171020
0.05	0.371216	0.2500	0.121216
0.02	0.306422	0.2500	0.056422
0.01	0.267036	0.2500	0.017036
0.005	0.256971	0.2500	0.006971
0.002	0.251475	0.2500	0.001475
0.001	0.250049	0.2500	0.000049

On four prescribed sinusoidal radial excursions peaking at s = 0.70, 0.93, 0.86, 0.99 with r_max=0.9, exact loss fraction 0.5, the per-particle gradient at width=0.01 is [3.1e-08, 1.757, 0.271, 5.1e-03] against [0, 0, 0, 0] for the exact one. The near-miss particle at 0.86 carries real weight, which is what lets an optimizer see a boundary before losses appear.

Objective

loss_soft_lost_fraction(field, particles, timestep=1.e-8, maxtime=1e-5, num_steps=300,
                        trace_tolerance=1e-5, model='GuidingCenterAdaptative',
                        boundary=None, r_max=0.99, width=0.02)

Placed with the other particle confinement losses in essos/objective_functions.py.

The exact diagnostic is unchanged

Tracing.loss_fraction is untouched. It remains the physically correct reported quantity and still returns (loss_fractions, total_particles_lost, lost_times) exactly as before. soft_loss_fraction sits alongside it and is only for driving gradients.

Tests

Added to the existing files, no new test module:

• tests/test_fields.py — from_arrays reproduces the file-built field; jax.grad flows through from_arrays into AbsB.
• tests/test_dynamics.py — surrogate converges to the exact value as width shrinks; exact gradient is zero where the surrogate's is not; from_arrays traces identically to the file route; soft_loss_fraction differentiates VMEC coefficients end to end.
• tests/test_objective_functions.py — loss_soft_lost_fraction in the mocked particle-loss test.

pytest tests/test_fields.py tests/test_dynamics.py tests/test_objective_functions.py -q
38 passed in 43.30s

pytest tests/ -q
85 passed, 1 xfailed in 49.59s

[codecov]

Codecov Report

✅ All modified and coverable lines are covered by tests.

Files with missing lines	Coverage Δ
essos/dynamics.py	46.81% <100.00%> (+3.10%)	⬆️
essos/fields.py	75.39% <100.00%> (+37.28%)	⬆️
essos/objective_functions.py	98.75% <100.00%> (+0.01%)	⬆️

... and 1 file with indirect coverage changes

[rogeriojorge]

@eduardolneto @EstevaoMGomes I'll be asking for your reviews in a few ESSOS PRs, this one is the first one, as it will unblock a VMEX PR and merge https://github.com/uwplasma/VMEX to do stellarator optimization

[Copilot]

Pull request overview

Adds in-memory VMEC construction from live JAX arrays and a differentiable soft alpha-particle loss objective while preserving the exact diagnostic.

Changes:

• Adds Vmec.from_arrays and shared VMEC array metadata.
• Adds Tracing.soft_loss_fraction and loss_soft_lost_fraction.
• Adds reconstruction, differentiation, and objective integration tests.

Review findings:

• essos/dynamics.py:1326 — Moderate (2 votes): non-finite solver failures may be misclassified as particle losses.
• essos/fields.py:229 — Moderate (3 votes): concrete NumPy/JAX integer metadata should be normalized to Python integers.

Reviewed changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 2 comments.

Show a summary per file

File	Description
tests/test_objective_functions.py	Tests objective integration.
tests/test_fields.py	Tests VMEC reconstruction and gradient flow.
tests/test_dynamics.py	Tests surrogate convergence and end-to-end differentiation.
essos/objective_functions.py	Exposes the differentiable loss objective.
essos/fields.py	Supports in-memory VMEC construction; see finding above.
essos/dynamics.py	Implements the differentiable loss surrogate; see finding above.

Suppressed comments (4)

essos/dynamics.py:1313

• width is used as a divisor in both the softmax temperature and sigmoid scale, so width=0 produces NaNs and a negative width reverses the surrogate. Reject non-positive widths before this computation so an invalid optimization hyperparameter cannot silently produce an invalid objective.

    def soft_loss_fraction(self, r_max=0.99, width=0.02):

essos/dynamics.py:1327

• This method assumes trajectories[:, :, 0] is VMEC's flux coordinate s, but the public Tracing API also supports Cartesian guiding-center/full-orbit traces and boundary=SurfaceClassifier losses. For those calls the first component is x, while the exact diagnostic uses boundary distance, so soft_loss_fraction silently optimizes sigmoid(x) and ignores the supplied boundary. Add a boundary-distance surrogate or reject unsupported field/model combinations instead of returning an incorrect objective.

        trajectories_r = self.trajectories[:, :, 0]
        trajectories_r = jnp.where(jnp.isfinite(trajectories_r), trajectories_r, 1.0)
        r_soft = jnp.sum(trajectories_r * jax.nn.softmax(trajectories_r / width, axis=1), axis=1)

essos/dynamics.py:1326

• Please add a regression test for the non-finite radial-save branch. VMEC boundary-event traces can retain non-finite post-event saves, and this replacement is what gives the surrogate a loss signal for them; the current tests only exercise finite prescribed trajectories, so this behavior could regress unnoticed.

        trajectories_r = jnp.where(jnp.isfinite(trajectories_r), trajectories_r, 1.0)

essos/fields.py:239

• This constructor creates a fresh Vmec instance for every call, but all VMEC field kernels (B, AbsB, etc.) are JIT-compiled with self as a static argument. Since the default object hash/equality are identity-based, rebuilding the field from live arrays on each objective/gradient evaluation gives each kernel a new static cache key and forces recompilation (and retains the old compiled executable). In a vmex optimization loop this can dominate runtime and grow memory; use a stable static wrapper or make the coefficient state dynamic (for example, a pytree-based field) so coefficient updates reuse the compiled kernels.

        self = cls.__new__(cls)

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.