❯ ESSOS: e-Stellarator Simulation and Optimization Suite

Table of Contents

• Overview
• Features
• Project Structure
• Getting Started
  • Prerequisites
  • Installation
    • From PyPI
    • From Source
• Usage
  • VMEC magnetic-axis handling
• Testing
• Project Roadmap
• Contributing
• License
• Acknowledgments

Overview

ESSOS is an open-source project in Python that uses JAX to optimize stellarator coils. Optimization can be applied to several objectives, such as alpha particle confinement, plasma boundaries and magnetic field equilibria (including near-axis expansions). It leverages automatic differentiation and efficient numerical methods to streamline optimization efforts, creating a specialized and fast numerical tool for optimizing force-free stellarator equilibria. It is parallelized using JAX's sharding tools. It can be imported in a Python script using the essos package, or run directly in the command line as essos. To install it, use

pip install essos

Alternatively, you can download it and run the example scripts in the repository after downloading it as

git clone https://github.com/uwplasma/ESSOS
cd ESSOS
pip install .
python examples/trace_particles_from_coils.py

The project can be downloaded in its GitHub repository.

Features

• JAX Integration: Utilizes JAX for automatic differentiation and efficient numerical computations.
• Optimization: Implements optimization routines for stellarator coil design.
• Particle Tracing: Traces alpha particles in magnetic fields generated by coils.
• Fieldline Tracing: Traces magnetic field lines.
• Coils and Near-Axis Fields: Models and optimizes electromagnetic coils and near-axis magnetic fields.
• Stellarator Surfaces: Representation of toroidal surfaces using a spectral Fourier decomposition.

Project Structure

ESSOS/
├── essos/
│   ├── __init__.py
│   ├── __main__.py
│   ├── coils.py
│   ├── constants.py
│   ├── dynamics.py
│   ├── fields.py
│   ├── objective_functions.py
│   ├── optimization.py
│   ├── plot.py
│   └── surfaces.py
├── examples/
│   ├── create_stellarator_coils.py
│   ├── optimize_coils_and_nearaxis.py
│   ├── optimize_coils_for_nearaxis.py
│   ├── optimize_coils_particle_confinement_fullorbit.py
│   ├── optimize_coils_particle_confinement_guidingcenter.py
│   ├── optimize_coils_vmec_surface.py
│   ├── trace_fieldlines_coils.py
│   ├── trace_particles_coils_fullorbit.py
│   ├── trace_particles_coils_guidingcenter.py
│   ├── trace_particles_vmec.py
│   └── comparisons_SIMSOPT/
│   └── inputs/
│       ├── ESSOS_bio_savart_LandremanPaulQA.json
│       ├── SIMSOPT_bio_savart_LandremanPaulQA.json
│       ├── wout_n3are_R7.75B5.7.nc
│       └── wout_LandremanPaul2021_QA_reactorScale_lowres.nc
├── tests/
│   ├── test_coils.py
│   ├── test_constants.py
│   ├── test_dynamics.py
│   ├── test_fields.py
├── README.md
├── LICENSE.md
├── CODE_OF_CONDUCT.md
├── CONTRIBUTING.md
├── setup.py
├── pyproject.toml
└── requirements.txt

Getting Started

Prerequisites

• Python 3.8 or higher

Installation

From PyPI

To install ESSOS from PyPI, run:

pip install essos

From Source

To install ESSOS from source, clone the repository and install the package:

git clone https://github.com/uwplasma/ESSOS
cd ESSOS
pip install .

Usage

ESSOS can be run directly in the command line as essos, or by following one of the examples in the examples folder. For example, to trace particles in a magnetic field generated from coils, run:

python examples/trace_particles_coils_guidingcenter.py

VMEC magnetic-axis handling

VMEC guiding-center trajectories use flux coordinates (s, theta, phi, ...), where s is normalized toroidal flux. The poloidal angle is not defined on the magnetic axis, and continuing the VMEC representation to negative s can make a trace stiff or non-finite. ESSOS therefore stops a VMEC guiding-center trace when s <= axis_threshold (default 1e-6) and records that numerical termination separately from a loss at s >= 1:

tracing = Tracing(
    field=vmec,
    model="GuidingCenterAdaptative",
    particles=particles,
    axis_threshold=1e-6,
)

print(tracing.axis_hits)                  # one boolean per particle
print(tracing.boundary_hits)              # LCFS termination mask
print(tracing.total_particles_unresolved) # number of axis terminations

Post-event output slots for axis-terminated trajectories are filled with the last finite in-domain saved state so plotting and diagnostics do not receive Diffrax's inf padding or a saved state beyond the axis threshold. Axis terminations are not counted as particle losses. This safeguard is applied to fixed- and adaptive-step VMEC guiding-center models, including the collision/SDE variants. It is not applied to full-orbit/Boris models or coil-field traces, whose positions are Cartesian and can cross the physical magnetic axis without this flux-coordinate singularity. Supplying a custom condition replaces the automatic VMEC axis and boundary events.

This is a numerical safeguard, not a physical continuation through the axis. A trajectory that must continue across the axis requires a regular coordinate chart (for example, pseudo-Cartesian flux coordinates) or a Cartesian/full-orbit handoff.

Stopping coil-field traces

For Cartesian coil fields, stop trajectories after they leave a prescribed distance from a reference surface. This keeps lost or chaotic lines from dominating Poincare axes while retaining a per-line termination mask:

surface_distance = SurfaceClassifier(surface, padding=0.2)
escape = LevelsetStoppingCriterion(surface_distance, maximum_distance=0.1)
tracing = Tracing(field=coils, model="FieldLineAdaptative",
                  initial_conditions=seeds, stopping_criteria=escape)
print(tracing.boundary_hits)

Progress bars are off by default because Diffrax updates can overwhelm batch logs; pass progress=True when interactive per-solve progress is useful. Toroidal Poincare sections unwrap the Cartesian azimuth before extracting crossings, including the branch at phi=0. Tracing uses the largest compatible subset of visible devices; pass devices=jax.devices("gpu")[:1] (or another explicit list) to select them.

Use model="FieldLineArclength" for Cartesian magnetic fields when fields with different strengths should be traced for the same physical length. It has the same field-line topology as FieldLineAdaptative but avoids runtime changes caused only by rescaling B. For flux-coordinate fields with nonzero toroidal field, use model="FieldLineToroidal" to set the coverage directly in toroidal angle. The convenience API makes that choice explicit and reports compilation and integration progress:

from essos.dynamics import trace_field_lines
trace = trace_field_lines(vmec_field, seeds, toroidal_turns=200,
                          samples=10000, progress=True)
coil_trace = trace_field_lines(coil_field, xyz_seeds, length=1000.0,
                               stopping_criteria=escape, progress=True)

The finite-beta coil example optimize_coils_finite_beta_vmex.py uses a VMEX equilibrium and virtual casing to optimize the external ESSOS coil field directly, with exact reverse-mode gradients. BiotSavart.b_cyl(R, phi, Z) also supplies VMEX/NESTOR with the traceable coil field on a moving free boundary, retaining coil-shape and current derivatives without an mgrid file.

Testing

To run the tests, use pytest:

pytest .

Project Roadmap

• Allow several optimization algorithms
• Allow plotly and/or Mayavi visualization
• Add DESC and SPEC equilibria for tracing
• Add beam injection examples
• Add plotting for near-axis expansion
• Add particle collisions

Contributing

Contributions are welcome! Please fork the repository and submit a pull request. For major changes, please open an issue first to discuss what you would like to change.

• 💬 Join the Discussions: Share your insights, provide feedback, or ask questions.
• 🐛 Report Issues: Submit bugs found or log feature requests for the ESSOS project.
• 💡 Submit Pull Requests: Review open PRs, and submit your own PRs.

Contributing Guidelines

1. Fork the Repository: Start by forking the project repository to your github account.
2. Clone Locally: Clone the forked repository to your local machine using a git client.

   git clone https://github.com/uwplasma/ESSOS

3. Create a New Branch: Always work on a new branch, giving it a descriptive name.

   git checkout -b new-feature-x

4. Make Your Changes: Develop and test your changes locally.
5. Commit Your Changes: Commit with a clear message describing your updates.

   git commit -m 'Implemented new feature x.'

6. Push to github: Push the changes to your forked repository.

   git push origin new-feature-x

7. Submit a Pull Request: Create a PR against the original project repository. Clearly describe the changes and their motivations.
8. Review: Once your PR is reviewed and approved, it will be merged into the main branch. Congratulations on your contribution!

Contributor Graph

---

License

This project is protected under the MIT License. For more details, refer to the LICENSE file.

---

Acknowledgments

• This project was developed as part of the New Talents in Physics Fellowship, awarded by the Calouste Gulbenkian Foundation.
• We acknowledge the help of the whole UWPlasma plasma group.

---