Ohm solver: Electron energy equation

In these examples the electron temperature \(T_e\) used for the electron pressure term of the generalized Ohm’s law is evolved with the electron energy equation

\[\frac{\partial U_e}{\partial t} + \nabla\cdot(U_e \mathbf{V}_e) + P_e \nabla\cdot\mathbf{V}_e = \eta J^2 - Q_{ei},\]

where \(U_e = n_e k_B T_e/(\gamma_e - 1)\) is the electron internal energy density, solved with the QDSMC kinetic-enslaving scheme of Belyaev et al. [18] (hybrid_pic_model.solve_electron_energy_equation). Each of the four tests below isolates one piece of the equation with an exact analytic solution: the transport terms on the left-hand side (adiabatic compression, and slab transport through a below-floor halo), the Joule-heating source (force-free field decay), and the electron-ion temperature-relaxation sink \(Q_{ei}\).

Adiabatic compression

With all sources off, entropy-conserving transport of an initially uniform entropy requires the pointwise adiabat

\[T_e(x, t) = T_{e0} \left( \frac{n(x,t)}{n_0} \right)^{\gamma_e - 1}\]

at every cell and time, independent of the flow. A uniform, unmagnetized, zero-resistivity plasma is given a sinusoidal ion velocity perturbation which drives an electron-pressure ion-acoustic compression/rarefaction wave, and the measured \(T_e\) is compared against the adiabat.

Run

Script inputs_test_2d_ohm_solver_electron_energy_picmi.py
Listing 75 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/inputs_test_2d_ohm_solver_electron_energy_picmi.py. One script covers all four cases; select this one with --case adiabat.
#!/usr/bin/env python3
#
# --- Test suite for the hybrid-PIC (Ohm's law) electron energy equation.
# --- One script, four cases (selected with --case), each isolating one
# --- term of the equation in a 2D Cartesian (x,z) periodic box:
# ---
# ---   adiabat : transport terms only (B=0, eta=0, all sources off),
# ---                 dU_e/dt + div(U_e V_e) + P_e div(V_e) = 0,
# ---             solved by the QDSMC scheme advecting the electron entropy
# ---             K_e = T_e n_e^(1-gamma) with Lagrangian markers moving at
# ---             V_e.  A sinusoidal ion velocity perturbation v_x = V0
# ---             sin(kx) drives an ion-acoustic compression, and entropy
# ---             conservation gives the pointwise check
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1).
# ---             Analyse with analysis_adiabat.py.
# ---
# ---   joule   : eta*J^2 source only.  A linear force-free field
# ---                 B(x) = B0 [0, sin(kx), cos(kx)],  k = 2 pi/Lx,
# ---             carries the uniform, parallel current J = curl(B)/mu0
# ---             (J x B = 0): no bulk motion, uniform heating, transport
# ---             terms identically zero, so
# ---                 dT_e/dt = (gamma_e - 1) eta J^2 / (n_e k_B),
# ---             a linear ramp turning sub-linear on the resistive-decay
# ---             time tau_R = mu0/(eta k^2).  Analyse with
# ---             analysis_joule.py (B-energy decay + T_e ramp).
# ---
# ---   qei     : electron-ion thermal-equilibration sink only (B=0, eta=0),
# ---                 dU_e/dt = -Q_ei,  Q_ei = 3 n_e k_B nu_ei (T_e - T_i),
# ---             enabled by the rate parser
# ---             hybrid_pic_model.electron_ion_relaxation_rate(rho,Te,Ti,t).
# ---             The single parameter enables BOTH the electron-side sink
# ---             AND the conjugate ion heating, so the exchange conserves
# ---             energy.  For constant nu_ei (single proton species, Z=1),
# ---                 (T_e - T_i)(t) = (T_e0 - T_i0) exp(-rate t),
# ---                 rate = [3(gamma_e-1) + 2] nu_ei,
# ---             and C_e T_e + C_i T_i is conserved.  Analyse with
# ---             analysis_qei.py (difference rate + budget).
# ---
# ---   vacuum  : transport through a below-floor halo (B=0, eta=0, all
# ---             sources off).  A slab (n = n0) drifts at c_s through a
# ---             halo with n = 0.02 n0, below the solver's n_floor,
# ---             starting from uniform entropy K_e, so entropy-conserving
# ---             transport must keep
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1)
# ---             pointwise at all times: the below-floor halo must act as
# ---             an insulating boundary, not a heat sink.  Analyse with
# ---             analysis_vacuum.py.

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np
from mpi4py import MPI as mpi

from pywarpx import picmi

constants = picmi.constants
comm = mpi.COMM_WORLD

simulation = None


class ElectronEnergyCase(object):
    """Shared 2D periodic-box setup; the case subclasses supply the physics
    parameters, the solver sources and the diagnostic field list."""

    # ---- Common plasma parameters ------------------------------------------
    gamma_e = 5.0 / 3.0  # electron adiabatic index
    n0 = 2.0e20  # uniform density (m^-3)
    Lx = 0.5  # domain length in x (m)

    # ---- Case hooks (overridden where a case differs) -----------------------
    include_joule_heating = False
    relaxation_rate = None  # qei only: nu_ei expression (str)
    eta_h = None  # joule only: hyper-resistivity
    # qei leaves do_temperature_deposition unset on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is configured,
    # which that case also exercises.
    set_temperature_deposition = True
    load_B = False  # joule only: force-free initial field
    reduced_diags = False  # joule only: field/particle energy

    def __init__(self, test, verbose):
        self.test = test
        self.verbose = verbose or test

        self.configure()
        self.get_plasma_quantities()
        if comm.rank == 0:
            self._print_params()
        self.setup_run()

    def momentum_expressions(self):
        return ["0", "0", "0"]

    def density_expression(self):
        return "n0"

    def setup_run(self):
        global simulation

        self.grid = picmi.Cartesian2DGrid(
            number_of_cells=[self.NX, self.NZ],
            lower_bound=[0.0, -self.Lz / 2.0],
            upper_bound=[self.Lx, self.Lz / 2.0],
            lower_boundary_conditions=["periodic", "periodic"],
            upper_boundary_conditions=["periodic", "periodic"],
            lower_boundary_conditions_particles=["periodic", "periodic"],
            upper_boundary_conditions_particles=["periodic", "periodic"],
            warpx_max_grid_size=self.NZ,
        )

        # Electron energy equation ON; each case turns on exactly one source
        # (or none, for the pure-transport adiabat case).
        solver_kwargs = {}
        if self.eta_h is not None:
            solver_kwargs["plasma_hyper_resistivity"] = self.eta_h
        if self.relaxation_rate is not None:
            solver_kwargs["electron_ion_relaxation_rate"] = self.relaxation_rate
        self.solver = picmi.HybridPICSolver(
            grid=self.grid,
            gamma=self.gamma_e,
            Te=self.te_eV,
            n0=self.n0,
            n_floor=0.05 * self.n0,
            plasma_resistivity=self.eta,
            substeps=self.substeps,
            solve_electron_energy_equation=True,
            include_joule_heating=self.include_joule_heating,
            **solver_kwargs,
        )

        simulation = picmi.Simulation(
            solver=self.solver,
            time_step_size=self.dt,
            max_steps=self.total_steps,
            verbose=self.verbose,
            particle_shape=1,
            warpx_serialize_initial_conditions=True,
            warpx_current_deposition_algo="direct",
            warpx_use_filter=True,
        )

        if self.load_B:
            B_init = picmi.LoadInitialFieldFromPython(
                load_from_python=self.load_initial_B,
                load_B=True,
                load_E=False,
            )
            simulation.add_applied_field(B_init)

        species_kwargs = {}
        if self.set_temperature_deposition:
            species_kwargs["warpx_do_temperature_deposition"] = True
        self.ions = picmi.Species(
            name="ions",
            charge="q_e",
            mass=constants.m_p,
            initial_distribution=picmi.AnalyticDistribution(
                density_expression=self.density_expression(),
                momentum_expressions=self.momentum_expressions(),
                warpx_momentum_spread_expressions=[str(self.vi_th)] * 3,
                n0=self.n0,
            ),
            **species_kwargs,
        )
        simulation.add_species(
            self.ions,
            layout=picmi.PseudoRandomLayout(
                grid=self.grid, n_macroparticles_per_cell=self.NPPC
            ),
        )

        # Remove any diags from a previous run in the same directory, so
        # stale openPMD dumps (one file per iteration) cannot mix into the
        # analysis of this run.
        if comm.rank == 0 and Path("diags").exists():
            shutil.rmtree("diags")
        comm.Barrier()

        field_diag = picmi.FieldDiagnostic(
            name="field_diag",
            grid=self.grid,
            period=self.diag_steps,
            data_list=self.diag_data_list,
            write_dir="diags",
            warpx_file_prefix="field_diags",
            warpx_format="openpmd",
            warpx_openpmd_backend="h5",
        )
        simulation.add_diagnostic(field_diag)

        if self.reduced_diags:
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="FieldEnergy",
                    name="field_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="ParticleEnergy",
                    name="part_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )

        simulation.initialize_inputs()
        simulation.initialize_warpx()


class AdiabaticCompression(ElectronEnergyCase):
    """Transport-terms (LHS) test: entropy-conserving compression."""

    te_eV = 100.0  # initial (uniform) electron temperature (eV)
    ti_eV = 10.0  # ion temperature (eV); cold vs Te for a clean,
    #   electron-pressure-driven acoustic wave

    # ---- Perturbation -------------------------------------------------------
    pert_frac = 0.30  # ion velocity amplitude V0 = pert_frac * c_s
    n_wave = 1  # wavelengths across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    periods = 2.0  # acoustic periods to simulate
    steps_per_period = 400
    substeps = 10

    diag_data_list = ["rho", "Te", "J", "B"]

    def configure(self):
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self._steps_override = 60
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 40

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Electron-pressure sound speed (cold-ion limit) sets the wave period.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.omega = self.k * self.c_s  # acoustic angular frequency
        self.T_period = 2.0 * np.pi / self.omega  # = Lx / c_s for n_wave=1
        self.V0 = self.pert_frac * self.c_s  # velocity perturbation amplitude

        self.dt = self.T_period / self.steps_per_period
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.periods * self.steps_per_period)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def momentum_expressions(self):
        # Sinusoidal x-velocity perturbation v_x = V0 sin(kx); uniform n0 and
        # uniform Te0 -> uniform initial entropy.
        return [f"({self.V0})*sin(({self.k})*x)", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Adiabatic-compression (electron-energy-equation LHS) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  c_s       = {self.c_s:.3e} m/s   (electron-pressure sound speed)\n"
            f"  V0        = {self.V0:.3e} m/s   (= {self.pert_frac:.2f} c_s)\n"
            f"  k         = {self.k:.4e} 1/m   ({self.n_wave} wavelength(s))\n"
            f"  T_period  = {self.T_period:.3e} s   (acoustic)\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   ({self.steps_per_period}/period)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, pure advection+compression\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise\n"
        )


class VacuumSlabTransport(ElectronEnergyCase):
    """Insulating-halo transport test: a plasma slab drifting through a
    below-floor halo must keep its electron entropy.

    A slab (n = n0) drifts at v0 = c_s through a tenuous halo whose density
    sits BELOW the solver's n_floor. The initial T_e is the floored adiabat
    (uniform entropy K_e), and there are no sources (B = 0, eta = 0), so
    entropy-conserving transport requires T_e = Te0 (n/n0)^(gamma-1)
    pointwise at all times -- exactly, even through the CIC-mixed slab edge,
    because a uniform K_e is invariant under any mass-weighted mixing.

    This guards the insulating treatment of below-floor cells: if the QDSMC
    transport left K_e = 0 there (instead of flooring the density in the
    K_e <-> T_e conversion), the halo would dilute and erase the slab's
    entropy at the drifting edge and T_e would fall off the adiabat within
    tens of steps.
    """

    te_eV = 100.0  # slab electron temperature (eV) at n0
    ti_eV = 10.0  # ion temperature (eV); cold, so the slab holds together

    # ---- Slab / halo geometry -----------------------------------------------
    halo_frac = 0.02  # halo density fraction of n0; BELOW the 0.05 n_floor
    slab_frac = 0.5  # slab width as a fraction of Lx
    cfl_marker = 0.2  # QDSMC marker displacement per step, v0 dt / dx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    substeps = 10

    diag_data_list = ["rho", "Te"]

    def configure(self):
        if self.test:
            self.NX = 64
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 8
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Drift at the electron-pressure sound speed: markers stream through
        # the slab edge at cfl_marker cells per step.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.v0 = self.c_s

        self.dt = self.cfl_marker * self.dx / self.v0
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = 400
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def density_expression(self):
        x0 = 0.5 * self.Lx
        hw = 0.5 * self.slab_frac * self.Lx
        return f"n0*({self.halo_frac} + (1 - {self.halo_frac})*(abs(x - {x0}) < {hw}))"

    def momentum_expressions(self):
        # Uniform drift: the slab translates without compression.
        return [f"{self.v0}", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Vacuum-slab (insulating-halo) transport test\n"
            f"  Te0 = {self.te_eV:.1f} eV (at n0),  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3,  halo = {self.halo_frac:.3f} n0 (below the 0.05 n0 floor)\n"
            f"  slab      = {self.slab_frac:.2f} Lx wide, drifting at v0 = c_s = {self.v0:.3e} m/s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (marker CFL {self.cfl_marker:.2f} cells/step)\n"
            f"  steps     = {self.total_steps} (slab travels "
            f"{self.cfl_marker * self.total_steps / self.NX:.2f} Lx),  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  no sources, pure transport through the halo\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise (uniform K_e)\n"
        )


class ForceFreeJoule(ElectronEnergyCase):
    """eta*J^2 source test: force-free field, uniform Joule ramp."""

    ti_eV = 500.0  # ion temperature (eV)
    te_eV = 500.0  # initial electron temperature (eV)

    # ---- Force-free field ---------------------------------------------------
    B0 = 0.1  # field magnitude (T); |B| is uniform
    n_wave = 1  # number of full wavelengths of B across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128  # cells in x (full run)
    NZ = 16  # cells in z (field is z-independent; periodic)
    NPPC = 800  # particles per cell; the T_e ramp sits on an ion shot-noise
    #   heating floor that converges as 1/NPPC
    DT = 0.0025  # timestep as a fraction of the ion cyclotron period; small
    #   enough for the forward-Euler Joule deposit to be converged
    TOTAL_STEPS = 3000  # full run
    DIAG_EVERY = 150  # diagnostic cadence (steps)
    substeps = 20

    include_joule_heating = True
    load_B = True
    reduced_diags = True
    diag_data_list = ["B", "E", "rho", "J", "Te", "T_ions"]

    def configure(self):
        self.eta_scale = self.args.eta_scale
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self.DT = 0.01
            self.total_steps = 50
            self.diag_steps = 10
        else:
            self.total_steps = self.TOTAL_STEPS
            self.diag_steps = self.DIAG_EVERY

    def get_plasma_quantities(self):
        mi = constants.m_p

        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ  # square cells
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Uniform plasma current magnitude from curl(B) = k B.
        self.J0 = self.k * self.B0 / constants.mu0  # A/m^2
        # Electron drift carrying it (ions start at rest): V_e = J/(e n0).
        self.v_drift = self.J0 / (constants.q_e * self.n0)

        # Ion cyclotron period at B0 sets the timestep scale.
        self.w_ci = constants.q_e * self.B0 / mi
        self.t_ci = 2.0 * np.pi / self.w_ci
        self.dt = self.DT * self.t_ci

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)

        # Constant resistivity (Ohm*m), scaled for the heating-signal amplitude.
        self.eta = 1.0e-5 * self.eta_scale
        # Resistive decay time of the force-free current: tau_R = mu0/(eta k^2).
        self.tau_R = constants.mu0 / (self.eta * self.k**2)

        # Hyper-resistivity off: grid-scale damping is not needed for a smooth,
        # single-wavelength field and would complicate the eta*J^2 budget.
        self.eta_h = 0.0

        # Analytic prediction (for the printout / cross-check).
        self.dTe_dt_pred = (
            (self.gamma_e - 1.0) * self.eta * self.J0**2 / (self.n0 * constants.kb)
        )  # K/s

    def _print_params(self):
        print(
            f"\n[setup] Force-free Joule-heating test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  B0        = {self.B0:.3e} T  (|B| uniform)\n"
            f"  k         = {self.k:.4e} 1/m  ({self.n_wave} wavelength(s) across Lx)\n"
            f"  |J|       = {self.J0:.3e} A/m^2  (uniform, force-free)\n"
            f"  V_e drift = {self.v_drift:.3e} m/s\n"
            f"  eta       = {self.eta:.3e} Ohm*m  (1e-5 x scale {self.eta_scale:g})\n"
            f"  tau_R     = {self.tau_R:.3e} s  (current resistive-decay time)\n"
            f"  Grid      = {self.NX} x {self.NZ}  (x x z),  dx = {self.dx:.3e} m\n"
            f"  t_ci      = {self.t_ci:.3e} s,  dt = {self.dt:.3e} s ({self.DT:.4f} t_ci)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  ----\n"
            f"  PREDICTED dTe/dt = (gamma_e-1) eta J^2 / (n0 kB) = {self.dTe_dt_pred:.4e} K/s\n"
            f"                   = {self.dTe_dt_pred * constants.kb / constants.q_e:.4e} eV/s\n"
        )

    def load_initial_B(self):
        """Set the linear force-free field B(x) = B0[0, sin(kx), cos(kx)].

        WarpX folds Bfield_fp_external into Bfield_fp at initialization, after
        which it evolves self-consistently. div(B) = 0 analytically (B has no
        x-component and the others depend only on x), so no cleaning needed.
        """
        Bx = simulation.fields.get("Bfield_fp_external", dir="x", level=0)
        By = simulation.fields.get("Bfield_fp_external", dir="y", level=0)
        Bz = simulation.fields.get("Bfield_fp_external", dir="z", level=0)

        Bx[:, :] = 0.0
        # Each component on its own (possibly staggered) mesh.
        XBy, _ = np.meshgrid(By.mesh("x"), By.mesh("z"), indexing="ij")
        XBz, _ = np.meshgrid(Bz.mesh("x"), Bz.mesh("z"), indexing="ij")
        By[:, :] = self.B0 * np.sin(self.k * XBy)
        Bz[:, :] = self.B0 * np.cos(self.k * XBz)
        comm.Barrier()


class QeiRelaxation(ElectronEnergyCase):
    """Q_ei electron-ion thermal-equilibration test: pure exponential."""

    te_eV = 300.0  # initial (uniform) electron temperature (eV), hot
    ti_eV = 50.0  # ion temperature (eV); the relaxation target

    # ---- Relaxation ---------------------------------------------------------
    nu_ei = 1.0e6  # electron-ion relaxation rate (1/s), constant;
    #   Te sink rate = 3(gamma_e-1)*nu_ei = 2e6 1/s -> tau = 0.5 us

    # ---- Geometry (small; the physics is 0-D / uniform) / numerics ----------
    NX = 32
    NZ = 8
    NPPC = 400
    n_tau = 3.0  # number of relaxation times to simulate
    steps_per_tau = 100  # rate*dt = 0.01 -> forward-Euler ~ exponential
    substeps = 10

    # do_temperature_deposition is NOT set on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is
    # configured, which this test also exercises (T_ions is dumped below).
    set_temperature_deposition = False
    diag_data_list = ["rho", "Te", "T_ions"]

    def configure(self):
        if self.test:
            self.NX = 16
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Analytic electron-sink rate and e-folding time.
        self.rate = 3.0 * (self.gamma_e - 1.0) * self.nu_ei
        self.tau = 1.0 / self.rate

        self.dt = self.tau / self.steps_per_tau
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.n_tau * self.steps_per_tau)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure Q_ei).
        self.eta = 0.0
        # Constant relaxation rate so the relaxation is a pure exponential.
        self.relaxation_rate = f"{self.nu_ei}"

    def _print_params(self):
        print(
            f"\n[setup] Electron-ion relaxation (Q_ei) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti0 = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  nu_ei     = {self.nu_ei:.3e} 1/s   (constant)\n"
            f"  rate      = 3(gamma-1)nu_ei = {self.rate:.3e} 1/s\n"
            f"  tau       = 1/rate = {self.tau:.3e} s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (rate*dt = {self.rate * self.dt:.3f})\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, Q_ei ON (e-sink + conjugate ion heating)\n"
            f"  CHECK:  (Te-Ti)(t) = (Te0-Ti0) exp(-[3(g-1)+2]nu t),  energy conserved\n"
        )


CASES = {
    "adiabat": AdiabaticCompression,
    "joule": ForceFreeJoule,
    "qei": QeiRelaxation,
    "vacuum": VacuumSlabTransport,
}

parser = argparse.ArgumentParser()
parser.add_argument(
    "--case",
    required=True,
    choices=sorted(CASES.keys()),
    help="which electron-energy-equation term to test",
)
parser.add_argument(
    "-t",
    "--test",
    help="toggle whether this script is run as a short CI test",
    action="store_true",
)
parser.add_argument(
    "-v",
    "--verbose",
    help="Verbose output",
    action="store_true",
)
parser.add_argument(
    "--eta-scale",
    type=float,
    default=1.0,
    help="joule case only: multiplier on the base resistivity eta=1e-5 "
    "(amplifies the eta*J^2 heating signal; the CI test uses 100)",
)
args, left = parser.parse_known_args()
sys.argv = sys.argv[:1] + left

case_class = CASES[args.case]
case_class.args = args
run = case_class(test=args.test, verbose=args.verbose)
simulation.step()

Execute:

python3 inputs_test_2d_ohm_solver_electron_energy_picmi.py --case adiabat

Analyze

Script analysis_adiabat.py
Listing 76 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/analysis_adiabat.py.
#!/usr/bin/env python3
"""Validate the electron-energy-equation LHS (advection + compression) via
the adiabat.

For a sources-off, uniform-initial-entropy compression run, entropy-conserving
transport requires, at every cell and time,

    T_e(x,t) = T_e0 * ( n(x,t) / n0 )^(gamma - 1).

Low-density cells (below the n_floor used by the solver) are masked, since
T_e is gated there.

Produces:
  * left  : T_e(x) measured (solid) vs Te0 (n/n0)^(gamma-1) (dashed) at
            several times;
  * right : a scatter of T_e/Te0 vs n/n0 (all cells & times) that must
            collapse onto the single adiabat curve y = x^(gamma-1),
and checks the pointwise relative error against --tol-median / --tol-max.
"""

import argparse
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries

Q_E = 1.602176634e-19
K_B = 1.380649e-23
K_PER_EV = K_B / Q_E


def main(argv=None):
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--diag-dir", default="diags/field_diags")
    ap.add_argument("--gamma", type=float, default=5.0 / 3.0)
    ap.add_argument(
        "--n-floor-frac",
        type=float,
        default=0.05,
        help="mask cells with n < this fraction of n0 (matches the solver n_floor)",
    )
    ap.add_argument(
        "--tol-median",
        type=float,
        default=0.005,
        help="allowed median pointwise relative error on the adiabat",
    )
    ap.add_argument(
        "--tol-max",
        type=float,
        default=0.02,
        help="allowed max pointwise relative error on the adiabat",
    )
    ap.add_argument("--out", default="adiabat_check.png")
    args = ap.parse_args(argv)

    ts = OpenPMDTimeSeries(args.diag_dir)
    its = list(ts.iterations)
    times = np.asarray(ts.t, dtype=float)
    if len(its) < 2:
        raise SystemExit(f"Need >=2 dumps in {args.diag_dir}")
    g1 = args.gamma - 1.0

    def zavg(name, it):
        arr, info = ts.get_field(name, iteration=it)
        return np.asarray(info.x, dtype=float), np.asarray(arr, dtype=float).mean(
            axis=0
        )

    coord = None
    Te_x, n_x = [], []
    for it in its:
        cc, Te = zavg("Te", it)
        _, rho = zavg("rho", it)
        if coord is None:
            coord = cc
        Te_x.append(Te * K_PER_EV)
        n_x.append(rho / Q_E)
    Te_x = np.array(Te_x)  # (nt, nx)
    n_x = np.array(n_x)

    # Reference state = median of the first dump (uniform initial fill).
    Te0 = float(np.median(Te_x[0]))
    n0 = float(np.median(n_x[0]))
    Te_pred = Te0 * (n_x / n0) ** g1

    valid = n_x > args.n_floor_frac * n0
    rel = np.abs(Te_x - Te_pred) / np.maximum(Te_x, 1e-30)
    # Score only meaningfully compressed cells above the density floor.
    sig = valid & (np.abs(n_x / n0 - 1.0) > 0.03)
    med = float(np.median(rel[sig])) if np.any(sig) else float("nan")
    mx = float(np.max(rel[sig])) if np.any(sig) else float("nan")

    dn = float(np.max(np.abs((n_x / n0 - 1.0)[valid])))
    print("=" * 62)
    print("Adiabatic-compression check  Te = Te0 (n/n0)^(gamma-1)")
    print(f"  gamma = {args.gamma:.5f}   Te0 = {Te0:.2f} eV   n0 = {n0:.3e} m^-3")
    print(f"  peak density swing |n/n0 - 1| = {dn:.1%}  (cells above n_floor)")
    print(
        f"  relative error (compressed, above floor): median {med:.2%} "
        f"(tol {args.tol_median:.2%}), max {mx:.2%} (tol {args.tol_max:.2%})"
    )
    print("=" * 62)

    c_cm = coord * 100.0
    fig, (axP, axS) = plt.subplots(1, 2, figsize=(13, 5.0))

    nt = len(its)
    idxs = sorted(set(np.linspace(0, nt - 1, 5).astype(int)))
    for j in idxs:
        c = plt.cm.viridis(j / max(nt - 1, 1))
        axP.plot(
            c_cm,
            Te_x[j],
            "-",
            color=c,
            lw=1.8,
            label=f"t={times[j] * 1e6:.2f}" + r" $\mu$s",
        )
        axP.plot(c_cm, Te_pred[j], "--", color=c, lw=1.0)
    axP.set_xlabel("x (cm)")
    axP.set_ylabel("$T_e$ (eV)")
    axP.set_title("solid: measured $T_e$   dashed: $T_{e0}(n/n_0)^{\\gamma-1}$")
    axP.legend(fontsize=8, ncol=2)
    axP.grid(alpha=0.3)

    nn = (n_x / n0)[valid].ravel()
    tt = (Te_x / Te0)[valid].ravel()
    tcol = np.broadcast_to(times[:, None] * 1e6, n_x.shape)[valid].ravel()
    sc = axS.scatter(nn, tt, c=tcol, s=6, cmap="plasma", alpha=0.5)
    xs = np.linspace(nn.min(), nn.max(), 200)
    axS.plot(xs, xs**g1, "k-", lw=2, label=r"adiabat $(n/n_0)^{\gamma-1}$")
    fig.colorbar(sc, ax=axS, label=r"time ($\mu$s)")
    axS.set_xlabel("$n / n_0$")
    axS.set_ylabel("$T_e / T_{e0}$")
    axS.set_title(f"adiabat collapse  (median err {med:.2%}, max {mx:.2%})")
    axS.legend()
    axS.grid(alpha=0.3)

    fig.tight_layout()
    fig.savefig(args.out, dpi=150)
    print(f"[saved] {args.out}")

    ok = np.any(sig) and med <= args.tol_median and mx <= args.tol_max
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())
Electron temperature collapsing onto the adiabat

Fig. 19 Measured \(T_e\) profiles against the analytic adiabat (left) and the collapse of all cells and times onto \(T_e/T_{e0} = (n/n_0)^{\gamma_e-1}\) (right).

Vacuum-slab transport

A plasma slab (\(n = n_0\)) drifts at the electron-pressure sound speed \(c_s\) through a tenuous halo whose density lies below the solver’s density floor hybrid_pic_model.n_floor. The run starts from uniform electron entropy (\(T_e\) on the floored adiabat) with no sources (\(\mathbf{B} = 0\), \(\eta = 0\)), so entropy-conserving transport must keep the same pointwise adiabat as in the adiabatic-compression case at every cell and time – here even through the below-floor halo, which must act as an insulating boundary rather than a heat sink for the drifting slab’s electron thermal energy.

Run

Script inputs_test_2d_ohm_solver_electron_energy_picmi.py
Listing 77 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/inputs_test_2d_ohm_solver_electron_energy_picmi.py. One script covers all four cases; select this one with --case vacuum.
#!/usr/bin/env python3
#
# --- Test suite for the hybrid-PIC (Ohm's law) electron energy equation.
# --- One script, four cases (selected with --case), each isolating one
# --- term of the equation in a 2D Cartesian (x,z) periodic box:
# ---
# ---   adiabat : transport terms only (B=0, eta=0, all sources off),
# ---                 dU_e/dt + div(U_e V_e) + P_e div(V_e) = 0,
# ---             solved by the QDSMC scheme advecting the electron entropy
# ---             K_e = T_e n_e^(1-gamma) with Lagrangian markers moving at
# ---             V_e.  A sinusoidal ion velocity perturbation v_x = V0
# ---             sin(kx) drives an ion-acoustic compression, and entropy
# ---             conservation gives the pointwise check
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1).
# ---             Analyse with analysis_adiabat.py.
# ---
# ---   joule   : eta*J^2 source only.  A linear force-free field
# ---                 B(x) = B0 [0, sin(kx), cos(kx)],  k = 2 pi/Lx,
# ---             carries the uniform, parallel current J = curl(B)/mu0
# ---             (J x B = 0): no bulk motion, uniform heating, transport
# ---             terms identically zero, so
# ---                 dT_e/dt = (gamma_e - 1) eta J^2 / (n_e k_B),
# ---             a linear ramp turning sub-linear on the resistive-decay
# ---             time tau_R = mu0/(eta k^2).  Analyse with
# ---             analysis_joule.py (B-energy decay + T_e ramp).
# ---
# ---   qei     : electron-ion thermal-equilibration sink only (B=0, eta=0),
# ---                 dU_e/dt = -Q_ei,  Q_ei = 3 n_e k_B nu_ei (T_e - T_i),
# ---             enabled by the rate parser
# ---             hybrid_pic_model.electron_ion_relaxation_rate(rho,Te,Ti,t).
# ---             The single parameter enables BOTH the electron-side sink
# ---             AND the conjugate ion heating, so the exchange conserves
# ---             energy.  For constant nu_ei (single proton species, Z=1),
# ---                 (T_e - T_i)(t) = (T_e0 - T_i0) exp(-rate t),
# ---                 rate = [3(gamma_e-1) + 2] nu_ei,
# ---             and C_e T_e + C_i T_i is conserved.  Analyse with
# ---             analysis_qei.py (difference rate + budget).
# ---
# ---   vacuum  : transport through a below-floor halo (B=0, eta=0, all
# ---             sources off).  A slab (n = n0) drifts at c_s through a
# ---             halo with n = 0.02 n0, below the solver's n_floor,
# ---             starting from uniform entropy K_e, so entropy-conserving
# ---             transport must keep
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1)
# ---             pointwise at all times: the below-floor halo must act as
# ---             an insulating boundary, not a heat sink.  Analyse with
# ---             analysis_vacuum.py.

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np
from mpi4py import MPI as mpi

from pywarpx import picmi

constants = picmi.constants
comm = mpi.COMM_WORLD

simulation = None


class ElectronEnergyCase(object):
    """Shared 2D periodic-box setup; the case subclasses supply the physics
    parameters, the solver sources and the diagnostic field list."""

    # ---- Common plasma parameters ------------------------------------------
    gamma_e = 5.0 / 3.0  # electron adiabatic index
    n0 = 2.0e20  # uniform density (m^-3)
    Lx = 0.5  # domain length in x (m)

    # ---- Case hooks (overridden where a case differs) -----------------------
    include_joule_heating = False
    relaxation_rate = None  # qei only: nu_ei expression (str)
    eta_h = None  # joule only: hyper-resistivity
    # qei leaves do_temperature_deposition unset on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is configured,
    # which that case also exercises.
    set_temperature_deposition = True
    load_B = False  # joule only: force-free initial field
    reduced_diags = False  # joule only: field/particle energy

    def __init__(self, test, verbose):
        self.test = test
        self.verbose = verbose or test

        self.configure()
        self.get_plasma_quantities()
        if comm.rank == 0:
            self._print_params()
        self.setup_run()

    def momentum_expressions(self):
        return ["0", "0", "0"]

    def density_expression(self):
        return "n0"

    def setup_run(self):
        global simulation

        self.grid = picmi.Cartesian2DGrid(
            number_of_cells=[self.NX, self.NZ],
            lower_bound=[0.0, -self.Lz / 2.0],
            upper_bound=[self.Lx, self.Lz / 2.0],
            lower_boundary_conditions=["periodic", "periodic"],
            upper_boundary_conditions=["periodic", "periodic"],
            lower_boundary_conditions_particles=["periodic", "periodic"],
            upper_boundary_conditions_particles=["periodic", "periodic"],
            warpx_max_grid_size=self.NZ,
        )

        # Electron energy equation ON; each case turns on exactly one source
        # (or none, for the pure-transport adiabat case).
        solver_kwargs = {}
        if self.eta_h is not None:
            solver_kwargs["plasma_hyper_resistivity"] = self.eta_h
        if self.relaxation_rate is not None:
            solver_kwargs["electron_ion_relaxation_rate"] = self.relaxation_rate
        self.solver = picmi.HybridPICSolver(
            grid=self.grid,
            gamma=self.gamma_e,
            Te=self.te_eV,
            n0=self.n0,
            n_floor=0.05 * self.n0,
            plasma_resistivity=self.eta,
            substeps=self.substeps,
            solve_electron_energy_equation=True,
            include_joule_heating=self.include_joule_heating,
            **solver_kwargs,
        )

        simulation = picmi.Simulation(
            solver=self.solver,
            time_step_size=self.dt,
            max_steps=self.total_steps,
            verbose=self.verbose,
            particle_shape=1,
            warpx_serialize_initial_conditions=True,
            warpx_current_deposition_algo="direct",
            warpx_use_filter=True,
        )

        if self.load_B:
            B_init = picmi.LoadInitialFieldFromPython(
                load_from_python=self.load_initial_B,
                load_B=True,
                load_E=False,
            )
            simulation.add_applied_field(B_init)

        species_kwargs = {}
        if self.set_temperature_deposition:
            species_kwargs["warpx_do_temperature_deposition"] = True
        self.ions = picmi.Species(
            name="ions",
            charge="q_e",
            mass=constants.m_p,
            initial_distribution=picmi.AnalyticDistribution(
                density_expression=self.density_expression(),
                momentum_expressions=self.momentum_expressions(),
                warpx_momentum_spread_expressions=[str(self.vi_th)] * 3,
                n0=self.n0,
            ),
            **species_kwargs,
        )
        simulation.add_species(
            self.ions,
            layout=picmi.PseudoRandomLayout(
                grid=self.grid, n_macroparticles_per_cell=self.NPPC
            ),
        )

        # Remove any diags from a previous run in the same directory, so
        # stale openPMD dumps (one file per iteration) cannot mix into the
        # analysis of this run.
        if comm.rank == 0 and Path("diags").exists():
            shutil.rmtree("diags")
        comm.Barrier()

        field_diag = picmi.FieldDiagnostic(
            name="field_diag",
            grid=self.grid,
            period=self.diag_steps,
            data_list=self.diag_data_list,
            write_dir="diags",
            warpx_file_prefix="field_diags",
            warpx_format="openpmd",
            warpx_openpmd_backend="h5",
        )
        simulation.add_diagnostic(field_diag)

        if self.reduced_diags:
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="FieldEnergy",
                    name="field_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="ParticleEnergy",
                    name="part_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )

        simulation.initialize_inputs()
        simulation.initialize_warpx()


class AdiabaticCompression(ElectronEnergyCase):
    """Transport-terms (LHS) test: entropy-conserving compression."""

    te_eV = 100.0  # initial (uniform) electron temperature (eV)
    ti_eV = 10.0  # ion temperature (eV); cold vs Te for a clean,
    #   electron-pressure-driven acoustic wave

    # ---- Perturbation -------------------------------------------------------
    pert_frac = 0.30  # ion velocity amplitude V0 = pert_frac * c_s
    n_wave = 1  # wavelengths across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    periods = 2.0  # acoustic periods to simulate
    steps_per_period = 400
    substeps = 10

    diag_data_list = ["rho", "Te", "J", "B"]

    def configure(self):
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self._steps_override = 60
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 40

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Electron-pressure sound speed (cold-ion limit) sets the wave period.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.omega = self.k * self.c_s  # acoustic angular frequency
        self.T_period = 2.0 * np.pi / self.omega  # = Lx / c_s for n_wave=1
        self.V0 = self.pert_frac * self.c_s  # velocity perturbation amplitude

        self.dt = self.T_period / self.steps_per_period
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.periods * self.steps_per_period)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def momentum_expressions(self):
        # Sinusoidal x-velocity perturbation v_x = V0 sin(kx); uniform n0 and
        # uniform Te0 -> uniform initial entropy.
        return [f"({self.V0})*sin(({self.k})*x)", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Adiabatic-compression (electron-energy-equation LHS) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  c_s       = {self.c_s:.3e} m/s   (electron-pressure sound speed)\n"
            f"  V0        = {self.V0:.3e} m/s   (= {self.pert_frac:.2f} c_s)\n"
            f"  k         = {self.k:.4e} 1/m   ({self.n_wave} wavelength(s))\n"
            f"  T_period  = {self.T_period:.3e} s   (acoustic)\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   ({self.steps_per_period}/period)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, pure advection+compression\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise\n"
        )


class VacuumSlabTransport(ElectronEnergyCase):
    """Insulating-halo transport test: a plasma slab drifting through a
    below-floor halo must keep its electron entropy.

    A slab (n = n0) drifts at v0 = c_s through a tenuous halo whose density
    sits BELOW the solver's n_floor. The initial T_e is the floored adiabat
    (uniform entropy K_e), and there are no sources (B = 0, eta = 0), so
    entropy-conserving transport requires T_e = Te0 (n/n0)^(gamma-1)
    pointwise at all times -- exactly, even through the CIC-mixed slab edge,
    because a uniform K_e is invariant under any mass-weighted mixing.

    This guards the insulating treatment of below-floor cells: if the QDSMC
    transport left K_e = 0 there (instead of flooring the density in the
    K_e <-> T_e conversion), the halo would dilute and erase the slab's
    entropy at the drifting edge and T_e would fall off the adiabat within
    tens of steps.
    """

    te_eV = 100.0  # slab electron temperature (eV) at n0
    ti_eV = 10.0  # ion temperature (eV); cold, so the slab holds together

    # ---- Slab / halo geometry -----------------------------------------------
    halo_frac = 0.02  # halo density fraction of n0; BELOW the 0.05 n_floor
    slab_frac = 0.5  # slab width as a fraction of Lx
    cfl_marker = 0.2  # QDSMC marker displacement per step, v0 dt / dx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    substeps = 10

    diag_data_list = ["rho", "Te"]

    def configure(self):
        if self.test:
            self.NX = 64
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 8
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Drift at the electron-pressure sound speed: markers stream through
        # the slab edge at cfl_marker cells per step.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.v0 = self.c_s

        self.dt = self.cfl_marker * self.dx / self.v0
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = 400
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def density_expression(self):
        x0 = 0.5 * self.Lx
        hw = 0.5 * self.slab_frac * self.Lx
        return f"n0*({self.halo_frac} + (1 - {self.halo_frac})*(abs(x - {x0}) < {hw}))"

    def momentum_expressions(self):
        # Uniform drift: the slab translates without compression.
        return [f"{self.v0}", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Vacuum-slab (insulating-halo) transport test\n"
            f"  Te0 = {self.te_eV:.1f} eV (at n0),  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3,  halo = {self.halo_frac:.3f} n0 (below the 0.05 n0 floor)\n"
            f"  slab      = {self.slab_frac:.2f} Lx wide, drifting at v0 = c_s = {self.v0:.3e} m/s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (marker CFL {self.cfl_marker:.2f} cells/step)\n"
            f"  steps     = {self.total_steps} (slab travels "
            f"{self.cfl_marker * self.total_steps / self.NX:.2f} Lx),  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  no sources, pure transport through the halo\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise (uniform K_e)\n"
        )


class ForceFreeJoule(ElectronEnergyCase):
    """eta*J^2 source test: force-free field, uniform Joule ramp."""

    ti_eV = 500.0  # ion temperature (eV)
    te_eV = 500.0  # initial electron temperature (eV)

    # ---- Force-free field ---------------------------------------------------
    B0 = 0.1  # field magnitude (T); |B| is uniform
    n_wave = 1  # number of full wavelengths of B across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128  # cells in x (full run)
    NZ = 16  # cells in z (field is z-independent; periodic)
    NPPC = 800  # particles per cell; the T_e ramp sits on an ion shot-noise
    #   heating floor that converges as 1/NPPC
    DT = 0.0025  # timestep as a fraction of the ion cyclotron period; small
    #   enough for the forward-Euler Joule deposit to be converged
    TOTAL_STEPS = 3000  # full run
    DIAG_EVERY = 150  # diagnostic cadence (steps)
    substeps = 20

    include_joule_heating = True
    load_B = True
    reduced_diags = True
    diag_data_list = ["B", "E", "rho", "J", "Te", "T_ions"]

    def configure(self):
        self.eta_scale = self.args.eta_scale
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self.DT = 0.01
            self.total_steps = 50
            self.diag_steps = 10
        else:
            self.total_steps = self.TOTAL_STEPS
            self.diag_steps = self.DIAG_EVERY

    def get_plasma_quantities(self):
        mi = constants.m_p

        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ  # square cells
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Uniform plasma current magnitude from curl(B) = k B.
        self.J0 = self.k * self.B0 / constants.mu0  # A/m^2
        # Electron drift carrying it (ions start at rest): V_e = J/(e n0).
        self.v_drift = self.J0 / (constants.q_e * self.n0)

        # Ion cyclotron period at B0 sets the timestep scale.
        self.w_ci = constants.q_e * self.B0 / mi
        self.t_ci = 2.0 * np.pi / self.w_ci
        self.dt = self.DT * self.t_ci

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)

        # Constant resistivity (Ohm*m), scaled for the heating-signal amplitude.
        self.eta = 1.0e-5 * self.eta_scale
        # Resistive decay time of the force-free current: tau_R = mu0/(eta k^2).
        self.tau_R = constants.mu0 / (self.eta * self.k**2)

        # Hyper-resistivity off: grid-scale damping is not needed for a smooth,
        # single-wavelength field and would complicate the eta*J^2 budget.
        self.eta_h = 0.0

        # Analytic prediction (for the printout / cross-check).
        self.dTe_dt_pred = (
            (self.gamma_e - 1.0) * self.eta * self.J0**2 / (self.n0 * constants.kb)
        )  # K/s

    def _print_params(self):
        print(
            f"\n[setup] Force-free Joule-heating test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  B0        = {self.B0:.3e} T  (|B| uniform)\n"
            f"  k         = {self.k:.4e} 1/m  ({self.n_wave} wavelength(s) across Lx)\n"
            f"  |J|       = {self.J0:.3e} A/m^2  (uniform, force-free)\n"
            f"  V_e drift = {self.v_drift:.3e} m/s\n"
            f"  eta       = {self.eta:.3e} Ohm*m  (1e-5 x scale {self.eta_scale:g})\n"
            f"  tau_R     = {self.tau_R:.3e} s  (current resistive-decay time)\n"
            f"  Grid      = {self.NX} x {self.NZ}  (x x z),  dx = {self.dx:.3e} m\n"
            f"  t_ci      = {self.t_ci:.3e} s,  dt = {self.dt:.3e} s ({self.DT:.4f} t_ci)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  ----\n"
            f"  PREDICTED dTe/dt = (gamma_e-1) eta J^2 / (n0 kB) = {self.dTe_dt_pred:.4e} K/s\n"
            f"                   = {self.dTe_dt_pred * constants.kb / constants.q_e:.4e} eV/s\n"
        )

    def load_initial_B(self):
        """Set the linear force-free field B(x) = B0[0, sin(kx), cos(kx)].

        WarpX folds Bfield_fp_external into Bfield_fp at initialization, after
        which it evolves self-consistently. div(B) = 0 analytically (B has no
        x-component and the others depend only on x), so no cleaning needed.
        """
        Bx = simulation.fields.get("Bfield_fp_external", dir="x", level=0)
        By = simulation.fields.get("Bfield_fp_external", dir="y", level=0)
        Bz = simulation.fields.get("Bfield_fp_external", dir="z", level=0)

        Bx[:, :] = 0.0
        # Each component on its own (possibly staggered) mesh.
        XBy, _ = np.meshgrid(By.mesh("x"), By.mesh("z"), indexing="ij")
        XBz, _ = np.meshgrid(Bz.mesh("x"), Bz.mesh("z"), indexing="ij")
        By[:, :] = self.B0 * np.sin(self.k * XBy)
        Bz[:, :] = self.B0 * np.cos(self.k * XBz)
        comm.Barrier()


class QeiRelaxation(ElectronEnergyCase):
    """Q_ei electron-ion thermal-equilibration test: pure exponential."""

    te_eV = 300.0  # initial (uniform) electron temperature (eV), hot
    ti_eV = 50.0  # ion temperature (eV); the relaxation target

    # ---- Relaxation ---------------------------------------------------------
    nu_ei = 1.0e6  # electron-ion relaxation rate (1/s), constant;
    #   Te sink rate = 3(gamma_e-1)*nu_ei = 2e6 1/s -> tau = 0.5 us

    # ---- Geometry (small; the physics is 0-D / uniform) / numerics ----------
    NX = 32
    NZ = 8
    NPPC = 400
    n_tau = 3.0  # number of relaxation times to simulate
    steps_per_tau = 100  # rate*dt = 0.01 -> forward-Euler ~ exponential
    substeps = 10

    # do_temperature_deposition is NOT set on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is
    # configured, which this test also exercises (T_ions is dumped below).
    set_temperature_deposition = False
    diag_data_list = ["rho", "Te", "T_ions"]

    def configure(self):
        if self.test:
            self.NX = 16
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Analytic electron-sink rate and e-folding time.
        self.rate = 3.0 * (self.gamma_e - 1.0) * self.nu_ei
        self.tau = 1.0 / self.rate

        self.dt = self.tau / self.steps_per_tau
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.n_tau * self.steps_per_tau)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure Q_ei).
        self.eta = 0.0
        # Constant relaxation rate so the relaxation is a pure exponential.
        self.relaxation_rate = f"{self.nu_ei}"

    def _print_params(self):
        print(
            f"\n[setup] Electron-ion relaxation (Q_ei) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti0 = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  nu_ei     = {self.nu_ei:.3e} 1/s   (constant)\n"
            f"  rate      = 3(gamma-1)nu_ei = {self.rate:.3e} 1/s\n"
            f"  tau       = 1/rate = {self.tau:.3e} s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (rate*dt = {self.rate * self.dt:.3f})\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, Q_ei ON (e-sink + conjugate ion heating)\n"
            f"  CHECK:  (Te-Ti)(t) = (Te0-Ti0) exp(-[3(g-1)+2]nu t),  energy conserved\n"
        )


CASES = {
    "adiabat": AdiabaticCompression,
    "joule": ForceFreeJoule,
    "qei": QeiRelaxation,
    "vacuum": VacuumSlabTransport,
}

parser = argparse.ArgumentParser()
parser.add_argument(
    "--case",
    required=True,
    choices=sorted(CASES.keys()),
    help="which electron-energy-equation term to test",
)
parser.add_argument(
    "-t",
    "--test",
    help="toggle whether this script is run as a short CI test",
    action="store_true",
)
parser.add_argument(
    "-v",
    "--verbose",
    help="Verbose output",
    action="store_true",
)
parser.add_argument(
    "--eta-scale",
    type=float,
    default=1.0,
    help="joule case only: multiplier on the base resistivity eta=1e-5 "
    "(amplifies the eta*J^2 heating signal; the CI test uses 100)",
)
args, left = parser.parse_known_args()
sys.argv = sys.argv[:1] + left

case_class = CASES[args.case]
case_class.args = args
run = case_class(test=args.test, verbose=args.verbose)
simulation.step()

Execute:

python3 inputs_test_2d_ohm_solver_electron_energy_picmi.py --case vacuum

Analyze

Script analysis_vacuum.py
Listing 78 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/analysis_vacuum.py.
#!/usr/bin/env python3
"""Validate entropy-conserving transport through a below-floor halo.

The vacuum case drifts a plasma slab (n = n0) through a halo whose density is
below the solver's n_floor. The run starts from the floored adiabat (uniform
electron entropy K_e = T_e n_e^(1-gamma)), and has no sources (B = 0, eta = 0),
so entropy-conserving transport requires, at every cell and time,

    T_e(x,t) = T_e0 * ( n(x,t) / n0 )^(gamma - 1),

exactly -- a uniform K_e is invariant under any mass-weighted mixing, so the
check holds even through the CIC-mixed drifting slab edge. If the transport
instead left K_e = 0 in below-floor cells (an absorbing halo), the halo would
dilute and erase the slab's entropy at the edge: T_e would fall off the
adiabat within tens of steps and the slab's electron thermal energy would
drain away.

Scored on slab cells (n > 0.5 n0). Also reports the slab-mean T_e retention
between the first and last dump.
"""

import argparse
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries

Q_E = 1.602176634e-19
K_B = 1.380649e-23
K_PER_EV = K_B / Q_E


def main(argv=None):
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--diag-dir", default="diags/field_diags")
    ap.add_argument("--gamma", type=float, default=5.0 / 3.0)
    ap.add_argument(
        "--slab-frac",
        type=float,
        default=0.5,
        help="score cells with n > this fraction of n0 (the slab interior)",
    )
    ap.add_argument(
        "--tol-median",
        type=float,
        default=0.01,
        help="allowed median pointwise relative error on the adiabat",
    )
    ap.add_argument(
        "--tol-max",
        type=float,
        default=0.02,
        help="allowed max pointwise relative error on the adiabat "
        "(correct transport stays below ~0.2%%; an absorbing K_e = 0 halo "
        "reaches ~8%% within the 80-step CI run)",
    )
    ap.add_argument("--out", default="vacuum_check.png")
    args = ap.parse_args(argv)

    ts = OpenPMDTimeSeries(args.diag_dir)
    # Skip the iteration-0 dump: it is written before the floored-adiabat
    # T_e seed runs (T_e is still the uniform InitData fill there).
    its = [it for it in ts.iterations if it > 0]
    times = np.asarray(ts.t, dtype=float)[-len(its) :]
    if len(its) < 2:
        raise SystemExit(f"Need >=2 post-step dumps in {args.diag_dir}")
    g1 = args.gamma - 1.0

    def zavg(name, it):
        arr, info = ts.get_field(name, iteration=it)
        return np.asarray(info.x, dtype=float), np.asarray(arr, dtype=float).mean(
            axis=0
        )

    coord = None
    Te_x, n_x = [], []
    for it in its:
        cc, Te = zavg("Te", it)
        _, rho = zavg("rho", it)
        if coord is None:
            coord = cc
        Te_x.append(Te * K_PER_EV)
        n_x.append(rho / Q_E)
    Te_x = np.array(Te_x)  # (nt, nx)
    n_x = np.array(n_x)

    # Reference state from the first dump: the slab is the high-density mode.
    n0 = float(np.percentile(n_x[0], 90))
    slab0 = n_x[0] > args.slab_frac * n0
    Te0 = float(np.median(Te_x[0][slab0]))

    Te_pred = Te0 * (n_x / n0) ** g1
    slab = n_x > args.slab_frac * n0

    rel = np.abs(Te_x - Te_pred) / np.maximum(Te_pred, 1e-30)
    med = float(np.median(rel[slab]))
    mx = float(np.max(rel[slab]))

    # Slab-mean retention (density-weighted), reported for context only (not
    # asserted): edge rarefaction moves scored cells down the adiabat, so it
    # sits below 1 even for exact transport (~0.82 over the 80-step CI run),
    # and an absorbing halo shows up in the pointwise max error long before
    # it moves this mean.
    slab_end = n_x[-1] > args.slab_frac * n0
    Te_mean0 = float(np.sum((Te_x[0] * n_x[0])[slab0]) / np.sum(n_x[0][slab0]))
    Te_mean1 = float(np.sum((Te_x[-1] * n_x[-1])[slab_end]) / np.sum(n_x[-1][slab_end]))
    retention = Te_mean1 / Te_mean0

    print("=" * 62)
    print("Vacuum-slab (insulating halo) check  Te = Te0 (n/n0)^(gamma-1)")
    print(f"  gamma = {args.gamma:.5f}   Te0 = {Te0:.2f} eV   n0 = {n0:.3e} m^-3")
    print(f"  slab cells scored: n > {args.slab_frac:.2f} n0")
    print(
        f"  relative error on the adiabat: median {med:.2%} "
        f"(tol {args.tol_median:.2%}), max {mx:.2%} (tol {args.tol_max:.2%})"
    )
    print(f"  slab-mean Te retention (last/first dump): {retention:.4f}")
    print("=" * 62)

    c_cm = coord * 100.0
    fig, (axP, axS) = plt.subplots(1, 2, figsize=(13, 5.0))

    nt = len(its)
    idxs = sorted(set(np.linspace(0, nt - 1, 5).astype(int)))
    for j in idxs:
        c = plt.cm.viridis(j / max(nt - 1, 1))
        axP.plot(
            c_cm,
            Te_x[j],
            "-",
            color=c,
            lw=1.8,
            label=f"t={times[j] * 1e6:.2f}" + r" $\mu$s",
        )
        axP.plot(c_cm, Te_pred[j], "--", color=c, lw=1.0)
    axP.set_xlabel("x (cm)")
    axP.set_ylabel("$T_e$ (eV)")
    axP.set_title("solid: measured $T_e$   dashed: $T_{e0}(n/n_0)^{\\gamma-1}$")
    axP.legend(fontsize=8, ncol=2)
    axP.grid(alpha=0.3)

    nn = (n_x / n0)[slab].ravel()
    tt = (Te_x / Te0)[slab].ravel()
    tcol = np.broadcast_to(times[:, None] * 1e6, n_x.shape)[slab].ravel()
    sc = axS.scatter(nn, tt, c=tcol, s=6, cmap="plasma", alpha=0.5)
    xs = np.linspace(nn.min(), nn.max(), 200)
    axS.plot(xs, xs**g1, "k-", lw=2, label=r"adiabat $(n/n_0)^{\gamma-1}$")
    fig.colorbar(sc, ax=axS, label=r"time ($\mu$s)")
    axS.set_xlabel("$n / n_0$")
    axS.set_ylabel("$T_e / T_{e0}$")
    axS.set_title(f"adiabat collapse  (median err {med:.2%}, max {mx:.2%})")
    axS.legend()
    axS.grid(alpha=0.3)

    fig.tight_layout()
    fig.savefig(args.out, dpi=150)
    print(f"[saved] {args.out}")

    ok = med <= args.tol_median and mx <= args.tol_max
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())

Joule heating

A linear force-free field \(\mathbf{B}(x) = B_0[0, \sin kx, \cos kx]\) satisfies \(\nabla\times\mathbf{B} = k\mathbf{B}\), so the current is parallel to the field (no \(\mathbf{J}\times\mathbf{B}\) force) with uniform magnitude \(|J| = k B_0/\mu_0\). Nothing moves, the transport terms vanish identically, and the electron temperature obeys the pure Joule-heating ramp

\[\frac{d T_e}{d t} = (\gamma_e - 1)\, \frac{\eta J^2}{n_e k_B},\]

while the magnetic field energy decays resistively as \(E_B(t) = E_B(0)\, e^{-2t/\tau_R}\) with \(\tau_R = \mu_0/(\eta k^2)\). The analysis fits the input resistivity from both signatures independently.

Run

Script inputs_test_2d_ohm_solver_electron_energy_picmi.py
Listing 79 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/inputs_test_2d_ohm_solver_electron_energy_picmi.py. One script covers all four cases; select this one with --case joule.
#!/usr/bin/env python3
#
# --- Test suite for the hybrid-PIC (Ohm's law) electron energy equation.
# --- One script, four cases (selected with --case), each isolating one
# --- term of the equation in a 2D Cartesian (x,z) periodic box:
# ---
# ---   adiabat : transport terms only (B=0, eta=0, all sources off),
# ---                 dU_e/dt + div(U_e V_e) + P_e div(V_e) = 0,
# ---             solved by the QDSMC scheme advecting the electron entropy
# ---             K_e = T_e n_e^(1-gamma) with Lagrangian markers moving at
# ---             V_e.  A sinusoidal ion velocity perturbation v_x = V0
# ---             sin(kx) drives an ion-acoustic compression, and entropy
# ---             conservation gives the pointwise check
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1).
# ---             Analyse with analysis_adiabat.py.
# ---
# ---   joule   : eta*J^2 source only.  A linear force-free field
# ---                 B(x) = B0 [0, sin(kx), cos(kx)],  k = 2 pi/Lx,
# ---             carries the uniform, parallel current J = curl(B)/mu0
# ---             (J x B = 0): no bulk motion, uniform heating, transport
# ---             terms identically zero, so
# ---                 dT_e/dt = (gamma_e - 1) eta J^2 / (n_e k_B),
# ---             a linear ramp turning sub-linear on the resistive-decay
# ---             time tau_R = mu0/(eta k^2).  Analyse with
# ---             analysis_joule.py (B-energy decay + T_e ramp).
# ---
# ---   qei     : electron-ion thermal-equilibration sink only (B=0, eta=0),
# ---                 dU_e/dt = -Q_ei,  Q_ei = 3 n_e k_B nu_ei (T_e - T_i),
# ---             enabled by the rate parser
# ---             hybrid_pic_model.electron_ion_relaxation_rate(rho,Te,Ti,t).
# ---             The single parameter enables BOTH the electron-side sink
# ---             AND the conjugate ion heating, so the exchange conserves
# ---             energy.  For constant nu_ei (single proton species, Z=1),
# ---                 (T_e - T_i)(t) = (T_e0 - T_i0) exp(-rate t),
# ---                 rate = [3(gamma_e-1) + 2] nu_ei,
# ---             and C_e T_e + C_i T_i is conserved.  Analyse with
# ---             analysis_qei.py (difference rate + budget).
# ---
# ---   vacuum  : transport through a below-floor halo (B=0, eta=0, all
# ---             sources off).  A slab (n = n0) drifts at c_s through a
# ---             halo with n = 0.02 n0, below the solver's n_floor,
# ---             starting from uniform entropy K_e, so entropy-conserving
# ---             transport must keep
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1)
# ---             pointwise at all times: the below-floor halo must act as
# ---             an insulating boundary, not a heat sink.  Analyse with
# ---             analysis_vacuum.py.

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np
from mpi4py import MPI as mpi

from pywarpx import picmi

constants = picmi.constants
comm = mpi.COMM_WORLD

simulation = None


class ElectronEnergyCase(object):
    """Shared 2D periodic-box setup; the case subclasses supply the physics
    parameters, the solver sources and the diagnostic field list."""

    # ---- Common plasma parameters ------------------------------------------
    gamma_e = 5.0 / 3.0  # electron adiabatic index
    n0 = 2.0e20  # uniform density (m^-3)
    Lx = 0.5  # domain length in x (m)

    # ---- Case hooks (overridden where a case differs) -----------------------
    include_joule_heating = False
    relaxation_rate = None  # qei only: nu_ei expression (str)
    eta_h = None  # joule only: hyper-resistivity
    # qei leaves do_temperature_deposition unset on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is configured,
    # which that case also exercises.
    set_temperature_deposition = True
    load_B = False  # joule only: force-free initial field
    reduced_diags = False  # joule only: field/particle energy

    def __init__(self, test, verbose):
        self.test = test
        self.verbose = verbose or test

        self.configure()
        self.get_plasma_quantities()
        if comm.rank == 0:
            self._print_params()
        self.setup_run()

    def momentum_expressions(self):
        return ["0", "0", "0"]

    def density_expression(self):
        return "n0"

    def setup_run(self):
        global simulation

        self.grid = picmi.Cartesian2DGrid(
            number_of_cells=[self.NX, self.NZ],
            lower_bound=[0.0, -self.Lz / 2.0],
            upper_bound=[self.Lx, self.Lz / 2.0],
            lower_boundary_conditions=["periodic", "periodic"],
            upper_boundary_conditions=["periodic", "periodic"],
            lower_boundary_conditions_particles=["periodic", "periodic"],
            upper_boundary_conditions_particles=["periodic", "periodic"],
            warpx_max_grid_size=self.NZ,
        )

        # Electron energy equation ON; each case turns on exactly one source
        # (or none, for the pure-transport adiabat case).
        solver_kwargs = {}
        if self.eta_h is not None:
            solver_kwargs["plasma_hyper_resistivity"] = self.eta_h
        if self.relaxation_rate is not None:
            solver_kwargs["electron_ion_relaxation_rate"] = self.relaxation_rate
        self.solver = picmi.HybridPICSolver(
            grid=self.grid,
            gamma=self.gamma_e,
            Te=self.te_eV,
            n0=self.n0,
            n_floor=0.05 * self.n0,
            plasma_resistivity=self.eta,
            substeps=self.substeps,
            solve_electron_energy_equation=True,
            include_joule_heating=self.include_joule_heating,
            **solver_kwargs,
        )

        simulation = picmi.Simulation(
            solver=self.solver,
            time_step_size=self.dt,
            max_steps=self.total_steps,
            verbose=self.verbose,
            particle_shape=1,
            warpx_serialize_initial_conditions=True,
            warpx_current_deposition_algo="direct",
            warpx_use_filter=True,
        )

        if self.load_B:
            B_init = picmi.LoadInitialFieldFromPython(
                load_from_python=self.load_initial_B,
                load_B=True,
                load_E=False,
            )
            simulation.add_applied_field(B_init)

        species_kwargs = {}
        if self.set_temperature_deposition:
            species_kwargs["warpx_do_temperature_deposition"] = True
        self.ions = picmi.Species(
            name="ions",
            charge="q_e",
            mass=constants.m_p,
            initial_distribution=picmi.AnalyticDistribution(
                density_expression=self.density_expression(),
                momentum_expressions=self.momentum_expressions(),
                warpx_momentum_spread_expressions=[str(self.vi_th)] * 3,
                n0=self.n0,
            ),
            **species_kwargs,
        )
        simulation.add_species(
            self.ions,
            layout=picmi.PseudoRandomLayout(
                grid=self.grid, n_macroparticles_per_cell=self.NPPC
            ),
        )

        # Remove any diags from a previous run in the same directory, so
        # stale openPMD dumps (one file per iteration) cannot mix into the
        # analysis of this run.
        if comm.rank == 0 and Path("diags").exists():
            shutil.rmtree("diags")
        comm.Barrier()

        field_diag = picmi.FieldDiagnostic(
            name="field_diag",
            grid=self.grid,
            period=self.diag_steps,
            data_list=self.diag_data_list,
            write_dir="diags",
            warpx_file_prefix="field_diags",
            warpx_format="openpmd",
            warpx_openpmd_backend="h5",
        )
        simulation.add_diagnostic(field_diag)

        if self.reduced_diags:
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="FieldEnergy",
                    name="field_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="ParticleEnergy",
                    name="part_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )

        simulation.initialize_inputs()
        simulation.initialize_warpx()


class AdiabaticCompression(ElectronEnergyCase):
    """Transport-terms (LHS) test: entropy-conserving compression."""

    te_eV = 100.0  # initial (uniform) electron temperature (eV)
    ti_eV = 10.0  # ion temperature (eV); cold vs Te for a clean,
    #   electron-pressure-driven acoustic wave

    # ---- Perturbation -------------------------------------------------------
    pert_frac = 0.30  # ion velocity amplitude V0 = pert_frac * c_s
    n_wave = 1  # wavelengths across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    periods = 2.0  # acoustic periods to simulate
    steps_per_period = 400
    substeps = 10

    diag_data_list = ["rho", "Te", "J", "B"]

    def configure(self):
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self._steps_override = 60
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 40

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Electron-pressure sound speed (cold-ion limit) sets the wave period.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.omega = self.k * self.c_s  # acoustic angular frequency
        self.T_period = 2.0 * np.pi / self.omega  # = Lx / c_s for n_wave=1
        self.V0 = self.pert_frac * self.c_s  # velocity perturbation amplitude

        self.dt = self.T_period / self.steps_per_period
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.periods * self.steps_per_period)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def momentum_expressions(self):
        # Sinusoidal x-velocity perturbation v_x = V0 sin(kx); uniform n0 and
        # uniform Te0 -> uniform initial entropy.
        return [f"({self.V0})*sin(({self.k})*x)", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Adiabatic-compression (electron-energy-equation LHS) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  c_s       = {self.c_s:.3e} m/s   (electron-pressure sound speed)\n"
            f"  V0        = {self.V0:.3e} m/s   (= {self.pert_frac:.2f} c_s)\n"
            f"  k         = {self.k:.4e} 1/m   ({self.n_wave} wavelength(s))\n"
            f"  T_period  = {self.T_period:.3e} s   (acoustic)\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   ({self.steps_per_period}/period)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, pure advection+compression\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise\n"
        )


class VacuumSlabTransport(ElectronEnergyCase):
    """Insulating-halo transport test: a plasma slab drifting through a
    below-floor halo must keep its electron entropy.

    A slab (n = n0) drifts at v0 = c_s through a tenuous halo whose density
    sits BELOW the solver's n_floor. The initial T_e is the floored adiabat
    (uniform entropy K_e), and there are no sources (B = 0, eta = 0), so
    entropy-conserving transport requires T_e = Te0 (n/n0)^(gamma-1)
    pointwise at all times -- exactly, even through the CIC-mixed slab edge,
    because a uniform K_e is invariant under any mass-weighted mixing.

    This guards the insulating treatment of below-floor cells: if the QDSMC
    transport left K_e = 0 there (instead of flooring the density in the
    K_e <-> T_e conversion), the halo would dilute and erase the slab's
    entropy at the drifting edge and T_e would fall off the adiabat within
    tens of steps.
    """

    te_eV = 100.0  # slab electron temperature (eV) at n0
    ti_eV = 10.0  # ion temperature (eV); cold, so the slab holds together

    # ---- Slab / halo geometry -----------------------------------------------
    halo_frac = 0.02  # halo density fraction of n0; BELOW the 0.05 n_floor
    slab_frac = 0.5  # slab width as a fraction of Lx
    cfl_marker = 0.2  # QDSMC marker displacement per step, v0 dt / dx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    substeps = 10

    diag_data_list = ["rho", "Te"]

    def configure(self):
        if self.test:
            self.NX = 64
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 8
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Drift at the electron-pressure sound speed: markers stream through
        # the slab edge at cfl_marker cells per step.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.v0 = self.c_s

        self.dt = self.cfl_marker * self.dx / self.v0
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = 400
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def density_expression(self):
        x0 = 0.5 * self.Lx
        hw = 0.5 * self.slab_frac * self.Lx
        return f"n0*({self.halo_frac} + (1 - {self.halo_frac})*(abs(x - {x0}) < {hw}))"

    def momentum_expressions(self):
        # Uniform drift: the slab translates without compression.
        return [f"{self.v0}", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Vacuum-slab (insulating-halo) transport test\n"
            f"  Te0 = {self.te_eV:.1f} eV (at n0),  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3,  halo = {self.halo_frac:.3f} n0 (below the 0.05 n0 floor)\n"
            f"  slab      = {self.slab_frac:.2f} Lx wide, drifting at v0 = c_s = {self.v0:.3e} m/s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (marker CFL {self.cfl_marker:.2f} cells/step)\n"
            f"  steps     = {self.total_steps} (slab travels "
            f"{self.cfl_marker * self.total_steps / self.NX:.2f} Lx),  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  no sources, pure transport through the halo\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise (uniform K_e)\n"
        )


class ForceFreeJoule(ElectronEnergyCase):
    """eta*J^2 source test: force-free field, uniform Joule ramp."""

    ti_eV = 500.0  # ion temperature (eV)
    te_eV = 500.0  # initial electron temperature (eV)

    # ---- Force-free field ---------------------------------------------------
    B0 = 0.1  # field magnitude (T); |B| is uniform
    n_wave = 1  # number of full wavelengths of B across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128  # cells in x (full run)
    NZ = 16  # cells in z (field is z-independent; periodic)
    NPPC = 800  # particles per cell; the T_e ramp sits on an ion shot-noise
    #   heating floor that converges as 1/NPPC
    DT = 0.0025  # timestep as a fraction of the ion cyclotron period; small
    #   enough for the forward-Euler Joule deposit to be converged
    TOTAL_STEPS = 3000  # full run
    DIAG_EVERY = 150  # diagnostic cadence (steps)
    substeps = 20

    include_joule_heating = True
    load_B = True
    reduced_diags = True
    diag_data_list = ["B", "E", "rho", "J", "Te", "T_ions"]

    def configure(self):
        self.eta_scale = self.args.eta_scale
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self.DT = 0.01
            self.total_steps = 50
            self.diag_steps = 10
        else:
            self.total_steps = self.TOTAL_STEPS
            self.diag_steps = self.DIAG_EVERY

    def get_plasma_quantities(self):
        mi = constants.m_p

        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ  # square cells
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Uniform plasma current magnitude from curl(B) = k B.
        self.J0 = self.k * self.B0 / constants.mu0  # A/m^2
        # Electron drift carrying it (ions start at rest): V_e = J/(e n0).
        self.v_drift = self.J0 / (constants.q_e * self.n0)

        # Ion cyclotron period at B0 sets the timestep scale.
        self.w_ci = constants.q_e * self.B0 / mi
        self.t_ci = 2.0 * np.pi / self.w_ci
        self.dt = self.DT * self.t_ci

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)

        # Constant resistivity (Ohm*m), scaled for the heating-signal amplitude.
        self.eta = 1.0e-5 * self.eta_scale
        # Resistive decay time of the force-free current: tau_R = mu0/(eta k^2).
        self.tau_R = constants.mu0 / (self.eta * self.k**2)

        # Hyper-resistivity off: grid-scale damping is not needed for a smooth,
        # single-wavelength field and would complicate the eta*J^2 budget.
        self.eta_h = 0.0

        # Analytic prediction (for the printout / cross-check).
        self.dTe_dt_pred = (
            (self.gamma_e - 1.0) * self.eta * self.J0**2 / (self.n0 * constants.kb)
        )  # K/s

    def _print_params(self):
        print(
            f"\n[setup] Force-free Joule-heating test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  B0        = {self.B0:.3e} T  (|B| uniform)\n"
            f"  k         = {self.k:.4e} 1/m  ({self.n_wave} wavelength(s) across Lx)\n"
            f"  |J|       = {self.J0:.3e} A/m^2  (uniform, force-free)\n"
            f"  V_e drift = {self.v_drift:.3e} m/s\n"
            f"  eta       = {self.eta:.3e} Ohm*m  (1e-5 x scale {self.eta_scale:g})\n"
            f"  tau_R     = {self.tau_R:.3e} s  (current resistive-decay time)\n"
            f"  Grid      = {self.NX} x {self.NZ}  (x x z),  dx = {self.dx:.3e} m\n"
            f"  t_ci      = {self.t_ci:.3e} s,  dt = {self.dt:.3e} s ({self.DT:.4f} t_ci)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  ----\n"
            f"  PREDICTED dTe/dt = (gamma_e-1) eta J^2 / (n0 kB) = {self.dTe_dt_pred:.4e} K/s\n"
            f"                   = {self.dTe_dt_pred * constants.kb / constants.q_e:.4e} eV/s\n"
        )

    def load_initial_B(self):
        """Set the linear force-free field B(x) = B0[0, sin(kx), cos(kx)].

        WarpX folds Bfield_fp_external into Bfield_fp at initialization, after
        which it evolves self-consistently. div(B) = 0 analytically (B has no
        x-component and the others depend only on x), so no cleaning needed.
        """
        Bx = simulation.fields.get("Bfield_fp_external", dir="x", level=0)
        By = simulation.fields.get("Bfield_fp_external", dir="y", level=0)
        Bz = simulation.fields.get("Bfield_fp_external", dir="z", level=0)

        Bx[:, :] = 0.0
        # Each component on its own (possibly staggered) mesh.
        XBy, _ = np.meshgrid(By.mesh("x"), By.mesh("z"), indexing="ij")
        XBz, _ = np.meshgrid(Bz.mesh("x"), Bz.mesh("z"), indexing="ij")
        By[:, :] = self.B0 * np.sin(self.k * XBy)
        Bz[:, :] = self.B0 * np.cos(self.k * XBz)
        comm.Barrier()


class QeiRelaxation(ElectronEnergyCase):
    """Q_ei electron-ion thermal-equilibration test: pure exponential."""

    te_eV = 300.0  # initial (uniform) electron temperature (eV), hot
    ti_eV = 50.0  # ion temperature (eV); the relaxation target

    # ---- Relaxation ---------------------------------------------------------
    nu_ei = 1.0e6  # electron-ion relaxation rate (1/s), constant;
    #   Te sink rate = 3(gamma_e-1)*nu_ei = 2e6 1/s -> tau = 0.5 us

    # ---- Geometry (small; the physics is 0-D / uniform) / numerics ----------
    NX = 32
    NZ = 8
    NPPC = 400
    n_tau = 3.0  # number of relaxation times to simulate
    steps_per_tau = 100  # rate*dt = 0.01 -> forward-Euler ~ exponential
    substeps = 10

    # do_temperature_deposition is NOT set on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is
    # configured, which this test also exercises (T_ions is dumped below).
    set_temperature_deposition = False
    diag_data_list = ["rho", "Te", "T_ions"]

    def configure(self):
        if self.test:
            self.NX = 16
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Analytic electron-sink rate and e-folding time.
        self.rate = 3.0 * (self.gamma_e - 1.0) * self.nu_ei
        self.tau = 1.0 / self.rate

        self.dt = self.tau / self.steps_per_tau
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.n_tau * self.steps_per_tau)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure Q_ei).
        self.eta = 0.0
        # Constant relaxation rate so the relaxation is a pure exponential.
        self.relaxation_rate = f"{self.nu_ei}"

    def _print_params(self):
        print(
            f"\n[setup] Electron-ion relaxation (Q_ei) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti0 = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  nu_ei     = {self.nu_ei:.3e} 1/s   (constant)\n"
            f"  rate      = 3(gamma-1)nu_ei = {self.rate:.3e} 1/s\n"
            f"  tau       = 1/rate = {self.tau:.3e} s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (rate*dt = {self.rate * self.dt:.3f})\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, Q_ei ON (e-sink + conjugate ion heating)\n"
            f"  CHECK:  (Te-Ti)(t) = (Te0-Ti0) exp(-[3(g-1)+2]nu t),  energy conserved\n"
        )


CASES = {
    "adiabat": AdiabaticCompression,
    "joule": ForceFreeJoule,
    "qei": QeiRelaxation,
    "vacuum": VacuumSlabTransport,
}

parser = argparse.ArgumentParser()
parser.add_argument(
    "--case",
    required=True,
    choices=sorted(CASES.keys()),
    help="which electron-energy-equation term to test",
)
parser.add_argument(
    "-t",
    "--test",
    help="toggle whether this script is run as a short CI test",
    action="store_true",
)
parser.add_argument(
    "-v",
    "--verbose",
    help="Verbose output",
    action="store_true",
)
parser.add_argument(
    "--eta-scale",
    type=float,
    default=1.0,
    help="joule case only: multiplier on the base resistivity eta=1e-5 "
    "(amplifies the eta*J^2 heating signal; the CI test uses 100)",
)
args, left = parser.parse_known_args()
sys.argv = sys.argv[:1] + left

case_class = CASES[args.case]
case_class.args = args
run = case_class(test=args.test, verbose=args.verbose)
simulation.step()

Execute:

python3 inputs_test_2d_ohm_solver_electron_energy_picmi.py --case joule --eta-scale 20

Analyze

Script analysis_joule.py
Listing 80 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/analysis_joule.py.
#!/usr/bin/env python3
"""Validate the eta*J^2 Joule source of the electron energy equation with two
independent measurements of the resistivity from the force-free run:

1. FIELD DECAY (primary): the force-free mode decays resistively,
       E_B(t) = E_B(0) exp(-2 t / tau_R),   tau_R = mu0 / (eta k^2),
   so eta = mu0 * rate / (2 k^2) from the FieldEnergy reduced diagnostic.
   This measures the Ohm's-law friction directly and is immune to PIC-noise
   heating of T_e.

2. Te RAMP (secondary): the Joule source gives
       dTe(t) = (gamma-1) eta J0^2/(n0 kB) * (tau_R/2)(1 - e^{-2t/tau_R}),
   fitted for eta by a 1-parameter least-squares scan.  This checks that the
   heating deposited into T_e uses the same eta.  It sits on a small
   ion-current shot-noise heating floor (~1/N_ppc), hence the looser
   tolerance.

The figure additionally shows the cumulative energy budget: the electron
thermal gain Delta E_e tracks the magnetic-field loss Delta E_B plus the
(small, shot-noise driven) ion kinetic drain Delta E_ion, with the total
conserved.

PASS if the field-decay fit is within --tol-field and the Te fit within
--tol-te of the input resistivity.
"""

import argparse
import sys
from pathlib import Path

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries

Q_E = 1.602176634e-19
K_B = 1.380649e-23
MU0 = 4.0e-7 * np.pi
K_PER_EV = Q_E / K_B

# Must match the input deck.
N0, B0, LX = 2.0e20, 0.1, 0.5
GAMMA = 5.0 / 3.0
KWAVE = 2.0 * np.pi / LX
J0 = KWAVE * B0 / MU0


def read_te_series(diag_dir):
    """Read the openPMD field diags.

    Returns (t[s], density-weighted <Te>[eV], electron thermal energy E_e[J])
    with E_e = kB/(gamma-1) * sum(n_e Te dV) integrated over the domain
    (per meter in the ignored y direction, consistent with the reduced
    diagnostics in 2D).
    """
    ts = OpenPMDTimeSeries(str(diag_dir))
    t = np.asarray(ts.t, dtype=float)
    Te_m, E_e = [], []
    for it in ts.iterations:
        Te, info = ts.get_field("Te", iteration=it)
        rho, _ = ts.get_field("rho", iteration=it)
        Te = np.asarray(Te, dtype=float)  # K
        ne = np.asarray(rho, dtype=float) / Q_E  # m^-3
        Te_m.append(float(np.sum(Te * ne) / np.sum(ne)) / K_PER_EV)
        dV = info.dx * info.dz  # (x 1 m in y)
        E_e.append(K_B / (GAMMA - 1.0) * float(np.sum(ne * Te)) * dV)
    return t, np.asarray(Te_m), np.asarray(E_e)


def eta_from_field_decay(t, E_B):
    """eta from the exponential decay of the magnetic field energy."""
    # E_B ~ exp(-2t/tau_R): linear fit of log E_B.
    rate = -np.polyfit(t, np.log(E_B), 1)[0]  # = 2/tau_R
    return MU0 * rate / (2.0 * KWAVE**2)


def model_dTe(t, eta):
    """Joule Te ramp [eV] with the resistive J decay folded in."""
    tau = MU0 / (eta * KWAVE**2)
    pref = (GAMMA - 1.0) * eta * J0**2 / (N0 * K_B)  # K/s at t=0
    return pref * (tau / 2.0) * (1.0 - np.exp(-2.0 * t / tau)) / K_PER_EV


def eta_from_te_ramp(t, dTe, eta_input):
    """1-parameter least-squares fit of eta (coarse scan)."""
    grid = np.linspace(0.1 * eta_input, 3.0 * eta_input, 4001)
    ssr = [np.sum((model_dTe(t, e) - dTe) ** 2) for e in grid]
    return grid[int(np.argmin(ssr))]


def main(argv=None):
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument("--diag-dir", default="diags/field_diags")
    ap.add_argument("--reduced-dir", default="diags")
    ap.add_argument(
        "--eta-scale",
        type=float,
        default=1.0,
        help="multiplier on the base resistivity eta=1e-5; must match the run",
    )
    ap.add_argument(
        "--tol-field",
        type=float,
        default=0.05,
        help="allowed relative error on the field-decay eta fit",
    )
    ap.add_argument(
        "--tol-te",
        type=float,
        default=0.20,
        help="allowed relative error on the Te-ramp eta fit",
    )
    ap.add_argument("--out", default="joule_check.png")
    args = ap.parse_args(argv)

    eta_input = 1.0e-5 * args.eta_scale

    # Reduced diagnostics: magnetic field and ion kinetic energies.
    fdata = np.loadtxt(Path(args.reduced_dir) / "field_energy.txt", skiprows=1)
    t_r, E_B = fdata[:, 1], fdata[:, 4]  # column 4 = B_lev0 (J)
    pdata = np.loadtxt(Path(args.reduced_dir) / "part_energy.txt", skiprows=1)
    E_ion = pdata[:, 2]  # column 2 = total (J)

    eta_B = eta_from_field_decay(t_r, E_B)
    err_B = abs(eta_B - eta_input) / eta_input

    t, Te, E_e = read_te_series(args.diag_dir)
    # Skip the iteration-0 dump in the Te-ramp fit (fields are written
    # before the first current deposition, so J = 0 there).
    t_fit, dTe = t[1:] - t[1], Te[1:] - Te[1]
    eta_T = eta_from_te_ramp(t_fit, dTe, eta_input)
    err_T = abs(eta_T - eta_input) / eta_input

    # Cumulative energy budget (each series relative to its first sample).
    dE_e = E_e - E_e[0]
    dE_B = E_B - E_B[0]
    dE_i = E_ion - E_ion[0]
    n_b = min(t_r.size, t.size)
    dE_tot = dE_e[:n_b] + dE_B[:n_b] + dE_i[:n_b]
    noncons = dE_tot[-1] / dE_e[n_b - 1] if dE_e[n_b - 1] != 0.0 else 0.0

    print("=" * 66)
    print("Force-free Joule-heating check, dTe/dt = (gamma-1) eta J^2/(n kB)")
    print(f"  eta (input)      = {eta_input:.4e} Ohm*m")
    print(
        f"  eta (field decay)= {eta_B:.4e} Ohm*m  "
        f"({100 * err_B:+.2f}%, tol {100 * args.tol_field:.1f}%)"
    )
    print(
        f"  eta (Te ramp)    = {eta_T:.4e} Ohm*m  "
        f"({100 * err_T:+.2f}%, tol {100 * args.tol_te:.1f}%)"
    )
    print(
        f"  energy budget: dE_e = {dE_e[n_b - 1]:+.3f} J, dE_B = {dE_B[n_b - 1]:+.3f} J, "
        f"dE_ion = {dE_i[n_b - 1]:+.3f} J"
    )
    print(f"  final non-conservation    = {100 * noncons:+.2f}% of dE_e")
    print("=" * 66)

    fig, (axE, axT) = plt.subplots(1, 2, figsize=(12, 4.6))

    tus_r = t_r * 1e6
    axE.plot(t * 1e6, dE_e, "o-", ms=4, label=r"$\Delta E_e$ (electron thermal)")
    axE.plot(tus_r, dE_B, "s-", ms=4, label=r"$\Delta E_B$ (magnetic)")
    axE.plot(tus_r, dE_i, "^-", ms=4, label=r"$\Delta E_{ion}$")
    axE.plot(tus_r[:n_b], dE_tot, "k-", lw=2.5, label=r"$\Delta E_{tot}$ (should be 0)")
    axE.axhline(0.0, color="gray", lw=0.8, ls=":")
    axE.set_xlabel(r"time ($\mu$s)")
    axE.set_ylabel("cumulative energy change (J)")
    axE.set_title(
        f"energy budget (non-conservation {100 * noncons:+.2f}% of "
        r"$\Delta E_e$)"
    )
    axE.legend(fontsize=9)
    axE.grid(alpha=0.3)

    tm = np.linspace(0.0, t_fit[-1], 200)
    axT.plot(t_fit * 1e6, dTe, "o", ms=5, label="measured")
    axT.plot(
        tm * 1e6, model_dTe(tm, eta_input), "-", lw=1.5, label=r"analytic, input $\eta$"
    )
    axT.plot(
        tm * 1e6,
        model_dTe(tm, eta_T),
        "--",
        lw=1.2,
        label=rf"fit, $\eta$ = {eta_T:.3e}",
    )
    axT.set_xlabel(r"time ($\mu$s)")
    axT.set_ylabel(r"$\Delta\langle T_e\rangle_n$ (eV)")
    axT.set_title("electron temperature ramp")
    axT.legend(fontsize=9)
    axT.grid(alpha=0.3)

    fig.suptitle("Joule heating of the force-free equilibrium")
    fig.tight_layout(rect=[0, 0, 1, 0.94])
    fig.savefig(args.out, dpi=150)
    print(f"[saved] {args.out}")

    ok = err_B <= args.tol_field and err_T <= args.tol_te
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())

Execute:

python3 analysis_joule.py --eta-scale 20
Joule heating energy budget and electron temperature ramp

Fig. 20 Cumulative energy budget (left): the electron thermal gain tracks the magnetic-field loss plus a small ion kinetic drain, with the total conserved. Density-weighted mean electron temperature rise against the analytic Joule-heating ramp, with the resistive current decay folded in (right).

Electron-ion temperature relaxation

A uniform, unmagnetized, zero-resistivity plasma with hot electrons (\(T_{e0} \gg T_{i0}\)) relaxes purely through the electron-ion thermal-equilibration exchange \(Q_{ei} = 3 n_e k_B \nu_{ei} (T_e - T_i)\), which cools the electron fluid and heats the kinetic ions by exactly the same amount. For a constant \(\nu_{ei}\) the temperature difference decays exponentially at the rate \([3(\gamma_e - 1) + 2]\,\nu_{ei}\) while the total thermal energy is conserved; for \(\gamma_e = 5/3\) both species meet at \((T_{e0} + T_{i0})/2\).

Run

Script inputs_test_2d_ohm_solver_electron_energy_picmi.py
Listing 81 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/inputs_test_2d_ohm_solver_electron_energy_picmi.py. One script covers all four cases; select this one with --case qei.
#!/usr/bin/env python3
#
# --- Test suite for the hybrid-PIC (Ohm's law) electron energy equation.
# --- One script, four cases (selected with --case), each isolating one
# --- term of the equation in a 2D Cartesian (x,z) periodic box:
# ---
# ---   adiabat : transport terms only (B=0, eta=0, all sources off),
# ---                 dU_e/dt + div(U_e V_e) + P_e div(V_e) = 0,
# ---             solved by the QDSMC scheme advecting the electron entropy
# ---             K_e = T_e n_e^(1-gamma) with Lagrangian markers moving at
# ---             V_e.  A sinusoidal ion velocity perturbation v_x = V0
# ---             sin(kx) drives an ion-acoustic compression, and entropy
# ---             conservation gives the pointwise check
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1).
# ---             Analyse with analysis_adiabat.py.
# ---
# ---   joule   : eta*J^2 source only.  A linear force-free field
# ---                 B(x) = B0 [0, sin(kx), cos(kx)],  k = 2 pi/Lx,
# ---             carries the uniform, parallel current J = curl(B)/mu0
# ---             (J x B = 0): no bulk motion, uniform heating, transport
# ---             terms identically zero, so
# ---                 dT_e/dt = (gamma_e - 1) eta J^2 / (n_e k_B),
# ---             a linear ramp turning sub-linear on the resistive-decay
# ---             time tau_R = mu0/(eta k^2).  Analyse with
# ---             analysis_joule.py (B-energy decay + T_e ramp).
# ---
# ---   qei     : electron-ion thermal-equilibration sink only (B=0, eta=0),
# ---                 dU_e/dt = -Q_ei,  Q_ei = 3 n_e k_B nu_ei (T_e - T_i),
# ---             enabled by the rate parser
# ---             hybrid_pic_model.electron_ion_relaxation_rate(rho,Te,Ti,t).
# ---             The single parameter enables BOTH the electron-side sink
# ---             AND the conjugate ion heating, so the exchange conserves
# ---             energy.  For constant nu_ei (single proton species, Z=1),
# ---                 (T_e - T_i)(t) = (T_e0 - T_i0) exp(-rate t),
# ---                 rate = [3(gamma_e-1) + 2] nu_ei,
# ---             and C_e T_e + C_i T_i is conserved.  Analyse with
# ---             analysis_qei.py (difference rate + budget).
# ---
# ---   vacuum  : transport through a below-floor halo (B=0, eta=0, all
# ---             sources off).  A slab (n = n0) drifts at c_s through a
# ---             halo with n = 0.02 n0, below the solver's n_floor,
# ---             starting from uniform entropy K_e, so entropy-conserving
# ---             transport must keep
# ---                 T_e(x,t) = T_e0 (n(x,t)/n0)^(gamma_e-1)
# ---             pointwise at all times: the below-floor halo must act as
# ---             an insulating boundary, not a heat sink.  Analyse with
# ---             analysis_vacuum.py.

import argparse
import shutil
import sys
from pathlib import Path

import numpy as np
from mpi4py import MPI as mpi

from pywarpx import picmi

constants = picmi.constants
comm = mpi.COMM_WORLD

simulation = None


class ElectronEnergyCase(object):
    """Shared 2D periodic-box setup; the case subclasses supply the physics
    parameters, the solver sources and the diagnostic field list."""

    # ---- Common plasma parameters ------------------------------------------
    gamma_e = 5.0 / 3.0  # electron adiabatic index
    n0 = 2.0e20  # uniform density (m^-3)
    Lx = 0.5  # domain length in x (m)

    # ---- Case hooks (overridden where a case differs) -----------------------
    include_joule_heating = False
    relaxation_rate = None  # qei only: nu_ei expression (str)
    eta_h = None  # joule only: hyper-resistivity
    # qei leaves do_temperature_deposition unset on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is configured,
    # which that case also exercises.
    set_temperature_deposition = True
    load_B = False  # joule only: force-free initial field
    reduced_diags = False  # joule only: field/particle energy

    def __init__(self, test, verbose):
        self.test = test
        self.verbose = verbose or test

        self.configure()
        self.get_plasma_quantities()
        if comm.rank == 0:
            self._print_params()
        self.setup_run()

    def momentum_expressions(self):
        return ["0", "0", "0"]

    def density_expression(self):
        return "n0"

    def setup_run(self):
        global simulation

        self.grid = picmi.Cartesian2DGrid(
            number_of_cells=[self.NX, self.NZ],
            lower_bound=[0.0, -self.Lz / 2.0],
            upper_bound=[self.Lx, self.Lz / 2.0],
            lower_boundary_conditions=["periodic", "periodic"],
            upper_boundary_conditions=["periodic", "periodic"],
            lower_boundary_conditions_particles=["periodic", "periodic"],
            upper_boundary_conditions_particles=["periodic", "periodic"],
            warpx_max_grid_size=self.NZ,
        )

        # Electron energy equation ON; each case turns on exactly one source
        # (or none, for the pure-transport adiabat case).
        solver_kwargs = {}
        if self.eta_h is not None:
            solver_kwargs["plasma_hyper_resistivity"] = self.eta_h
        if self.relaxation_rate is not None:
            solver_kwargs["electron_ion_relaxation_rate"] = self.relaxation_rate
        self.solver = picmi.HybridPICSolver(
            grid=self.grid,
            gamma=self.gamma_e,
            Te=self.te_eV,
            n0=self.n0,
            n_floor=0.05 * self.n0,
            plasma_resistivity=self.eta,
            substeps=self.substeps,
            solve_electron_energy_equation=True,
            include_joule_heating=self.include_joule_heating,
            **solver_kwargs,
        )

        simulation = picmi.Simulation(
            solver=self.solver,
            time_step_size=self.dt,
            max_steps=self.total_steps,
            verbose=self.verbose,
            particle_shape=1,
            warpx_serialize_initial_conditions=True,
            warpx_current_deposition_algo="direct",
            warpx_use_filter=True,
        )

        if self.load_B:
            B_init = picmi.LoadInitialFieldFromPython(
                load_from_python=self.load_initial_B,
                load_B=True,
                load_E=False,
            )
            simulation.add_applied_field(B_init)

        species_kwargs = {}
        if self.set_temperature_deposition:
            species_kwargs["warpx_do_temperature_deposition"] = True
        self.ions = picmi.Species(
            name="ions",
            charge="q_e",
            mass=constants.m_p,
            initial_distribution=picmi.AnalyticDistribution(
                density_expression=self.density_expression(),
                momentum_expressions=self.momentum_expressions(),
                warpx_momentum_spread_expressions=[str(self.vi_th)] * 3,
                n0=self.n0,
            ),
            **species_kwargs,
        )
        simulation.add_species(
            self.ions,
            layout=picmi.PseudoRandomLayout(
                grid=self.grid, n_macroparticles_per_cell=self.NPPC
            ),
        )

        # Remove any diags from a previous run in the same directory, so
        # stale openPMD dumps (one file per iteration) cannot mix into the
        # analysis of this run.
        if comm.rank == 0 and Path("diags").exists():
            shutil.rmtree("diags")
        comm.Barrier()

        field_diag = picmi.FieldDiagnostic(
            name="field_diag",
            grid=self.grid,
            period=self.diag_steps,
            data_list=self.diag_data_list,
            write_dir="diags",
            warpx_file_prefix="field_diags",
            warpx_format="openpmd",
            warpx_openpmd_backend="h5",
        )
        simulation.add_diagnostic(field_diag)

        if self.reduced_diags:
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="FieldEnergy",
                    name="field_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )
            simulation.add_diagnostic(
                picmi.ReducedDiagnostic(
                    diag_type="ParticleEnergy",
                    name="part_energy",
                    period=self.diag_steps,
                    path="diags/",
                )
            )

        simulation.initialize_inputs()
        simulation.initialize_warpx()


class AdiabaticCompression(ElectronEnergyCase):
    """Transport-terms (LHS) test: entropy-conserving compression."""

    te_eV = 100.0  # initial (uniform) electron temperature (eV)
    ti_eV = 10.0  # ion temperature (eV); cold vs Te for a clean,
    #   electron-pressure-driven acoustic wave

    # ---- Perturbation -------------------------------------------------------
    pert_frac = 0.30  # ion velocity amplitude V0 = pert_frac * c_s
    n_wave = 1  # wavelengths across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    periods = 2.0  # acoustic periods to simulate
    steps_per_period = 400
    substeps = 10

    diag_data_list = ["rho", "Te", "J", "B"]

    def configure(self):
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self._steps_override = 60
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 40

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Electron-pressure sound speed (cold-ion limit) sets the wave period.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.omega = self.k * self.c_s  # acoustic angular frequency
        self.T_period = 2.0 * np.pi / self.omega  # = Lx / c_s for n_wave=1
        self.V0 = self.pert_frac * self.c_s  # velocity perturbation amplitude

        self.dt = self.T_period / self.steps_per_period
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.periods * self.steps_per_period)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def momentum_expressions(self):
        # Sinusoidal x-velocity perturbation v_x = V0 sin(kx); uniform n0 and
        # uniform Te0 -> uniform initial entropy.
        return [f"({self.V0})*sin(({self.k})*x)", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Adiabatic-compression (electron-energy-equation LHS) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  c_s       = {self.c_s:.3e} m/s   (electron-pressure sound speed)\n"
            f"  V0        = {self.V0:.3e} m/s   (= {self.pert_frac:.2f} c_s)\n"
            f"  k         = {self.k:.4e} 1/m   ({self.n_wave} wavelength(s))\n"
            f"  T_period  = {self.T_period:.3e} s   (acoustic)\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   ({self.steps_per_period}/period)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, pure advection+compression\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise\n"
        )


class VacuumSlabTransport(ElectronEnergyCase):
    """Insulating-halo transport test: a plasma slab drifting through a
    below-floor halo must keep its electron entropy.

    A slab (n = n0) drifts at v0 = c_s through a tenuous halo whose density
    sits BELOW the solver's n_floor. The initial T_e is the floored adiabat
    (uniform entropy K_e), and there are no sources (B = 0, eta = 0), so
    entropy-conserving transport requires T_e = Te0 (n/n0)^(gamma-1)
    pointwise at all times -- exactly, even through the CIC-mixed slab edge,
    because a uniform K_e is invariant under any mass-weighted mixing.

    This guards the insulating treatment of below-floor cells: if the QDSMC
    transport left K_e = 0 there (instead of flooring the density in the
    K_e <-> T_e conversion), the halo would dilute and erase the slab's
    entropy at the drifting edge and T_e would fall off the adiabat within
    tens of steps.
    """

    te_eV = 100.0  # slab electron temperature (eV) at n0
    ti_eV = 10.0  # ion temperature (eV); cold, so the slab holds together

    # ---- Slab / halo geometry -----------------------------------------------
    halo_frac = 0.02  # halo density fraction of n0; BELOW the 0.05 n_floor
    slab_frac = 0.5  # slab width as a fraction of Lx
    cfl_marker = 0.2  # QDSMC marker displacement per step, v0 dt / dx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128
    NZ = 16
    NPPC = 800
    substeps = 10

    diag_data_list = ["rho", "Te"]

    def configure(self):
        if self.test:
            self.NX = 64
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 8
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Drift at the electron-pressure sound speed: markers stream through
        # the slab edge at cfl_marker cells per step.
        self.c_s = np.sqrt(self.gamma_e * constants.q_e * self.te_eV / mi)
        self.v0 = self.c_s

        self.dt = self.cfl_marker * self.dx / self.v0
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = 400
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure LHS).
        self.eta = 0.0

    def density_expression(self):
        x0 = 0.5 * self.Lx
        hw = 0.5 * self.slab_frac * self.Lx
        return f"n0*({self.halo_frac} + (1 - {self.halo_frac})*(abs(x - {x0}) < {hw}))"

    def momentum_expressions(self):
        # Uniform drift: the slab translates without compression.
        return [f"{self.v0}", "0", "0"]

    def _print_params(self):
        print(
            f"\n[setup] Vacuum-slab (insulating-halo) transport test\n"
            f"  Te0 = {self.te_eV:.1f} eV (at n0),  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3,  halo = {self.halo_frac:.3f} n0 (below the 0.05 n0 floor)\n"
            f"  slab      = {self.slab_frac:.2f} Lx wide, drifting at v0 = c_s = {self.v0:.3e} m/s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (marker CFL {self.cfl_marker:.2f} cells/step)\n"
            f"  steps     = {self.total_steps} (slab travels "
            f"{self.cfl_marker * self.total_steps / self.NX:.2f} Lx),  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  no sources, pure transport through the halo\n"
            f"  CHECK:  T_e(x,t) = Te0 (n/n0)^(gamma_e-1)  pointwise (uniform K_e)\n"
        )


class ForceFreeJoule(ElectronEnergyCase):
    """eta*J^2 source test: force-free field, uniform Joule ramp."""

    ti_eV = 500.0  # ion temperature (eV)
    te_eV = 500.0  # initial electron temperature (eV)

    # ---- Force-free field ---------------------------------------------------
    B0 = 0.1  # field magnitude (T); |B| is uniform
    n_wave = 1  # number of full wavelengths of B across Lx

    # ---- Geometry / numerics ------------------------------------------------
    NX = 128  # cells in x (full run)
    NZ = 16  # cells in z (field is z-independent; periodic)
    NPPC = 800  # particles per cell; the T_e ramp sits on an ion shot-noise
    #   heating floor that converges as 1/NPPC
    DT = 0.0025  # timestep as a fraction of the ion cyclotron period; small
    #   enough for the forward-Euler Joule deposit to be converged
    TOTAL_STEPS = 3000  # full run
    DIAG_EVERY = 150  # diagnostic cadence (steps)
    substeps = 20

    include_joule_heating = True
    load_B = True
    reduced_diags = True
    diag_data_list = ["B", "E", "rho", "J", "Te", "T_ions"]

    def configure(self):
        self.eta_scale = self.args.eta_scale
        if self.test:
            self.NX = 32
            self.NZ = 8
            self.NPPC = 64
            self.DT = 0.01
            self.total_steps = 50
            self.diag_steps = 10
        else:
            self.total_steps = self.TOTAL_STEPS
            self.diag_steps = self.DIAG_EVERY

    def get_plasma_quantities(self):
        mi = constants.m_p

        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ  # square cells
        self.k = 2.0 * np.pi * self.n_wave / self.Lx

        # Uniform plasma current magnitude from curl(B) = k B.
        self.J0 = self.k * self.B0 / constants.mu0  # A/m^2
        # Electron drift carrying it (ions start at rest): V_e = J/(e n0).
        self.v_drift = self.J0 / (constants.q_e * self.n0)

        # Ion cyclotron period at B0 sets the timestep scale.
        self.w_ci = constants.q_e * self.B0 / mi
        self.t_ci = 2.0 * np.pi / self.w_ci
        self.dt = self.DT * self.t_ci

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)

        # Constant resistivity (Ohm*m), scaled for the heating-signal amplitude.
        self.eta = 1.0e-5 * self.eta_scale
        # Resistive decay time of the force-free current: tau_R = mu0/(eta k^2).
        self.tau_R = constants.mu0 / (self.eta * self.k**2)

        # Hyper-resistivity off: grid-scale damping is not needed for a smooth,
        # single-wavelength field and would complicate the eta*J^2 budget.
        self.eta_h = 0.0

        # Analytic prediction (for the printout / cross-check).
        self.dTe_dt_pred = (
            (self.gamma_e - 1.0) * self.eta * self.J0**2 / (self.n0 * constants.kb)
        )  # K/s

    def _print_params(self):
        print(
            f"\n[setup] Force-free Joule-heating test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  B0        = {self.B0:.3e} T  (|B| uniform)\n"
            f"  k         = {self.k:.4e} 1/m  ({self.n_wave} wavelength(s) across Lx)\n"
            f"  |J|       = {self.J0:.3e} A/m^2  (uniform, force-free)\n"
            f"  V_e drift = {self.v_drift:.3e} m/s\n"
            f"  eta       = {self.eta:.3e} Ohm*m  (1e-5 x scale {self.eta_scale:g})\n"
            f"  tau_R     = {self.tau_R:.3e} s  (current resistive-decay time)\n"
            f"  Grid      = {self.NX} x {self.NZ}  (x x z),  dx = {self.dx:.3e} m\n"
            f"  t_ci      = {self.t_ci:.3e} s,  dt = {self.dt:.3e} s ({self.DT:.4f} t_ci)\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  ----\n"
            f"  PREDICTED dTe/dt = (gamma_e-1) eta J^2 / (n0 kB) = {self.dTe_dt_pred:.4e} K/s\n"
            f"                   = {self.dTe_dt_pred * constants.kb / constants.q_e:.4e} eV/s\n"
        )

    def load_initial_B(self):
        """Set the linear force-free field B(x) = B0[0, sin(kx), cos(kx)].

        WarpX folds Bfield_fp_external into Bfield_fp at initialization, after
        which it evolves self-consistently. div(B) = 0 analytically (B has no
        x-component and the others depend only on x), so no cleaning needed.
        """
        Bx = simulation.fields.get("Bfield_fp_external", dir="x", level=0)
        By = simulation.fields.get("Bfield_fp_external", dir="y", level=0)
        Bz = simulation.fields.get("Bfield_fp_external", dir="z", level=0)

        Bx[:, :] = 0.0
        # Each component on its own (possibly staggered) mesh.
        XBy, _ = np.meshgrid(By.mesh("x"), By.mesh("z"), indexing="ij")
        XBz, _ = np.meshgrid(Bz.mesh("x"), Bz.mesh("z"), indexing="ij")
        By[:, :] = self.B0 * np.sin(self.k * XBy)
        Bz[:, :] = self.B0 * np.cos(self.k * XBz)
        comm.Barrier()


class QeiRelaxation(ElectronEnergyCase):
    """Q_ei electron-ion thermal-equilibration test: pure exponential."""

    te_eV = 300.0  # initial (uniform) electron temperature (eV), hot
    ti_eV = 50.0  # ion temperature (eV); the relaxation target

    # ---- Relaxation ---------------------------------------------------------
    nu_ei = 1.0e6  # electron-ion relaxation rate (1/s), constant;
    #   Te sink rate = 3(gamma_e-1)*nu_ei = 2e6 1/s -> tau = 0.5 us

    # ---- Geometry (small; the physics is 0-D / uniform) / numerics ----------
    NX = 32
    NZ = 8
    NPPC = 400
    n_tau = 3.0  # number of relaxation times to simulate
    steps_per_tau = 100  # rate*dt = 0.01 -> forward-Euler ~ exponential
    substeps = 10

    # do_temperature_deposition is NOT set on purpose -- it is enabled
    # automatically on charged species when the Q_ei relaxation is
    # configured, which this test also exercises (T_ions is dumped below).
    set_temperature_deposition = False
    diag_data_list = ["rho", "Te", "T_ions"]

    def configure(self):
        if self.test:
            self.NX = 16
            self.NZ = 8
            self.NPPC = 200
            self._steps_override = 80
            self.ndiag = 10
        else:
            self._steps_override = None
            self.ndiag = 20

    def get_plasma_quantities(self):
        mi = constants.m_p
        self.dx = self.Lx / self.NX
        self.Lz = self.dx * self.NZ

        # Analytic electron-sink rate and e-folding time.
        self.rate = 3.0 * (self.gamma_e - 1.0) * self.nu_ei
        self.tau = 1.0 / self.rate

        self.dt = self.tau / self.steps_per_tau
        if self._steps_override is not None:
            self.total_steps = self._steps_override
        else:
            self.total_steps = int(self.n_tau * self.steps_per_tau)
        self.diag_steps = max(1, self.total_steps // self.ndiag)

        self.vi_th = np.sqrt(constants.q_e * self.ti_eV / mi)
        # No applied B (B=0). No resistivity (eta=0 -> no Joule, pure Q_ei).
        self.eta = 0.0
        # Constant relaxation rate so the relaxation is a pure exponential.
        self.relaxation_rate = f"{self.nu_ei}"

    def _print_params(self):
        print(
            f"\n[setup] Electron-ion relaxation (Q_ei) test\n"
            f"  Te0 = {self.te_eV:.1f} eV,  Ti0 = {self.ti_eV:.1f} eV,  gamma_e = {self.gamma_e:.4f}\n"
            f"  n0        = {self.n0:.3e} m^-3\n"
            f"  nu_ei     = {self.nu_ei:.3e} 1/s   (constant)\n"
            f"  rate      = 3(gamma-1)nu_ei = {self.rate:.3e} 1/s\n"
            f"  tau       = 1/rate = {self.tau:.3e} s\n"
            f"  Grid      = {self.NX} x {self.NZ},  Lx x Lz = {self.Lx:.3f} x {self.Lz:.4f} m\n"
            f"  dt        = {self.dt:.3e} s   (rate*dt = {self.rate * self.dt:.3f})\n"
            f"  steps     = {self.total_steps},  diag every {self.diag_steps}\n"
            f"  B = 0,  eta = 0  ->  Joule OFF, Q_ei ON (e-sink + conjugate ion heating)\n"
            f"  CHECK:  (Te-Ti)(t) = (Te0-Ti0) exp(-[3(g-1)+2]nu t),  energy conserved\n"
        )


CASES = {
    "adiabat": AdiabaticCompression,
    "joule": ForceFreeJoule,
    "qei": QeiRelaxation,
    "vacuum": VacuumSlabTransport,
}

parser = argparse.ArgumentParser()
parser.add_argument(
    "--case",
    required=True,
    choices=sorted(CASES.keys()),
    help="which electron-energy-equation term to test",
)
parser.add_argument(
    "-t",
    "--test",
    help="toggle whether this script is run as a short CI test",
    action="store_true",
)
parser.add_argument(
    "-v",
    "--verbose",
    help="Verbose output",
    action="store_true",
)
parser.add_argument(
    "--eta-scale",
    type=float,
    default=1.0,
    help="joule case only: multiplier on the base resistivity eta=1e-5 "
    "(amplifies the eta*J^2 heating signal; the CI test uses 100)",
)
args, left = parser.parse_known_args()
sys.argv = sys.argv[:1] + left

case_class = CASES[args.case]
case_class.args = args
run = case_class(test=args.test, verbose=args.verbose)
simulation.step()

Execute:

python3 inputs_test_2d_ohm_solver_electron_energy_picmi.py --case qei

Analyze

Script analysis_qei.py
Listing 82 You can copy this file from Examples/Tests/ohm_solver_electron_energy_eq/analysis_qei.py.
#!/usr/bin/env python3
"""Validate the electron-ion temperature relaxation (Q_ei), both the
electron-side sink AND the conjugate ion heating -- i.e. that the exchange is
energy-conserving.

The companion deck evolves a uniform, unmagnetized, zero-resistivity plasma
with the ions at rest and hot electrons (Te0 >> Ti0), with ONLY the Q_ei
exchange active:

    dU_e/dt = -Q_ei,   Q_ei = 3 n_e k_B nu_ei (T_e - T_i),     (electron sink)
    ions GAIN exactly Q_ei via a thermal-velocity rescale.     (ion source)

With U_e = n_e k_B T_e/(gamma_e-1), the (3/2) n_i k_B T_i ion thermal energy,
a single proton species (Z=1, n_e=n_i) and constant nu_ei, the two
temperatures relax toward a common value:

    dT_e/dt = -3(gamma_e-1) nu_ei (T_e - T_i)        [electron side]
    dT_i/dt = +2 nu_ei (T_e - T_i)                   [ion side, gamma-indep.]

so the difference decays exponentially,

    (T_e - T_i)(t) = (T_e0 - T_i0) exp(-rate t),  rate = [3(gamma_e-1) + 2] nu_ei,

(= 4 nu_ei for gamma_e = 5/3), and the total thermal energy is conserved:

    C_e T_e + C_i T_i = const,   C_e = n_e k_B/(gamma_e-1),  C_i = (3/2) n_i k_B.

For gamma_e=5/3, C_e=C_i so T_e and T_i meet at (T_e0+T_i0)/2.

This script reads domain-mean T_e(t) (Kelvin->eV) and T_i(t) (eV) and checks
  (1) the difference-decay rate vs [3(gamma_e-1)+2] nu_ei, and
  (2) energy conservation: C_e T_e + C_i T_i constant over the run.
"""

import argparse
import sys

import matplotlib

matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from openpmd_viewer import OpenPMDTimeSeries

Q_E = 1.602176634e-19
K_B = 1.380649e-23
K_PER_EV = Q_E / K_B  # T[eV] * this = T[K];  T[K] / this = T[eV]


def domain_means(diag_dir):
    """Return (t[s], <Te>[eV], <Ti>[eV]) density-weighted domain means."""
    ts = OpenPMDTimeSeries(str(diag_dir))
    t = np.asarray(ts.t, dtype=float)

    Te_m, Ti_m = [], []
    for it in ts.iterations:
        Te, _ = ts.get_field("Te", iteration=it)
        Ti, _ = ts.get_field("T_ions", iteration=it)
        rho, _ = ts.get_field("rho", iteration=it)
        w = np.asarray(rho, dtype=float) / Q_E
        wsum = float(np.sum(w))
        Te_m.append(float(np.sum(np.asarray(Te, float) * w) / wsum) / K_PER_EV)
        Ti_m.append(float(np.sum(np.asarray(Ti, float) * w) / wsum))  # already eV
    return t, np.array(Te_m), np.array(Ti_m)


def main(argv=None):
    ap = argparse.ArgumentParser(
        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
    )
    ap.add_argument(
        "--diag-dir",
        default="diags/field_diags",
        help="openPMD field-diagnostics directory",
    )
    ap.add_argument(
        "--nu-ei",
        type=float,
        default=1.0e6,
        help="constant relaxation rate used in the run (1/s); must match the deck",
    )
    ap.add_argument(
        "--gamma", type=float, default=5.0 / 3.0, help="electron adiabatic index"
    )
    ap.add_argument(
        "--rtol",
        type=float,
        default=0.05,
        help="allowed relative error on the fitted difference-rate",
    )
    ap.add_argument(
        "--etol",
        type=float,
        default=0.02,
        help="allowed relative drift of total thermal energy",
    )
    ap.add_argument("--out", default="qei_check.png")
    args = ap.parse_args(argv)

    t, Te, Ti = domain_means(args.diag_dir)
    if t.size < 3:
        print(f"ERROR: need >=3 dumps, found {t.size} in {args.diag_dir}")
        return 1

    g = args.gamma
    # rate at which (Te - Ti) decays = [3(g-1) + 2] nu_ei.
    rate_pred = (3.0 * (g - 1.0) + 2.0) * args.nu_ei
    Te0, Ti0 = Te[0], Ti[0]

    # (1) fit ln((Te-Ti)/(Te0-Ti0)) = -rate t.
    d = (Te - Ti) / (Te0 - Ti0)
    good = d > 1e-3
    rate_fit = -np.polyfit(t[good], np.log(d[good]), 1)[0]
    rel_err = abs(rate_fit - rate_pred) / rate_pred

    # (2) energy conservation: heat capacities (per n k_B; n_e=n_i cancels).
    ce = 1.0 / (g - 1.0)  # C_e / (n k_B)
    ci = 1.5  # C_i / (n k_B)
    E = ce * Te + ci * Ti  # total thermal "energy" per n k_B (eV units)
    E_drift = (E - E[0]) / E[0]
    e_max = float(np.max(np.abs(E_drift)))
    T_eq_pred = (ce * Te0 + ci * Ti0) / (ce + ci)

    print("=" * 66)
    print("Electron-ion relaxation (Q_ei), energy-conserving exchange")
    print(f"  Te0 = {Te0:.2f} eV,  Ti0 = {Ti0:.2f} eV,  gamma = {g:.4f}")
    print(f"  nu_ei (input)             = {args.nu_ei:.4e} 1/s")
    print(f"  diff-rate predicted [3(g-1)+2]nu_ei = {rate_pred:.4e} 1/s")
    print(f"  diff-rate fitted          = {rate_fit:.4e} 1/s")
    print(
        f"  relative error            = {rel_err * 100:.2f}%   (tol {args.rtol * 100:.1f}%)"
    )
    print(f"  equilibrium T predicted   = {T_eq_pred:.2f} eV")
    print(f"  Te_end / Ti_end (meet?)   = {Te[-1]:.2f} / {Ti[-1]:.2f} eV")
    print(
        f"  total-energy max drift    = {e_max * 100:.3f}%   (tol {args.etol * 100:.2f}%)"
    )
    print("=" * 66)

    tus = t * 1e6
    fig, ax = plt.subplots(1, 3, figsize=(15, 4.4))
    ax[0].plot(tus, Te, "o-", ms=4, label=r"$\langle T_e\rangle$")
    ax[0].plot(tus, Ti, "s-", ms=4, color="C3", label=r"$\langle T_i\rangle$")
    ax[0].axhline(T_eq_pred, color="gray", lw=0.9, ls=":", label=r"$T_{eq}$ pred")
    ax[0].set_xlabel(r"time ($\mu$s)")
    ax[0].set_ylabel("temperature (eV)")
    ax[0].set_title("e-i relaxation to common T")
    ax[0].legend()
    ax[0].grid(alpha=0.3)

    ax[1].semilogy(tus[good], d[good], "o", ms=5, label="measured")
    ax[1].semilogy(
        tus, np.exp(-rate_fit * t), "-", lw=2, label=f"fit  rate={rate_fit:.2e}"
    )
    ax[1].semilogy(
        tus, np.exp(-rate_pred * t), "--", lw=2, label=f"pred rate={rate_pred:.2e}"
    )
    ax[1].set_xlabel(r"time ($\mu$s)")
    ax[1].set_ylabel(r"$(T_e-T_i)/(T_{e0}-T_{i0})$")
    ax[1].set_title(f"difference decay (err {rel_err * 100:.1f}%)")
    ax[1].legend()
    ax[1].grid(alpha=0.3, which="both")

    ax[2].plot(tus, E_drift * 100, "o-", ms=4, color="C2")
    ax[2].axhline(0.0, color="gray", lw=0.8, ls=":")
    ax[2].set_xlabel(r"time ($\mu$s)")
    ax[2].set_ylabel(r"$(E-E_0)/E_0$ (%)")
    ax[2].set_title(f"total thermal energy (max {e_max * 100:.2f}%)")
    ax[2].grid(alpha=0.3)

    fig.suptitle("$Q_{ei}$ energy-conserving electron-ion relaxation")
    fig.tight_layout(rect=[0, 0, 1, 0.95])
    fig.savefig(args.out, dpi=150)
    print(f"[saved] {args.out}")

    ok = (rel_err <= args.rtol) and (e_max <= args.etol)
    print("PASS" if ok else "FAIL")
    return 0 if ok else 1


if __name__ == "__main__":
    sys.exit(main())

Execute:

python3 analysis_qei.py --nu-ei 1e6
Electron-ion temperature relaxation

Fig. 21 Electron and ion temperatures relaxing to the common equilibrium value (left), the exponential decay of the temperature difference against the analytic rate (center), and the drift of the total thermal energy (right).