spectrax 0.2.1

❯ SPECTRAX: Hermite-Fourier Vlasov-Maxwell solver in JAX for plasma physics simulations

Table of Contents

• Table of Contents
• Overview
• Mathematical Method
• Features
• Getting Started
  • Prerequisites
  • Installation
  • Usage
    • Command‑line Interface
  • Testing
• Input File Format
• Contributing
• License
• Acknowledgments

Overview

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. This approach was previously implemented in SpectralPlasmaSolver (SPS), developed at Los Alamos National Laboratory, and described in Delzanno (2025), Vencels et al. (2016) and Roytershteyn & Delzanno (2018), where the one particle distribution is expanded in asymmetrically weighted (AW) Hermite functions in velocity space and Fourier modes in configuration space. By performing an AW 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 implements a Hermite–Fourier spectral approach for Vlasov–Maxwell simulations in a JAX framework. It uses just‑in‑time compilation to run efficiently on CPUs and GPUs, 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.

Mathematical Method

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 moments of the one-particle probability density function and the electric and magnetic fields. The Hermite-Fourier expansion is truncated for closure and a hypercollisional operator is introduced in velocity space to suppress recurrence. 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.

Features

• JAX‑based spectral solver – all core operations are implemented in JAX and compiled with jit, enabling efficient execution on CPUs and GPUs.

• 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 as ImplicitMidpoint) to advance the Hermite–Fourier coefficients in time. The simulation function assembles the right‑hand‑side, applies a 2⁄3 de‑aliasing mask on Fourier modes in the nonlinear term, and integrates the system until t_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. Resolution in spatial dimensions is controlled via Nx, Ny, Nz, and velocity Hermite orders via Nn, Nm, Np.

• Diagnostics – after each simulation the diagnostics function computes the Debye length, kinetic energies of each species, electromagnetic energy and total energy and stores them in the output dictionary.

• Plotting utilities – the plot function 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 the inverse_HF_transform function, which evaluates Hermite polynomials and sums over all modes. The phase-space plots assume a 1D simulation.

• Flexible initialization – simulation parameters may be provided through simple TOML files or directly in Python. The load_parameters function reads a TOML file and merges it with sensible defaults that initialize a two‑stream instability. Users can also initialize their own spectral coefficients, 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.

Getting Started

Prerequisites

• 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.

Installation

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.

1. Clone this repository:

git clone https://github.com/uwplasma/SPECTRAX.git
cd SPECTRAX

2. (Optional) create a virtual environment and activate it.

3. Install the package in editable mode:

pip install -r requirements.txt
pip install -e .

JAX will automatically select the available hardware (CPU or GPU). For GPU support you may need the appropriate CUDA-enabled version of jaxlib; consult the JAX installation guide.

Usage

SPECTRAX can be used either via a command‑line interface or directly as a Python module.

Command‑line Interface

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.

spectrax

To run it with different input parameters, us a TOML file like those in the Examples directory.

spectrax example_input.toml

Other examples written in Python scripts, like those in the Examples directory, can be executed from the terminal as follows:

python example_script.py

The simulation function returns a dictionary containing the evolved Hermite coefficients Ck, electromagnetic coefficients Fk, time array, the input parameters and diagnostic quantities.

Testing

Run the test suite using the following command:

pytest .

Input File Format

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.

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.

Contributing

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

License

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

Acknowledgments

• We acknowledge the help of the whole UWPlasma plasma group.