Add arbitrary-harmonic 2D3P VFP solver

adept/__init__.py

@@ -1,3 +1,3 @@
 from ._base_ import ADEPTModule, ergoExo  # noqa: I001
 from .mlflow_logging import MlflowLoggingModule
-from . import hermite_legendre_1d, hermite_poisson_1d, lpse2d, vlasov1d, vlasov2d
+from . import hermite_legendre_1d, hermite_poisson_1d, lpse2d, vfp2d, vlasov1d, vlasov2d

adept/_base_.py

@@ -321,6 +321,9 @@ def _get_adept_module_(self, cfg: dict) -> ADEPTModule:
         elif cfg["solver"] == "vfp-1d":
             from adept.vfp1d.base import BaseVFP1D as this_module
 
+        elif cfg["solver"] == "vfp-2d":
+            from adept.vfp2d import BaseVFP2D as this_module
+
         elif cfg["solver"] == "spectrax-1d":
             from adept.spectrax1d import BaseSpectrax1D as this_module
 

adept/driftdiffusion.py

@@ -667,8 +667,12 @@ class LogMeanFlux(AbstractDriftDiffusionDifferencingScheme):
     from the Peclet number C·dv/D.
 
     This ensures the edge-interpolated f matches the log-mean used in
-    kernel-based compute_D, preserving the bilinear symmetry needed
-    for energy conservation.
+    kernel-based compute_D, preserving the bilinear symmetry needed for a
+    zero semidiscrete energy derivative at the frozen distribution. A finite
+    implicit step freezes these nonlinear coefficients at the old state and
+    therefore does not conserve energy exactly away from a Maxwellian. Use
+    timestep convergence tests; Chang-Cooper is generally the safer option
+    for strongly non-Maxwellian distributions.
 
     The implicit system uses the same tridiagonal structure as Chang-Cooper;
     only the delta computation differs.

adept/vfp1d/fokker_planck.py

@@ -15,7 +15,7 @@
 import equinox as eqx
 import lineax as lx
 import numpy as np
-from jax import Array, vmap
+from jax import Array, lax, vmap
 from jax import numpy as jnp
 
 from adept.driftdiffusion import (
@@ -30,6 +30,54 @@
 from adept.vfp1d.grid import Grid
 
 
+def _linear_solve_value_or_nan(op: lx.AbstractLinearOperator, rhs: Array) -> Array:
+    """Solve without executing Lineax's host error callback during shape tracing.
+
+    Diffrax evaluates a vector field with ``filter_eval_shape`` before compiling a
+    solve. With JAX 0.10, Lineax's default ``throw=True`` callback can inspect the
+    abstract placeholder values and report a spurious non-finite-input failure at
+    ``t=0``. ``throw=False`` keeps the result code in the traced computation. Map a
+    genuine runtime failure to NaNs so downstream finite checks cannot mistake it
+    for a valid collision update.
+    """
+
+    if isinstance(op, lx.TridiagonalLinearOperator):
+        zero = jnp.zeros((1,), dtype=op.diagonal.dtype)
+        lower = jnp.concatenate((zero, op.lower_diagonal))
+        upper = jnp.concatenate((op.upper_diagonal, zero))
+        return lax.linalg.tridiagonal_solve(lower, op.diagonal, upper, rhs[:, None])[:, 0]
+
+    solution = lx.linear_solve(op, rhs, solver=lx.AutoLinearSolver(well_posed=True), throw=False)
+    return jnp.where(
+        solution.result == lx.RESULTS.successful,
+        solution.value,
+        jnp.full_like(solution.value, jnp.nan),
+    )
+
+
+def inverse_bremsstrahlung_resonance_ratio(
+    Z: float | Array,
+    ni: float | Array,
+    nuee_coeff: float,
+    logLam_ratio: float,
+    w0_norm: float,
+) -> Array:
+    """Return the coefficient in ``nu_ei(v) / omega_0 = ratio / v**3``.
+
+    ``nuee_coeff`` contains the normalization-time and reference-density
+    factors, while ``logLam_ratio`` converts the electron-electron reference
+    rate to the electron-ion rate.
+    """
+
+    return (
+        jnp.asarray(nuee_coeff)
+        * jnp.asarray(logLam_ratio)
+        * jnp.asarray(Z) ** 2
+        * jnp.asarray(ni)
+        / jnp.asarray(w0_norm)
+    )
+
+
 class FastVFP(AbstractBetaBasedModel):
     """
     FastVFP model: D = 1/(2β·v).
@@ -248,6 +296,11 @@ class F0Collisions(eqx.Module):
     Uses a positive-only velocity grid (0 to vmax) with zero-flux boundary conditions.
     At v=0, where the drift coefficient C=v=0, zero-flux is equivalent to a reflective
     boundary condition, correctly representing the physics of the isotropic distribution.
+
+    Zero-flux boundaries conserve density. Energy conservation is timestep dependent
+    because the nonlinear drift and diffusion coefficients are frozen during each
+    implicit solve; neither differencing scheme should be treated as an exact finite-step
+    nonlinear energy projection.
     """
 
     nuee_coeff: float
@@ -273,7 +326,9 @@ def _solve_one_vslice_(
         :param dt: time step
         :param D0_heating: Maxwellian heating rate D₀ (>0 heats, <0 cools). None to skip.
         :param ib_vosc2: IB quiver velocity squared v_osc². None to skip.
-        :param ib_Z2ni_w0: IB parameter Z²nᵢ/ω₀. Required when ib_vosc2 is not None.
+        :param ib_Z2ni_w0: Coefficient in νₑᵢ(v)/ω₀ = ib_Z2ni_w0/v³.
+            The legacy name is retained for API compatibility. Required when
+            ib_vosc2 is not None.
 
         :return: updated distribution function (nv,)
         """
@@ -301,7 +356,7 @@ def _solve_one_vslice_(
             D = D + D0_heating * v_edge**2
         if ib_vosc2 is not None:
             # IB heating (Ridgers eq 4.39): D̄ = D + (v_osc²/(6v))·g(v)
-            # g(v) = [1 + (Z²nᵢ/(ω₀v³))²]⁻¹
+            # g(v) = [1 + (νₑᵢ(v)/ω₀)²]⁻¹
             ib_arg = ib_Z2ni_w0 / v_edge**3
             D = D + (ib_vosc2 / (6.0 * v_edge)) / (1.0 + ib_arg**2)
 
@@ -313,7 +368,7 @@ def _solve_one_vslice_(
 
         # Delta formulation: solve for increment δ = f^{n+1} - f^n to reduce
         # floating-point error in density conservation (cf. diffrax)
-        delta = lx.linear_solve(op, f0 - op.mv(f0), solver=lx.AutoLinearSolver(well_posed=True)).value
+        delta = _linear_solve_value_or_nan(op, f0 - op.mv(f0))
         return f0 + delta
 
     def __call__(
@@ -334,7 +389,8 @@ def __call__(
         :param dt: time step
         :param D0_heating: Maxwellian heating rate D₀. Scalar or (nx,). None to skip.
         :param ib_vosc2: IB quiver velocity squared. Scalar or (nx,). None to skip.
-        :param ib_Z2ni_w0: IB parameter Z²nᵢ/ω₀. Scalar or (nx,). Required with ib_vosc2.
+        :param ib_Z2ni_w0: Coefficient in νₑᵢ(v)/ω₀ = ib_Z2ni_w0/v³.
+            Scalar or (nx,). The legacy name is retained for API compatibility.
 
         :return: updated distribution function (nx, nv)
         """
@@ -352,7 +408,8 @@ def __call__(
 
 class FLMCollisions:
     """
-    The FLM collision operator is as described in Tzoufras2014.
+    The linearized anisotropic FLM collision operator is described by
+    Tzoufras et al., JCP 230 (2011), Eq. (41).
 
     It also has an implementation of electron-electron hack
     where the off-diagonal terms in the electron-electron collision
@@ -488,11 +545,23 @@ def _solve_one_x_tridiag_(self, diag: Array, upper: Array, lower: Array, f10: Ar
         Solves a tridiagonal system of equations.
         """
         op = lx.TridiagonalLinearOperator(diagonal=diag, upper_diagonal=upper, lower_diagonal=lower)
-        return lx.linear_solve(op, f10, solver=lx.AutoLinearSolver(well_posed=True)).value
-
-    def __call__(self, Z, ni, f0, f10, dt, include_ee_offdiag_explicitly=True):
+        if jnp.iscomplexobj(f10):
+            # Lineax requires the operator and vector PyTree structures/dtypes
+            # to agree. The collision matrix is real, so solve the two complex
+            # components independently without constructing a complex matrix.
+            real = _linear_solve_value_or_nan(op, jnp.real(f10))
+            imag = _linear_solve_value_or_nan(op, jnp.imag(f10))
+            return real + 1j * imag
+        return _linear_solve_value_or_nan(op, f10)
+
+    def solve_harmonic(self, Z, ni, f0, flm, dt, il: int, include_ee_offdiag_explicitly=True):
         """
-        Solves the FLM collision operator for all l and m.
+        Solve the collision operator for one harmonic order ``il``.
+
+        Collisions are diagonal in ``m`` for the linearized isotropic
+        background used here, so every ``m`` at a given ``il`` uses the same
+        radial operator. Inputs are flattened over configuration space with
+        shape ``(nspace, nv)``.
 
         The solve has two options:
 
@@ -502,32 +571,45 @@ def __call__(self, Z, ni, f0, f10, dt, include_ee_offdiag_explicitly=True):
         2. The ee collision operator is ignored and the Z* scaling is used instead
 
         """
+        if not 1 <= il <= self.grid.nl:
+            raise ValueError(f"il must satisfy 1 <= il <= {self.grid.nl}; got {il}")
+
         v = self.grid.v[None, :]
         dv = self.grid.dv
-        for il in range(1, self.grid.nl + 1):
-            ei_diag = -il * (il + 1) / 2.0 * (Z[:, None] ** 2.0) * ni[:, None] / v**3.0
+        ei_diag = -il * (il + 1) / 2.0 * (Z[:, None] ** 2.0) * ni[:, None] / v**3.0
 
-            if self.full_aniso_ee:
-                ee_diag, ee_lower, ee_upper = self.get_ee_diagonal_contrib(f0)
-                pad_f0 = jnp.concatenate([f0[:, 1::-1], f0], axis=1)
-                #
-                d2dv2 = 0.5 / v * jnp.gradient(jnp.gradient(pad_f0, dv, axis=1), dv, axis=1)[:, 2:]
+        if self.full_aniso_ee:
+            ee_diag, ee_lower, ee_upper = self.get_ee_diagonal_contrib(f0)
+            pad_f0 = jnp.concatenate([f0[:, 1::-1], f0], axis=1)
+            d2dv2 = 0.5 / v * jnp.gradient(jnp.gradient(pad_f0, dv, axis=1), dv, axis=1)[:, 2:]
 
-                ddv = v**-2.0 * jnp.gradient(pad_f0, dv, axis=1)[:, 2:]
+            ddv = v**-2.0 * jnp.gradient(pad_f0, dv, axis=1)[:, 2:]
 
-                diag = 1 - dt * (self.nuei_coeff * ei_diag + self.nuee_coeff * ee_diag)
-                lower = -dt * self.nuee_coeff * ee_lower
-                upper = -dt * self.nuee_coeff * ee_upper
+            diag = 1 - dt * (self.nuei_coeff * ei_diag + self.nuee_coeff * ee_diag)
+            lower = -dt * self.nuee_coeff * ee_lower
+            upper = -dt * self.nuee_coeff * ee_upper
 
-                new_f10 = vmap(self._solve_one_x_tridiag_, in_axes=(0, 0, 0, 0))(diag, upper, lower, f10)
+            new_flm = vmap(self._solve_one_x_tridiag_, in_axes=(0, 0, 0, 0))(diag, upper, lower, flm)
 
-                if include_ee_offdiag_explicitly:
-                    new_f10 = new_f10 + dt * self.nuee_coeff * self.get_ee_offdiagonal_contrib(
-                        None, f10, {"ddvf0": ddv, "d2dv2f0": d2dv2, "il": il}
-                    )
+            if include_ee_offdiag_explicitly:
+                new_flm = new_flm + dt * self.nuee_coeff * self.get_ee_offdiagonal_contrib(
+                    None, flm, {"ddvf0": ddv, "d2dv2f0": d2dv2, "il": il}
+                )
+        else:
+            # Epperlein--Haines Z* approximation to electron--electron scattering.
+            new_flm = flm / (1 - dt * self.nuei_coeff * self.Z_nuei_scaling * ei_diag)
 
-            else:
-                # only uses the Z* epperlein haines scaling instead of solving the ee collisions
-                new_f10 = f10 / (1 - dt * self.nuei_coeff * self.Z_nuei_scaling * ei_diag)
+        return new_flm
 
-        return new_f10
+    def __call__(self, Z, ni, f0, f10, dt, include_ee_offdiag_explicitly=True):
+        """Backward-compatible ``f10`` solve for the VFP-1D ``l=1`` state."""
+
+        return self.solve_harmonic(
+            Z=Z,
+            ni=ni,
+            f0=f0,
+            flm=f10,
+            dt=dt,
+            il=1,
+            include_ee_offdiag_explicitly=include_ee_offdiag_explicitly,
+        )

adept/vfp1d/vector_field.py

@@ -3,7 +3,14 @@
 from jax import Array
 from jax import numpy as jnp
 
-from adept.vfp1d.fokker_planck import F0Collisions, FLMCollisions, SelfConsistentBetaConfig, get_model, get_scheme
+from adept.vfp1d.fokker_planck import (
+    F0Collisions,
+    FLMCollisions,
+    SelfConsistentBetaConfig,
+    get_model,
+    get_scheme,
+    inverse_bremsstrahlung_resonance_ratio,
+)
 
 
 class OSHUN1D:
@@ -51,6 +58,8 @@ def __init__(self, cfg: dict, grid):
         vosc2_per_intensity = cfg["units"]["derived"].get("vosc2_per_intensity", 0.0)
         self.ib_vosc2 = vosc2_per_intensity * ib_intensity
         self.ib_w0 = cfg["units"]["derived"].get("w0_norm", 1.0)
+        self.ib_nuee_coeff = nuee_coeff
+        self.ib_logLam_ratio = cfg["units"]["derived"].get("logLam_ratio", 1.0)
         self._ib_enabled = self.ib_vosc2 > 0.0
 
         self.solve_aniso = FLMCollisions(
@@ -393,7 +402,13 @@ def __call__(self, t, y, args) -> dict:
             heating_kwargs["D0_heating"] = self.D0_heating
         if self._ib_enabled:
             heating_kwargs["ib_vosc2"] = self.ib_vosc2
-            heating_kwargs["ib_Z2ni_w0"] = Z**2 * ni / self.ib_w0
+            heating_kwargs["ib_Z2ni_w0"] = inverse_bremsstrahlung_resonance_ratio(
+                Z,
+                ni,
+                self.ib_nuee_coeff,
+                self.ib_logLam_ratio,
+                self.ib_w0,
+            )
         f0_star = self.solve_Cee0(None, f0_star, self.grid.dt, **heating_kwargs)
 
         # Interpolate center quantities to edges for FLM collision operator

adept/vfp2d/__init__.py

@@ -0,0 +1,55 @@
+"""Two-dimensional spherical-harmonic Vlasov--Fokker--Planck solver."""
+
+from adept.vfp2d.base import BaseVFP2D
+from adept.vfp2d.collisions import AnisotropicCollisions, CollisionStep
+from adept.vfp2d.grid import Grid
+from adept.vfp2d.harmonics import (
+    HarmonicLayout,
+    HouLiFilter2D,
+    TzoufrasVlasov,
+    cartesian_l2,
+    complex_to_real,
+    conservative_f00_positivity,
+    current,
+    density,
+    nernst_velocity,
+    real_to_complex,
+    scalar_velocity_moment,
+    tensor_velocity_moment,
+    vector_velocity_moment,
+)
+from adept.vfp2d.ohm import KineticOhm2D, project_current_moment
+from adept.vfp2d.vector_field import (
+    KineticOhmStep,
+    Maxwell2D,
+    SpectralPoisson2D,
+    SplitStepVFP2D,
+    VlasovMaxwell,
+)
+
+__all__ = [
+    "AnisotropicCollisions",
+    "BaseVFP2D",
+    "CollisionStep",
+    "Grid",
+    "HarmonicLayout",
+    "HouLiFilter2D",
+    "KineticOhm2D",
+    "KineticOhmStep",
+    "Maxwell2D",
+    "SpectralPoisson2D",
+    "SplitStepVFP2D",
+    "TzoufrasVlasov",
+    "VlasovMaxwell",
+    "cartesian_l2",
+    "complex_to_real",
+    "conservative_f00_positivity",
+    "current",
+    "density",
+    "nernst_velocity",
+    "project_current_moment",
+    "real_to_complex",
+    "scalar_velocity_moment",
+    "tensor_velocity_moment",
+    "vector_velocity_moment",
+]