spectrax

SPECTRAX: a Hermite–Fourier spectral Vlasov–Maxwell solver.

This package re-exports the main user-facing functions from internal modules for convenient access (e.g. from spectrax import simulation, load_parameters).

Submodules

• spectrax.__main__
• spectrax._diagnostics
• spectrax._initialization
• spectrax._initialize_maxwellian
• spectrax._inverse_transform
• spectrax._model
• spectrax._plot
• spectrax._simulation
• spectrax.midpoint_solver

Functions

diagnostics(→ None)	Compute and attach common diagnostics to output (mutates in-place).
load_parameters(input_file)	Load simulation input parameters and solver configuration from a TOML file.
initialize_simulation_parameters([user_parameters, ...])	Assemble the parameter dictionary used to run a Hermite-Fourier Vlasov-Maxwell
compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns)	Build the Hermite-Fourier coefficients for a drifting Maxwellian distribution function.
inverse_HF_transform(Ck, Nn, Nm, Np, xi_x, xi_y, xi_z)	Reconstruct f(x, v) from Hermite–Fourier coefficients.
plasma_current(qs, alpha_s, u_s, Ck, Nn, Nm, Np, Ns)	Compute the spectral Ampère-Maxwell current from Hermite-Fourier coefficients.
Hermite_Fourier_system(Ck, C, F, kx_grid, ky_grid, ...)	Evaluate the right-hand side of the coupled Hermite-Fourier moment equations.
plot(output)	Plot common diagnostics from a completed simulation output assuming two species in a 1D setup.
cross_product(k_vec, F_vec)	Compute the cross product k × F for length-3 vectors or broadcastable arrays.
ode_system(Nx, Ny, Nz, Nn, Nm, Np, Ns, t, Ck_Fk, args)	Right-hand side for the coupled Vlasov-Maxwell system expressed in spectral form.
simulation([input_parameters, Nx, Ny, Nz, Nn, Nm, Np, ...])	Run a spectral Vlasov-Maxwell simulation and return the solution together with

Package Contents

spectrax.diagnostics(output: dict) → None

Compute and attach common diagnostics to output (mutates in-place).

Parameters:

• output (The following keys are added to) –

  Dictionary holding simulation results and parameters. This function expects at least the following keys:

  • alpha_s: array of shape (3*Ns,) (thermal scales per species)

  • u_s: array of shape (3*Ns,) (drift velocities per species)

  • Ck: Hermite-Fourier coefficients of shape (Nt, Ns*Hs, Ny, Nx, Nz)

  • Fk: field Fourier coefficients of shape (Nt, 6, Ny, Nx, Nz)

  • Omega_cs: array of shape (Ns,) (cyclotron frequencies)

  • Lx: domain length in x (used for k_norm)

  Optionally, it may use: - Nn, Nm, Np to decode flattened Hermite indices - mi_me for legacy 2-species Debye length and mass scaling - ms or masses for general per-species mass scaling

• keys (Added)

• ----------

• output

• Ns (-)

• lambda_D (-)

• k_norm (-)

• kinetic_energy_species (-)

• kinetic_energy (-)

• EM_energy (-)

• total_energy (-)

• keys) (- kinetic_energy_species{j} for j=1..Ns (back-compat convenience)

Returns:

output is mutated in-place.

Return type:

None

spectrax.load_parameters(input_file)

Load simulation input parameters and solver configuration from a TOML file.

Parameters:

input_file (str or pathlib.Path) – Path to the TOML file containing simulation parameters.

Returns:

A pair (input_parameters, solver_parameters) where solver_parameters includes an instantiated Diffrax solver ready for diffeqsolve.

Return type:

tuple[dict, dict]

spectrax.initialize_simulation_parameters(user_parameters={}, Nx=33, Ny=1, Nz=1, Nn=50, Nm=1, Np=1, Ns=2, timesteps=500, dt=0.01)

Assemble the parameter dictionary used to run a Hermite-Fourier Vlasov-Maxwell simulation, starting from library defaults and overriding them with user input. The defaults include a two-stream perturbation, precomputed spectral grids, and helper tables required by the RHS evaluation. Derived quantities are evaluated after merging with any user-provided overrides so that dependent fields remain consistent.

Parameters:

• user_parameters (Mapping, optional) – Optional dictionary of parameter overrides. Any key present here replaces the corresponding default before derived quantities are computed.

• Nx (int, optional) – Number of Fourier modes along each spatial direction.

• Ny (int, optional) – Number of Fourier modes along each spatial direction.

• Nz (int, optional) – Number of Fourier modes along each spatial direction.

• Nn (int, optional) – Number of Hermite modes along each velocity-space axis.

• Nm (int, optional) – Number of Hermite modes along each velocity-space axis.

• Np (int, optional) – Number of Hermite modes along each velocity-space axis.

• Ns (int, optional) – Number of particle species represented in the simulation.

• timesteps (int, optional) – Number of time samples to store in the solution.

• dt (float, optional) – Initial guess for the integrator time step.

Returns:

Dictionary containing the merged parameters, derived helper arrays, and initial spectral coefficients such as Ck_0 and Fk_0.

Return type:

dict

spectrax.compute_C_nmp(Us_grid, alpha_s, u_s, Nn, Nm, Np, Ns)

Build the Hermite-Fourier coefficients for a drifting Maxwellian distribution function.

Parameters:

• Us_grid (jnp.ndarray) – Velocity grid for each species with shape (Ns, 3, Ny, Nx, Nz), containing the phase-space coordinates at which the Maxwellian is sampled.

• alpha_s (array-like) – Sequence of thermal scaling parameters flattened as (3 * Ns,), grouped as (alpha_x, alpha_y, alpha_z) per species.

• u_s (array-like) – Sequence of drift velocities flattened as (3 * Ns,), grouped as (u_x, u_y, u_z) per species.

• Nn (int) – Number of Hermite modes retained along the x, y, and z velocity axes.

• Nm (int) – Number of Hermite modes retained along the x, y, and z velocity axes.

• Np (int) – Number of Hermite modes retained along the x, y, and z velocity axes.

• Ns (int) – Number of species.

Returns:

Complex Hermite-Fourier coefficients with shape (Ns, Np, Nm, Nn, Ny, Nx, Nz) corresponding to the Maxwellian evaluated on the supplied grid.

Return type:

jnp.ndarray

spectrax.inverse_HF_transform(Ck, Nn, Nm, Np, xi_x, xi_y, xi_z)

Reconstruct f(x, v) from Hermite–Fourier coefficients.

Parameters:

• Ck (jnp.ndarray) – Hermite-Fourier coefficients for (one or more) species/moments. The inverse FFT is applied along the last three axes.

• Nn (int) – Number of Hermite modes retained along each velocity-space axis.

• Nm (int) – Number of Hermite modes retained along each velocity-space axis.

• Np (int) – Number of Hermite modes retained along each velocity-space axis.

• xi_x (jnp.ndarray) – Normalized velocity-space coordinates (typically (v - u)/alpha) used to evaluate Hermite polynomials and the Gaussian weight.

• xi_y (jnp.ndarray) – Normalized velocity-space coordinates (typically (v - u)/alpha) used to evaluate Hermite polynomials and the Gaussian weight.

• xi_z (jnp.ndarray) – Normalized velocity-space coordinates (typically (v - u)/alpha) used to evaluate Hermite polynomials and the Gaussian weight.

Returns:

The reconstructed distribution function evaluated on (t, x, y, z, xi_x, xi_y, xi_z).

Return type:

jnp.ndarray

spectrax.plasma_current(qs, alpha_s, u_s, Ck, Nn, Nm, Np, Ns)

Compute the spectral Ampère-Maxwell current from Hermite-Fourier coefficients.

Parameters:

• qs (jnp.ndarray, shape (Ns,)) – Charges of the species.

• alpha_s (jnp.ndarray, shape (3 * Ns,)) – Velocity scaling factors for each species.

• u_s (jnp.ndarray, shape (3 * Ns,)) – Velocity shift for each species.

• Ck (jnp.ndarray, shape (Ns * Np * Nm * Nn, Ny, Nx, Nz)) – Hermite-Fourier coefficients for all species stacked along the first axis.

• Nn (int) – Number of Hermite modes in x, y, and z respectively.

• Nm (int) – Number of Hermite modes in x, y, and z respectively.

• Np (int) – Number of Hermite modes in x, y, and z respectively.

• Ns (int) – Number of species.

Returns:

The total Ampère-Maxwell current components (Jx, Jy, Jz).

Return type:

jnp.ndarray, shape (3, Ny, Nx, Nz)

spectrax.Hermite_Fourier_system(Ck, C, F, kx_grid, ky_grid, kz_grid, k2_grid, col, sqrt_n_plus, sqrt_n_minus, sqrt_m_plus, sqrt_m_minus, sqrt_p_plus, sqrt_p_minus, Lx, Ly, Lz, nu, D, alpha_s, u_s, qs, Omega_cs, Nn, Nm, Np, Ns, mask23)

Evaluate the right-hand side of the coupled Hermite-Fourier moment equations.

Parameters:

• Ck (jnp.ndarray) – Hermite-Fourier coefficients with shape (Ns * Np * Nm * Nn, Ny, Nx, Nz).

• C (jnp.ndarray) – Configuration-space Hermite coefficients (i.e., inverse FFT of Ck), with shape (Ns * Np * Nm * Nn, Ny, Nx, Nz). The array is reshaped internally to separate the species and Hermite indices.

• F (jnp.ndarray) – Configuration-space electromagnetic fields with shape (6, Ny, Nx, Nz) ordered as (Ex, Ey, Ez, Bx, By, Bz).

• kx_grid (jnp.ndarray) – Fourier wave-number grids scaled to the physical domain length.

• ky_grid (jnp.ndarray) – Fourier wave-number grids scaled to the physical domain length.

• kz_grid (jnp.ndarray) – Fourier wave-number grids scaled to the physical domain length.

• k2_grid (jnp.ndarray) – Squared magnitude of the wave number.

• col (jnp.ndarray) – Precomputed collision coefficients.

• sqrt_* (jnp.ndarray) – Square-root ladder coefficients for the Hermite recurrences along each axis.

• Lx (float) – Domain lengths in each spatial direction.

• Ly (float) – Domain lengths in each spatial direction.

• Lz (float) – Domain lengths in each spatial direction.

• nu (float) – Collision frequency.

• D (float) – Hyper-diffusion coefficient.

• alpha_s (jnp.ndarray) – Thermal scaling parameters and drift velocities per species.

• u_s (jnp.ndarray) – Thermal scaling parameters and drift velocities per species.

• qs (jnp.ndarray) – Species charges.

• Omega_cs (jnp.ndarray) – Cyclotron frequencies per species.

• Nn (int) – Number of Hermite modes and species.

• Nm (int) – Number of Hermite modes and species.

• Np (int) – Number of Hermite modes and species.

• Ns (int) – Number of Hermite modes and species.

• mask23 (jnp.ndarray) – Boolean mask implementing the 2/3 de-aliasing rule in Fourier space.

Returns:

Time derivative dCk/dt with shape (Ns, Np, Nm, Nn, Ny, Nx, Nz).

Return type:

jnp.ndarray

spectrax.plot(output)

Plot common diagnostics from a completed simulation output assuming two species in a 1D setup.

Parameters:

output (dict) – Simulation output dictionary containing keys like time, Ck, Fk, energies computed by spectrax._diagnostics.diagnostics(), and basic grid parameters.

spectrax.cross_product(k_vec, F_vec)

Compute the cross product k × F for length-3 vectors or broadcastable arrays.

Parameters:

• k_vec (array-like) – First vector with leading dimension 3.

• F_vec (array-like) – Second vector with leading dimension 3.

Returns:

Array representing the cross product with the same trailing shape as the inputs.

Return type:

jnp.ndarray

spectrax.ode_system(Nx, Ny, Nz, Nn, Nm, Np, Ns, t, Ck_Fk, args)

Right-hand side for the coupled Vlasov-Maxwell system expressed in spectral form.

Parameters:

• Nx (int) – Number of Fourier modes per spatial dimension.

• Ny (int) – Number of Fourier modes per spatial dimension.

• Nz (int) – Number of Fourier modes per spatial dimension.

• Nn (int) – Number of Hermite modes per velocity-space dimension.

• Nm (int) – Number of Hermite modes per velocity-space dimension.

• Np (int) – Number of Hermite modes per velocity-space dimension.

• Ns (int) – Number of species.

• t (float) – Integration time (unused but required by Diffrax interface).

• Ck_Fk (jnp.ndarray) – Flattened state vector containing concatenated Hermite coefficients followed by electromagnetic field coefficients.

• args (tuple) – Cached parameter tuple produced in simulation providing physical constants, grids, and helper arrays.

Returns:

Flattened derivative vector matching the shape of Ck_Fk.

Return type:

jnp.ndarray

spectrax.simulation(input_parameters={}, Nx=33, Ny=1, Nz=1, Nn=20, Nm=1, Np=1, Ns=2, timesteps=200, dt=0.01, solver=Dopri8(), adaptive_time_step=True)

Run a spectral Vlasov-Maxwell simulation and return the solution together with the parameter dictionary used to produce it.

Parameters:

• input_parameters (dict, optional) – User-specified overrides passed to initialize_simulation_parameters.

• Nx (int, optional) – Number of retained Fourier modes per spatial direction.

• Ny (int, optional) – Number of retained Fourier modes per spatial direction.

• Nz (int, optional) – Number of retained Fourier modes per spatial direction.

• Nn (int, optional) – Number of Hermite modes per velocity-space axis.

• Nm (int, optional) – Number of Hermite modes per velocity-space axis.

• Np (int, optional) – Number of Hermite modes per velocity-space axis.

• Ns (int, optional) – Number of species.

• timesteps (int, optional) – Number of solution snapshots to save between t=0 and t_max.

• dt (float, optional) – Initial integration step size.

• solver (diffrax.AbstractSolver, optional) – Diffrax solver instance controlling the time integration.

Returns:

Dictionary containing the evolved coefficients (Ck, Fk), time samples, perturbation diagnostics, and all simulation parameters.

Return type:

dict