Skip to content

Fix VFP1D anisotropic collision matrix #348

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
joglekara merged 1 commit into from
Aug 22, 2026
Merged

Fix VFP1D anisotropic collision matrix #348

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
29 changes: 21 additions & 8 deletions adept/vfp1d/fokker_planck.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -439,7 +439,7 @@ def __init__(self, Z: float, nuee_coeff: float, grid, logLam_ratio: float = 1.0,
b_coeffs = np.stack(
[
(-ll - (il + 1)) / denom_plus,
(ll + (il + 2)) / denom_plus,
(-ll + (il + 2)) / denom_plus,
(ll + (il - 1)) / denom_minus,
(ll - il) / denom_minus,
]
Expand Down Expand Up @@ -503,14 +503,18 @@ def get_ee_offdiagonal_contrib(self, t, y: Array, args: dict) -> Array:

return contrib

def get_ee_diagonal_contrib(self, f0: Array) -> Array:
def get_ee_diagonal_contrib(self, f0: Array, il: int = 1) -> Array:
"""
Returns the tridiagonal operator for the electron-electron collision operator.

:param f0: the distribution function (nx, nv)
:param il: spherical-harmonic order

:return: tuple(diagonal, lower diagonal, upper diagonal) of shape (nx, nv), (nx, nv-1), (nx, nv-1)
"""
if not 1 <= il <= self.grid.nl:
raise ValueError(f"il must satisfy 1 <= il <= {self.grid.nl}; got {il}")

i0 = self.calc_ros_i(f0, power=0.0)
jm1 = self.calc_ros_j(f0, power=-1.0)
i2 = self.calc_ros_i(f0, power=2.0)
Expand All @@ -524,10 +528,12 @@ def get_ee_diagonal_contrib(self, f0: Array) -> Array:
diag_d2dv2 = (i2 + jm1) / (3.0 * v) / dv**2.0
upper_d2dv2 = (i2 + jm1) / (3.0 * v) / dv**2.0

diag_angular = -(-i2 + 2 * jm1 + 3 * i0) / (3.0 * v**3.0)
tri_i1 = (-i2 + 2 * jm1 + 3 * i0) / 3.0
angular_eigenvalue = il * (il + 1) / 2.0
diag_angular = -angular_eigenvalue * tri_i1 / v**3.0

lower_ddv = (-i2 + 2 * jm1 + 3 * i0) / (3.0 * v**2.0) / 2 / dv
upper_ddv = (-i2 + 2 * jm1 + 3 * i0) / (3.0 * v**2.0) / 2 / dv
lower_ddv = tri_i1 / v**2.0 / 2 / dv
upper_ddv = tri_i1 / v**2.0 / 2 / dv

# adding spatial differencing coefficients here
# 1 -2 1 for d2dv2
Expand All @@ -536,9 +542,16 @@ def get_ee_diagonal_contrib(self, f0: Array) -> Array:
diag = diag_term1 - 2.0 * diag_d2dv2 + diag_angular
upper = upper_d2dv2 + upper_ddv

diag = diag.at[:, 0].add(lower[:, 0])
# Regular spherical harmonics satisfy f_l(-v) = (-1)^l f_l(v).
# Fold the origin ghost cell into the first diagonal with that parity.
origin_parity = -1.0 if il % 2 else 1.0
diag = diag.at[:, 0].add(origin_parity * lower[:, 0])

return diag, lower[:, :-1], upper[:, 1:]
# Lineax stores A[i + 1, i] in lower_diagonal[i] and A[i, i + 1]
# in upper_diagonal[i]. The finite-difference coefficients above are
# indexed by the matrix row, so the lower diagonal starts at row 1
# while the upper diagonal ends at row nv - 2.
return diag, lower[:, 1:], upper[:, :-1]

def _solve_one_x_tridiag_(self, diag: Array, upper: Array, lower: Array, f10: Array) -> Array:
"""
Expand Down Expand Up @@ -579,7 +592,7 @@ def solve_harmonic(self, Z, ni, f0, flm, dt, il: int, include_ee_offdiag_explici
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)
ee_diag, ee_lower, ee_upper = self.get_ee_diagonal_contrib(f0, il=il)
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:]

Expand Down
110 changes: 110 additions & 0 deletions tests/test_vfp1d/test_flm_collisions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,110 @@
"""Equation-level regression tests for the anisotropic FLM collision operator."""

import lineax as lx
import numpy as np
import pytest
from jax import numpy as jnp

from adept.vfp1d.fokker_planck import FLMCollisions
from adept.vfp1d.grid import Grid


def _make_collisions(nv: int = 32, nl: int = 3) -> FLMCollisions:
grid = Grid(
xmin=0.0,
xmax=1.0,
nx=1,
tmin=0.0,
tmax=0.1,
dt=0.01,
nv=nv,
vmax=6.0,
nl=nl,
)
return FLMCollisions(Z=1.0, nuee_coeff=1.0, grid=grid, full_aniso_ee=True)


def _raw_tridiagonal_coefficients(collisions: FLMCollisions, f0: jnp.ndarray, il: int):
"""Return row-indexed coefficients before packing the two off-diagonals."""
i0 = collisions.calc_ros_i(f0, power=0.0)
jm1 = collisions.calc_ros_j(f0, power=-1.0)
i2 = collisions.calc_ros_i(f0, power=2.0)

v = collisions.grid.v[None, :]
dv = collisions.grid.dv
tri_i1 = (-i2 + 2.0 * jm1 + 3.0 * i0) / 3.0
tri_i2 = (i2 + jm1) / 3.0

diffusion = tri_i2 / v / dv**2
drift = tri_i1 / v**2 / (2.0 * dv)
lower = diffusion - drift
upper = diffusion + drift
diag = 8.0 * jnp.pi * f0 - 2.0 * diffusion - 0.5 * il * (il + 1) * tri_i1 / v**3
return diag, lower, upper


def test_flm_b_coefficients_match_tzoufras_operator():
collisions = _make_collisions(nl=4)
ell = np.arange(1, 5)
ll = ell * (ell + 1) / 2.0
denom_plus = (2 * ell + 1) * (2 * ell + 3)
denom_minus = (2 * ell + 1) * (2 * ell - 1)

np.testing.assert_allclose(collisions.b1[1:], (-ll - (ell + 1)) / denom_plus)
np.testing.assert_allclose(collisions.b2[1:], (-ll + (ell + 2)) / denom_plus)
np.testing.assert_allclose(collisions.b3[1:], (ll + (ell - 1)) / denom_minus)
np.testing.assert_allclose(collisions.b4[1:], (ll - ell) / denom_minus)

# In particular, the f10 J_{-2} coefficient is 2/15, not 4/15.
np.testing.assert_allclose(collisions.b2[1], 2.0 / 15.0)


@pytest.mark.parametrize("il", [1, 2, 3])
def test_ee_tridiagonal_packs_row_coefficients_and_origin_parity(il):
collisions = _make_collisions(nv=24, nl=3)
f0 = jnp.exp(-(collisions.grid.v[None, :] ** 2))
diag, lower, upper = collisions.get_ee_diagonal_contrib(f0, il=il)
raw_diag, raw_lower, raw_upper = _raw_tridiagonal_coefficients(collisions, f0, il)

expected_diag = raw_diag.at[:, 0].add((-1.0 if il % 2 else 1.0) * raw_lower[:, 0])
np.testing.assert_allclose(diag, expected_diag)
np.testing.assert_allclose(lower, raw_lower[:, 1:])
np.testing.assert_allclose(upper, raw_upper[:, :-1])

matrix = np.asarray(
lx.TridiagonalLinearOperator(diagonal=diag[0], lower_diagonal=lower[0], upper_diagonal=upper[0]).as_matrix()
)
np.testing.assert_allclose(np.diag(matrix, k=-1), np.asarray(raw_lower[0, 1:]))
np.testing.assert_allclose(np.diag(matrix, k=1), np.asarray(raw_upper[0, :-1]))


def _drifting_maxwellian_collision_residual(nv: int) -> float:
"""Maximum C_ee[f0, f10] for the infinitesimal-drift Maxwellian mode."""
collisions = _make_collisions(nv=nv, nl=1)
v = collisions.grid.v
dv = collisions.grid.dv
f0 = (jnp.exp(-(v**2)) / jnp.pi**1.5)[None, :]
f10 = 2.0 * v[None, :] * f0

diag, lower, upper = collisions.get_ee_diagonal_contrib(f0, il=1)
operator = lx.TridiagonalLinearOperator(diagonal=diag[0], lower_diagonal=lower[0], upper_diagonal=upper[0])

padded_f0 = jnp.concatenate([f0[:, 1::-1], f0], axis=1)
d2dv2 = 0.5 / v * jnp.gradient(jnp.gradient(padded_f0, dv, axis=1), dv, axis=1)[:, 2:]
ddv = v**-2 * jnp.gradient(padded_f0, dv, axis=1)[:, 2:]
offdiagonal = collisions.get_ee_offdiagonal_contrib(
None,
f10,
{"ddvf0": ddv, "d2dv2f0": d2dv2, "il": 1},
)
residual = operator.mv(f10[0]) + offdiagonal[0]
return float(jnp.max(jnp.abs(residual)))


def test_drifting_maxwellian_collision_residual_converges_at_velocity_origin():
"""The regular l=1 mode must not develop a 1/dv collision singularity."""
coarse = _drifting_maxwellian_collision_residual(64)
fine = _drifting_maxwellian_collision_residual(256)

assert fine < 0.5 * coarse, (coarse, fine)
assert fine < 0.02, fine
Loading