❯ SPECTRAX: Hermite-Fourier Vlasov-Maxwell solver in JAX for plasma physics simulations
- Overview
- Mathematical Method
- Features
- Project Structure
- Getting Started
- Input File Format
- Project Roadmap
- How to Cite
- Contributing
- License
- Acknowledgments
SPECTRAX is an open-source spectral kinetic plasma solver written in Python with the JAX ecosystem. It solves the collisionless Vlasov–Maxwell equations by evolving the Hermite–Fourier coefficients of the particle distribution function and the electromagnetic fields. The approach builds on the SpectralPlasmaSolver (SPS) algorithm developed at Los Alamos National Laboratory and described by Vencels et al. (2016) and Roytershteyn & Delzanno (2018), where the particle distribution is expanded in Hermite functions in velocity space and Fourier modes in configuration space. By performing a Hermite expansion in velocity space, the method naturally couples fluid and kinetic physics—the lowest‐order Hermite coefficients correspond to fluid moments and higher modes capture kinetic corrections.
SPECTRAX re‑implements this algorithm in a JAX framework. It uses just‑in‑time compilation to run efficiently on CPUs, GPUs or TPUs, adopts state‑of‑the‑art ODE solvers from the Diffrax library, and includes utilities for diagnostics and plotting. The code supports multi‑species plasmas and arbitrary spatial dimensionality (1D to 3D) and can serve as a test bed for studying kinetic instabilities, turbulence, and the transition between fluid and kinetic regimes.
The Hermite–Fourier spectral method replaces the Vlasov equation in the 6-dimensional phase-space with a hierarchy of coupled ordinary differential equations (ODEs) for the Hermite–Fourier coefficients. In the SPS formulation the velocity part of the distribution function is represented by Hermite functions while the spatial dependence is captured by Fourier modes. The expansion is truncated for closure. The resulting system of ODEs is integrated in time using solvers from the Diffrax library. Diffrax's implementation of the Dormand-Prince’s 8/7 method, Dopri8, proved to be notoriously fast and stable, and is set as the default solver.
-
JAX‑based spectral solver – all core operations are implemented in JAX and compiled with
jit, enabling execution on CPUs, GPUs or TPUs. -
Efficient time integration – SPECTRAX uses ODE solvers from the Diffrax library (e.g.,
Dopri5,Dopri8,Tsit5; a Diffrax-based, custom-made implicit midpoint solver is also available asImplicitMidpoint) to advance the Hermite–Fourier coefficients in time. Thesimulationfunction assembles the right‑hand‑side, applies a 2⁄3 de‑aliasing mask on Fourier modes in the nonlinear term, and integrates the system untilt_max, returning the time‑evolved coefficients. -
Multi‑species and multi‑dimensional – the code supports multiple particle species with distinct mass ratios, temperatures and drift velocities. Spatial dimensions are controlled via
Nx,Ny,Nz, and velocity Hermite orders viaNn,Nm,Np. -
Diagnostics – after each simulation the
diagnosticsfunction computes the Debye length, normalized wavenumber, kinetic energies of each species, electromagnetic energy and total energy and stores them in the output dictionary. -
Plotting utilities – the
plotfunction produces a multi‑panel figure showing energy evolution, relative energy error, density fluctuations and phase‑space distributions for each species. It reconstructs the distribution function by performing an inverse Fourier transform followed by an inverse Hermite transform. The phase‑space reconstruction uses theinverse_HF_transformfunction, which evaluates Hermite polynomials and sums over all modes. -
Flexible initialization – initial conditions may be provided through simple TOML files or directly in Python. The
load_parametersfunction reads a TOML file and merges it with sensible defaults that initialize a two‑stream instability. Users can also construct the spectral coefficients manually, as shown in the example scripts. -
Open source and extensible – SPECTRAX is released under the MIT License. Its modular structure allows researchers to experiment with new closures, collision operators or boundary conditions.
- Programming Language: Python
Besides Python, SPECTRAX has minimum requirements. These are stated in requirements.txt, and consist of the Python libraries jax, jax_tqdm and matplotlib.
SPECTRAX is a standard Python package that may be installed from a local checkout. The project depends
on jax, jaxlib, jax_tqdm, diffrax, orthax, and matplotlib.
- Clone this repository:
git clone https://github.com/uwplasma/SPECTRAX.git
cd SPECTRAX-
(Optional) create a virtual environment and activate it.
-
Install the package in editable mode:
pip install -r requirements.txt
pip install -e .JAX will automatically select the available hardware (CPU, GPU or TPU). For GPU support you may need the
appropriate CUDA-enabled version of jaxlib; consult the JAX installation guide.
SPECTRAX can be used either via a command‑line interface or directly as a Python module.
After installation, the spectrax CLI entry point is available. You can run an instance of the 1D two-stream instability by simply calling spectrax from the terminal.
spectraxTo run it with different input parameters, us a TOML file like those in the Examples directory.
spectrax example_input.tomlOther examples written in Python scripts, like those in the Examples directory, can be run from the terminal as follows:
python example_script.pyThe simulation function returns a dictionary containing the evolved Hermite coefficients Ck, electromagnetic coefficients Fk, time array, the input parameters and diagnostic quantities.
Run the test suite using the following command:
pytest .Input files are written in TOML and define both physical and solver parameters. Below is a summary of the most important keys. Keys absent from the file fall back to sensible defaults specified in the code.
| Parameter | Description |
|---|---|
Lx, Ly, Lz |
Domain lengths in the spatial directions (periodic boundary condition). |
mi_me |
Ion‑to‑electron mass ratio. |
Ti_Te |
Ion‑to‑electron temperature ratio. |
qs |
Array of species charges (electron negative). |
alpha_s |
Thermal scales (inverse thermal velocities) for each species. |
u_s |
Drift velocities for each species (packed as [u_x,u_y,u_z] ). |
Omega_cs |
Cyclotron frequencies for each species. |
nu |
Artificial collision frequency to damp recurrence. |
D |
Hyper‑diffusion coefficient (stabilization). |
t_max |
Final simulation time. |
nx, ny, nz |
Mode numbers used to seed sinusoidal perturbations. |
dn1, dn2, dE |
Amplitudes of initial density or field perturbations (see examples). |
ode_tolerance |
Relative/absolute tolerance for adaptive solvers. |
Nx, Ny, Nz |
Number of retained Fourier modes per spatial dimension. |
Nn, Nm, Np |
Number of Hermite modes per velocity dimension (x, y, z). |
Ns |
Number of species (two by default). |
timesteps |
Number of solution snapshots to store between t=0 and t_max. |
dt |
Initial step size provided to the ODE solver. |
solver |
Name of Diffrax solver (e.g., Tsit5, Dopri5, Dopri8, ImplicitMidpoint). |
adaptive_time_step |
Timestep adaptability (true by default). |
Many of these parameters can be arrays; for example alpha_s must contain three values per species (one
for each velocity dimension) and can be used to represent anisotropic plasmas. (11)
- 💬 Join the Discussions: Share your insights, provide feedback, or ask questions.
- 🐛 Report Issues: Submit bugs found or log feature requests for the
SPECTRAXproject. - 💡 Submit Pull Requests: Review open PRs, and submit your own PRs.
Contributing Guidelines
- Fork the Repository: Start by forking the project repository to your github account.
- Clone Locally: Clone the forked repository to your local machine using a git client.
git clone https://github.com/uwplasma/SPECTRAX
- Create a New Branch: Always work on a new branch, giving it a descriptive name.
git checkout -b new-feature-x
- Make Your Changes: Develop and test your changes locally.
- Commit Your Changes: Commit with a clear message describing your updates.
git commit -m 'Implemented new feature x.' - Push to github: Push the changes to your forked repository.
git push origin new-feature-x
- Submit a Pull Request: Create a PR against the original project repository. Clearly describe the changes and their motivations.
- Review: Once your PR is reviewed and approved, it will be merged into the main branch. Congratulations on your contribution!
Contributor Graph
This project is protected under the MIT License. For more details, refer to the LICENSE file.
- We acknowledge the help of the whole UWPlasma plasma group.