[rogeriojorge]

What this adds

Field algebra on MagneticField.

• The base class now defaults to Cartesian coordinates (to_xyz is the identity and sqrtg = 1), so a Cartesian field only defines B.
• f + g, f - g, a * f and sum([...]) build a CombinedField. Nested sums are flattened, and the weights are static.
• B_xyz(xyz) evaluates any field at a Cartesian point.
• f.compare(g, points) returns |B_f - B_g| / |B_f| at the same physical points, even when the two fields use different coordinates.

ToroidalField(MagneticField) is the base for fields in (s, theta, phi).

• flux_coordinates (the inversion of to_xyz) moves up from Vmec.
• boundary_distance also moves up and is now built on to_xyz alone. Vmec keeps its analytic Fourier version of the boundary curve.
• Vmec subclasses it and drops its copies of the base-class curl_B, curl_b, kappa and dB_by_dX.
• The tracer checks is_toroidal(field) instead of isinstance(field, Vmec). A scaled field or a sum led by a toroidal field therefore keeps the axis-regular chart, the boundary stop and the exterior-field handling.

MRXField(ToroidalField) wraps a state of MRX (arXiv:2510.26986).

• MRX stores B as a 2-form: spline coefficients of Bhat = J DPhi^-1 B on a spline map Phi of the unit cube.
• With (s, theta, phi) = (r^2, 2 pi theta_l, 2 pi zeta / nfp) and scale factors c = (2r, 2pi, 2pi/nfp), the field is:
  • B = DPhi Bhat / J,
  • B^i = c_i Bhat^i / J,
  • B_i = (DPhi^T DPhi Bhat)_i / (J c_i),
  • sqrt(g) = J / (c_1 c_2 c_3).
• Field periods beyond the first are rotations of it.
• Evaluation is MRX's own local de Boor contraction, so there is no resampling and div B = 0 holds exactly.
• The field is a pytree and is differentiable with respect to its coefficients.
• MRXField.from_mrx(geometry, checkpoint=None) loads a relaxed checkpoint, or computes the vacuum (harmonic) field of a VMEC or DESC domain. MRX is imported lazily, so it is not a dependency.

PointDipoleField(MagneticField) is a sum of point dipoles (permanent magnets), with the moments as dofs. It is named so it can coexist with the permanent-magnet DipoleField of #32, which keeps orientations fixed and optimizes scalar strengths.

Examples.

• examples/particle_tracing/trace_particles_mrx.py: MRX vs VMEC, field lines and alpha particles. It also takes an MRX checkpoint, for example an island state.
• examples/coil_optimization/optimize_coils_mrx.py: coils that reproduce an MRX field.

essos/fields.py changes by +315/-99 lines: three new classes, the algebra, and the generic inversion and boundary distance, net of the code removed from Vmec.

Benchmarks

All on CPU (Apple M4), float64.

Consistency of MRXField

On the MRX vacuum field of Landreman-Paul QA, at 200 random points:

identity	max relative error
B = e_i B^i	1.4e-15
B_i = e_i . B	1.8e-15
sqrt(g) = det(e)	3.4e-14
B_xyz(to_xyz(p)) = B(p)	4.5e-14
div B (from sqrt(g) B^i)	1.9e-15

Here e = d to_xyz / d(s, theta, phi), computed by automatic differentiation.

MRX vs VMEX on Landreman-Paul QA

Setup:

• input.LandremanPaul2021_QA_lowres was solved with VMEX to fsqr = 9.9e-14.
• MRX vacuum fields were computed in the domain of that wout at p = 2.
• Fields were compared at 2000 random points of the MRX domain, after one scale factor, because MRX normalizes ||B||, not the flux.
• iota was traced in ESSOS through both fields (FieldLineToroidal, 60 field periods).

MRX mesh	rms, s >= 0.1	max, s >= 0.1	rms, s < 0.1	max |iota_MRX - iota_VMEX| (both traced), s >= 0.3
(8, 12, 12)	1.5e-2	6.3e-2	2.3e-2	1.3e-3
(12, 16, 16)	7.9e-3	3.2e-2	2.1e-2	3.0e-4
(16, 20, 20)	5.0e-3	2.1e-2	2.0e-2	8.9e-5
(20, 24, 24)	3.4e-3	1.6e-2	2.0e-2	4.9e-5

For reference, VMEC2000's wout of the same deck differs from VMEX by 3.5e-4 rms for s >= 0.1 and 1.0e-2 for s < 0.1. The difference near the axis is mostly VMEC's: its half-grid B is least accurate there. MRX's iota sign is opposite because its poloidal angle is right-handed.

Evaluation cost

QA mesh (12, 16, 16), p = 2.

	VMEC	MRX
B, batched (per point)	2.1 us	0.28 us
gc_quantities, batched (per point)	2.7 us	1.1 us
gc_quantities, single jitted call	38 us	71 us

Adaptive guiding-center tracing in an MRX field takes more steps than in VMEC. At p = 2, grad|B| is only continuous, so the steps shrink at every knot. At p = 3 the same orbits trace 3.2x faster than at p = 2 and conserve energy better (2e-7 vs 4e-7). from_mrx therefore builds vacuum fields at p = 3 by default.

Coils for an MRX field

optimize_coils_mrx.py uses the reactor-scale QA domain, 5.7 T on axis, and 4 coils per half period (16 coils in total). It matches the full MRX vector field on the boundary: in vacuum that fixes the field inside.

	value
optimization	1000 evaluations, 18.8 s
boundary rms of |dB|/|B|	1.5e-2
volume |B_coils - B_MRX|/|B_MRX|, mean / max	7.1e-3 / 2.8e-2

The coil field lines lie on the MRX flux surfaces in the Poincare section of the example.

DESC, VMEX and MRX on the same boundary (precise QA, vacuum)

Setup:

• DESC: the precise-QA equilibrium shipped with MRX (data/desc_LandremanPaul2021_QA.h5).
• VMEX: DESC's boundary, written out with VMECIO.write_vmec_input, solved to fsqr < 1e-13 at ns = 33, 65, 129 and 257.
• MRX: the vacuum field of the DESC domain at p = 3.
• B is compared at 2000 random points of DESC's domain, with DESC's own values as the reference. VMEX and MRX are evaluated through B_xyz.
• iota is traced in ESSOS for 600 field periods.

field	rms, rho >= 0.3	rms, rho < 0.3	max, rho < 0.3	max |iota - iota_DESC|, rho >= 0.3 / rho < 0.3
VMEX ns = 33	8.4e-4	2.9e-2	9.4e-2
VMEX ns = 65	3.3e-4	1.9e-2	9.7e-2
VMEX ns = 129	1.7e-4	1.5e-2	8.2e-2
VMEX ns = 257	1.4e-4	7.3e-3	5.9e-2	2.4e-4 / 8.8e-4
MRX (8, 12, 12)	3.8e-3	1.2e-3	3.5e-3
MRX (12, 16, 16)	1.3e-3	4.3e-4	1.3e-3
MRX (16, 20, 20)	5.5e-4	2.0e-4	7.6e-4	1.2e-4 / 3.7e-4

• Core. VMEX converges fastest in the core.
• Axis. Near the axis VMEX is still converging at ns = 257, and its axis iota moves with ns: -0.4210, -0.4195, -0.4186 and -0.4182 for ns = 33 to 257, against DESC's -0.4199. MRX's polar splines are smooth through the axis, and at (8, 12, 12) it already beats VMEX at ns = 257 there by 6x.
• Cross sections. DESC's surfaces, VMEX's surfaces and the Poincare points of MRX field lines overlay at phi = 0, pi/4 and pi/2.

Finite beta (2.5%, bootstrap current)

Setup:

• VMEX solves input.LandremanPaul2021_QA_beta2p5_bootstrap at ns = 33, 65, 129 and 257, all to fsqr < 1e-13.
• DESC re-solves the same equilibrium with the VMEX boundary, pressure and toroidal current (M = 12, N = 10). Its normalized force error is 1.4e-4.
• MRX relaxes the VMEX ns = 257 field: scripts/relax.py, (12, 16, 16) p = 3, 20 Newton steps. The force residual ends at 4.5e-9, div B at 1e-15, and the helicity drifts by 7e-6.
• The toroidal current I(s) of VMEX and MRX comes from Ampere's law around the poloidal loop at fixed s, evaluated through their ESSOS fields. J_phi comes from curl_B.

field	rms |B - B_DESC|/|B_DESC|, rho >= 0.3 / rho < 0.3	iota at rho = 0.1 (DESC 0.1659)	max |iota - iota_DESC|, rho >= 0.3	max |I - I_DESC| / I_edge
VMEX ns = 33	1.3e-3 / 1.6e-2	0.2001	9.6e-3	1.3e-3
VMEX ns = 65	1.2e-3 / 1.2e-2	0.1877	4.5e-3	8.8e-4
VMEX ns = 129	1.2e-3 / 7.4e-3	0.1763	1.3e-3	8.7e-4
VMEX ns = 257	1.2e-3 / 6.0e-3	0.1664	1.3e-3	8.8e-4
MRX (12, 16, 16)	2.3e-3 / 1.1e-3	0.1688	2.8e-3	8.2e-3

With the bootstrap current, VMEX's iota near the axis needs ns = 257 to reach DESC's: ns = 33 is 20% off at rho = 0.1. Near the axis, MRX's field is 5x closer to DESC's than VMEX's at ns = 257. The toroidal current profiles of all three overlay, and so do the J_phi contours on phi = 0.

Field lines and guiding centers through MRX islands

The README now shows six 10 keV protons started across the (6,1) chain, in the MRX island state and in VMEX li383 (ns = 65). In MRX their crossings of phi = 0 jump around the island chain; in VMEX each stays on its own nested ring. The movie is docs/readme_mrx_islands.mp4 (330 kB), made by docs/make_mrx_islands_movie.py.

The states are from MRX's scripts/tutorials/3_li383_newton.py and 4_li383_island_seed.py:

• the li383 Newton state ("nested"),
• the state relaxed after seeding the (5,1), (6,1) and (7,1) chains ("islands"),
• the VMEC li383 equilibrium they start from.

For each field:

• 51 field lines for 400 field periods, seeded across every chain,
• 128 guiding-center protons of 5 keV in a field scaled to 1 T on the axis ((1 / B_axis) * field), traced for 0.2 ms.

Run on CPU (Apple M4, 10 devices).

field	field lines	guiding centers	lost	max energy error, confined orbits
VMEC	6 s	35 s	10.9%	1.6e-2
MRX nested	122 s	2120 s	18.8%	8.5e-6
MRX islands	66 s	1709 s	17.2%	6.8e-6

• The Poincare sections show the seeded chains in the relaxed MRX state.
• The VMEC energy error comes from the wout, not the tracer. This li383 wout has ns = 16, and its |B| table differs from sqrt(B_i B^i) built from its component tables by 0.55% rms and 3% at most. The guiding-center equations mix the two tables. Every MRX quantity comes from one B.
• MRX evaluation costs more per step at p = 2 (these MRX states are p = 2).

Performance

Guiding centers stalled at VMEC grid points

_radial_interp was linear in s, so grad|B| jumped at every wout grid point. A guiding center whose radial drift changes sign at such a point chatters across it while the adaptive step shrinks to nothing, and runs until max_steps (1e6). Under vmap the whole batch waits for it. On li383 at 1 T, one 5 keV proton of 128 did this at s = 14/15. A 128-particle run then took more than 4 h on CPU and never finished on a GPU.

The modes are now interpolated with cubic Hermite polynomials and centred-difference slopes. The first interval stays linear, so the axis extrapolation is unchanged. Grid values are kept and the s derivative is continuous. The stalled orbit now finishes in under 2000 steps. All batches of the same 128 particles take 3-6 s, and 8192 particles on an A4000 run in 8 even batches of 36-66 s.

Compilation

Every Tracing compiles a new solve, about 4 s on a CPU and 20 s on a GPU. dynamics.py now enables XLA's persistent compilation cache:

• The directory is ESSOS_XLA_CACHE, default ~/.cache/essos/xla. An empty value disables the cache, and a directory the user already set is kept.
• On an A4000, repeated 64-particle calls drop from 23-26 s to 6-7 s.
• A new process starts at 18 s instead of 34 s.

Inverting to_xyz

B_xyz and compare invert to_xyz by Newton iterations. The 12 x 32 seed search and 40 damped steps cost 66x one evaluation of B. Newton now starts from the point's angle about the axis (measured from theta = 0, in the sense theta turns) and its distance from the axis relative to the boundary's. Twelve steps are enough.

On 2000 points from the axis to s = 1:

field	max residual	cost per point
VMEX QA	6.7e-16 m	92 us (was 656 us)
VMEC li383	1.1e-15 m	34 us
MRX li383 islands	1.1e-15 m	19 us

CPU cores

On a single CPU device, a vmapped trace runs its small kernels on about one core. essos/__init__.py now asks JAX for os.cpu_count() CPU devices, and Tracing already shards particles over devices:

• ESSOS_CPU_DEVICES overrides the count. An XLA_FLAGS or jax_num_cpu_devices setting already made by the user is kept.
• Tracing pads the batch to a multiple of the device count instead of falling back to fewer devices.

On an Apple M4 (10 cores), VMEC li383 guiding centers:

• 64 particles over 20 us trace in about 1 s instead of 8.5 s, with results equal to 7e-9.
• 512 particles run 6.6x faster.
• Field lines in a coil field run 2-3x faster. Biot-Savart already used several cores.

B_on_surface sharded every leaf of the surface pytree, and failed whenever one did not divide by the device count. It now shards only the points, and only when they divide evenly.

The examples set XLA_FLAGS to a fixed device count, usually 1, which turned the sharding off. They now use the default.

A differentiated trace (jax.grad through Tracing) stays on one device. Compiling a gradient through a while loop sharded over several CPU devices aborts in XLA with JAX 0.6, the newest JAX for Python 3.10. That was the first CI failure of this PR. I reproduced it on Linux with Python 3.10 and JAX 0.6.2, and confirmed the fix there.

Defaults on a GPU workstation

Two RTX A4000 and 36 CPU cores, with nothing configured: 1024 guiding centers in VMEC li383 over 100 us.

devices (default)	first call	second call (compiled)
2 GPUs (automatic)	30.6 s	8.4 s
1 GPU	43.8 s	29.5 s
36 CPU devices	369.5 s	350.6 s

The results agree to 4e-7, within the default solver tolerance. Two GPUs are 3.5x faster than one, because each shard stops at its own slowest orbit.

Tests

tests/test_fields.py adds tests for:

• the field algebra and compare,
• dipoles against the on-axis formula, plus curl, divergence and moment gradients,
• the generic boundary curve against the analytic VMEC series,
• MRXField on an analytic torus map, which needs no MRX install: the identities above, field-period rotation, the is_toroidal sum, and gradients in the coefficients.

Also tested: the radial interpolation keeps grid values and has a continuous s derivative (this test fails on the old linear interpolation). The suite has 323 passing tests. The two test_survivor_compaction_preserves_outputs[adaptive8-...] failures also fail on main on this machine; they are 1-ulp differences.

Not in this PR

Near-axis fields stay in pyQSC_JAX. Its near_axis gives B in Boozer coordinates, while its position map is in the axis's cylindrical angle, so making it a ToroidalField needs a root solve per point. Having it subclass an ESSOS class would also make pyQSC_JAX depend on ESSOS.

The benchmark scripts are not part of the PR. Two examples cover the API: trace_particles_mrx.py and optimize_coils_mrx.py.

[rogeriojorge]

Split into focused PRs, each reviewable on its own:

• Vmec: C1 radial interpolation (guiding centers stalled at grid points) #150 Vmec C1 radial interpolation (guiding-center stall)
• Use every CPU core and GPU by default; cache compiled traces #151 CPU/GPU defaults and compile cache
• Magnetic field classes: Cartesian defaults, field algebra, ToroidalField #152 Magnetic field classes (base of the next five)
• PointDipoleField: point dipoles with the moments as dofs #153 PointDipoleField
• MRXField: trace MRX states (vacuum or relaxed, with islands) #154 MRXField, examples and islands movie
• DescField: trace DESC equilibria as DESC writes them #155 DescField
• NearAxisField: trace pyQSC_JAX near-axis fields in Boozer coordinates #156 NearAxisField
• Full orbits in toroidal fields #157 Full orbits in toroidal fields