diff --git a/adept/_vlasov1d/datamodel.py b/adept/_vlasov1d/datamodel.py index 13010629..75d92293 100644 --- a/adept/_vlasov1d/datamodel.py +++ b/adept/_vlasov1d/datamodel.py @@ -169,6 +169,9 @@ class FokkerPlanckConfig(BaseModel): type: str time: EnvelopeConfig space: EnvelopeConfig + # Super-Gaussian exponent of the operator's equilibrium (only used by + # type: super_gaussian; m=2 is Maxwellian) + m: float = Field(default=2.0, ge=1.0) class KrookConfig(BaseModel): diff --git a/adept/_vlasov1d/solvers/pushers/fokker_planck.py b/adept/_vlasov1d/solvers/pushers/fokker_planck.py index dfb2a5b1..dec5fcee 100644 --- a/adept/_vlasov1d/solvers/pushers/fokker_planck.py +++ b/adept/_vlasov1d/solvers/pushers/fokker_planck.py @@ -13,12 +13,19 @@ import jax import lineax as lx import numpy as np +import optimistix as optx from jax import Array, shard_map, vmap from jax import numpy as jnp from jax.sharding import Mesh from jax.sharding import PartitionSpec as P -from adept.driftdiffusion import AbstractBetaBasedModel, CentralDifferencing, ChangCooper +from adept.driftdiffusion import ( + AbstractBetaBasedModel, + CentralDifferencing, + ChangCooper, + analytic_supergaussian_temperature_ratio, + chang_cooper_delta, +) class LenardBernstein(AbstractBetaBasedModel): @@ -96,6 +103,158 @@ def compute_D(self, f: Array, beta: Array) -> Array: return 1.0 / (2.0 * beta) +class SuperGaussianDougherty(Dougherty): + """ + Dougherty-type collision model whose equilibrium is a super-Gaussian. + + Physics: + Equilibrium: f₀ ∝ exp(-β·|v - vbar|^m) (m = 2 recovers Dougherty) + C = D · Δφ/dv with φ = β·|v - vbar|^m = -ln f₀ (exact finite + difference of the potential across each edge; the continuum limit + of C/D = -∂v ln f₀) + D = β^(-2/m)·Γ(3/m)/Γ(1/m) (the analytic temperature of the target + super-Gaussian; reduces to D = 1/(2β) = T at m = 2) + + Here β is the super-Gaussian shape parameter (units of v^-m), computed by + compute_beta from the energy-conservation closure β = n/(m·⟨|v-vbar|^m⟩) + (the generalization of the Lenard-Bernstein/Dougherty β = 1/(2T), which it + equals at m = 2). Paired with the Chang-Cooper scheme and the exact-Δφ + drift below, the sampled discrete super-Gaussian is the exact fixed point. + + This maintains a prescribed super-Gaussian order m — e.g. a Langdon/DLM + inverse-bremsstrahlung-heated distribution — against collisional + relaxation to a Maxwellian. Conserves density exactly (zero-flux BC) and + energy up to discretization error (via the β closure). Momentum is exact + only for f symmetric about vbar; skewed transients exchange momentum at + O(skewness) — unlike m = 2, where C ∝ (v - vbar) makes it exact. + """ + + m: float + + def __init__(self, v: Array, dv: float, m: float): + """Initialize model with velocity grid and super-Gaussian exponent m.""" + super().__init__(v=v, dv=dv) + self.m = m + + def compute_beta(self, f: Array, rtol: float = 1e-8, atol: float = 1e-12, max_steps: int = 3) -> Array: + """ + Compute the energy-conserving shape parameter β. + + The continuum operator conserves energy instantaneously iff + m·β·⟨|v-vbar|^m⟩ = ⟨1⟩ (from ∫v²·∂v[D·(φ'f + ∂vf)]dv = 0), giving the + closure β = n/(m·⟨|v-vbar|^m⟩). At m = 2 this reduces to + β = n/(2·⟨(v-vbar)²⟩) = 1/(2T), the Buet beta computed from the + discrete temperature — where it coincides with temperature matching, + which is why the Maxwellian operators get both for free. + + With max_steps > 0, β is refined by a Newton solve of the DISCRETE + energy-flux condition consistent with the Chang-Cooper weighting, + + h(β) = Σ_edges v_e·[w·f̃(w) + Δf] = 0, w = β·Δψ, + f̃(w) = δ(w)·f_i + (1-δ(w))·f_{i+1}, ψ = |v - vbar|^m, + + which zeroes the discrete energy flux of the operator exactly for the + current f. At a sampled super-Gaussian every edge term vanishes + identically (the Chang-Cooper equilibrium property), so β₀ is an + exact root: the discrete equilibrium is exactly stationary and there + is no secular drift. With max_steps = 0 (self_consistent_beta + disabled) the continuum closure is used directly; its O(dv²) + quadrature residual causes a slow secular temperature drift + (measured ~5e-7 per collision time at nv=128 for m=3). The continuum + closure lands within O(dv²) of the root, so a SINGLE Newton step + (max_steps: 1) already reduces the drift to machine level; use it for + long collisional runs. + + Transient (off-equilibrium) energy conservation is additionally + limited to O(nu·dt) by the operator splitting: β is frozen from f^n + while the implicit step advances to f^{n+1}. This is a one-time + offset proportional to the shape change, negligible for the small + nu·dt of production runs. + + Args: + f: Distribution function, shape (..., nv) + rtol: Relative tolerance for the Newton solve + atol: Absolute tolerance for the Newton solve + max_steps: Maximum Newton iterations (0 = continuum closure only) + + Returns: + beta: Shape parameter (units of v^-m), shape (...) + """ + m = self.m + v = self.v + v_edge = 0.5 * (v[1:] + v[:-1]) + + def beta_single(f_v: Array) -> Array: + vbar = jnp.sum(f_v * v) / jnp.sum(f_v) + psi = jnp.abs(v - vbar) ** m + # Continuum closure: β = n / (m·⟨ψ⟩); dv cancels in the ratio + beta_init = jnp.sum(f_v) / (m * jnp.sum(f_v * psi)) + if max_steps == 0: + return beta_init + + d_psi = psi[1:] - psi[:-1] + d_f = f_v[1:] - f_v[:-1] + + def residual(beta, args): + del args + w = beta * d_psi + delta = chang_cooper_delta(w) + f_tilde = delta * f_v[:-1] + (1.0 - delta) * f_v[1:] + return jnp.sum(v_edge * (w * f_tilde + d_f)) + + solver = optx.Newton(rtol=rtol, atol=atol) + sol = optx.root_find(fn=residual, solver=solver, y0=beta_init, args=None, max_steps=max_steps, throw=False) + return sol.value + + return vmap(beta_single)(f) + + def compute_D(self, f: Array, beta: Array) -> Array: + """ + Compute diffusion coefficient from the super-Gaussian shape parameter. + + D = β^(-2/m)·Γ(3/m)/Γ(1/m) is the analytic temperature of the target + shape, so that the m=2 case reduces to the Dougherty D = 1/(2β) = T. + + Args: + f: Distribution function, shape (..., nv) - unused + beta: Super-Gaussian shape parameter, shape (...) + + Returns: + D: Diffusion coefficient, shape (...) + """ + return beta ** (-2.0 / self.m) * analytic_supergaussian_temperature_ratio(self.m) + + def compute_C_and_D(self, f: Array, beta: Array) -> tuple[Array, Array]: + """ + Compute C and D from the super-Gaussian equilibrium condition. + + C is the exact finite difference of the equilibrium potential + φ(v) = β·|v - vbar|^m = -ln f₀ across each cell edge: + + C_edge = D · (φ(v_{i+1}) - φ(v_i)) / dv + + rather than the midpoint derivative m·β·|v_eff|^(m-1)·sgn(v_eff). + The Chang-Cooper fixed point satisfies f_{i+1}/f_i = exp(-C·dv/D), + so this choice makes the SAMPLED super-Gaussian the exact discrete + equilibrium (the midpoint form is only exact for m=2 and causes + secular temperature drift for m≠2). For m=2 the two coincide: + Δ(v²)/dv = 2·v_edge exactly. + + Args: + f: Distribution function, shape (..., nv) + beta: Super-Gaussian shape parameter, shape (...) + + Returns: + C_edge: Drift coefficient at cell edges, shape (..., nv-1) + D: Diffusion coefficient, shape (...) + """ + D = self.compute_D(f, beta) + vbar = self.compute_vbar(f) + phi = beta[..., None] * jnp.abs(self.v - vbar[..., None]) ** self.m + C_edge = D[..., None] * (phi[..., 1:] - phi[..., :-1]) / self.dv + return C_edge, D + + class Collisions: """High-level collision operator that wraps Fokker-Planck and Krook terms.""" @@ -150,6 +309,10 @@ def __init_fp_operator__(self): elif fp_type == "dougherty": model = Dougherty(v=v, dv=dv) return model, CentralDifferencing(dv=dv) + elif fp_type in ("super_gaussian", "super_gaussian_chang_cooper"): + m = float(self.cfg["terms"]["fokker_planck"].get("m", 2.0)) + model = SuperGaussianDougherty(v=v, dv=dv, m=m) + return model, ChangCooper(dv=dv) else: raise NotImplementedError(f"Unknown Fokker-Planck type: {fp_type}") @@ -200,17 +363,25 @@ def _apply_collisions(self, nu_fp: jnp.ndarray, nu_K: jnp.ndarray, f: jnp.ndarra def _collide(f_shard, nu_fp_shard, nu_K_shard): if self.cfg["terms"]["fokker_planck"]["is_on"]: - vbar = self.fp_model.compute_vbar(f_shard) - beta = find_self_consistent_beta( - f_shard, - v, - dv, - spherical=False, - vbar=vbar, - rtol=self._sc_rtol, - atol=self._sc_atol, - max_steps=self._sc_max_steps, - ) + if hasattr(self.fp_model, "compute_beta"): + # Model defines its own beta closure (e.g. the energy-conserving + # super-Gaussian beta); the self_consistent_beta knobs control + # its Newton refinement of the discrete energy-flux condition. + beta = self.fp_model.compute_beta( + f_shard, rtol=self._sc_rtol, atol=self._sc_atol, max_steps=self._sc_max_steps + ) + else: + vbar = self.fp_model.compute_vbar(f_shard) + beta = find_self_consistent_beta( + f_shard, + v, + dv, + spherical=False, + vbar=vbar, + rtol=self._sc_rtol, + atol=self._sc_atol, + max_steps=self._sc_max_steps, + ) C_edge, D = self.fp_model.compute_C_and_D(f_shard, beta) f_shard = vmap(self._solve_one_x, in_axes=(0, 0, 0, 0, None))(C_edge, D, nu_fp_shard, f_shard, dt) diff --git a/adept/driftdiffusion.py b/adept/driftdiffusion.py index a8a224d5..a2482ac8 100644 --- a/adept/driftdiffusion.py +++ b/adept/driftdiffusion.py @@ -36,7 +36,7 @@ └── AbstractKernelBasedModel (D, C via kernel integrals) Concrete models live alongside their solvers: -- LenardBernstein, Dougherty → adept._vlasov1d.solvers.pushers.fokker_planck +- LenardBernstein, Dougherty, SuperGaussianDougherty → adept._vlasov1d.solvers.pushers.fokker_planck - FastVFP → adept.vfp1d.fokker_planck - CoulombianKernel, AsymptoticLocal → adept.vfp1d.fokker_planck @@ -71,6 +71,7 @@ import optimistix as optx from jax import Array from jax import numpy as jnp +from jax.scipy.special import gammaln def chang_cooper_delta(w: Array) -> Array: @@ -136,6 +137,27 @@ def discrete_temperature( return v2_moment / norm +def analytic_supergaussian_temperature_ratio(m: float, spherical: bool = False) -> Array: + """ + Analytic ratio T/v_m² for a super-Gaussian f ∝ exp(-(|v|/v_m)^m). + + Cartesian (1D symmetric grid): T = ⟨v²⟩ = v_m²·Γ(3/m)/Γ(1/m) + Spherical (speed grid, v² weight): T = ⟨v⁴⟩/(3⟨v²⟩) = v_m²·Γ(5/m)/(3·Γ(3/m)) + + Both reduce to 1/2 at m=2, recovering the Maxwellian T = v_m²/2 (β = 1/(2T)). + + Args: + m: Super-Gaussian exponent (m=2 is Maxwellian) + spherical: If True, use the spherical moment definition of T + + Returns: + T/v_m², scalar + """ + if spherical: + return jnp.exp(gammaln(5.0 / m) - gammaln(3.0 / m)) / 3.0 + return jnp.exp(gammaln(3.0 / m) - gammaln(1.0 / m)) + + def _find_self_consistent_beta_single( f: Array, v: Array, @@ -145,14 +167,17 @@ def _find_self_consistent_beta_single( rtol: float = 1e-8, atol: float = 1e-12, max_steps: int = 3, + m: float = 2.0, ) -> Array: """ Find beta for a single spatial point (internal, not vmapped). - Given a distribution f, finds β* such that a Maxwellian with that β* - has the same discrete temperature as f. + Given a distribution f, finds β* such that the target equilibrium shape + exp(-β·|v-vbar|^m) has the same discrete temperature as f. - Uses Buet notation: beta = 1/(2T) so f = exp(-beta * (v-vbar)²). + For m=2 (default) this is the Maxwellian in Buet notation: beta = 1/(2T) + so f = exp(-beta * (v-vbar)²). For m>2 the target is a super-Gaussian and + β is the shape parameter (units of v^-m). Args: f: Distribution function, shape (nv,) @@ -163,27 +188,37 @@ def _find_self_consistent_beta_single( rtol: Relative tolerance for Newton solver (default: 1e-8) atol: Absolute tolerance for Newton solver (default: 1e-12) max_steps: Maximum Newton iterations (default: 3) + m: Super-Gaussian exponent of the target shape (default 2.0 = Maxwellian) Returns: - beta_star: Beta value where Maxwellian has T_discrete = T_f, scalar + beta_star: Beta value where the target shape has T_discrete = T_f, scalar """ # Compute target temperature from f (using vbar if provided) T_target = discrete_temperature(f, v, dv, spherical, vbar) - # Initial guess: Buet beta = 1/(2*T_f) - beta_init = 1.0 / (2.0 * T_target) + # Initial guess from the analytic moment relation. m is a static Python + # float, so the m=2 branch keeps the historical bitwise-exact expressions. + if m == 2.0: + # Buet beta = 1/(2*T_f) + beta_init = 1.0 / (2.0 * T_target) + else: + # T = β^(-2/m) · (T/v_m² ratio) → β = (T / ratio)^(-m/2) + beta_init = (T_target / analytic_supergaussian_temperature_ratio(m, spherical)) ** (-m / 2.0) - # Short-circuit: no SC iterations means just use the discrete-T beta + # Short-circuit: no SC iterations means just use the analytic-moment beta if max_steps == 0: return beta_init def residual(beta, args): v, dv, T_target, spherical, vbar = args - # Maxwellian: f = exp(-beta * (v-vbar)²) + # Target shape: f = exp(-beta * |v-vbar|^m) v_shifted = v if vbar is None else (v - vbar) - f_maxwellian = jnp.exp(-beta * v_shifted**2) - T_maxwellian = discrete_temperature(f_maxwellian, v, dv, spherical, vbar) - return T_maxwellian - T_target + if m == 2.0: + f_model = jnp.exp(-beta * v_shifted**2) + else: + f_model = jnp.exp(-beta * jnp.abs(v_shifted) ** m) + T_model = discrete_temperature(f_model, v, dv, spherical, vbar) + return T_model - T_target solver = optx.Newton(rtol=rtol, atol=atol) sol = optx.root_find( @@ -202,7 +237,7 @@ def residual(beta, args): # without needing a separate vmap definition. _find_self_consistent_beta_vmapped = eqx.filter_vmap( _find_self_consistent_beta_single, - in_axes=(0, None, None, None, 0, None, None, None), + in_axes=(0, None, None, None, 0, None, None, None, None), ) @@ -215,16 +250,19 @@ def find_self_consistent_beta( rtol: float = 1e-8, atol: float = 1e-12, max_steps: int = 3, + m: float = 2.0, ) -> Array: """ - Find beta such that a Maxwellian has the same discrete temperature as f. + Find beta such that the target equilibrium shape has the same discrete temperature as f. - Given a distribution f, finds β* such that a Maxwellian with that β* - has the same discrete temperature as f. This eliminates equilibrium drift - in Chang-Cooper schemes caused by the mismatch between analytical and + Given a distribution f, finds β* such that exp(-β*·|v-vbar|^m) has the + same discrete temperature as f. This eliminates equilibrium drift in + Chang-Cooper schemes caused by the mismatch between analytical and discrete temperatures. - Uses Buet notation: beta = 1/(2T) so f = exp(-beta * (v-vbar)²). + For m=2 (default, Maxwellian) this uses Buet notation: beta = 1/(2T) so + f = exp(-beta * (v-vbar)²). For m≠2 the target is a super-Gaussian and β + is its shape parameter (units of v^-m). Vmapped over f (and vbar if provided) for batching over spatial points. @@ -237,11 +275,12 @@ def find_self_consistent_beta( rtol: Relative tolerance for Newton solver (default: 1e-8) atol: Absolute tolerance for Newton solver (default: 1e-12) max_steps: Maximum Newton iterations (default: 3) + m: Super-Gaussian exponent of the target shape (default 2.0 = Maxwellian) Returns: - beta_star: Beta values where Maxwellian has T_discrete = T_f, shape (nx,) + beta_star: Beta values where the target shape has T_discrete = T_f, shape (nx,) """ - return _find_self_consistent_beta_vmapped(f, v, dv, spherical, vbar, rtol, atol, max_steps) + return _find_self_consistent_beta_vmapped(f, v, dv, spherical, vbar, rtol, atol, max_steps, m) class AbstractMaxwellianPreservingModel(eqx.Module): diff --git a/docs/source/solvers/vlasov1d/config.md b/docs/source/solvers/vlasov1d/config.md index a7e38f63..01e809f6 100644 --- a/docs/source/solvers/vlasov1d/config.md +++ b/docs/source/solvers/vlasov1d/config.md @@ -537,10 +537,94 @@ Fokker-Planck collision operator. | Field | Type | Description | |-------|------|-------------| | `is_on` | bool | Enable/disable | -| `type` | string | Collision type: `"Dougherty"` | +| `type` | string | Collision type (see table below) | +| `m` | float | Super-Gaussian exponent of the equilibrium (only used by `type: super_gaussian`; default `2.0` = Maxwellian) | +| `self_consistent_beta` | object | Optional Newton refinement of the equilibrium shape parameter (see below) | | `time` | object | Temporal profile | | `space` | object | Spatial profile | +Available `type` values (case-insensitive): + +| `type` | Model | Scheme | Equilibrium | +|--------|-------|--------|-------------| +| `lenard_bernstein` | Lenard-Bernstein | central differencing | Maxwellian at v=0 | +| `chang_cooper` | Lenard-Bernstein | Chang-Cooper | Maxwellian at v=0 | +| `dougherty` | Dougherty | central differencing | Maxwellian at $\bar v$ | +| `chang_cooper_dougherty` | Dougherty | Chang-Cooper | Maxwellian at $\bar v$ | +| `super_gaussian` | Super-Gaussian Dougherty | Chang-Cooper | Super-Gaussian of order `m` at $\bar v$ | + +The Chang-Cooper scheme is positivity-preserving and conserves density exactly; central differencing is provided for comparison only. + +#### super_gaussian + +Drift-diffusion operator whose equilibrium is a super-Gaussian +$f_0 \propto \exp(-\beta\,|v-\bar v|^m)$ instead of a Maxwellian +($m=2$ recovers `chang_cooper_dougherty`). Use this to *maintain* a prescribed +non-Maxwellian order — e.g. a Langdon/DLM inverse-bremsstrahlung-heated +distribution — against collisional relaxation during Landau-damping or LPI +runs. Note it does not *generate* super-Gaussians dynamically (there is no +competition between heating and e-e collisions as in the VFP-1D +inverse-bremsstrahlung operator); it relaxes toward the prescribed order `m`. + +The drift coefficient is the exact finite difference of the equilibrium +potential $\varphi = \beta|v-\bar v|^m$, so the sampled super-Gaussian is the +exact fixed point of the Chang-Cooper discretization. The shape parameter +$\beta$ is set by the energy-conservation closure +$\beta = n/(m\,\langle|v-\bar v|^m\rangle)$ (the generalization of the +Lenard-Bernstein/Dougherty $\beta = 1/(2T)$), refined by a Newton solve of the +discrete energy-flux condition when `self_consistent_beta` is enabled. + +Conservation properties: + +- **Density**: exact (zero-flux boundary conditions). +- **Energy**: no secular drift at equilibrium when `self_consistent_beta` is + enabled (recommended for long collisional runs; without it the continuum + closure leaves an O(dv²) quadrature residual that drifts T by ~5e-7 per + collision time at nv=128 for m=3 — already negligible unless many collision + times elapse). A single Newton step (`max_steps: 1`) reduces the drift to + machine level. Off-equilibrium transients carry a one-time O(nu·dt) offset + from operator splitting, negligible at production nu·dt. +- **Momentum**: exact for distributions symmetric about $\bar v$; skewed + transients exchange momentum at O(skewness), unlike $m=2$ where + $C \propto (v-\bar v)$ makes it exact. + +Note on initialization: the super-Gaussian initializer (`m` in the species +config) uses a width convention that fixes $\langle v^4\rangle/\langle v^2\rangle$, +while this operator preserves the super-Gaussian carrying the distribution's +own variance. Any super-Gaussian of order `m` is (approximately) a fixed +point regardless of width convention, so the two compose without drift. + +```yaml +terms: + fokker_planck: + is_on: True + type: super_gaussian + m: 3.0 + self_consistent_beta: + enabled: True + max_steps: 1 # one Newton step suffices (the closure lands within O(dv²) of the root) + time: + baseline: 1.0e-5 + # ... + space: + baseline: 1.0 + # ... +``` + +#### self_consistent_beta + +Optional sub-object controlling the Newton refinement of the equilibrium shape +parameter β. For the Maxwellian operators it matches the *discrete* temperature +of the equilibrium to that of f (eliminating equilibrium drift in Chang-Cooper +schemes); for `super_gaussian` it solves the discrete energy-flux condition. + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `enabled` | bool | `False` | Enable the Newton solve | +| `max_steps` | int | `3` | Maximum Newton iterations | +| `rtol` | float | `1e-8` | Relative tolerance | +| `atol` | float | `1e-12` | Absolute tolerance | + ### krook Krook collision operator (relaxation towards Maxwellian). diff --git a/tests/test_vlasov1d/test_super_gaussian_fp.py b/tests/test_vlasov1d/test_super_gaussian_fp.py new file mode 100644 index 00000000..d7b4f0c1 --- /dev/null +++ b/tests/test_vlasov1d/test_super_gaussian_fp.py @@ -0,0 +1,231 @@ +# Copyright (c) Ergodic LLC 2026 +# research@ergodic.io +""" +Tests for the super-Gaussian-preserving Fokker-Planck operator (vlasov1d). + +The SuperGaussianDougherty model relaxes toward f ∝ exp(-β·|v-vbar|^m) instead +of a Maxwellian, e.g. to maintain a Langdon/DLM inverse-bremsstrahlung-heated +distribution against collisional relaxation. These tests verify: + +- a super-Gaussian of the configured order is a fixed point (no secular drift), +- a Maxwellian relaxes TO the super-Gaussian (kurtosis moves to the SG value), +- density is conserved exactly; the transient energy error is O(dv²), +- m=2 reduces to the Dougherty (Chang-Cooper) operator. +""" + +from typing import NamedTuple + +import jax +import jax.numpy as jnp +import numpy as np +import pytest +from scipy.special import gamma + +from adept._vlasov1d.solvers.pushers.fokker_planck import Collisions + +VMAX = 6.0 + + +class Grid(NamedTuple): + v: np.ndarray + dv: float + + +def make_grid(nv: int) -> Grid: + dv = 2.0 * VMAX / nv + return Grid(v=np.linspace(-VMAX + dv / 2.0, VMAX - dv / 2.0, nv), dv=dv) + + +GRID = make_grid(128) + + +# ============================================================================= +# Helpers +# ============================================================================= + + +def make_collisions(fp_type: str, m: float | None = None, sc_steps: int = 3, grid: Grid = GRID) -> Collisions: + """Build a production Collisions instance for a single-x-point test grid.""" + fp_cfg = { + "type": fp_type, + "is_on": True, + "self_consistent_beta": {"enabled": sc_steps > 0, "max_steps": sc_steps}, + } + if m is not None: + fp_cfg["m"] = m + cfg = { + "grid": {"species_grids": {"electron": {"v": grid.v, "dv": grid.dv}}}, + "terms": {"fokker_planck": fp_cfg, "krook": {"is_on": False}}, + } + return Collisions(cfg) + + +def step_n(collisions: Collisions, f: jnp.ndarray, nsteps: int, nu: float = 1.0, dt: float = 0.1) -> jnp.ndarray: + """Apply nsteps implicit collision steps to f of shape (nx, nv).""" + nu_fp = jnp.full(f.shape[0], nu) + nu_k = jnp.zeros(f.shape[0]) + + @jax.jit + def body(f_carry, _): + return collisions(nu_fp, nu_k, f_carry, dt), None + + f_final, _ = jax.lax.scan(body, f, length=nsteps) + return f_final + + +def supergaussian(m: float, T: float, v0: float = 0.0, grid: Grid = GRID) -> np.ndarray: + """Unit-density super-Gaussian with variance T: f ∝ exp(-(|v-v0|/vm)^m).""" + vm = np.sqrt(T * gamma(1.0 / m) / gamma(3.0 / m)) + f = np.exp(-(np.abs((grid.v - v0) / vm) ** m)) + return f / np.sum(f * grid.dv) + + +def maxwellian(T: float, v0: float = 0.0, grid: Grid = GRID) -> np.ndarray: + return supergaussian(2.0, T, v0, grid) + + +def moments(f: np.ndarray, grid: Grid = GRID) -> tuple[float, float, float, float]: + """Return (density, vbar, T, kurtosis) of f(v); kurtosis = ⟨(v-vbar)⁴⟩/T².""" + f = np.asarray(f) + n = np.sum(f * grid.dv) + vbar = np.sum(f * grid.v * grid.dv) / n + T = np.sum(f * (grid.v - vbar) ** 2 * grid.dv) / n + kurt = np.sum(f * (grid.v - vbar) ** 4 * grid.dv) / n / T**2 + return n, vbar, T, kurt + + +def sg_kurtosis(m: float) -> float: + """Analytic kurtosis ⟨v⁴⟩/⟨v²⟩² of a 1D super-Gaussian of order m.""" + return gamma(5.0 / m) * gamma(1.0 / m) / gamma(3.0 / m) ** 2 + + +# ============================================================================= +# Tests +# ============================================================================= + + +@pytest.mark.parametrize("m", [3.0, 4.0]) +def test_supergaussian_is_fixed_point(m): + """A super-Gaussian of the configured order stays put; conservation holds.""" + collisions = make_collisions("super_gaussian", m=m) + f0 = jnp.asarray(supergaussian(m, T=1.0))[None, :] + f1 = step_n(collisions, f0, nsteps=100) # 10 collision times + + n0, _, T0, _ = moments(f0[0]) + n1, _, T1, k1 = moments(f1[0]) + + assert abs(n1 / n0 - 1.0) < 1e-12, f"density drift: {n1 / n0 - 1.0:.2e}" + assert abs(T1 / T0 - 1.0) < 5e-3, f"temperature drift: {T1 / T0 - 1.0:.2e}" + # Shape is maintained: kurtosis stays at the super-Gaussian value + assert abs(k1 - sg_kurtosis(m)) < 0.02, f"kurtosis {k1:.4f} left SG value {sg_kurtosis(m):.4f}" + # And f itself barely moves (analytic vs self-consistent discrete equilibrium) + rel_l2 = float(jnp.linalg.norm(f1 - f0) / jnp.linalg.norm(f0)) + assert rel_l2 < 5e-3, f"super-Gaussian not a fixed point: rel L2 change {rel_l2:.2e}" + assert float(jnp.min(f1)) > -1e-20, "positivity violated" + + +def test_no_secular_drift_at_equilibrium(): + """Once settled on the discrete equilibrium, a long run does not move it.""" + m = 3.0 + collisions = make_collisions("super_gaussian", m=m) + f0 = jnp.asarray(supergaussian(m, T=1.0))[None, :] + # Settle onto the self-consistent discrete equilibrium first + f_eq = step_n(collisions, f0, nsteps=100) + # Then run 10x longer: a secular drift would accumulate, a fixed point won't + f_end = step_n(collisions, f_eq, nsteps=1000) + + rel_l2 = float(jnp.linalg.norm(f_end - f_eq) / jnp.linalg.norm(f_eq)) + assert rel_l2 < 1e-8, f"secular drift at equilibrium: rel L2 change {rel_l2:.2e} over 100 tau" + _, _, T_eq, _ = moments(f_eq[0]) + _, _, T_end, _ = moments(f_end[0]) + assert abs(T_end / T_eq - 1.0) < 1e-8, f"secular T drift: {T_end / T_eq - 1.0:.2e}" + + +def test_dougherty_destroys_supergaussian_control(): + """Control: the standard Dougherty operator relaxes the same IC to Maxwellian.""" + collisions = make_collisions("chang_cooper_dougherty") + f0 = jnp.asarray(supergaussian(3.0, T=1.0))[None, :] + f1 = step_n(collisions, f0, nsteps=100) + _, _, _, k1 = moments(f1[0]) + assert abs(k1 - 3.0) < 0.02, f"Dougherty control did not maxwellianize: kurtosis {k1:.4f}" + + +def test_maxwellian_relaxes_to_supergaussian(): + """A Maxwellian IC relaxes to the super-Gaussian shape; energy error is O(nu·dt). + + Transient energy conservation for m≠2 is limited by operator splitting + (see SuperGaussianDougherty.compute_beta): β is frozen from f^n during the + implicit step, so a full Maxwellian→SG conversion accumulates a one-time + O(nu·dt) temperature offset. Verify the shape conversion and the + first-order convergence of the offset in dt (nu·dt is small in production). + """ + m = 3.0 + T_drift = {} + for dt in (0.1, 0.05): + collisions = make_collisions("super_gaussian", m=m) + f0 = jnp.asarray(maxwellian(T=1.0))[None, :] + f1 = step_n(collisions, f0, nsteps=round(20.0 / dt), dt=dt) # 20 collision times + + n0, _, T0, _ = moments(f0[0]) + n1, _, T1, k1 = moments(f1[0]) + T_drift[dt] = abs(T1 / T0 - 1.0) + + assert abs(n1 / n0 - 1.0) < 1e-12 + assert abs(k1 - sg_kurtosis(m)) < 0.02, f"dt={dt}: kurtosis {k1:.4f} != SG target {sg_kurtosis(m):.4f}" + + # Compare against the analytic super-Gaussian with the realized (n, T) + f_target = n1 * supergaussian(m, T=T1) + rel_l2 = float(np.linalg.norm(np.asarray(f1[0]) - f_target) / np.linalg.norm(f_target)) + assert rel_l2 < 1e-2, f"dt={dt}: did not relax to super-Gaussian: rel L2 {rel_l2:.2e}" + + assert T_drift[0.1] < 2e-2, f"transient energy error too large: {T_drift[0.1]:.2e}" + # O(nu·dt) convergence: halving dt should roughly halve the offset (allow 1.7x) + assert T_drift[0.05] < T_drift[0.1] / 1.7, f"energy error not O(dt): {T_drift[0.1]:.2e} -> {T_drift[0.05]:.2e}" + + +def test_shifted_supergaussian_conserves_momentum(): + """A drifting IC keeps its momentum and relaxes to a shifted super-Gaussian.""" + m = 3.0 + v_shift = 1.5 + collisions = make_collisions("super_gaussian", m=m) + f0 = jnp.asarray(maxwellian(T=0.5, v0=v_shift))[None, :] + f1 = step_n(collisions, f0, nsteps=200) + + _, vbar0, T0, _ = moments(f0[0]) + _, vbar1, T1, k1 = moments(f1[0]) + + assert abs(vbar1 - vbar0) < 5e-5, f"momentum drift: {vbar1 - vbar0:.2e}" + assert abs(T1 / T0 - 1.0) < 2e-2, f"transient energy error too large: {T1 / T0 - 1.0:.2e}" + assert abs(k1 - sg_kurtosis(m)) < 0.02, f"kurtosis {k1:.4f} != SG target {sg_kurtosis(m):.4f}" + + +def test_m2_reduces_to_dougherty(): + """super_gaussian with m=2 matches chang_cooper_dougherty step for step.""" + # sc_steps=0 so the Dougherty beta is exactly 1/(2·T_discrete), matching + # the super-Gaussian energy-conserving closure n/(2·⟨(v-vbar)²⟩) at m=2 + sg = make_collisions("super_gaussian", m=2.0, sc_steps=0) + dough = make_collisions("chang_cooper_dougherty", sc_steps=0) + + # Non-equilibrium two-temperature IC exercises the full operator + f0 = 0.7 * maxwellian(T=0.5) + 0.3 * maxwellian(T=2.0, v0=0.5) + f0 = jnp.asarray(f0)[None, :] + + f_sg = step_n(sg, f0, nsteps=10) + f_dough = step_n(dough, f0, nsteps=10) + + np.testing.assert_allclose(np.asarray(f_sg), np.asarray(f_dough), rtol=1e-10, atol=1e-14) + + +def test_self_consistent_beta_supergaussian_shape(): + """find_self_consistent_beta(m=3) returns β whose discrete SG matches T of f.""" + from adept.driftdiffusion import discrete_temperature, find_self_consistent_beta + + m = 3.0 + f = jnp.asarray(0.6 * maxwellian(T=0.8) + 0.4 * supergaussian(4.0, T=1.5))[None, :] + v = jnp.asarray(GRID.v) + beta = find_self_consistent_beta(f, v, GRID.dv, spherical=False, m=m) + + f_sg = jnp.exp(-beta[:, None] * jnp.abs(v[None, :]) ** m) + T_sg = discrete_temperature(f_sg, v, GRID.dv) + T_f = discrete_temperature(f, v, GRID.dv) + np.testing.assert_allclose(np.asarray(T_sg), np.asarray(T_f), rtol=1e-8)