ADEPT

ADEPT is an Automatic Differentiation Enabled Plasma Transport code.

The solvers are written in JAX and time-integrated with diffrax, so every simulation is differentiable end-to-end and runs on CPU or GPU. Runs are logged to MLflow.

Solvers

solver: key	Description	Docs
vlasov-1d	1D1V Vlasov-Poisson/Maxwell with Fokker-Planck collisions (vlasov-1d-iaw adds kinetic ions with Boltzmann electrons)	overview · config
vlasov-1d2v	1D2V in cylindrical velocity space, with a full-geometry Coulomb collision operator	overview · config
vlasov-2d	2D2V Vlasov-Maxwell	overview · config
vfp-1d	Vlasov-Fokker-Planck electron transport (spherical harmonic expansion)	overview · config
envelope-2d	2D laser-plasma envelope solver (TPD, SRS)	overview · config
lpse-2d	JAX port of C++ LPSE: TPD or SRS separately, or combined TPD + SRS; independent of envelope-2d	overview · config
spectrax-1d	1D Hermite-Fourier Vlasov-Maxwell (hermite-epw-1d and hermite-maxwell-1d are reduced entry points)	config
hermite-legendre-1d	1D1V mixed Hermite-Legendre electrostatic Vlasov-Poisson	config
pic-1d	1D electrostatic particle-in-cell	—
tf-1d	1D warm two-fluid with kinetic closures	overview · config

Example configs for each live in configs/.

Installation

ADEPT uses uv. Install it first:

curl -LsSf https://astral.sh/uv/install.sh | sh

As a dependency

uv add "adept @ git+https://github.com/ergodicio/adept.git"

or, into an existing environment:

uv pip install "git+https://github.com/ergodicio/adept.git"

For development

git clone https://github.com/ergodicio/adept.git
cd adept
uv sync --extra dev
uv run pre-commit install

uv sync creates .venv/ and installs from uv.lock, so you get the exact versions CI uses. Python 3.11+ is required. Prefix commands with uv run and you never need to activate the environment.

For an NVIDIA GPU, sync the gpu extra instead, which pulls jax[cuda12]:

uv sync --extra gpu

Usage

uv run run.py --cfg configs/vlasov-1d/epw

The --cfg path is given without the .yaml extension. This runs the simulation defined in the config and stores the output to an mlflow server.

Unless you have separately deployed an mlflow server somewhere (or set MLFLOW_TRACKING_URI), it simply writes files using the mlflow specification to the current working directory. To access and visualize the results, it is easiest to use the UI from the browser by running uv run mlflow ui from the same directory and opening http://localhost:5000.

There are other ways to use ADEPT, notably as part of a neural network training pipeline that leverages differentiable simulation. In reference [1], neural networks are trained to learn forcing functions that drive the system towards previously unseen behavior. In reference [2], neural networks are trained to help bridge the micro-macro physics gap in multiphysics simulations.

Tests

uv run pytest

Long-running parameter sweeps are marked slow and skipped by default; run them with uv run pytest -m slow.

Docs

https://ergodicio.github.io/adept/

Examples

https://github.com/ergodicio/adept-notebooks

Contributing guide

The contributing guide is in development but for now, just make an issue / pull request and we can go from there :)

Linting and formatting are handled by ruff via pre-commit. If you have run uv run pre-commit install, it runs on every commit; otherwise, run it manually before pushing:

uv run pre-commit run --all-files

If you change dependencies in pyproject.toml, refresh the lockfile with uv lock and commit it — CI installs with --frozen.

Citation

If you are using this package for your research, please cite the following

A. Joglekar and A. Thomas, “ADEPT - automatic differentiation enabled plasma transport,”
ICML - SynS & ML Workshop (https://syns-ml.github.io/2023/contributions/), 2023

References

[1] A. S. Joglekar & A. G. R. Thomas. "Unsupervised discovery of nonlinear plasma physics using differentiable kinetic simulations." J. Plasma Phys. 88, 905880608 (2022).

[2] A. S. Joglekar and A. G. R. Thomas, “Machine learning of hidden variables in multiscale fluid simulation,” Mach. Learn.: Sci. Technol., vol. 4, no. 3, p. 035049, Sep. 2023, doi: 10.1088/2632-2153/acf81a.