diff --git a/Examples/1D_bump_on_tail.py b/Examples/1D_bump_on_tail.py
new file mode 100644
index 0000000..988ba39
--- /dev/null
+++ b/Examples/1D_bump_on_tail.py
@@ -0,0 +1,59 @@
+"""Example: 1D bump-on-tail instability, with the beam expanded in the bulk's Hermite basis.
+
+A bulk Maxwellian (density 0.9, thermal speed 1) and a warm beam (density 0.1, thermal speed 0.5,
+drift 4.5) are shifted by -0.45 to the frame with no net current, so that SPECTRAX's Vlasov-Ampere
+system has no uniform plasma oscillation. Each population is its own species.
+
+The beam is expanded in the bulk's Hermite width (alpha = sqrt(2)), not its own (sqrt(2) * 0.5):
+`compute_C_nmp(..., vth_s=...)` gives the closed-form coefficients of a Maxwellian narrower than its
+basis. In its own narrow basis the trapped beam, which spreads over 2 < v < 6, has a divergent
+expansion: its coefficients grow by 1e2-1e7 and the adaptive step collapses (near t = 28 with nu = 3).
+Set BEAM_BASIS_VT = 0.5 to see it. Linear theory: omega = 0.866 + 0.198 i at k = 0.3.
+"""
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+
+from spectrax import compute_C_nmp, inverse_HF_transform, simulation
+
+jax.config.update("jax_enable_x64", True)
+BEAM_BASIS_VT = 1.0 # 0.5: the beam in its own basis (fails)
+species = [(0.9, -0.45, 1.0), (0.1, 4.05, 0.5)] # (density, drift, thermal speed)
+k, eps, t_max, Nx, Nn = 0.3, 0.04, 100.0, 33, 128
+alpha = np.array([np.sqrt(2.0), np.sqrt(2.0) * BEAM_BASIS_VT])
+
+Ck_0 = []
+for (n, u, vth), a in zip(species, alpha):
+ C = n * compute_C_nmp(jnp.full((1, 3, 1, Nx, 1), u), [a] * 3, [0.0] * 3, Nn, 1, 1, 1,
+ vth_s=[vth, a / np.sqrt(2), a / np.sqrt(2)])[0, 0, 0] # (Nn, 1, Nx//2+1, 1)
+ Ck_0.append(C.at[:, :, 1].set(0.5 * eps * C[:, :, 0])) # density n (1 + eps cos kx)
+parameters = {
+ "Lx": 2 * np.pi / k, "Ly": 1.0, "Lz": 1.0, "mi_me": 1.0, "Ti_Te": 1.0, "nx": 1, "t_max": t_max,
+ "qs": jnp.array([-1.0, -1.0]), "Omega_cs": jnp.array([1.0, 1.0]), "nu": 3.0, "D": 0.0, "ode_tolerance": 1e-8,
+ "alpha_s": jnp.repeat(jnp.asarray(alpha), 3), "u_s": jnp.array([0.0] * 6), "Ck_0": jnp.concatenate(Ck_0),
+ # Gauss's law for the density perturbation: E = -(eps / k) sin(kx).
+ "Fk_0": jnp.zeros((6, 1, Nx // 2 + 1, 1), dtype=jnp.complex128).at[0, 0, 1, 0].set(0.5j * eps / k),
+}
+output = simulation(parameters, Nx=Nx, Nn=Nn, Ns=2, timesteps=501, dt=0.01)
+
+t, E1 = output["time"], 2 * jnp.abs(output["Fk"][:, 0, 0, 1, 0])
+x, v = np.linspace(0, 2 * np.pi / k, Nx, endpoint=False), np.linspace(-8, 10, 241)
+f = 0
+for s, a in enumerate(alpha): # 1D marginal: f(v_x, v_y = v_z = 0) times the transverse integral a^2 pi
+ xi_y, xi_x, xi_z = np.meshgrid([0.0], v / a, [0.0], indexing="ij")
+ f += np.pi * a**2 * inverse_HF_transform(output["Ck"][-1:, s * Nn:(s + 1) * Nn], Nn, 1, 1, Nx, 1, 1,
+ xi_x, xi_y, xi_z)[0, 0, :, 0, 0, :, 0]
+
+fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(11, 3.8))
+ax1.semilogy(t, E1, label="SPECTRAX")
+linear = (t > 8) & (t < 20)
+ax1.semilogy(t[linear], E1[linear][0] * np.exp(0.198 * (t[linear] - 8)), "k--", label=r"$\gamma = 0.198$")
+ax1.set(xlabel=r"$t\,\omega_{pe}$", ylabel=r"$|\hat E_{k}|$", title="bump-on-tail, k = 0.3")
+ax1.legend()
+mesh = ax2.pcolormesh(x, v, np.asarray(f).T, shading="auto", cmap="magma")
+ax2.set(xlabel=r"$x/\lambda_D$", ylabel=r"$v/v_t$", title=f"f(x, v, t = {t_max:g})")
+fig.colorbar(mesh, ax=ax2)
+plt.tight_layout()
+plt.show()
diff --git a/Examples/2D_magnetic_reconnection.py b/Examples/2D_magnetic_reconnection.py
new file mode 100644
index 0000000..512c86a
--- /dev/null
+++ b/Examples/2D_magnetic_reconnection.py
@@ -0,0 +1,228 @@
+"""Periodic double-Harris-sheet magnetic reconnection.
+
+The short defaults are a smoke test. For a clear nonlinear island movie, run
+
+ python Examples/2D_magnetic_reconnection.py --nx 65 --ny 17 \
+ --t-max 500 --chunk-time 50 --snapshots 51 --perturbation 0.1
+
+Use ``--perturbation 1e-4`` instead for a linear tearing-growth diagnostic.
+The parameters follow Camporeale et al. (2006), section IV.C; the periodic
+double-sheet geometry and half-sheet magnetic diagnostic follow Koshkarov et
+al. (2021).
+"""
+
+import argparse
+import json
+import time
+from pathlib import Path
+
+import jax
+import jax.numpy as jnp
+import matplotlib.pyplot as plt
+import numpy as np
+from diffrax import Dopri8
+from matplotlib.animation import FuncAnimation, PillowWriter
+
+from spectrax import simulation
+
+
+def initial_conditions(nx, ny, hermite, perturbation, nu):
+ a = jnp.sqrt(2.0)
+ lx, ly = 32.0 * a, 4.0 * jnp.pi * a
+ x = jnp.linspace(0.0, lx, nx, endpoint=False)
+ y = jnp.linspace(0.0, ly, ny, endpoint=False)
+ X, Y, Z = jnp.meshgrid(x, y, jnp.zeros(1), indexing="xy")
+ sheets = (
+ jnp.cosh((X - lx / 4.0) / a) ** -2,
+ jnp.cosh((X - 3.0 * lx / 4.0) / a) ** -2,
+ )
+ density = jnp.stack((sheets[0], sheets[1], sheets[0], sheets[1]))
+ alpha = a * jnp.array(([1.0 / 6.0] * 6) + ([1.0 / 24.0] * 6))
+ coefficients = jnp.zeros(
+ (4, hermite, hermite, hermite, ny, nx, 1), dtype=jnp.float64
+ )
+ coefficients = coefficients.at[:, 0, 0, 0].set(
+ density / jnp.prod(alpha.reshape(4, 3), axis=1)[:, None, None, None]
+ )
+
+ kx, ky = 2.0 * jnp.pi / lx, 2.0 * jnp.pi / ly
+ magnetic = jnp.array(
+ [
+ -perturbation * ky * jnp.sin(ky * Y) * jnp.sin(kx * X),
+ jnp.tanh((X - lx / 4.0) / a)
+ - jnp.tanh((X - 3.0 * lx / 4.0) / a)
+ - 1.0
+ - perturbation * kx * jnp.cos(ky * Y) * jnp.cos(kx * X),
+ 5.0 * jnp.ones_like(X),
+ ]
+ )
+ transform = lambda value: jnp.fft.rfftn(
+ value, axes=(-1, -3, -2), norm="forward"
+ )
+ return {
+ "Lx": lx,
+ "Ly": ly,
+ "Lz": 1.0,
+ "mi_me": 16.0,
+ "qs": jnp.array((-1.0, -1.0, 1.0, 1.0)),
+ "alpha_s": alpha,
+ "u_s": jnp.array(
+ (0.0, 0.0, -1.0 / 6.0, 0.0, 0.0, 1.0 / 6.0,
+ 0.0, 0.0, 1.0 / 6.0, 0.0, 0.0, -1.0 / 6.0)
+ ) / a,
+ "Omega_cs": jnp.array((1.0 / 3.0, 1.0 / 3.0, 1.0 / 48.0, 1.0 / 48.0)),
+ "nu": nu,
+ "D": 0.0,
+ "ode_tolerance": 1e-7,
+ "Ck_0": transform(coefficients),
+ "Fk_0": transform(jnp.concatenate((jnp.zeros_like(magnetic), magnetic))),
+ }
+
+
+def fields_and_flux(fk, nx, lx, ly):
+ ny = fk.shape[-3]
+ fields = np.fft.irfftn(
+ fk, s=(1, ny, nx), axes=(-1, -3, -2), norm="forward"
+ ).real
+ kx = 2.0 * np.pi * np.fft.rfftfreq(nx, d=lx / nx)
+ ky = 2.0 * np.pi * np.fft.fftfreq(ny, d=ly / ny)
+ KX, KY = np.meshgrid(kx, ky)
+ bx, by = fk[:, 3, :, :, 0], fk[:, 4, :, :, 0]
+ k2 = KX**2 + KY**2
+ az_k = np.zeros_like(bx)
+ np.divide(-1j * KY * bx + 1j * KX * by, k2, out=az_k, where=k2 != 0.0)
+ jz_k = 1j * KX * by - 1j * KY * bx
+ az = np.fft.irfftn(az_k, s=(ny, nx), axes=(-2, -1), norm="forward").real
+ jz = np.fft.irfftn(jz_k, s=(ny, nx), axes=(-2, -1), norm="forward").real
+ return fields, az, jz
+
+
+def save_diagnostics(output_dir, t, fk, nx, lx, ly):
+ fields, az, jz = fields_and_flux(fk, nx, lx, ly)
+ half = (nx + 1) // 2
+ magnetic = fields[:, 3:, :, :half, 0]
+ phase = np.exp(2j * np.pi * np.arange(magnetic.shape[2]) / magnetic.shape[2])
+ amplitude = np.linalg.norm(
+ np.sum(magnetic * phase[None, None, :, None], axis=(-2, -1)), axis=1
+ ) / np.sqrt(magnetic.shape[2] * magnetic.shape[3])
+ peaks = np.flatnonzero(
+ (amplitude[1:-1] > amplitude[:-2]) & (amplitude[1:-1] > amplitude[2:])
+ ) + 1
+ peaks = peaks[t[peaks] >= t[-1] / 4.0]
+ fit = None
+ if len(peaks) >= 3:
+ gamma, intercept = np.polyfit(t[peaks], np.log(amplitude[peaks]), 1)
+ fit = np.exp(intercept + gamma * t)
+ else:
+ gamma = np.nan
+
+ sheet = np.argmin(abs(np.linspace(0.0, lx, nx, endpoint=False) - lx / 4.0))
+ flux = np.ptp(az[:, :, sheet], axis=1)
+ fig, axes = plt.subplots(2, 1, figsize=(6.4, 6.0), sharex=True, constrained_layout=True)
+ axes[0].semilogy(t, amplitude, color="0.65", label="half-sheet magnetic mode")
+ if fit is not None:
+ axes[0].semilogy(t[peaks], amplitude[peaks], "o", ms=3)
+ axes[0].semilogy(t, fit, "k--", label=rf"envelope $\gamma={gamma:.5f}$")
+ axes[0].set_ylabel("magnetic-mode amplitude")
+ axes[0].legend(fontsize=8)
+ axes[1].plot(t, flux)
+ axes[1].set(xlabel=r"$\omega_{pe}t$", ylabel=r"reconnected flux $\Delta A_z$")
+ fig.savefig(output_dir / "reconnection-growth.png", dpi=180)
+ plt.close(fig)
+
+ x = np.linspace(0.0, lx, nx, endpoint=False)[:half]
+ y = np.linspace(-ly / 2.0, ly / 2.0, fk.shape[-3], endpoint=False)
+ az = np.roll(az, fk.shape[-3] // 2, axis=1)[:, :, :half]
+ jz = np.roll(jz, fk.shape[-3] // 2, axis=1)[:, :, :half]
+ window = (x >= lx / 8.0) & (x <= 3.0 * lx / 8.0)
+ x, az, jz = x[window], az[:, :, window], jz[:, :, window]
+ jz -= jz.mean(axis=1, keepdims=True)
+ limit = max(np.quantile(abs(jz[-1]), 0.99), 1e-12)
+ sheet = np.argmin(abs(x - lx / 4.0))
+ fig, ax = plt.subplots(figsize=(6.4, 3.8), constrained_layout=True)
+ image = ax.pcolormesh(x, y, jz[0], shading="auto", cmap="RdBu_r", vmin=-limit, vmax=limit)
+ contours = [None]
+ title = ax.set_title("")
+ ax.set(xlabel="x", ylabel="y")
+ fig.colorbar(image, ax=ax, label=r"$J_z-\langle J_z\rangle_y$")
+
+ def update(frame):
+ image.set_array(jz[frame].ravel())
+ if contours[0] is not None:
+ contours[0].remove()
+ center = az[frame, :, sheet].mean()
+ span = max(1.25 * np.ptp(az[frame, :, sheet]), 1e-12)
+ contours[0] = ax.contour(
+ x, y, az[frame], center + np.linspace(-span, span, 19),
+ colors="k", linewidths=0.55,
+ )
+ title.set_text(rf"$\omega_{{pe}}t={t[frame]:.1f}$")
+ return image, title, contours[0]
+
+ frames = range(0, len(t), max(1, len(t) // 100))
+ animation = FuncAnimation(fig, update, frames=frames, interval=80, blit=False)
+ animation.save(output_dir / "reconnection-topology.gif", PillowWriter(fps=12), dpi=110)
+ plt.close(fig)
+ return amplitude, gamma
+
+
+def main():
+ parser = argparse.ArgumentParser()
+ parser.add_argument("--nx", type=int, default=33)
+ parser.add_argument("--ny", type=int, default=9)
+ parser.add_argument("--hermite", type=int, default=4)
+ parser.add_argument("--nu", type=float, default=1.0)
+ parser.add_argument("--perturbation", type=float, default=1e-4)
+ parser.add_argument("--t-max", type=float, default=10.0)
+ parser.add_argument("--chunk-time", type=float, default=10.0)
+ parser.add_argument("--snapshots", type=int, default=21)
+ parser.add_argument("--output-dir", type=Path, default=Path("reconnection_output"))
+ args = parser.parse_args()
+ args.output_dir.mkdir(parents=True, exist_ok=True)
+
+ values = initial_conditions(args.nx, args.ny, args.hermite, args.perturbation, args.nu)
+ times, fields, offset = [], [], 0.0
+ start = time.perf_counter()
+ while offset < args.t_max:
+ duration = min(args.chunk_time, args.t_max - offset)
+ values["t_max"] = duration
+ output = simulation(
+ values, Nx=args.nx, Ny=args.ny, Nz=1,
+ Nn=args.hermite, Nm=args.hermite, Np=args.hermite, Ns=4,
+ timesteps=args.snapshots, dt=1e-3, solver=Dopri8(),
+ )
+ jax.block_until_ready(output["Ck"])
+ first = 0 if not times else 1
+ times.append(np.asarray(output["time"])[first:] + offset)
+ fields.append(np.asarray(output["Fk"])[first:])
+ values["Ck_0"] = jnp.asarray(output["Ck"][-1])
+ values["Fk_0"] = jnp.asarray(output["Fk"][-1])
+ del output
+ offset += duration
+ print(f"t={offset:g}")
+
+ t, fk = np.concatenate(times), np.concatenate(fields)
+ amplitude, gamma = save_diagnostics(
+ args.output_dir, t, fk, args.nx, float(values["Lx"]), float(values["Ly"])
+ )
+ coefficients = np.asarray(values["Ck_0"]).reshape(
+ 4, args.hermite, args.hermite, args.hermite, args.ny, args.nx // 2 + 1, 1
+ )
+ power = abs(coefficients) ** 2
+ tails = {
+ name: float(np.take(power, -1, axis=axis).sum() / power.sum())
+ for name, axis in (("p", 1), ("m", 2), ("n", 3))
+ }
+ np.savez(args.output_dir / "reconnection.npz", time=t, amplitude=amplitude, Fk=fk)
+ linear = args.perturbation <= 1e-3
+ print(json.dumps({
+ "runtime_s": time.perf_counter() - start,
+ "regime": "linear diagnostic" if linear else "finite-seed nonlinear",
+ "gamma": None if not np.isfinite(gamma) else float(gamma),
+ "expected_gamma": 0.256 / 48.0 if linear else None,
+ "hermite_tail_fraction": tails,
+ }, sort_keys=True))
+
+
+if __name__ == "__main__":
+ main()
diff --git a/README.md b/README.md
index 9d50368..2e594f5 100644
--- a/README.md
+++ b/README.md
@@ -1,272 +1,272 @@
-
- 
-
-
-
- ❯ SPECTRAX: Hermite-Fourier Vlasov-Maxwell solver in JAX for plasma physics simulations
-
-
- 
- 
- 
-
- 
-
-
- 
-
-
- 
-
-
-
-
-
-
-
-
-
-
-
-## Table of Contents
-
-- [Table of Contents](#table-of-contents)
-- [Overview](#overview)
-- [Mathematical Method](#mathematical-method)
-- [Features](#features)
-- [Getting Started](#getting-started)
- - [Prerequisites](#prerequisites)
- - [Installation](#installation)
- - [Usage](#usage)
- - [Command‑line Interface](#commandline-interface)
- - [Testing](#testing)
-- [Input File Format](#input-file-format)
-- [Contributing](#contributing)
-- [License](#license)
-- [Acknowledgments](#acknowledgments)
-
----
-
-## Overview
-
-**SPECTRAX** is an open-source spectral kinetic plasma solver written in Python with the [JAX](https://github.com/jax-ml/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)](https://www.sciencedirect.com/science/article/pii/S0021999115004738), [Vencels et al. (2016)](https://iopscience.iop.org/article/10.1088/1742-6596/719/1/012022) and [Roytershteyn & Delzanno (2018)](https://www.frontiersin.org/journals/astronomy-and-space-sciences/articles/10.3389/fspas.2018.00027/full), 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](https://github.com/patrick-kidger/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](https://github.com/patrick-kidger/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](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:
-
-```sh
-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:
-
-```sh
-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](https://docs.jax.dev/en/latest/installation.html).
-
----
-
-### 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.
-
-```sh
-spectrax
-```
-
-To run it with different input parameters, us a TOML file like those in the `Examples` directory.
-
-```sh
-spectrax example_input.toml
-```
-
-Other examples written in Python scripts, like those in the `Examples` directory, can be executed from the terminal as follows:
-
-```sh
-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:
-```sh
-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.
-
-| 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. |
-| `alpha_s` | Thermal scales for each species for Hermite basis. |
-| `u_s` | Drift velocities for each species for Hermite basis (packed as `[u_x,u_y,u_z]` ). |
-| `Omega_cs` | Cyclotron frequencies for each species. |
-| `nu` | Hyper-collision frequency to damp recurrence. |
-| `D` | Hyper‑diffusion coefficient. |
-| `t_max` | Final simulation time. |
-| `nx, ny, nz` | Mode numbers used to seed sinusoidal perturbations (see examples). |
-| `dn1, dn2` | Amplitudes of initial density 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. |
-| `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). |
-| `dtmin` | Smallest adaptive step; below it the solve stops with an error instead of stalling. No floor by default. A collapsing step usually signals an under-resolved run (Hermite cutoff or basis width); it is reported, not fixed. |
-| `max_steps` | Step budget (`1000000` by default); exhausting it is reported as `max_steps_reached`, distinct from `dt_min_reached`. |
-| `throw` | `True` (default) raises on a failed solve; `False` returns and reports it in `output["solver_result"]`; only saves with `output["valid_times"]` (the first `num_valid_times`) hold a solution. |
-
-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](https://github.com/uwplasma/SPECTRAX/discussions)**: Share your insights, provide feedback, or ask questions.
-- **🐛 [Report Issues](https://github.com/uwplasma/SPECTRAX/issues)**: Submit bugs found or log feature requests for the `SPECTRAX` project.
-- **💡 [Submit Pull Requests](https://github.com/uwplasma/SPECTRAX/blob/main/CONTRIBUTING.md)**: 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.
- ```sh
- git clone https://github.com/uwplasma/SPECTRAX
- ```
-3. **Create a New Branch**: Always work on a new branch, giving it a descriptive name.
- ```sh
- 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.
- ```sh
- git commit -m 'Implemented new feature x.'
- ```
-6. **Push to github**: Push the changes to your forked repository.
- ```sh
- 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](LICENSE) file.
-
----
-
-## Acknowledgments
-
-- We acknowledge the help of the whole [UWPlasma](https://rogerio.physics.wisc.edu/) plasma group.
-
----
-
-
-
-
-
+
+
+
+
+
+ ❯ SPECTRAX: Hermite-Fourier Vlasov-Maxwell solver in JAX for plasma physics simulations
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+## Table of Contents
+
+- [Table of Contents](#table-of-contents)
+- [Overview](#overview)
+- [Mathematical Method](#mathematical-method)
+- [Features](#features)
+- [Getting Started](#getting-started)
+ - [Prerequisites](#prerequisites)
+ - [Installation](#installation)
+ - [Usage](#usage)
+ - [Command‑line Interface](#commandline-interface)
+ - [Testing](#testing)
+- [Input File Format](#input-file-format)
+- [Contributing](#contributing)
+- [License](#license)
+- [Acknowledgments](#acknowledgments)
+
+---
+
+## Overview
+
+**SPECTRAX** is an open-source spectral kinetic plasma solver written in Python with the [JAX](https://github.com/jax-ml/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)](https://www.sciencedirect.com/science/article/pii/S0021999115004738), [Vencels et al. (2016)](https://iopscience.iop.org/article/10.1088/1742-6596/719/1/012022) and [Roytershteyn & Delzanno (2018)](https://www.frontiersin.org/journals/astronomy-and-space-sciences/articles/10.3389/fspas.2018.00027/full), 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](https://github.com/patrick-kidger/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](https://github.com/patrick-kidger/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](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:
+
+```sh
+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:
+
+```sh
+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](https://docs.jax.dev/en/latest/installation.html).
+
+---
+
+### 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.
+
+```sh
+spectrax
+```
+
+To run it with different input parameters, us a TOML file like those in the `Examples` directory.
+
+```sh
+spectrax example_input.toml
+```
+
+Other examples written in Python scripts, like those in the `Examples` directory, can be executed from the terminal as follows:
+
+```sh
+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:
+```sh
+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.
+
+| 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. |
+| `alpha_s` | Thermal scales for each species for Hermite basis. |
+| `u_s` | Drift velocities for each species for Hermite basis (packed as `[u_x,u_y,u_z]` ). |
+| `Omega_cs` | Cyclotron frequencies for each species. |
+| `nu` | Hyper-collision frequency to damp recurrence. |
+| `D` | Hyper‑diffusion coefficient. |
+| `t_max` | Final simulation time. |
+| `nx, ny, nz` | Mode numbers used to seed sinusoidal perturbations (see examples). |
+| `dn1, dn2` | Amplitudes of initial density 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. |
+| `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). |
+| `dtmin` | Smallest adaptive step; below it the solve stops with an error instead of stalling. No floor by default. A collapsing step usually signals an under-resolved run (Hermite cutoff or basis width); it is reported, not fixed. |
+| `max_steps` | Step budget (`1000000` by default); exhausting it is reported as `max_steps_reached`, distinct from `dt_min_reached`. |
+| `throw` | `True` (default) raises on a failed solve; `False` returns and reports it in `output["solver_result"]`; only saves with `output["valid_times"]` (the first `num_valid_times`) hold a solution. |
+
+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](https://github.com/uwplasma/SPECTRAX/discussions)**: Share your insights, provide feedback, or ask questions.
+- **🐛 [Report Issues](https://github.com/uwplasma/SPECTRAX/issues)**: Submit bugs found or log feature requests for the `SPECTRAX` project.
+- **💡 [Submit Pull Requests](https://github.com/uwplasma/SPECTRAX/blob/main/CONTRIBUTING.md)**: 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.
+ ```sh
+ git clone https://github.com/uwplasma/SPECTRAX
+ ```
+3. **Create a New Branch**: Always work on a new branch, giving it a descriptive name.
+ ```sh
+ 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.
+ ```sh
+ git commit -m 'Implemented new feature x.'
+ ```
+6. **Push to github**: Push the changes to your forked repository.
+ ```sh
+ 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](LICENSE) file.
+
+---
+
+## Acknowledgments
+
+- We acknowledge the help of the whole [UWPlasma](https://rogerio.physics.wisc.edu/) plasma group.
+
+---
+
+
+
+
+
diff --git a/pyproject.toml b/pyproject.toml
index 6796111..72129f2 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -29,7 +29,7 @@ keywords = ["Plasma",
"Simulation",
"JAX"]
-dependencies = [ "jax", "jaxlib", "jax_tqdm", "matplotlib", "diffrax", "orthax"]
+dependencies = [ "jax", "jaxlib", "jax_tqdm", "matplotlib", "diffrax", "orthax", "solvax>=0.8.0"]
requires-python = ">=3.11"
diff --git a/requirements.txt b/requirements.txt
index d727662..91c9821 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -3,5 +3,6 @@ jaxlib
jax_tqdm
diffrax
orthax
+solvax>=0.8.0
matplotlib
pandas
diff --git a/spectrax/_diagnostics.py b/spectrax/_diagnostics.py
index 8fdb788..3d18bc6 100644
--- a/spectrax/_diagnostics.py
+++ b/spectrax/_diagnostics.py
@@ -19,7 +19,13 @@
import jax.numpy as jnp
-__all__ = ["diagnostics"]
+__all__ = ["diagnostics", "hermite_tail_fraction"]
+
+
+def _rfft_weights(Nx: int, Nx_kept: int) -> jnp.ndarray:
+ """Return Parseval weights for an RFFT of length ``Nx``."""
+ weights = jnp.full(Nx_kept, 2.0).at[0].set(1.0)
+ return weights.at[-1].set(1.0 + (Nx % 2) * (Nx_kept > 1))
def _infer_Ns(output: Mapping[str, Any]) -> int:
@@ -68,6 +74,50 @@ def _infer_masses(output: Mapping[str, Any], Ns: int) -> jnp.ndarray:
return jnp.ones((Ns,))
+def hermite_tail_fraction(output: Mapping[str, Any], width: int = 1) -> jnp.ndarray:
+ """Return the perturbation norm fraction at each Hermite boundary.
+
+ The tail is the union of the outer ``width`` modes on every Hermite axis
+ with more than one retained mode. Squared ``dCk`` coefficients are summed
+ with Parseval weights over Fourier space. For SPECTRAX's matched-Maxwellian
+ basis this norm is proportional to perturbed-distribution free energy, but
+ it is not the total plasma energy.
+
+ Returns an array with shape ``(Nt, Ns)``. A zero perturbation, or a layout
+ with no active Hermite axis, has zero tail fraction.
+ """
+ if not isinstance(width, int) or width < 1:
+ raise ValueError("width must be a positive integer.")
+
+ dCk = jnp.asarray(output["dCk"])
+ Ns = int(output["Ns"]) if "Ns" in output else _infer_Ns(output)
+ Nn, Nm, Np = (int(output[key]) for key in ("Nn", "Nm", "Np"))
+ expected_modes = Ns * Nn * Nm * Np
+ if dCk.ndim != 5 or dCk.shape[1] != expected_modes:
+ raise ValueError(f"dCk must have shape (Nt, {expected_modes}, Ny, Nx//2+1, Nz).")
+
+ n = jnp.arange(Nn)[None, None, :]
+ m = jnp.arange(Nm)[None, :, None]
+ p = jnp.arange(Np)[:, None, None]
+ tail = jnp.zeros((Np, Nm, Nn), dtype=bool)
+ if Nn > 1:
+ tail |= n >= max(Nn - width, 0)
+ if Nm > 1:
+ tail |= m >= max(Nm - width, 0)
+ if Np > 1:
+ tail |= p >= max(Np - width, 0)
+
+ Nx_kept = dCk.shape[-2]
+ weights = _rfft_weights(int(output.get("Nx", 2 * (Nx_kept - 1))), Nx_kept)
+ power = jnp.abs(dCk.reshape(dCk.shape[0], Ns, Np, Nm, Nn, *dCk.shape[-3:])) ** 2
+ power *= weights.reshape(1, 1, 1, 1, 1, 1, -1, 1)
+ total = jnp.sum(power, axis=(2, 3, 4, 5, 6, 7))
+ boundary = jnp.sum(power * tail.reshape(1, 1, Np, Nm, Nn, 1, 1, 1),
+ axis=(2, 3, 4, 5, 6, 7))
+ denominator = jnp.where(total > 0, total, 1.0)
+ return jnp.where(total > 0, boundary / denominator, 0.0)
+
+
def diagnostics(output: dict) -> None:
"""Compute and attach common diagnostics to ``output`` (mutates in-place).
@@ -198,9 +248,8 @@ def _take_mode(n: int, m: int, p: int) -> jnp.ndarray:
kinetic_energy_species = pref[None, :] * (term0[None, :] * C000 + term1 + term2) # (Nt, Ns)
kinetic_energy = jnp.sum(kinetic_energy_species, axis=1) # (Nt,)
# Field energy
- rfft_weights = jnp.full(Nx_kept, 2.0)
- rfft_weights = rfft_weights.at[0].set(1.0)
- rfft_weights = rfft_weights.at[-1].set(1.0)
+ Nx = output.get("Nx", 2 * (Nx_kept - 1))
+ rfft_weights = _rfft_weights(Nx, Nx_kept)
weight_grid = rfft_weights.reshape(1, 1, -1, 1)
EM_energy = 0.5 * jnp.sum((jnp.abs(Fk) ** 2) * weight_grid, axis=(-4, -3, -2, -1)) * Omega_cs[0] ** 2
diff --git a/spectrax/_initialize_maxwellian.py b/spectrax/_initialize_maxwellian.py
index 0c627db..e85916c 100644
--- a/spectrax/_initialize_maxwellian.py
+++ b/spectrax/_initialize_maxwellian.py
@@ -3,16 +3,30 @@
import jax
jax.config.update("jax_enable_x64", True)
import jax.numpy as jnp
-from jax import jit
+from jax import jit, lax
from jax.numpy.fft import rfftn
-from jax.scipy.special import factorial
from functools import partial
__all__ = ['compute_C_nmp']
+def _maxwellian_modes(d, c2, N):
+ """Normalised Hermite coefficients g_n = c^n H_n(d / c) / sqrt(2^n n!), n < N, of a 1D Maxwellian.
+
+ d is the drift in units of the Hermite width and c2 = 1 - 2 vth^2 / alpha^2; c2 = 0 when the
+ Maxwellian has the basis width (then g_n = (sqrt(2) d)^n / sqrt(n!)). The recurrence
+ g_{n+1} = (sqrt(2) d g_n - sqrt(n) c2 g_{n-1}) / sqrt(n + 1) avoids factorials.
+ """
+ def step(g, n): # g = (g_{n-1}, g_n) -> (g_n, g_{n+1}); a scan keeps compile time independent of N
+ g_next = (jnp.sqrt(2.0) * d * g[1] - jnp.sqrt(n) * c2 * g[0]) / jnp.sqrt(n + 1.0)
+ return (g[1], g_next), g_next
+ g0, g1 = jnp.ones_like(d), jnp.sqrt(2.0) * d
+ _, rest = lax.scan(step, (g0, g1), jnp.arange(1.0, max(N - 1, 1)))
+ return jnp.concatenate([g0[None], g1[None], rest])[:N]
+
+
@partial(jit, static_argnames=['Nn', 'Nm', 'Np', 'Ns'])
-def compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns):
+def compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns, vth_s=None):
"""
Build the Hermite-Fourier coefficients for a drifting Maxwellian distribution function.
@@ -31,6 +45,11 @@ def compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns):
Number of Hermite modes retained along the x, y, and z velocity axes.
Ns : int
Number of species.
+ vth_s : array-like, optional
+ Thermal speeds (standard deviations) of the Maxwellians, flattened as `(3 * Ns,)`.
+ Default: `alpha_s / sqrt(2)`, the basis width. A Maxwellian narrower than its basis
+ (`vth < alpha / sqrt(2)`) has a convergent expansion, so a thin beam can share a wide
+ basis; a wider one does not converge.
Returns
-------
@@ -38,29 +57,13 @@ def compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns):
Complex Hermite-Fourier coefficients with shape `(Ns, Np, Nm, Nn, Ny, Nx//2+1, Nz)`
corresponding to the Maxwellian evaluated on the supplied grid.
"""
-
- U_x = Us_grid[:, 0, None, None, None, :, :, :] # shape (Ns, 1, 1, 1, Ny, Nx, Nz)
- U_y = Us_grid[:, 1, None, None, None, :, :, :] # shape (Ns, 1, 1, 1, Ny, Nx, Nz)
- U_z = Us_grid[:, 2, None, None, None, :, :, :] # shape (Ns, 1, 1, 1, Ny, Nx, Nz)
-
- alpha = jnp.array(alpha_s).reshape(Ns, 3)
- alpha_x = alpha[:, 0, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
- alpha_y = alpha[:, 1, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
- alpha_z = alpha[:, 2, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
-
- u = jnp.array(u_s).reshape(Ns, 3)
- u_x = u[:, 0, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
- u_y = u[:, 1, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
- u_z = u[:, 2, None, None, None, None, None, None] # shape (Ns, 1, 1, 1, 1, 1, 1)
-
- p = jnp.arange(Np)[None, :, None, None, None, None, None] # shape (1, Np, 1, 1, 1, 1, 1)
- m = jnp.arange(Nm)[None, None, :, None, None, None, None] # shape (1, 1, Nm, 1, 1, 1, 1)
- n = jnp.arange(Nn)[None, None, None, :, None, None, None] # shape (1, 1, 1, Nn, 1, 1, 1)
-
- C = (jnp.sqrt(2 ** (n + m + p) / (factorial(n) * factorial(m) * factorial(p)))
- * (1 / (alpha_x ** (n + 1) * alpha_y ** (m + 1) * alpha_z ** (p + 1)))
- * (U_x - u_x) ** n * (U_y - u_y) ** m * (U_z - u_z) ** p)
-
+ alpha = jnp.array(alpha_s).reshape(Ns, 3)[:, :, None, None, None] # (Ns, 3, 1, 1, 1)
+ u = jnp.array(u_s).reshape(Ns, 3)[:, :, None, None, None]
+ vth = alpha / jnp.sqrt(2.0) if vth_s is None else jnp.array(vth_s).reshape(Ns, 3)[:, :, None, None, None]
+ c2 = 1 - 2 * (vth / alpha) ** 2 # 0 when the Maxwellian has the basis width
+ g = [_maxwellian_modes((Us_grid[:, i] - u[:, i]) / alpha[:, i], c2[:, i], N) / alpha[:, i]
+ for i, N in enumerate((Nn, Nm, Np))] # each (N, Ns, Ny, Nx, Nz)
+ C = jnp.einsum('psyxz,msyxz,nsyxz->spmnyxz', g[2], g[1], g[0]) # (Ns, Np, Nm, Nn, Ny, Nx, Nz)
Ck_0 = rfftn(C, axes=(-1, -3, -2), norm="forward") # shape (Ns, Np, Nm, Nn, Ny, Nx//2+1, Nz)
return Ck_0
diff --git a/spectrax/_simulation.py b/spectrax/_simulation.py
index f9dd349..611469d 100644
--- a/spectrax/_simulation.py
+++ b/spectrax/_simulation.py
@@ -9,6 +9,7 @@
from ._initialization import initialize_simulation_parameters
from ._model import plasma_current, Hermite_Fourier_system
from ._diagnostics import diagnostics
+from .midpoint_solver import ImplicitMidpoint
__all__ = ["cross_product", "ode_system", "simulation"]
@@ -151,7 +152,8 @@ def simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2,
-------
dict
Dictionary containing the evolved coefficients (`Ck`, `Fk`), time samples,
- solver statistics, perturbation diagnostics, and all simulation parameters.
+ solver statistics, perturbation diagnostics, all simulation parameters, and
+ `midpoint_stats` when using :class:`ImplicitMidpoint`.
"""
# **Initialize simulation parameters**
@@ -185,7 +187,8 @@ def simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2,
ODETerm(ode_system_partial), solver=solver,
stepsize_controller=stepsize_controller,
t0=0, t1=parameters["t_max"], dt0=dt,
- y0=initial_conditions, args=args, saveat=SaveAt(ts=time),
+ y0=initial_conditions, args=args,
+ saveat=SaveAt(ts=time, solver_state=isinstance(solver, ImplicitMidpoint)),
max_steps=max_steps, progress_meter=TqdmProgressMeter(), throw=throw)
# Reshape the solution to extract Ck and Fk
@@ -201,6 +204,8 @@ def simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2,
temporary_output = {"Ck": Ck, "Fk": Fk, "time": time, "dCk": dCk, "solver_stats": sol.stats,
"solver_result": sol.result, "valid_times": valid_times,
"num_valid_times": jnp.sum(valid_times)}
+ if sol.solver_state is not None:
+ temporary_output["midpoint_stats"] = sol.solver_state
output = {**temporary_output, **parameters}
diagnostics(output)
return output
diff --git a/spectrax/midpoint_solver.py b/spectrax/midpoint_solver.py
index 2b6bf06..fa25e29 100644
--- a/spectrax/midpoint_solver.py
+++ b/spectrax/midpoint_solver.py
@@ -1,10 +1,56 @@
"""Custom Diffrax solver: implicit midpoint with a Newton–GMRES nonlinear solve."""
+from collections.abc import Callable
+from typing import NamedTuple
+
import diffrax
import jax
import jax.numpy as jnp
-from jax import lax
-from jax.scipy.sparse.linalg import gmres
+import optimistix as optx
+from solvax import newton_krylov
+
+
+def collision_diffusion_x_streaming_preconditioner(args, dt):
+ """Invert collision, diffusion, and x-streaming Hermite line blocks."""
+ nu, diffusion = args[-20:-18]
+ alpha = args[-17].reshape(-1, 3)[:, 0]
+ drift = args[-16].reshape(-1, 3)[:, 0]
+ length, kx, k2 = args[-15], args[-12], args[-9]
+ collision, sqrt_plus, sqrt_minus = args[-7:-4]
+ phase = (0.5j * dt / length * kx)[None, None, None]
+ species = (slice(None),) + (None,) * 6
+ diagonal = (
+ 1
+ + 0.5 * dt * (
+ nu * collision[None, :, :, :, None, None, None]
+ + diffusion * k2[None, None, None, None]
+ )
+ + drift[species] * phase
+ )
+ coupling = alpha[species] * phase / jnp.sqrt(2)
+
+ def lines(value):
+ return jnp.moveaxis(jnp.broadcast_to(value, diagonal.shape), 3, -1)
+
+ main, lower = lines(diagonal), lines(coupling * sqrt_minus)
+ upper = lines(coupling * sqrt_plus).at[..., -1].set(0)
+
+ def apply(residual):
+ coefficients, fields = residual
+ rhs = jnp.moveaxis(coefficients, 3, -1)[..., None]
+ solution = jax.lax.linalg.tridiagonal_solve(lower, main, upper, rhs)[..., 0]
+ return jnp.moveaxis(solution, -1, 3), fields
+
+ return apply
+
+
+class MidpointSolverState(NamedTuple):
+ """Newton--GMRES iteration totals and per-step maxima."""
+
+ newton_iterations: jax.Array
+ linear_iterations: jax.Array
+ max_newton_iterations: jax.Array
+ max_linear_iterations: jax.Array
class ImplicitMidpoint(diffrax.AbstractSolver):
@@ -12,15 +58,26 @@ class ImplicitMidpoint(diffrax.AbstractSolver):
This solver implements the implicit midpoint rule:
- y_{n+1} = y_n + Δt * f(t_{n+1}, (y_n + y_{n+1}) / 2)
+ y_{n+1} = y_n + Δt * f(t_n + Δt / 2, (y_n + y_{n+1}) / 2)
The nonlinear equation for ``y_{n+1}`` is solved with Newton iterations,
where each linearized step is solved via GMRES using JAX's linearization.
+ ``rtol`` and ``atol`` control the nonlinear residual, while
+ ``linear_rtol`` and ``linear_atol`` control each GMRES solve.
+ Newton iteration stops when ``||F|| <= max(atol, rtol * ||F_initial||)``.
+ Accepted-step diagnostics are accumulated in :class:`MidpointSolverState`.
+ An optional ``preconditioner(args, dt)`` factory supplies the inverse action
+ used by GMRES.
"""
rtol: float = 1e-6
atol: float = 1e-8
max_iters: int = 300
+ linear_restart: int = 20
+ linear_rtol: float = 1e-4
+ linear_atol: float = 0.0
+ linear_max_restarts: int = 20
+ preconditioner: Callable | None = None
term_structure = diffrax.ODETerm
interpolation_cls = diffrax.LocalLinearInterpolation
@@ -29,70 +86,55 @@ def order(self, terms):
return 2
def init(self, terms, t0, t1, y0, args):
- return None
+ del terms, t0, t1, y0, args
+ zero = jnp.int32(0)
+ return MidpointSolverState(zero, zero, zero, zero)
def func(self, terms, t0, y0, args):
return terms.vf(t0, y0, args)
def step(self, terms, t0, t1, y0, args, solver_state, made_jump):
- del solver_state, made_jump
+ del made_jump
+ if solver_state is None:
+ solver_state = self.init(terms, t0, t1, y0, args)
δt = t1 - t0
+ t_mid = t0 + 0.5 * δt
f0 = terms.vf(t0, y0, args)
- y1_init = y0 + δt * f0
+ y1_init = jax.tree.map(lambda y, f: y + δt * f, y0, f0)
- # Define F(y1) = y1 - y0 - δt * f(t1, (y0 + y1)/2)
+ # Define F(y1) = y1 - y0 - δt * f(t_mid, (y0 + y1)/2)
def F_fn(y1):
- y_mid = 0.5 * (y0 + y1)
- return y1 - y0 - δt * terms.vf(t1, y_mid, args)
-
- y1 = _newton_gmres(F_fn, y0, y1_init, self.rtol, self.atol, self.max_iters)
-
- y_error = y1 - y1_init
+ y_mid = jax.tree.map(lambda a, b: 0.5 * (a + b), y0, y1)
+ f_mid = terms.vf(t_mid, y_mid, args)
+ return jax.tree.map(lambda a, b, f: a - b - δt * f, y1, y0, f_mid)
+
+ precond = None if self.preconditioner is None else self.preconditioner(args, δt)
+ solution = newton_krylov(
+ F_fn,
+ y1_init,
+ precond=precond,
+ rtol=self.rtol,
+ atol=self.atol,
+ max_steps=self.max_iters,
+ linear_restart=self.linear_restart,
+ linear_rtol=self.linear_rtol,
+ linear_atol=self.linear_atol,
+ linear_max_restarts=self.linear_max_restarts,
+ )
+ y1 = solution.x
+
+ y_error = jax.tree.map(lambda a, b: a - b, y1, y1_init)
dense_info = dict(y0=y0, y1=y1)
- return y1, y_error, dense_info, None, diffrax.RESULTS.successful
-
-
-def _newton_gmres(F_fn, y0, y_init, rtol, atol, max_iters):
- """Solve ``F(y)=0`` using Newton iterations with GMRES linear solves.
-
- Notes
- -----
- - The Jacobian-vector product is obtained via ``jax.linearize``.
- - The GMRES tolerance is chosen adaptively (Eisenstat–Walker style) based on
- the current scaled residual norm.
- """
-
- @jax.jit
- def loop_fn(y_init):
- def cond_fn(state):
- _, not_converged, i = state
- return (i < max_iters) & not_converged
-
- def body_fn(state):
- y1, _, i = state
-
- res, jvp = jax.linearize(F_fn, y1)
-
- scale = atol + rtol * jnp.maximum(jnp.abs(y0), jnp.abs(y1))
- norm = jnp.linalg.norm(res / scale)
-
- # Adaptive inner tolerance (Eisenstat–Walker).
- inner_tol = jnp.minimum(0.1, norm * 0.5)
-
- delta, _ = gmres(
- jvp,
- -res,
- tol=inner_tol,
- atol=atol,
- maxiter=min(20, max_iters // 2),
- )
-
- y1_next = jnp.where(norm < 1.0, y1, y1 + delta)
- return y1_next, norm >= 1.0, i + 1
-
- init_state = (y_init, True, 0)
- y1_final, _, _ = lax.while_loop(cond_fn, body_fn, init_state)
- return y1_final
-
- return loop_fn(y_init)
+ result = diffrax.RESULTS.where(
+ solution.converged,
+ diffrax.RESULTS.successful,
+ diffrax.RESULTS.promote(optx.RESULTS.nonlinear_max_steps_reached),
+ )
+ next_solver_state = MidpointSolverState(
+ solver_state.newton_iterations + solution.newton_iterations,
+ solver_state.linear_iterations + solution.linear_iterations,
+ jnp.maximum(solver_state.max_newton_iterations, solution.newton_iterations),
+ jnp.maximum(solver_state.max_linear_iterations, solution.linear_iterations),
+ )
+ return y1, y_error, dense_info, next_solver_state, result
diff --git a/tests/test_diagnostics.py b/tests/test_diagnostics.py
new file mode 100644
index 0000000..cac4234
--- /dev/null
+++ b/tests/test_diagnostics.py
@@ -0,0 +1,63 @@
+import jax.numpy as jnp
+import pytest
+
+from spectrax import diagnostics, hermite_tail_fraction
+
+
+@pytest.mark.parametrize(("Nx", "expected_energy"), [(1, 0.5), (4, 0.5), (5, 1.0)])
+def test_rfft_field_energy_weights_nyquist_only_for_even_grid(Nx, expected_energy):
+ Nx_kept = Nx // 2 + 1
+ output = {
+ "Nx": Nx,
+ "Nn": 3,
+ "Nm": 1,
+ "Np": 1,
+ "alpha_s": jnp.ones(6),
+ "u_s": jnp.zeros(6),
+ "Omega_cs": jnp.ones(2),
+ "mi_me": 1.0,
+ "Lx": 1.0,
+ "Ck": jnp.zeros((1, 6, 1, Nx_kept, 1), dtype=jnp.complex128),
+ "Fk": jnp.zeros((1, 6, 1, Nx_kept, 1), dtype=jnp.complex128)
+ .at[0, 0, 0, -1, 0]
+ .set(1.0),
+ }
+
+ diagnostics(output)
+
+ assert output["EM_energy"][0] == pytest.approx(expected_energy)
+
+
+def test_hermite_tail_fraction_uses_odd_rfft_weights():
+ dCk = jnp.zeros((1, 4, 1, 3, 1), dtype=jnp.complex128)
+ dCk = dCk.at[0, 0, 0, -1, 0].set(1.0)
+ dCk = dCk.at[0, 3, 0, 0, 0].set(1.0)
+ output = {"dCk": dCk, "Nx": 5, "Nn": 4, "Nm": 1, "Np": 1,
+ "Ns": 1, "u_s": jnp.zeros(3)}
+
+ assert hermite_tail_fraction(output)[0, 0] == pytest.approx(1 / 3)
+
+
+def test_hermite_tail_fraction_resolves_species_and_active_axes():
+ Nn, Nm, Np, Ns = 3, 2, 2, 2
+ modes = Nn * Nm * Np
+ dCk = jnp.zeros((1, Ns * modes, 1, 1, 1), dtype=jnp.complex128)
+ dCk = dCk.at[0, 0, 0, 0, 0].set(2.0)
+ dCk = dCk.at[0, Nn, 0, 0, 0].set(1.0)
+ dCk = dCk.at[0, modes, 0, 0, 0].set(1.0)
+ output = {"dCk": dCk, "Nx": 1, "Nn": Nn, "Nm": Nm, "Np": Np,
+ "Ns": Ns, "u_s": jnp.zeros(3 * Ns)}
+
+ assert hermite_tail_fraction(output)[0] == pytest.approx(jnp.array([0.2, 0.0]))
+
+
+def test_hermite_tail_fraction_ignores_inactive_axes_and_zero_norm():
+ output = {"dCk": jnp.zeros((1, 1, 1, 1, 1)), "Nx": 1,
+ "Nn": 1, "Nm": 1, "Np": 1, "Ns": 1}
+
+ assert hermite_tail_fraction(output)[0, 0] == 0.0
+
+
+def test_hermite_tail_fraction_rejects_invalid_width():
+ with pytest.raises(ValueError, match="positive integer"):
+ hermite_tail_fraction({}, width=0)
diff --git a/tests/test_initialize_maxwellian.py b/tests/test_initialize_maxwellian.py
new file mode 100644
index 0000000..13b1646
--- /dev/null
+++ b/tests/test_initialize_maxwellian.py
@@ -0,0 +1,56 @@
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from scipy.special import factorial
+
+from spectrax import compute_C_nmp
+
+jax.config.update("jax_enable_x64", True)
+
+
+def test_basis_width_maxwellian_keeps_the_closed_form():
+ """Default vth: C_nmp = sqrt(2^(n+m+p) / (n! m! p!)) (U - u)^n ... / alpha^(n+1) ... (the previous formula)."""
+ Nn, Nm, Np, Ns, alpha, u = 6, 3, 2, 2, np.array([1.0, 0.7, 1.3, 0.5, 0.9, 1.1]), np.array([0.1, -0.2, 0.3, 1.0, 0.0, -0.5])
+ U = np.random.default_rng(1).normal(size=(Ns, 3, 2, 5, 1))
+ a, w = alpha.reshape(Ns, 3), u.reshape(Ns, 3)
+ n, m, p = np.meshgrid(np.arange(Nn), np.arange(Nm), np.arange(Np), indexing="ij")
+ C = np.empty((Ns, Np, Nm, Nn, 2, 5, 1))
+ for s in range(Ns):
+ for idx in np.ndindex(Nn, Nm, Np):
+ i, j, k = idx
+ C[s, k, j, i] = (np.sqrt(2.0 ** (i + j + k) / (factorial(i) * factorial(j) * factorial(k)))
+ * (U[s, 0] - w[s, 0]) ** i * (U[s, 1] - w[s, 1]) ** j * (U[s, 2] - w[s, 2]) ** k
+ / (a[s, 0] ** (i + 1) * a[s, 1] ** (j + 1) * a[s, 2] ** (k + 1)))
+ expected = np.fft.rfftn(C, axes=(-1, -3, -2), norm="forward")
+ np.testing.assert_allclose(compute_C_nmp(U, alpha, u, Nn, Nm, Np, Ns), expected, rtol=1e-12, atol=1e-14)
+ np.testing.assert_allclose(compute_C_nmp(U, alpha, u, Nn, Nm, Np, Ns, vth_s=alpha / np.sqrt(2)), expected,
+ rtol=1e-12, atol=1e-14)
+
+
+@pytest.mark.parametrize("vth, drift", [(0.5, 0.0), (0.5, 0.8), (0.7, -0.4), (1.0, 0.3)])
+def test_narrow_maxwellian_in_a_wide_basis_is_reconstructed(vth, drift):
+ """A Maxwellian of width vth drifting by `drift` in a basis of width alpha = sqrt(2) (vth <= 1): the Hermite
+ series sum_n C_n H_n(xi) e^{-xi^2} / sqrt(2^n n! pi) must return the Gaussian (1D, x only). The series
+ converges like (1 - vth^2)^(n/2), so 200 modes reach round-off for vth >= 0.5."""
+ Nn, alpha = 200, np.sqrt(2.0)
+ U = np.full((1, 3, 1, 1, 1), 0.0)
+ U[0, 0] = drift
+ C = np.asarray(compute_C_nmp(U, [alpha, alpha, alpha], [0.0] * 3, Nn, 1, 1, 1,
+ vth_s=[vth, alpha / np.sqrt(2), alpha / np.sqrt(2)]))[0, 0, 0, :, 0, 0, 0].real
+ v = np.linspace(-4, 4, 81)
+ xi = v / alpha
+ h = np.zeros((Nn, v.size)) # orthonormal Hermite functions
+ h[0], h[1] = np.pi**-0.25 * np.exp(-xi**2 / 2), np.sqrt(2) * xi * np.pi**-0.25 * np.exp(-xi**2 / 2)
+ for n in range(2, Nn):
+ h[n] = np.sqrt(2 / n) * xi * h[n - 1] - np.sqrt((n - 1) / n) * h[n - 2]
+ f = alpha**2 * np.pi**-0.25 * np.exp(-xi**2 / 2) * (C @ h) # 1D marginal, as in inverse_HF_transform
+ exact = np.exp(-(v - drift) ** 2 / (2 * vth**2)) / (np.sqrt(2 * np.pi) * vth) # unit density
+ np.testing.assert_allclose(f, exact, atol=1e-12)
+
+
+def test_high_order_maxwellian_is_finite():
+ """Nn = 256 overflowed the factorial form (inf / inf = NaN) before the recurrence."""
+ U = jnp.linspace(-0.1, 0.1, 6).reshape(2, 3, 1, 1, 1)
+ alpha = jnp.array([0.25, 0.3, 0.35, 0.05, 0.06, 0.07])
+ assert jnp.all(jnp.isfinite(compute_C_nmp(U, alpha, jnp.zeros(6), 256, 4, 4, 2)))
diff --git a/tests/test_simulation_runs.py b/tests/test_simulation_runs.py
index d64c0aa..a8843a8 100644
--- a/tests/test_simulation_runs.py
+++ b/tests/test_simulation_runs.py
@@ -1,6 +1,31 @@
import pytest
+import diffrax
import jax.numpy as jnp
+import optimistix as optx
from spectrax import simulation
+from spectrax.midpoint_solver import (
+ ImplicitMidpoint,
+ collision_diffusion_x_streaming_preconditioner,
+)
+
+def test_multispecies_background_removed_from_perturbation():
+ """Remove every species' homogeneous density from dCk."""
+ Ns = 3
+ Ck_0 = jnp.zeros((Ns, 1, 2, 1), dtype=jnp.complex128)
+ Ck_0 = Ck_0.at[:, 0, 0, 0].set(jnp.arange(1, Ns + 1))
+ parameters = {
+ "Ck_0": Ck_0,
+ "Fk_0": jnp.zeros((6, 1, 2, 1), dtype=jnp.complex128),
+ "qs": jnp.array([-1.0, 1.0, 1.0]),
+ "alpha_s": jnp.ones(3 * Ns),
+ "u_s": jnp.zeros(3 * Ns),
+ "Omega_cs": jnp.ones(Ns),
+ "t_max": 0.001,
+ }
+
+ result = simulation(parameters, Nx=3, Nn=1, Ns=Ns, timesteps=2,
+ dt=0.001, adaptive_time_step=False)
+ assert jnp.all(result["dCk"][:, :, 0, 0, 0] == 0)
@pytest.mark.parametrize("Ns", [1, 2, 3, 4])
def test_perturbation_removes_every_background_density_and_keeps_mean_current(Ns):
@@ -41,6 +66,23 @@ def test_simulation_runs():
result["solver_stats"]["num_accepted_steps"]
+ result["solver_stats"]["num_rejected_steps"]
)
+ assert result["Ck"].shape[1:] == (result["Ns"] * result["Nn"] * result["Nm"] * result["Np"],
+ result["Ny"], result["Nx"] // 2 + 1, result["Nz"])
+ assert result["Fk"].shape[1:] == (6, result["Ny"], result["Nx"] // 2 + 1, result["Nz"])
+
+def test_solver_state_keeps_species_and_hermite_axes(monkeypatch):
+ """The integrator sees (Ck, Fk) with their species, Hermite and grid axes, not one flat vector."""
+ import spectrax._simulation as sim
+ captured = {}
+ def stop(term, **kwargs):
+ captured.update(kwargs, rhs=term.vector_field(0.0, kwargs["y0"], kwargs["args"]))
+ raise RuntimeError("captured")
+ monkeypatch.setattr(sim, "diffeqsolve", stop)
+ with pytest.raises(RuntimeError, match="captured"):
+ simulation(Nx=7, Nn=4, timesteps=3)
+ shapes = [(2, 1, 1, 4, 1, 4, 1), (6, 1, 4, 1)]
+ assert [y.shape for y in captured["y0"]] == shapes
+ assert [y.shape for y in captured["rhs"]] == shapes
def test_solver_state_keeps_species_and_hermite_axes(monkeypatch):
"""The integrator sees (Ck, Fk) with their species, Hermite and grid axes, not one flat vector."""
@@ -68,6 +110,72 @@ def test_step_limits_stop_the_solve_with_an_error(limits):
simulation(**limits)
def test_throw_false_reports_the_failure_instead_of_raising():
+ """With throw=False a stopped solve returns, and solver_result says why."""
+ from diffrax import RESULTS
+ assert simulation(throw=False)["solver_result"] == RESULTS.successful
+ assert simulation(max_steps=2, dtmin=0, throw=False)["solver_result"] == RESULTS.max_steps_reached
+
+def test_implicit_midpoint_step_structured_state():
+ term = diffrax.ODETerm(lambda t, y, args: (-y[0], 2 * y[1]))
+ y1, _, _, _, result = ImplicitMidpoint().step(
+ term, 0.0, 0.1, (jnp.ones(2), jnp.ones((2, 2))), None, None, False
+ )
+ assert result == diffrax.RESULTS.successful
+ assert jnp.allclose(y1[0], 0.95 / 1.05, rtol=1e-5)
+ assert jnp.allclose(y1[1], 1.1 / 0.9, rtol=1e-5)
+
+def test_implicit_midpoint_reports_newton_exhaustion():
+ term = diffrax.ODETerm(lambda t, y, args: y + 1)
+ result = ImplicitMidpoint(max_iters=0).step(
+ term, 0.0, 0.1, jnp.array(0.0), None, None, False
+ )[-1]
+ assert result == diffrax.RESULTS.promote(optx.RESULTS.nonlinear_max_steps_reached)
+
+def test_implicit_midpoint_evaluates_rhs_at_midpoint_time():
+ term = diffrax.ODETerm(lambda t, y, args: t)
+ y1 = ImplicitMidpoint(max_iters=2).step(
+ term, 0.0, 0.2, jnp.array(0.0), None, None, False
+ )[0]
+ assert jnp.allclose(y1, 0.02, rtol=0, atol=1e-12)
+
+def test_implicit_midpoint_simulation_reports_iterations():
+ result = simulation(
+ {"t_max": 0.01, "ode_tolerance": 1e-8}, Nx=5, Nn=3,
+ timesteps=2, dt=0.01, adaptive_time_step=False,
+ solver=ImplicitMidpoint(
+ max_iters=4, linear_restart=2, linear_max_restarts=2,
+ preconditioner=collision_diffusion_x_streaming_preconditioner,
+ ),
+ )
+ assert jnp.all(jnp.isfinite(result["Ck"]))
+ stats = result["midpoint_stats"]
+ assert stats.newton_iterations >= stats.max_newton_iterations > 0
+ assert stats.linear_iterations >= stats.max_linear_iterations > 0
+
+def test_collision_diffusion_x_streaming_line():
+ args = [None] * 21
+ args[-20], args[-19] = 2.0, 3.0
+ args[-17], args[-16], args[-15] = jnp.arange(4.0, 7.0), jnp.arange(7.0, 10.0), 10.0
+ args[-12], args[-9] = jnp.full((1, 1, 1), 2.0), jnp.full((1, 1, 1), 5.0)
+ args[-7] = jnp.arange(3.0).reshape(1, 1, 3)
+ args[-6] = jnp.sqrt(jnp.arange(1.0, 4.0)).reshape(1, 1, 1, 3, 1, 1, 1)
+ args[-5] = jnp.sqrt(jnp.arange(3.0)).reshape(1, 1, 1, 3, 1, 1, 1)
+ coefficients = jnp.arange(1.0, 4.0).astype(complex).reshape(1, 1, 1, 3, 1, 1, 1)
+ fields = jnp.ones((6, 1, 1, 1))
+
+ actual = collision_diffusion_x_streaming_preconditioner(tuple(args), 0.2)(
+ (coefficients, fields)
+ )
+ phase = 0.1j * 2 / 10
+ diagonal = 1 + 0.1 * (2 * jnp.arange(3.0) + 15) + phase * 7
+ coupling = phase * 4 / jnp.sqrt(2)
+ matrix = (jnp.diag(diagonal)
+ + jnp.diag(coupling * jnp.sqrt(jnp.arange(1.0, 3.0)), 1)
+ + jnp.diag(coupling * jnp.sqrt(jnp.arange(1.0, 3.0)), -1))
+ expected = jnp.linalg.solve(matrix, coefficients.reshape(3))
+ assert jnp.allclose(actual[0].reshape(3), expected, rtol=1e-6, atol=1e-6)
+ assert jnp.array_equal(actual[1], fields)
+
"""With throw=False a stopped solve returns, says why, and marks which saves hold a solution."""
from diffrax import RESULTS
ok = simulation(throw=False)