ESSOS

Stellarator coil optimization and particle tracing in JAX, on CPU or GPU. Coil objectives and adaptive tracing support automatic differentiation.

pip install essos

What it does

• Automatic differentiation. jax.grad works through coil geometry, Biot-Savart fields and adaptive tracing.
• Coil optimization. Fit coils to a plasma boundary under length, curvature, separation, coil-surface-distance and force constraints, with least_squares, an augmented Lagrangian, multi-objective (Pareto) search, or stochastic optimization over Gaussian coil perturbations. Coils can also target a near-axis field, particle confinement, or a finite-beta boundary (VMEX, or the plasma field from virtual_casing_jax in optimize_coils_virtual_casing.py).
• Particle tracing. Guiding-centre and full-orbit (Boris) models, with Monte Carlo collisions on background species with density and temperature profiles (NRL Coulomb logarithm per species pair), electric fields and alpha-loss diagnostics.
• Boozer-coordinate tracing. Guiding-centre diagnostics from the Boozer |B| spectrum and flux functions iota, G and I, with fixed-step (RK4, Dopri5, Dopri8) or adaptive Dormand-Prince stepping and optional collisions.
• Field-line tracing. Adaptive, arclength and toroidal-angle models, with Poincare sections.
• Fields. Biot-Savart from coils, VMEC equilibria (stellarator-symmetric or not, analytic derivatives, optional mode_tolerance truncation), near-axis expansions, and CombinedField to trace a sum of fields as one.
• Interpolated fields. InterpolatedField(source, R=..., Z=..., nr, nz, nphi, nfp) tabulates any Cartesian field (coils, dipoles, MGRID, combined) on a cylindrical grid and evaluates B, dB/dX and guiding-centre terms from a tricubic spline; it saves to .npz or MGRID and works under jit and grad. See trace_particles_interpolated_dipoles.py.
• VMEC MGRID. Export coil fields and load MGRID files as JAX-compatible three-dimensional magnetic fields.
• Surfaces. Fourier-represented toroidal surfaces, from a VMEC wout or built directly.
• Parallel by default. ESSOS asks JAX for one device per CPU core and shards particles over every CPU core or GPU (ESSOS_CPU_DEVICES overrides the count; devices= picks). Compiled traces are cached on disk (ESSOS_XLA_CACHE, default ~/.cache/essos/xla).
• Checked against SIMSOPT. examples/comparisons_simsopt compares coils, surfaces, VMEC import, field lines, guiding-centre and full-orbit tracing, and losses with SIMSOPT on the same inputs. Alpha losses are also compared with SIMSOPT, FIRM3D and SIMPLE (below).

Coil optimization

Fit coils to a VMEC boundary, trading normal-field error against coil length and curvature. Losses compose with +, and L.grad is the exact gradient.

from essos.coils import Coils, CreateEquallySpacedCurves
from essos.fields import BiotSavart
from essos.losses import custom_loss
from essos.surfaces import BdotN_over_B, SurfaceRZFourier
from scipy.optimize import least_squares

surface = SurfaceRZFourier.from_wout_file("wout.nc", s=1, ntheta=30, nphi=30,
                                          range_torus="half period")
coils = Coils(curves=CreateEquallySpacedCurves(3, 3, 10.0, 5.6, n_segments=45,
                                               nfp=2, stellsym=True),
              currents=[1.0] * 3)

L = (custom_loss(lambda field, surface: abs(BdotN_over_B(surface, field)).sum(),
                 "field", surface=surface)
     + custom_loss(lambda field: (field.coils.length - 32.0).clip(0).mean(), "field")
     + custom_loss(lambda field: (field.coils.curvature - 0.1).clip(0).mean(), "field"))
L.dependencies = {"field": BiotSavart(coils)}

result = least_squares(L, L.starting_dofs, L.grad, max_nfev=400)
optimized = L.dofs_to_pytree(result.x)["field"].coils

More in examples/coil_optimization.

Field-line tracing

import jax.numpy as jnp
from essos.coils import Coils
from essos.dynamics import Tracing
from essos.fields import BiotSavart

field = BiotSavart(Coils.from_json("coils.json"))
R0 = jnp.linspace(1.21, 1.40, 8)
seeds = jnp.array([R0, jnp.zeros_like(R0), jnp.zeros_like(R0)]).T

tracing = Tracing(field=field, model="FieldLineAdaptative", initial_conditions=seeds,
                  maxtime=8000, times_to_trace=40000, atol=1e-8, rtol=1e-8)
tracing.poincare_plot(shifts=[0.0])

More in examples/fieldline_tracing.

VMEC MGRID fields

Run examples/simple_examples/mgrid_from_coils.py to export the included Landreman-Paul QA coils, load the file as a magnetic field, and compare its interpolated field with direct Biot-Savart. It prints the differences and saves a plot; the maximum difference at its four sample points is about 0.2%. Accuracy depends on grid resolution and distance from the coils, so compare in your region of interest before using a grid.

The cylindrical grid covers one field period in phi. The field repeats across periods, and R and Z queries are clamped to the grid bounds, so set the bounds to cover the region you intend to evaluate.

Particle tracing

from essos.constants import ALPHA_PARTICLE_CHARGE, ALPHA_PARTICLE_MASS, ONE_EV
from essos.dynamics import Particles, Tracing

particles = Particles(initial_xyz=seeds, mass=ALPHA_PARTICLE_MASS,
                      charge=ALPHA_PARTICLE_CHARGE, energy=4000 * ONE_EV)
tracing = Tracing(field=field, model="GuidingCenterAdaptative", particles=particles,
                  maxtime=1e-4, times_to_trace=800, atol=1e-7, rtol=1e-7)
tracing.plot()
print(tracing.loss_fractions)

More in examples/particle_tracing, including full-orbit, collisional and electric-field variants.

Tracing to the wall

Inside the last closed flux surface (LCFS), guiding centres move in VMEC coordinates. An orbit that crosses the LCFS continues in Cartesian coordinates in a field you supply for the outside (coils, an mgrid, or a VMEX VmecExtender) until it strikes the wall or comes back in. Below, 3 keV protons in a VMEX free-boundary equilibrium of the Landreman-Paul QA coils, with a wall 3 cm outside the LCFS.

vmec = vmex.essos_vmec_field(wout)                     # or Vmec("wout_*.nc")
outside = vmex.VmecExtender.from_wout(wout, external_field=coil_grid)
tracing = Tracing(field=vmec, model="GuidingCenterAdaptative", particles=particles,
                  maxtime=4e-5, exterior_field=outside, wall=0.03)  # gap [m] or a surface
print(tracing.wall_hits.mean(), tracing.wall_positions[tracing.wall_hits])

The full script is trace_particles_vmec_to_wall.py; the figure comes from python docs/make_readme_figures.py wall.

Boozer-coordinate tracing

Transform a VMEC equilibrium with booz_xform_jax, then trace guiding centres in an axis-regular chart with optional Monte Carlo collisions. method="rk4" (default), "dopri5" or "dopri8" take fixed steps; "adaptive" (Dopri5) and "adaptive8" (Dopri8) control each particle's step to tolerance. This tracer provides forward diagnostics.

Install the transform extra with pip install 'essos[booz]'.

import numpy as np
from netCDF4 import Dataset
from booz_xform_jax import Booz_xform
from essos.boozer import BoozerField, trace_boozer
from essos.constants import ALPHA_PARTICLE_CHARGE, ALPHA_PARTICLE_MASS, ONE_EV

booz = Booz_xform(verbose=0, mboz=32, nboz=32)
booz.read_wout("wout.nc", flux=False)
booz.run()
with Dataset("wout.nc") as wout:
    phi_edge = float(wout.variables["phi"][-1])  # boundary toroidal flux
field = BoozerField.from_booz_xform(booz, psi0=-phi_edge / (2 * np.pi))

n = 1000
speed = np.sqrt(2 * 3.5e6 * ONE_EV / ALPHA_PARTICLE_MASS)
result = trace_boozer(field, s=np.full(n, 0.25), theta=np.random.uniform(0, 2*np.pi, n),
                      zeta=np.random.uniform(0, 2*np.pi, n), pitch=np.random.uniform(-1, 1, n),
                      speed=speed, mass=ALPHA_PARTICLE_MASS, charge=ALPHA_PARTICLE_CHARGE,
                      tmax=1e-2)      # species=BackgroundSpecies(...) adds collisions
print(result.lost.mean(), result.loss_fractions())

psi0 is minus VMEC's boundary toroidal flux over 2 pi. A particle is lost at s = 1; result.loss_times and result.states hold when and where.

Boozer tracing evaluates |B| = sum [bc_mn(s) cos(m theta - n zeta) + bs_mn(s) sin(m theta - n zeta)] and its derivatives, together with the flux functions G, I and iota.

Comparison with other codes

256 collisionless 3.5 MeV alphas per equilibrium, started at s = 0.5 with uniform Boozer angles and pitch, traced for 10 ms and lost at s = 1, in the Landreman-Paul QA reactor-scale wout, the ARIES-CS n3are wout, and the W7-X standard configuration scaled to the same minor radius (1.70 m) and volume-averaged |B| (5.86 T). Apple M4 (load about 3 from other processes), 8 CPU cores per code; compilation and field set-up excluded. DESC traces the first 32 alphas only, on VMECIO.load(wout, L=24, M=8, N=8) with the half-mesh lambda fix of PlasmaControl/DESC#2352, and its times were taken with other DESC runs sharing the machine. Run with examples/cross_code_benchmark/benchmark.py.

The fate agreement is the share of particles lost or confined as in SIMPLE and FIRM3D, over the particles on which those two agree. The bar plot scales DESC's time from 32 to 256 particles. ESSOS, SIMSOPT and FIRM3D use the same 32x32 booz_xform spectrum (ESSOS keeps modes above its default 1e-6 of B00 and, above 48 modes, evaluates them from a table in the Boozer angles); SIMPLE reads the wout directly and DESC fits it. ± is the binomial standard error.

Tracing notes

• VMEC magnetic axis. VMEC guiding centres are integrated in sqrt(s) (cos theta, sin theta), which is regular on the axis, so orbits cross it; trajectories are returned in (s, theta, phi, ...) with theta in [0, 2 pi).
• LCFS and wall. A VMEC guiding-centre trace stops at the exact LCFS crossing (tracing.lcfs_times, lcfs_positions, lcfs_energies), found by root finding. Pass exterior_field= (ESSOS coils, a VMEX VmecExtender or MgridField, or any batched xyz -> B) and wall= (a SurfaceClassifier or a signed-distance callable, positive inside) to follow the orbit outside in Cartesian coordinates until it strikes the wall (wall_times, wall_positions, wall_energies) or comes back inside, where it returns to flux coordinates (returns). tracing.status gives each outcome (essos.dynamics.VMEC_STATUS); trajectories_xyz and region hold the whole orbit. A non-finite step ends that orbit as failed, with its time. Loss fractions use the exact wall, or LCFS, times. Collisions are switched off outside the LCFS. Supplying condition replaces these events; it is evaluated on (s, theta, phi, ...). See examples/particle_tracing/trace_particles_vmec_to_wall.py.
• Stopping coil-field traces. Pass stopping_criteria=LevelsetStoppingCriterion(...) to end Cartesian traces once they leave a prescribed distance from a surface, and read the per-line mask from tracing.boundary_hits.
• Step budget. max_steps (default 1_000_000) bounds every Diffrax solve, so a trace that cannot finish returns instead of running unbounded.
• Progress bars are off by default; pass progress=True when interactive. With particle_batch_size=n particles are traced n at a time and the bar counts completed particles (batches are bypassed under jax.jit/grad).
• Model choice. FieldLineArclength traces a fixed physical length so rescaling B does not change the run; FieldLineToroidal sets coverage directly in toroidal angle for flux-coordinate fields.

Optional: near-axis fields

Near-axis expansions need pyQSC_JAX, which is not on PyPI and therefore cannot be a declared dependency:

pip install git+https://github.com/uwplasma/pyQSC_JAX.git

Everything else works without it; the near-axis entry points raise with this command if it is missing.

Examples

Testing

pytest

License

MIT, see LICENSE.

Acknowledgments

Developed by the UWPlasma group at the University of Wisconsin-Madison, with support from Simons Foundation grant 560651.