ESSOS

Stellarator coil and particle optimization in JAX. Everything ESSOS computes — coil geometry, Biot-Savart fields, guiding-centre orbits, field lines — is differentiable end to end and runs on CPU or GPU, so a design objective and its gradient come from the same code.

pip install essos

What it does

• Differentiable throughout. jax.grad works through coil geometry, the field, and the traced orbits, so objectives compose without finite differences.
• Coil optimization. Fit coils to a plasma boundary under length, curvature, separation and coil-surface-distance constraints.
• Particle tracing. Guiding-centre and full-orbit (Boris) models, with collisions, electric fields and alpha-loss diagnostics.
• Field-line tracing. Adaptive, arclength and toroidal-angle models, with Poincare sections.
• Fields. Biot-Savart from coils, VMEC equilibria, near-axis expansions, and CombinedField to trace a sum of fields as one.
• Surfaces. Fourier-represented toroidal surfaces, from a VMEC wout or built directly.
• Parallel. JAX sharding across the visible devices; pass devices= to pick.

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.

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 notes

• VMEC magnetic axis. The poloidal angle is undefined on axis, so a VMEC guiding-centre trace stops at s <= axis_threshold (default 1e-6) and reports it through tracing.axis_hits and tracing.total_particles_unresolved, separately from a loss at s >= 1. This is a numerical safeguard, not a continuation through the axis; a trajectory that must cross it needs a regular chart or a full-orbit handoff. Supplying condition replaces the automatic axis and boundary events.
• 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.
• 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.

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.