Vlasov 1D1V Solver

Example decks live in configs/vlasov-1d/. To run one:

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

The top-level solver: key selects this module — vlasov-1d, or vlasov-1d-iaw for the ion-acoustic turbulence variant.

Things You Might Care About

Infinite length (single mode) plasma waves — Landau damping, trapping
Finite length plasma waves — everything in 1. plus wavepackets
Wave dynamics on density gradients — 2. plus density gradients
Stimulated Raman Scattering — 3. plus light waves
Ion-acoustic turbulence driven by stochastic box-scale forcing (vlasov-1d-iaw)

Equations and Quantities

We solve the following coupled set of partial differential equations:

\[ \frac{\partial f}{\partial t} + v \frac{\partial f}{\partial x} + \frac{q}{m} (E + E_D) \frac{\partial f}{\partial v} = \nu \partial_v (v f + v_0^2 \partial_v f) \]

\[ \partial_x E = 1 - \int f \, dv \]

where \(f\) is the distribution function, \(E\) is the electric field, \(C(f)\) is the collision operator, \(q\) is the charge, \(m\) is the mass, and \(v\) is the velocity.

The distribution function is \(f = f(t, x, v)\) and the electric field is \(E = E(t, x)\).

Multispecies Support

The solver supports multiple particle species (e.g., electrons and ions) evolving self-consistently under a shared electric field. Each species can have:

Independent charge and mass (determining the species-specific \(q/m\) ratio)
Independent velocity grid parameters (vmax, nv)
Multiple density components

For multispecies simulations, the Poisson equation sums over all species:

\[ \partial_x E = \sum_s q_s \int f_s \, dv_s \]

See the Configuration Reference for details on configuring multispecies simulations.

These simulations can be initialized via perturbing the distribution function or the electric field. The electric field can be “driven” using \(E_D\) which is a user defined function of time and space.

Solver Options

As with all other solvers, the configuration is passed in via a yaml file. Below we describe the key solver options, and then link to the full configuration reference.

Velocity Advection

exponential - This solver (incorrectly) assumes periodic boundaries in the velocity direction and uses a direct exponential solve such that

\[ f^{n+1} = f^n \times \exp(A \cdot dt) \]

where \(A\) is the advection operator. This is a much faster solver than the cubic-spline solver, but is less accurate. Use this if you are confident that the distribution function will be well behaved in the tails.

cubic-spline - This is a semi-Lagrangian solver that uses a cubic-spline interpolator to advect the distribution function in velocity space. Use this if you have trouble with the exponential solver.

Spatial Advection

exponential - This is the only solver that is available. We only have periodic boundaries implemented in space (for the plasma) so this is perfectly fine. It is also very fast.

Field Solver

poisson - This is the standard spectral Poisson solver. This is the fastest and most accurate solver available.
hampere - This solver uses a Hamiltonian formulation of the Vlasov-Ampere system that conserves energy exactly. This is the 2nd most reliable solver.
ampere - This solver uses Ampere’s law to solve for the electric field.

Collisions

none - No collisions are included in the simulation.
lenard-bernstein - This solver uses the Lenard-Bernstein collision operator to include collisions in the simulation.
dougherty - This solver uses the Dougherty collision operator to include collisions in the simulation.

Boundary Conditions

Axis / quantity	Condition	Notes
\(x\) (distribution and electrostatic field)	Periodic	Spatial advection and the Poisson/Ampère solves are spectral, so periodicity is structural rather than a choice.
\(v\) (`edfdv: exponential`)	Periodic	An artifact of doing the velocity push spectrally. It wraps the forward tail onto the \(-v\) edge, so it is only safe when \(f \approx 0\) at both velocity edges.
\(v\) (`edfdv: cubic-spline`)	Semi-Lagrangian interpolation	Does not assume periodicity; use it when the tails are populated.
\(v\) (collision operator)	Zero-flux	Applied at both velocity edges, which is what makes the Fokker-Planck operators conserve density exactly.
\(x\) (transverse vector potential \(a\))	Absorbing, 2nd order	The transverse wave equation is solved by finite differences on a grid with two boundary cells, independently of the periodic plasma domain.

The velocity grid is uniform and cell-centered, and may be asymmetric (vmin != -vmax). Choose bounds wide enough that \(f \approx 0\) at both edges, or the periodic velocity push and the zero-flux collision stencil stop being accurate.

Forcing and Drivers

Four mechanisms can drive the system, all optional:

Block	What it does
`drivers.ex`	A prescribed longitudinal field \(E_D(x, t)\) added to the force term. Each pulse is a travelling wave \(a_0 \sin(k_0 x - (\omega_0 + \delta\omega_0) t)\) with independent tanh envelopes in space and time. It never enters the Poisson solve, so the self-consistent field-energy diagnostic excludes it.
`drivers.ex_stochastic`	Time-correlated forcing: a set of Fourier modes whose complex amplitudes evolve as independent Ornstein-Uhlenbeck processes with correlation time `tau` and a prescribed stationary RMS. This is the box-scale stirring used for ion-acoustic turbulence.
`drivers.ey`	A transverse EM driver. It sources the wave equation for the vector potential \(a\), and the plasma feels the resulting ponderomotive force \(-\tfrac{1}{2}\partial_x(a^2)\). Extended sources use \(S = -\omega^2 a_0\,\text{env}(x,t)\sin(kx - \omega t)\); point sources use a single-cell delta with amplitude scaled by the vacuum Green’s function, and radiate both ways — put them next to the absorbing boundary for a unidirectional wave.
`terms.krook`	Not a driver but a forcing term: relaxation toward a Maxwellian at rate \(\nu_K\), enveloped in space and time. Mostly useful as a hard thermalization layer at a boundary.

What Gets Saved

Every run writes to a temporary directory that ergoExo logs to MLflow as artifacts. Two trees:

binary/ — netCDF, one file per save stream:

File	Contents
`fields-shared-t=<t>.nc`	The EM fields on the requested \((t, x)\) grid
`fields-<species>-t=<t>.nc`	Per-species real-space moments (density, velocity, temperature, …)
`scalars-t=<t>.nc`	Scalar time series — field and kinetic energies, and the conservation diagnostics
`dist-<save_key>.nc`	The distribution function for each species save block you configured, on its requested axes

plots/ — PNGs generated from those datasets:

Directory	Contents
`plots/fields/`	Space-time plots of the shared EM fields, with `logplots/` and `lineouts/` variants
`plots/fields/<species>/`	The same for each species’ moments, again with `logplots/` and `lineouts/`
`plots/scalars/` and `plots/scalars/<species>/`	Scalar time-series plots
`plots/dists/<save_key>/`	Distribution-function snapshots, one subdirectory per distribution save block

With solver: vlasov-1d-iaw you additionally get plots/iaw/density_spectrum.png, nk_spectrogram.png, and phase_space_dfx.png, plus binary/nk.nc.

Timing metrics (run_time, postprocess_time_min, total_time) go to MLflow as metrics rather than files. Which streams exist, and at what cadence, is entirely determined by the save block — see the Configuration Reference.

Running on Multiple GPUs

grid.parallel splits the phase-space pushes across every GPU that the process can see. It is a deliberately naive scheme: one process, one node, no distributed memory. It buys throughput on a distribution function that already fits in a single GPU’s memory; it does not let you run a bigger one.

Set it to the list of axes to split over:

grid:
  nx: 17280
  parallel: ["x", "v"]

What gets split

Each pusher is wrapped in jax.shard_map over a one-dimensional mesh of jax.devices():

Axis	Operators	Why no halo is needed
`"x"`	`edfdv` (`exponential` and `cubic-spline`) and the collision operator (Fokker-Planck + Krook)	Both are pointwise in \(x\) — an independent velocity-space solve per spatial cell
`"v"`	`vdfdx` (`exponential`)	The spectral \(x\)-advection is an independent phase rotation per velocity

Because the two axes are different, ["x", "v"] makes XLA insert an all-to-all between the velocity and spatial pushes on every step. That transpose is the entire cost of the scheme, so it pays off only when \(f\) is large enough that the per-device push dominates the reshuffle. Splitting a single axis (["x"]) avoids the transpose but leaves the other push serial.

Everything else — the Poisson/Ampère solve, the Hou-Li filter, the drivers, the diagnostics, and the saves — operates on the global array and is gathered by XLA as needed. Nothing about save or post-processing changes: the state diffrax carries is an ordinary global array, so netCDF output is byte-for-byte the same as a serial run’s.

Requirements and limits

nx must be divisible by the device count for "x", and every species’ nv for "v". Otherwise shard_map raises at trace time, naming the offending axis and size.
One process must see all the GPUs. On a NERSC Perlmutter node that means requesting the four GPUs and launching one task — no srun -n 4:
```
srun -n 1 -c 32 -G 4 uv run run.py --cfg configs/vlasov-1d/my-deck
```
The full \(f\) is allocated on the default device at initialization, so it must fit on one GPU. This is the “naive” part, and the reason sharded initialization and sharded checkpointing are not involved.
Results match the serial path to round-off. The pushes are bitwise identical per step; over a run, reduction reordering inside the shards leaves a drift of order \(10^{-15}\) in \(E\).
"v" breaks reverse-mode AD. jax.grad through the \(v\)-sharded vdfdx fails on jax 0.9.0.1 with a cotangent-type mismatch from the FFT along the unsharded axis (an upstream shard_map limitation, reproducible in a few lines of pure JAX). Forward mode (jax.jvp) is fine, and the \(x\)-sharded operators differentiate correctly in both modes with gradients identical to serial. For gradient work, use parallel: ["x"].

For scale: an electron + ion deck at nx: 17280, nv: 2048 with parallel: ["x", "v"] ran at roughly 45 ms/step on four A100s.

Practical Notes

Density profile. Uniform is easy. For a non-uniform profile you have to specify the parameters of the profile, which can be a sinusoidal perturbation or a tanh flat top — see initialization for the tanh flat-top parameters.

Collision frequency. nu_ee will modify the dynamics substantially depending on how far the distribution is driven from Maxwellian. Its envelope is specified the same way as a driver’s.

Krook frequency. In terms of physical correspondence this mostly resembles sideloss. Use it as a hard thermalization operator, for instance at the boundaries as in the SRS example.

Configuration Reference

See the Configuration Reference for complete YAML schema documentation.