diff --git a/adept/__init__.py b/adept/__init__.py index 3b5effaf..f6281821 100644 --- a/adept/__init__.py +++ b/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 diff --git a/adept/_base_.py b/adept/_base_.py index 2e8bd570..f41e934e 100644 --- a/adept/_base_.py +++ b/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 diff --git a/adept/driftdiffusion.py b/adept/driftdiffusion.py index a2482ac8..c085af50 100644 --- a/adept/driftdiffusion.py +++ b/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. diff --git a/adept/vfp1d/fokker_planck.py b/adept/vfp1d/fokker_planck.py index 910a0975..9170a857 100644 --- a/adept/vfp1d/fokker_planck.py +++ b/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, + ) diff --git a/adept/vfp1d/vector_field.py b/adept/vfp1d/vector_field.py index 613759cd..a3d57b32 100644 --- a/adept/vfp1d/vector_field.py +++ b/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 diff --git a/adept/vfp2d/__init__.py b/adept/vfp2d/__init__.py new file mode 100644 index 00000000..36287fe5 --- /dev/null +++ b/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", +] diff --git a/adept/vfp2d/base.py b/adept/vfp2d/base.py new file mode 100644 index 00000000..726c6865 --- /dev/null +++ b/adept/vfp2d/base.py @@ -0,0 +1,571 @@ +"""ADEPTModule entry point for the arbitrary-harmonic VFP-2D solver.""" + +from __future__ import annotations + +from dataclasses import asdict + +import jax.numpy as jnp +import jax.tree_util as jtu +import numpy as np +import xarray as xr +from diffrax import ODETerm, SaveAt, diffeqsolve + +from adept._base_ import ADEPTModule, Stepper +from adept.normalization import UREG, laser_normalization, normalize +from adept.utils import filter_scalars +from adept.vfp1d.fokker_planck import ( + F0Collisions, + FLMCollisions, + SelfConsistentBetaConfig, + get_model, + get_scheme, + inverse_bremsstrahlung_resonance_ratio, +) +from adept.vfp1d.grid import Grid as CollisionGrid +from adept.vfp1d.helpers import _initialize_distribution_, calc_logLambda, load_profile_on_grid +from adept.vfp2d.collisions import AnisotropicCollisions, CollisionStep +from adept.vfp2d.distributed import create_spatial_sharding +from adept.vfp2d.grid import Grid +from adept.vfp2d.harmonics import ( + HarmonicLayout, + HouLiFilter2D, + TzoufrasVlasov, + complex_to_real, + current, + density, + nernst_velocity, + real_to_complex, + scalar_velocity_moment, + tensor_velocity_moment, +) +from adept.vfp2d.ohm import KineticOhm2D +from adept.vfp2d.plotting import add_reconnection_diagnostics, reconnection_metrics, save_artifacts +from adept.vfp2d.vector_field import ( + KineticOhmStep, + Maxwell2D, + SpectralPoisson2D, + SplitStepVFP2D, + VlasovMaxwell, +) + + +def _profile_1d(profile: dict, axis, norm, reference=None) -> jnp.ndarray: + basis = profile.get("basis", "uniform") + baseline = float(profile.get("baseline", profile.get("value", 1.0))) + if basis == "uniform": + return baseline * jnp.ones_like(axis) + if basis in ("sine", "cosine"): + amplitude = float(profile.get("amplitude", 0.0)) + wavelength = normalize(profile["wavelength"], norm, dim="x") + trig = jnp.sin if basis == "sine" else jnp.cos + return baseline * (1.0 + amplitude * trig(2.0 * jnp.pi * axis / wavelength)) + if basis == "tanh": + center = normalize(profile["center"], norm, dim="x") + width = normalize(profile["width"], norm, dim="x") + rise = normalize(profile["rise"], norm, dim="x") + left, right = center - 0.5 * width, center + 0.5 * width + envelope = 0.5 * (jnp.tanh((axis - left) / rise) - jnp.tanh((axis - right) / rise)) + if profile.get("bump_or_trough", "bump") == "trough": + envelope = 1.0 - envelope + return baseline + float(profile.get("bump_height", 0.0)) * envelope + if basis == "file": + loaded = load_profile_on_grid(profile, axis, norm) + if reference is None: + raise ValueError("A physical reference quantity is required for file profiles") + return jnp.asarray((loaded / reference).to("").magnitude) + raise NotImplementedError(f"Unsupported VFP-2D profile basis: {basis}") + + +def _profile_2d(profile: dict, grid: Grid, norm, reference=None) -> jnp.ndarray: + """Build a separable 2D profile while accepting VFP-1D profile syntax.""" + + if profile.get("basis") == "gaussian_spots": + x_center = normalize(profile.get("x_center", 0.0), norm, dim="x") + x_radius = normalize(profile["x_radius"], norm, dim="x") + y_radius = normalize(profile.get("y_radius", profile["x_radius"]), norm, dim="x") + y_centers = profile.get("y_centers", [profile.get("y_center", 0.0)]) + y_centers = jnp.asarray([normalize(center, norm, dim="x") for center in y_centers]) + x_envelope = jnp.exp(-(((grid.x - x_center) / x_radius) ** 2)) + y_envelope = jnp.sum(jnp.exp(-(((grid.y[:, None] - y_centers[None, :]) / y_radius) ** 2)), axis=1) + return float(profile.get("amplitude", 1.0)) * x_envelope[:, None] * y_envelope[None, :] + + if "x" in profile or "y" in profile: + px = _profile_1d(profile.get("x", {"basis": "uniform", "baseline": 1.0}), grid.x, norm, reference) + py = _profile_1d(profile.get("y", {"basis": "uniform", "baseline": 1.0}), grid.y, norm, reference) + return px[:, None] * py[None, :] + target_axis = profile.get("axis", "x") + if target_axis == "y": + return jnp.broadcast_to(_profile_1d(profile, grid.y, norm, reference)[None, :], (grid.nx, grid.ny)) + return jnp.broadcast_to(_profile_1d(profile, grid.x, norm, reference)[:, None], (grid.nx, grid.ny)) + + +class BaseVFP2D(ADEPTModule): + """2D3P VFP solver with arbitrary packed complex spherical harmonics.""" + + def __init__(self, cfg: dict): + super().__init__(cfg) + self.plasma_norm = laser_normalization( + cfg["units"]["laser_wavelength"], cfg["units"]["reference electron temperature"] + ) + g = cfg["grid"] + l_max = int(g.get("lmax", g.get("nl", 1))) + m_max = int(g.get("mmax", l_max)) + if g.get("vmax_is_normalized", False): + vmax = float(g["vmax"]) + else: + vmax = float(g.get("vmax", 8.0)) * self.plasma_norm.vth_norm() / np.sqrt(2.0) + self.grid = Grid( + xmin=normalize(g["xmin"], self.plasma_norm, dim="x"), + xmax=normalize(g["xmax"], self.plasma_norm, dim="x"), + nx=int(g["nx"]), + ymin=normalize(g["ymin"], self.plasma_norm, dim="x"), + ymax=normalize(g["ymax"], self.plasma_norm, dim="x"), + ny=int(g["ny"]), + vmax=vmax, + nv=int(g["nv"]), + dt=normalize(g["dt"], self.plasma_norm, dim="t"), + l_max=l_max, + m_max=m_max, + ) + self.layout = HarmonicLayout(l_max, m_max) + self.tmin = normalize(g.get("tmin", 0.0), self.plasma_norm, dim="t") + requested_tmax = normalize(g["tmax"], self.plasma_norm, dim="t") + self.nt = int(np.ceil((requested_tmax - self.tmin) / self.grid.dt)) + self.tmax = self.tmin + self.nt * self.grid.dt + self.max_steps = self.nt + 4 + self._density = None + field_cfg = cfg.get("terms", {}).get("field_solver", {}) + self.field_mode = field_cfg if isinstance(field_cfg, str) else field_cfg.get("mode", "maxwell") + self._kinetic_ohm = None + self._maxwell = None + self.spatial_sharding = create_spatial_sharding(g.get("sharding"), self.grid.nx) + + def write_units(self) -> dict: + norm = self.plasma_norm + z = self.cfg["units"]["Z"] + ne = UREG.Quantity(self.cfg["units"]["reference electron density"]).to("1/cc") + log_ei, log_ee = calc_logLambda( + self.cfg, ne, norm.T0.to("eV"), z, self.cfg["units"]["Ion"], force_ee_equal_ei=True + ) + r_e = 2.8179403205e-13 * UREG.cm + nuee_coeff = float( + (4 * jnp.pi * norm.n0 * r_e**2 * UREG.c**4 * log_ee * norm.tau / norm.v0**3).to("").magnitude + ) + lam0 = UREG.Quantity(self.cfg["units"]["laser_wavelength"]).to("um") + ib_cfg = self.cfg.get("drivers", {}).get("ib", {}) + polarisation = ib_cfg.get("polarisation", "linear") + if polarisation == "linear": + alpha_pol = 1.0 + elif polarisation == "circular": + alpha_pol = 0.5 + else: + alpha_pol = float(polarisation) + vosc2_per_intensity = float( + (0.093373 * (lam0 / UREG.um) ** 2 / (alpha_pol * (norm.T0 / UREG.keV))).to("").magnitude + ) + w0_norm = float((2 * np.pi * UREG.c / lam0 * norm.tau).to("")) + derived = { + "n0": norm.n0.to("1/cc"), + "T0": norm.T0.to("eV"), + "x0": norm.L0.to("nm"), + "t0": norm.tau.to("fs"), + "vth_norm": norm.vth_norm(), + "c_norm": norm.speed_of_light_norm(), + "logLambda_ei": log_ei, + "logLambda_ee": log_ee, + "nuee_coeff": nuee_coeff, + "logLam_ratio": log_ei / log_ee, + "vosc2_per_intensity": vosc2_per_intensity, + "w0_norm": w0_norm, + } + self.cfg["units"]["derived"] = derived + return {key: str(value) for key, value in derived.items()} + + def get_derived_quantities(self): + values = filter_scalars(asdict(self.grid)) + values.update({"tmin": self.tmin, "tmax": self.tmax, "nt": self.nt, "max_steps": self.max_steps}) + self.cfg["grid"].update(values) + + def get_solver_quantities(self): + self.cfg["grid"].update(asdict(self.grid)) + self.cfg["grid"].update({"harmonic_pairs": self.layout.pairs}) + + def init_state_and_args(self): + f00 = jnp.zeros((self.grid.nx, self.grid.ny, self.grid.nv)) + n_total = jnp.zeros((self.grid.nx, self.grid.ny)) + found = False + for name, component in self.cfg["density"].items(): + if not name.startswith("species-"): + continue + n_prof = _profile_2d( + component["n"], + self.grid, + self.plasma_norm, + reference=UREG.Quantity(self.cfg["units"]["reference electron density"]), + ) + t_prof = _profile_2d(component["T"], self.grid, self.plasma_norm, reference=self.plasma_norm.T0) + if self.cfg["grid"].get("relativistic", False): + theta0 = float((self.plasma_norm.T0 / (UREG.m_e * UREG.c**2)).to("").magnitude) + theta = theta0 * t_prof[..., None] + gamma = jnp.sqrt(1.0 + self.grid.v**2) + local_f = jnp.exp(-(gamma[None, None, :] - 1.0) / theta) + norm = 4.0 * jnp.pi * jnp.sum(local_f * self.grid.v**2, axis=-1) * self.grid.dv + local_f = n_prof[..., None] * local_f / norm[..., None] + else: + local_f, _ = _initialize_distribution_( + nv=self.grid.nv, + m=float(component.get("m", 2.0)), + vth=self.plasma_norm.vth_norm(), + vmax=self.grid.vmax, + n_prof=n_prof.reshape(-1), + T_prof=t_prof.reshape(-1), + ) + local_f = local_f.reshape((self.grid.nx, self.grid.ny, self.grid.nv)) + f00 = f00 + local_f + n_total = n_total + n_prof + found = True + if not found: + raise ValueError("VFP-2D density must contain at least one 'species-*' component") + + ne_over_n0 = float( + (UREG.Quantity(self.cfg["units"]["reference electron density"]) / self.plasma_norm.n0).to("").magnitude + ) + f00 = f00 * ne_over_n0 + n_total = n_total * ne_over_n0 + flm = ( + jnp.zeros((self.grid.nx, self.grid.ny, self.layout.size, self.grid.nv), dtype=jnp.complex128) + .at[..., self.layout.index(0, 0), :] + .set(f00) + ) + + zref = float(self.cfg["units"]["Z"]) + ion_charge = n_total if self.cfg["density"].get("quasineutrality", True) else jnp.mean(n_total) + charge_density = ion_charge - density(flm, self.layout, self.grid.v, self.grid.dv) + e = SpectralPoisson2D(self.grid.kx, self.grid.ky)(charge_density) + # Diffrax currently warns that complex state support is experimental. + # Keep its PyTree purely real while retaining complex arithmetic inside + # the harmonic operator. + self.state = {"flm": complex_to_real(flm), "e": e, "b": jnp.zeros_like(e)} + self._density = n_total + self.args = {"Z": jnp.ones_like(n_total), "ni": n_total / zref} + drivers = self.cfg.get("drivers", {}) + maxwellian = drivers.get("maxwellian_heating", {}) + if "D0" in maxwellian: + profile = _profile_2d( + maxwellian.get("profile", {"basis": "uniform", "baseline": 1.0}), + self.grid, + self.plasma_norm, + ) + self.args["D0_heating"] = float(maxwellian["D0"]) * profile + + ib = drivers.get("ib", {}) + intensity = float(ib.get("intensity_1e15_Wcm2", 0.0)) + if intensity > 0.0: + profile = _profile_2d( + ib.get("profile", {"basis": "uniform", "baseline": 1.0}), + self.grid, + self.plasma_norm, + ) + self.args["ib_vosc2"] = self.cfg["units"]["derived"]["vosc2_per_intensity"] * intensity * profile + derived = self.cfg["units"]["derived"] + self.args["ib_Z2ni_w0"] = inverse_bremsstrahlung_resonance_ratio( + self.args["Z"], + self.args["ni"], + derived["nuee_coeff"], + derived["logLam_ratio"], + derived["w0_norm"], + ) + for source_key, arg_key in ( + ("switch_on", "ib_t_on"), + ("switch_off", "ib_t_off"), + ("switch_width", "ib_switch_width"), + ): + if source_key in ib: + self.args[arg_key] = normalize(ib[source_key], self.plasma_norm, dim="t") + + field_cfg = self.cfg.get("terms", {}).get("field_solver", {}) + if isinstance(field_cfg, dict) and self.field_mode == "kinetic-ohm": + hidden = field_cfg.get("hidden_density_gradient", {}) + if hidden.get("active", False): + profile = _profile_2d( + hidden.get("profile", {"basis": "uniform", "baseline": 1.0}), + self.grid, + self.plasma_norm, + ) + scale_length = normalize(hidden["scale_length"], self.plasma_norm, dim="x") + reference_density = float( + (UREG.Quantity(self.cfg["units"]["reference electron density"]) / self.plasma_norm.n0) + .to("") + .magnitude + ) + self.args["hidden_dndz"] = reference_density * profile / scale_length + if "switch_off" in hidden: + self.args["hidden_gradient_t_off"] = normalize(hidden["switch_off"], self.plasma_norm, dim="t") + if "switch_width" in hidden: + self.args["hidden_gradient_switch_width"] = normalize( + hidden["switch_width"], self.plasma_norm, dim="t" + ) + + def _collision_step(self) -> CollisionStep | None: + fp = self.cfg.get("terms", {}).get("fokker_planck", {}) + if not fp.get("active", True): + return None + if self.cfg["grid"].get("relativistic", False): + raise NotImplementedError( + "The Tzoufras linearized collision operator is non-relativistic; " + "set grid.relativistic=false when Fokker-Planck collisions are active." + ) + collision_grid = CollisionGrid( + xmin=0.0, + xmax=1.0, + nx=self.grid.nx * self.grid.ny, + tmin=0.0, + tmax=self.grid.dt, + dt=self.grid.dt, + nv=self.grid.nv, + vmax=self.grid.vmax, + nl=self.layout.l_max, + ) + f00_cfg = fp.get("f00", {}) + model = get_model(f00_cfg.get("model", "CoulombianKernel"), collision_grid.v, collision_grid.dv) + scheme = get_scheme(f00_cfg.get("scheme", "central"), collision_grid.dv) + sc = fp.get("self_consistent_beta", {}) + isotropic = F0Collisions( + nuee_coeff=self.cfg["units"]["derived"]["nuee_coeff"], + grid=collision_grid, + model=model, + scheme=scheme, + sc_beta=SelfConsistentBetaConfig( + max_steps=sc.get("max_steps", 3) if sc.get("enabled", False) else 0, + rtol=sc.get("rtol", 1e-8), + atol=sc.get("atol", 1e-12), + ), + ) + flm_operator = FLMCollisions( + Z=float(self.cfg["units"]["Z"]), + nuee_coeff=self.cfg["units"]["derived"]["nuee_coeff"], + grid=collision_grid, + logLam_ratio=self.cfg["units"]["derived"]["logLam_ratio"], + full_aniso_ee=fp.get("flm", {}).get("ee", True), + ) + return CollisionStep( + self.layout, + isotropic, + AnisotropicCollisions(flm_operator, self.layout), + mesh=None if self.spatial_sharding is None else self.spatial_sharding.mesh, + ) + + def init_diffeqsolve(self): + if self.spatial_sharding is not None: + self.state = jtu.tree_map(self.spatial_sharding.put, self.state) + self.args = jtu.tree_map(self.spatial_sharding.put, self.args) + relativistic = bool(self.cfg["grid"].get("relativistic", False)) + streaming_speed = self.grid.v / jnp.sqrt(1.0 + self.grid.v**2) if relativistic else self.grid.v + partitioned_dx = self.grid.dx if self.spatial_sharding is not None else None + partitioned_dy = self.grid.dy if self.spatial_sharding is not None else None + vlasov = TzoufrasVlasov( + self.layout, + self.grid.v, + self.grid.dv, + self.grid.kx, + self.grid.ky, + streaming_speed=streaming_speed, + dx=partitioned_dx, + dy=partitioned_dy, + mesh=None if self.spatial_sharding is None else self.spatial_sharding.mesh, + ) + maxwell = Maxwell2D( + self.grid.kx, + self.grid.ky, + c=self.plasma_norm.speed_of_light_norm(), + dx=partitioned_dx, + dy=partitioned_dy, + mesh=None if self.spatial_sharding is None else self.spatial_sharding.mesh, + ) + self._maxwell = maxwell + collisions = self._collision_step() + if self.field_mode == "maxwell": + rhs = VlasovMaxwell( + vlasov, + maxwell, + self.layout, + self.grid.v, + self.grid.dv, + real_storage=True, + streaming_speed=streaming_speed, + ) + step = SplitStepVFP2D(rhs, self.grid.dt, collisions=collisions) + elif self.field_mode == "kinetic-ohm": + zref = float(self.cfg["units"]["Z"]) + resistivity_coefficient = ( + 0.5 * zref * self.cfg["units"]["derived"]["nuee_coeff"] * self.cfg["units"]["derived"]["logLam_ratio"] + ) + self._kinetic_ohm = KineticOhm2D( + self.layout, + self.grid.v, + self.grid.dv, + self.grid.kx, + self.grid.ky, + resistivity_coefficient=resistivity_coefficient, + dx=partitioned_dx, + dy=partitioned_dy, + mesh=None if self.spatial_sharding is None else self.spatial_sharding.mesh, + ) + filter_cfg = self.cfg.get("terms", {}).get("hou_li_filter", {}) + spatial_filter = None + if filter_cfg.get("is_on", False): + dimensions = set(filter_cfg.get("dimensions", ["x", "y"])) + if not dimensions or not dimensions <= {"x", "y"}: + raise ValueError("VFP-2D Hou-Li filtering dimensions must be a nonempty subset of [x, y]") + spatial_filter = HouLiFilter2D( + self.grid.nx, + self.grid.ny, + alpha=float(filter_cfg.get("alpha", 36.0)), + order=int(filter_cfg.get("order", 36)), + dimensions=tuple(sorted(dimensions)), + mesh=None if self.spatial_sharding is None else self.spatial_sharding.mesh, + ) + step = KineticOhmStep( + vlasov, + maxwell, + self._kinetic_ohm, + self.layout, + self.grid.v, + self.grid.dv, + self.grid.dt, + collisions=collisions, + real_storage=True, + enforce_f00_positivity=( + self.cfg.get("terms", {}).get("fokker_planck", {}).get("f00", {}).get("positivity", "none") + == "conservative" + ), + spatial_filter=spatial_filter, + ) + initial_flm = real_to_complex(self.state["flm"]) + initial_current = maxwell.c2 * maxwell.curl(self.state["b"]) + initial_hidden_dndz = KineticOhmStep._hidden_dndz(self.tmin, self.args, self.state["b"][..., 0]) + initial_e, _terms = self._kinetic_ohm( + initial_flm, + self.state["b"], + plasma_current=initial_current, + hidden_dndz=initial_hidden_dndz, + ) + self.state = {**self.state, "e": initial_e} + else: + raise ValueError( + f"Unsupported VFP-2D field solver mode {self.field_mode!r}; expected 'maxwell' or 'kinetic-ohm'" + ) + save_cfg = self.cfg.get("save", {}).get("t", {}) + save_tmin = normalize(save_cfg.get("tmin", self.tmin), self.plasma_norm, dim="t") + save_tmax = normalize(save_cfg.get("tmax", self.tmax), self.plasma_norm, dim="t") + save_nt = int(save_cfg.get("nt", min(self.nt + 1, 101))) + self.save_times = jnp.linspace(save_tmin, save_tmax, save_nt) + self.time_quantities = {"t0": self.tmin, "t1": self.tmax, "max_steps": self.max_steps} + if self.spatial_sharding is not None: + + def save_fn(_t, state, _args): + return jtu.tree_map(self.spatial_sharding.replicate, state) + + saveat = SaveAt(ts=self.save_times, fn=save_fn) + else: + saveat = SaveAt(ts=self.save_times) + self.diffeqsolve_quants = { + "terms": ODETerm(step), + "solver": Stepper(), + "saveat": saveat, + } + + def __call__(self, trainable_modules: dict | None, args: dict | None): + def solve(): + return diffeqsolve( + terms=self.diffeqsolve_quants["terms"], + solver=self.diffeqsolve_quants["solver"], + t0=self.tmin, + t1=self.tmax, + dt0=self.grid.dt, + max_steps=self.max_steps, + y0=self.state, + args=self.args if args is None else args, + saveat=self.diffeqsolve_quants["saveat"], + ) + + if self.spatial_sharding is not None: + with self.spatial_sharding.mesh: + result = solve() + else: + result = solve() + return {"solver result": result} + + def post_process(self, run_output: dict, td: str) -> dict: + result = run_output["solver result"] + flm_jax = real_to_complex(result.ys["flm"]) + flm = np.asarray(flm_jax) + ne = density(flm_jax, self.layout, self.grid.v, self.grid.dv) + plasma_current = current(flm_jax, self.layout, self.grid.v, self.grid.dv) + mean_v2 = scalar_velocity_moment(flm_jax, self.layout, self.grid.v, self.grid.dv, power=2) + temperature_normalized = (2.0 / 3.0) * mean_v2 / self.plasma_norm.vth_norm() ** 2 + pressure_anisotropy = tensor_velocity_moment(flm_jax, self.layout, self.grid.v, self.grid.dv, power=0) + v_nernst = nernst_velocity(flm_jax, self.layout, self.grid.v, self.grid.dv, plasma_current=plasma_current) + coords = { + "t": np.asarray(result.ts), + "x": np.asarray(self.grid.x), + "y": np.asarray(self.grid.y), + "harmonic": np.arange(self.layout.size), + "v": np.asarray(self.grid.v), + "component": ["x", "y", "z"], + "ell": ("harmonic", self.layout.ell), + "m": ("harmonic", self.layout.m), + } + data_vars = { + "flm_real": (("t", "x", "y", "harmonic", "v"), flm.real), + "flm_imag": (("t", "x", "y", "harmonic", "v"), flm.imag), + "e": (("t", "x", "y", "component"), np.asarray(result.ys["e"])), + "b": (("t", "x", "y", "component"), np.asarray(result.ys["b"])), + "ne": (("t", "x", "y"), np.asarray(ne)), + "temperature": (("t", "x", "y"), np.asarray(temperature_normalized)), + "current": (("t", "x", "y", "component"), np.asarray(plasma_current)), + "v_nernst": (("t", "x", "y", "component"), np.asarray(v_nernst)), + "pressure_anisotropy": ( + ("t", "x", "y", "component", "component_2"), + np.asarray(pressure_anisotropy), + ), + } + if self._kinetic_ohm is not None and self._maxwell is not None: + ohm_history = {key: [] for key in ("resistive", "hall", "nernst", "scalar_pressure", "tensor_pressure")} + for index, time in enumerate(np.asarray(result.ts)): + frame_flm = flm_jax[index] + frame_b = result.ys["b"][index] + if self.spatial_sharding is not None: + frame_flm = self.spatial_sharding.put(frame_flm) + frame_b = self.spatial_sharding.put(frame_b) + target_current = self._maxwell.c2 * self._maxwell.curl(frame_b) + hidden_dndz = KineticOhmStep._hidden_dndz(float(time), self.args, frame_b[..., 0]) + _electric, terms = self._kinetic_ohm( + frame_flm, + frame_b, + plasma_current=target_current, + hidden_dndz=hidden_dndz, + ) + for key, value in terms.items(): + ohm_history[key].append(value) + for key, values in ohm_history.items(): + data_vars[f"ohm_{key}"] = ( + ("t", "x", "y", "component"), + np.asarray(jnp.stack(values)), + ) + + ds = xr.Dataset( + data_vars, + coords={**coords, "component_2": ["x", "y", "z"]}, + attrs={ + "solver": "vfp-2d", + "harmonic_convention": "Tzoufras JCP 230 (2011)", + "length_unit_um": float(self.plasma_norm.L0.to("um").magnitude), + "time_unit_ps": float(self.plasma_norm.tau.to("ps").magnitude), + }, + ) + ds = add_reconnection_diagnostics(ds) + if td: + n_panels = int(self.cfg.get("output", {}).get("n_panels", 9)) + save_artifacts(ds, td, n_panels=n_panels) + return {"vfp2d": ds, "metrics": reconnection_metrics(ds)} diff --git a/adept/vfp2d/collisions.py b/adept/vfp2d/collisions.py new file mode 100644 index 00000000..929e9642 --- /dev/null +++ b/adept/vfp2d/collisions.py @@ -0,0 +1,143 @@ +"""Collision adapters for packed arbitrary-``f_lm`` VFP-2D states.""" + +from __future__ import annotations + +import jax.numpy as jnp +from jax import Array, shard_map +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P + +from adept.vfp1d.fokker_planck import F0Collisions, FLMCollisions +from adept.vfp2d.harmonics import HarmonicLayout + + +class AnisotropicCollisions: + """Apply the Tzoufras linearized anisotropic operator to every packed mode. + + The radial operator depends on ``l`` but is diagonal in ``m``. Spatial axes + are flattened into a single batch for the existing JAX/Lineax tridiagonal + solve, then restored. The isotropic ``f00`` mode is intentionally unchanged; + it is advanced by the conservative isotropic collision solver. + """ + + def __init__(self, operator: FLMCollisions, layout: HarmonicLayout): + if operator.grid.nl < layout.l_max: + raise ValueError("FLMCollisions grid.nl must be at least layout.l_max") + self.operator = operator + self.layout = layout + + @staticmethod + def _spatial_field(value: Array | float, shape: tuple[int, ...], dtype) -> Array: + return jnp.broadcast_to(jnp.asarray(value, dtype=dtype), shape).reshape(-1) + + def __call__(self, flm: Array, Z: Array | float, ni: Array | float, dt: float) -> Array: + spatial_shape = flm.shape[:-2] + nv = flm.shape[-1] + f0 = jnp.real(flm[..., self.layout.index(0, 0), :]).reshape((-1, nv)) + flat_Z = self._spatial_field(Z, spatial_shape, f0.dtype) + flat_ni = self._spatial_field(ni, spatial_shape, f0.dtype) + result = flm + + for i, (ell, _m) in enumerate(self.layout.pairs): + if ell == 0: + continue + mode = flm[..., i, :].reshape((-1, nv)) + updated = self.operator.solve_harmonic(flat_Z, flat_ni, f0, mode, dt, il=ell) + result = result.at[..., i, :].set(updated.reshape((*spatial_shape, nv))) + return result + + +class CollisionStep: + """Conservative ``f00`` plus arbitrary-``f_lm`` implicit collision step.""" + + def __init__( + self, + layout: HarmonicLayout, + isotropic: F0Collisions | None = None, + anisotropic: AnisotropicCollisions | None = None, + mesh: Mesh | None = None, + ): + self.layout = layout + self.isotropic = isotropic + self.anisotropic = anisotropic + self.mesh = mesh + + def _local_step( + self, + flm: Array, + Z: Array | float, + ni: Array | float, + dt: float, + *, + D0_heating: Array | float | None = None, + ib_vosc2: Array | float | None = None, + ib_Z2ni_w0: Array | float | None = None, + ) -> Array: + spatial_shape = flm.shape[:-2] + nv = flm.shape[-1] + result = flm + if self.isotropic is not None: + f00 = jnp.real(result[..., self.layout.index(0, 0), :]).reshape((-1, nv)) + heating = {} + if D0_heating is not None: + heating["D0_heating"] = jnp.broadcast_to(jnp.asarray(D0_heating), spatial_shape).reshape(-1) + if ib_vosc2 is not None: + if ib_Z2ni_w0 is None: + raise ValueError("ib_Z2ni_w0 is required when inverse-bremsstrahlung heating is enabled") + heating["ib_vosc2"] = jnp.broadcast_to(jnp.asarray(ib_vosc2), spatial_shape).reshape(-1) + heating["ib_Z2ni_w0"] = jnp.broadcast_to(jnp.asarray(ib_Z2ni_w0), spatial_shape).reshape(-1) + f00 = self.isotropic(None, f00, dt, **heating).reshape((*spatial_shape, nv)) + result = result.at[..., self.layout.index(0, 0), :].set(f00) + if self.anisotropic is not None: + result = self.anisotropic(result, Z=Z, ni=ni, dt=dt) + return result + + def __call__( + self, + flm: Array, + Z: Array | float, + ni: Array | float, + dt: float, + *, + D0_heating: Array | float | None = None, + ib_vosc2: Array | float | None = None, + ib_Z2ni_w0: Array | float | None = None, + ) -> Array: + if self.mesh is None: + return self._local_step( + flm, + Z, + ni, + dt, + D0_heating=D0_heating, + ib_vosc2=ib_vosc2, + ib_Z2ni_w0=ib_Z2ni_w0, + ) + + spatial_shape = flm.shape[:-2] + template = jnp.broadcast_to(jnp.asarray(ni), spatial_shape) + z_field = jnp.broadcast_to(jnp.asarray(Z), spatial_shape) + d0_field = jnp.zeros_like(template) if D0_heating is None else jnp.broadcast_to(D0_heating, spatial_shape) + ib_field = jnp.zeros_like(template) if ib_vosc2 is None else jnp.broadcast_to(ib_vosc2, spatial_shape) + ratio_field = jnp.zeros_like(template) if ib_Z2ni_w0 is None else jnp.broadcast_to(ib_Z2ni_w0, spatial_shape) + + def _mapped(local_flm, local_z, local_ni, local_d0, local_ib, local_ratio, local_dt): + return self._local_step( + local_flm, + local_z, + local_ni, + local_dt, + D0_heating=local_d0 if D0_heating is not None else None, + ib_vosc2=local_ib if ib_vosc2 is not None else None, + ib_Z2ni_w0=local_ratio if ib_Z2ni_w0 is not None else None, + ) + + spatial_spec = P("x", None) + flm_spec = P("x", None, None, None) + return shard_map( + _mapped, + mesh=self.mesh, + in_specs=(flm_spec, spatial_spec, spatial_spec, spatial_spec, spatial_spec, spatial_spec, P()), + out_specs=flm_spec, + check_vma=False, + )(flm, z_field, template, d0_field, ib_field, ratio_field, dt) diff --git a/adept/vfp2d/distributed.py b/adept/vfp2d/distributed.py new file mode 100644 index 00000000..494e5e5f --- /dev/null +++ b/adept/vfp2d/distributed.py @@ -0,0 +1,60 @@ +"""Named x-axis sharding for the VFP-2D spatial state.""" + +from __future__ import annotations + +from collections.abc import Sequence +from dataclasses import dataclass +from typing import Any + +import jax +import numpy as np +from jax.sharding import Mesh, NamedSharding +from jax.sharding import PartitionSpec as P + + +@dataclass(frozen=True) +class SpatialSharding: + """Shard arrays whose leading dimension is the VFP-2D x axis.""" + + mesh: Mesh + nx: int + + def for_array(self, value: Any) -> NamedSharding: + ndim = np.ndim(value) + partition = P("x", *(None for _ in range(ndim - 1))) if ndim else P() + return NamedSharding(self.mesh, partition) + + def put(self, value: Any): + array = jax.numpy.asarray(value) + if array.ndim and array.shape[0] == self.nx: + return jax.device_put(array, self.for_array(array)) + return jax.device_put(array, NamedSharding(self.mesh, P(*(None for _ in range(array.ndim))))) + + def replicate(self, value: Any): + """Replicate a saved snapshot so Diffrax can assemble its leading time axis.""" + + array = jax.numpy.asarray(value) + replicated = NamedSharding(self.mesh, P(*(None for _ in range(array.ndim)))) + return jax.device_put(array, replicated) + + +def create_spatial_sharding( + raw_cfg: Any, + nx: int, + devices: Sequence[jax.Device] | None = None, +) -> SpatialSharding | None: + """Create an all-visible-device x mesh when ``grid.sharding.enabled`` is true.""" + + cfg = dict(raw_cfg or {}) + if not cfg.get("enabled", False): + return None + if cfg.get("axis", "x") != "x": + raise ValueError("VFP-2D currently supports sharding only along the x axis") + + devices = tuple(jax.devices() if devices is None else devices) + if not devices: + raise ValueError("No JAX devices are available for VFP-2D sharding") + if nx % len(devices): + raise ValueError(f"VFP-2D nx={nx} must be divisible by the {len(devices)} visible devices") + mesh = jax.make_mesh((len(devices),), ("x",), devices=devices) + return SpatialSharding(mesh=mesh, nx=nx) diff --git a/adept/vfp2d/grid.py b/adept/vfp2d/grid.py new file mode 100644 index 00000000..5134f919 --- /dev/null +++ b/adept/vfp2d/grid.py @@ -0,0 +1,79 @@ +"""Periodic configuration-space and radial-momentum grid for VFP-2D.""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax.numpy as jnp +import numpy as np +from jax import Array + + +@dataclass(frozen=True) +class Grid: + xmin: float + xmax: float + nx: int + ymin: float + ymax: float + ny: int + vmax: float + nv: int + dt: float + l_max: int + m_max: int + dx: float + dy: float + dv: float + x: Array + y: Array + v: Array + kx: Array + ky: Array + + def __init__( + self, + *, + xmin: float, + xmax: float, + nx: int, + ymin: float, + ymax: float, + ny: int, + vmax: float, + nv: int, + dt: float, + l_max: int, + m_max: int | None = None, + ): + if nx < 1 or ny < 1 or nv < 2: + raise ValueError("nx and ny must be positive and nv must be at least 2") + if xmax <= xmin or ymax <= ymin or vmax <= 0 or dt <= 0: + raise ValueError("grid extents, vmax, and dt must be positive") + if m_max is None: + m_max = l_max + if not 0 <= m_max <= l_max: + raise ValueError("m_max must satisfy 0 <= m_max <= l_max") + + dx = (xmax - xmin) / nx + dy = (ymax - ymin) / ny + dv = vmax / nv + object.__setattr__(self, "xmin", xmin) + object.__setattr__(self, "xmax", xmax) + object.__setattr__(self, "nx", nx) + object.__setattr__(self, "ymin", ymin) + object.__setattr__(self, "ymax", ymax) + object.__setattr__(self, "ny", ny) + object.__setattr__(self, "vmax", vmax) + object.__setattr__(self, "nv", nv) + object.__setattr__(self, "dt", dt) + object.__setattr__(self, "l_max", l_max) + object.__setattr__(self, "m_max", m_max) + object.__setattr__(self, "dx", dx) + object.__setattr__(self, "dy", dy) + object.__setattr__(self, "dv", dv) + object.__setattr__(self, "x", jnp.linspace(xmin + dx / 2, xmax - dx / 2, nx)) + object.__setattr__(self, "y", jnp.linspace(ymin + dy / 2, ymax - dy / 2, ny)) + object.__setattr__(self, "v", jnp.linspace(dv / 2, vmax - dv / 2, nv)) + object.__setattr__(self, "kx", jnp.fft.fftfreq(nx, d=dx) * 2.0 * np.pi) + object.__setattr__(self, "ky", jnp.fft.fftfreq(ny, d=dy) * 2.0 * np.pi) diff --git a/adept/vfp2d/harmonics.py b/adept/vfp2d/harmonics.py new file mode 100644 index 00000000..7af5f548 --- /dev/null +++ b/adept/vfp2d/harmonics.py @@ -0,0 +1,560 @@ +"""Spherical-harmonic layout and Vlasov operators for VFP-2D. + +The normalization and complex ``f[l, m]`` convention follow Tzoufras et al., +J. Comput. Phys. 230 (2011), equations (5)--(24). Only non-negative ``m`` +are stored because a real distribution satisfies ``f[l, -m] = conj(f[l,m])``. + +Arrays use the compact layout ``(..., harmonic, speed)``. This is much more +amenable to JAX transformations than a nested ``dict[l][m]`` PyTree and allows +``m_max`` to be chosen independently of ``l_max``. +""" + +from __future__ import annotations + +from dataclasses import dataclass + +import jax.numpy as jnp +import numpy as np +from jax import Array, lax, shard_map +from jax.sharding import Mesh +from jax.sharding import PartitionSpec as P + + +@dataclass(frozen=True) +class HarmonicLayout: + """Packed indexing for ``0 <= l <= l_max`` and ``0 <= m <= min(l,m_max)``.""" + + l_max: int + m_max: int + pairs: tuple[tuple[int, int], ...] + _indices: tuple[tuple[int, ...], ...] + + def __init__(self, l_max: int, m_max: int | None = None): + if l_max < 0: + raise ValueError("l_max must be non-negative") + if m_max is None: + m_max = l_max + if not 0 <= m_max <= l_max: + raise ValueError("m_max must satisfy 0 <= m_max <= l_max") + + pairs = tuple((ell, m) for ell in range(l_max + 1) for m in range(min(ell, m_max) + 1)) + lookup = {pair: i for i, pair in enumerate(pairs)} + rows = tuple(tuple(lookup.get((ell, m), -1) for m in range(m_max + 1)) for ell in range(l_max + 1)) + object.__setattr__(self, "l_max", l_max) + object.__setattr__(self, "m_max", m_max) + object.__setattr__(self, "pairs", pairs) + object.__setattr__(self, "_indices", rows) + + @property + def size(self) -> int: + return len(self.pairs) + + @property + def ell(self) -> np.ndarray: + return np.asarray([ell for ell, _ in self.pairs], dtype=np.int32) + + @property + def m(self) -> np.ndarray: + return np.asarray([m for _, m in self.pairs], dtype=np.int32) + + def index(self, ell: int, m: int) -> int: + """Return the packed index, or ``-1`` when the mode is truncated.""" + + if ell < 0 or ell > self.l_max or m < 0 or m > self.m_max: + return -1 + return self._indices[ell][m] + + +def _spectral_derivative(a: Array, k: Array, axis: int) -> Array: + """Periodic spectral derivative, preserving complex-valued harmonics.""" + + shape = [1] * a.ndim + shape[axis] = k.size + multiplier = 1j * k.reshape(shape) + return jnp.fft.ifft(multiplier * jnp.fft.fft(a, axis=axis), axis=axis) + + +def periodic_central_derivative( + a: Array, + spacing: float, + axis: int, + mesh: Mesh | None = None, +) -> Array: + """Periodic finite-difference derivative suitable for a partitioned axis.""" + + if mesh is not None: + if axis != 0: + raise ValueError("The partitioned finite-difference axis must be leading axis 0") + mesh_size = int(mesh.shape["x"]) + local_size = a.shape[0] // mesh_size + halo = 2 if local_size >= 2 and a.shape[0] >= 5 else 1 + permutation_left = tuple((rank, (rank + 1) % mesh_size) for rank in range(mesh_size)) + permutation_right = tuple((rank, (rank - 1) % mesh_size) for rank in range(mesh_size)) + + def _local(local: Array) -> Array: + left = lax.ppermute(local[-halo:], "x", permutation_left) + right = lax.ppermute(local[:halo], "x", permutation_right) + padded = jnp.concatenate((left, local, right), axis=0) + if halo == 1: + return (padded[2:] - padded[:-2]) / (2.0 * spacing) + nlocal = local.shape[0] + return ( + padded[:nlocal] - 8.0 * padded[1 : nlocal + 1] + 8.0 * padded[3 : nlocal + 3] - padded[4 : nlocal + 4] + ) / (12.0 * spacing) + + spec = P("x", *(None for _ in range(a.ndim - 1))) + return shard_map(_local, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False)(a) + + if a.shape[axis] < 5: + return (jnp.roll(a, -1, axis=axis) - jnp.roll(a, 1, axis=axis)) / (2.0 * spacing) + return ( + jnp.roll(a, 2, axis=axis) + - 8.0 * jnp.roll(a, 1, axis=axis) + + 8.0 * jnp.roll(a, -1, axis=axis) + - jnp.roll(a, -2, axis=axis) + ) / (12.0 * spacing) + + +def partitioned_high_order_filter(a: Array, mesh: Mesh) -> Array: + """Remove the x-Nyquist mode with an eighth-difference shard-local filter.""" + + mesh_size = int(mesh.shape["x"]) + local_size = a.shape[0] // mesh_size + halo = 4 if local_size >= 4 else 1 + permutation_left = tuple((rank, (rank + 1) % mesh_size) for rank in range(mesh_size)) + permutation_right = tuple((rank, (rank - 1) % mesh_size) for rank in range(mesh_size)) + + def _local(local: Array) -> Array: + left = lax.ppermute(local[-halo:], "x", permutation_left) + right = lax.ppermute(local[:halo], "x", permutation_right) + padded = jnp.concatenate((left, local, right), axis=0) + if halo == 1: + return 0.25 * padded[:-2] + 0.5 * padded[1:-1] + 0.25 * padded[2:] + nlocal = local.shape[0] + coefficients = (1.0, -8.0, 28.0, -56.0, 70.0, -56.0, 28.0, -8.0, 1.0) + high_pass = sum( + coefficient * padded[offset : offset + nlocal] for offset, coefficient in enumerate(coefficients) + ) + return local - high_pass / 256.0 + + spec = P("x", *(None for _ in range(a.ndim - 1))) + return shard_map(_local, mesh=mesh, in_specs=spec, out_specs=spec, check_vma=False)(a) + + +class TzoufrasVlasov: + """Arbitrary-``f_lm`` 2D3P Vlasov operator from Tzoufras (2011). + + Configuration space is periodic in ``x`` and ``y``. Momentum space is + represented by a positive, cell-centred speed grid and a truncated complex + spherical-harmonic expansion. The third configuration-space derivative is + zero, while all three components of ``E`` and ``B`` are retained. + """ + + def __init__( + self, + layout: HarmonicLayout, + v: Array, + dv: float, + kx: Array, + ky: Array, + streaming_speed: Array | None = None, + dx: float | None = None, + dy: float | None = None, + mesh: Mesh | None = None, + ): + self.layout = layout + self.v = jnp.asarray(v) + self.dv = float(dv) + self.kx = jnp.asarray(kx) + self.ky = jnp.asarray(ky) + self.streaming_speed = self.v if streaming_speed is None else jnp.asarray(streaming_speed) + self.dx = None if dx is None else float(dx) + self.dy = None if dy is None else float(dy) + self.mesh = mesh + + def ddv(self, f: Array, ell: int) -> Array: + """Centred radial derivative with the regularity parity ``f_l(-v)=(-1)^l f_l(v)``.""" + + left = ((-1) ** ell) * f[..., :1] + right = jnp.zeros_like(f[..., :1]) + padded = jnp.concatenate((left, f, right), axis=-1) + return (padded[..., 2:] - padded[..., :-2]) / (2.0 * self.dv) + + def gh(self, f: Array) -> tuple[Array, Array]: + """Return the radial operators ``G_l`` and ``H_l`` from Eqs. (20)--(22).""" + + g = jnp.zeros_like(f) + h = jnp.zeros_like(f) + inv_v = 1.0 / self.v + for i, (ell, _m) in enumerate(self.layout.pairs): + derivative = self.ddv(f[..., i, :], ell) + g = g.at[..., i, :].set(derivative - ell * inv_v * f[..., i, :]) + h = h.at[..., i, :].set(derivative + (ell + 1) * inv_v * f[..., i, :]) + return g, h + + def streaming(self, f: Array, dfdz: Array | None = None) -> Array: + """Spatial-advection contribution from Tzoufras Eqs. (17)--(19). + + Configuration space is evolved in x and y. ``dfdz`` optionally supplies + a prescribed derivative in the unresolved direction, which is useful + for an integrable 2.5D density gradient without allocating a z grid. + """ + + dfdx = ( + _spectral_derivative(f, self.kx, axis=0) + if self.dx is None + else periodic_central_derivative(f, self.dx, axis=0, mesh=self.mesh) + ) + dfdy = ( + _spectral_derivative(f, self.ky, axis=1) + if self.dy is None + else periodic_central_derivative(f, self.dy, axis=1) + ) + if dfdz is None: + dfdz = jnp.zeros_like(f) + transverse_minus = dfdy - 1j * dfdz + transverse_plus = dfdy + 1j * dfdz + out = jnp.zeros_like(f) + v = self.streaming_speed + + for target, (ell, m) in enumerate(self.layout.pairs): + value = jnp.zeros_like(f[..., target, :]) + lower = self.layout.index(ell - 1, m) + upper = self.layout.index(ell + 1, m) + if lower >= 0: + value -= v * (ell - m) / (2 * ell - 1) * dfdx[..., lower, :] + if upper >= 0: + value -= v * (ell + m + 1) / (2 * ell + 3) * dfdx[..., upper, :] + + if m > 0: + lm = self.layout.index(ell - 1, m - 1) + lp = self.layout.index(ell - 1, m + 1) + um = self.layout.index(ell + 1, m - 1) + up = self.layout.index(ell + 1, m + 1) + if lm >= 0: + value -= 0.5 * v / (2 * ell - 1) * transverse_minus[..., lm, :] + if lp >= 0: + value += 0.5 * v * (ell - m) * (ell - m - 1) / (2 * ell - 1) * transverse_plus[..., lp, :] + if um >= 0: + value += 0.5 * v / (2 * ell + 3) * transverse_minus[..., um, :] + if up >= 0: + value -= 0.5 * v * (ell + m + 1) * (ell + m + 2) / (2 * ell + 3) * transverse_plus[..., up, :] + else: + lower1 = self.layout.index(ell - 1, 1) + upper1 = self.layout.index(ell + 1, 1) + transverse = jnp.zeros_like(value) + if lower1 >= 0: + transverse -= ell * (ell - 1) / (2 * ell - 1) * transverse_plus[..., lower1, :] + if upper1 >= 0: + transverse += (ell + 1) * (ell + 2) / (2 * ell + 3) * transverse_plus[..., upper1, :] + value -= v * jnp.real(transverse) + + out = out.at[..., target, :].set(value) + return out + + def electric(self, f: Array, electric_field: Array) -> Array: + """Electric-force contribution, Eqs. (20)--(22).""" + + g, h = self.gh(f) + ex = electric_field[..., 0, None] + ey_minus_iez = (electric_field[..., 1] - 1j * electric_field[..., 2])[..., None] + ey_plus_iez = (electric_field[..., 1] + 1j * electric_field[..., 2])[..., None] + out = jnp.zeros_like(f) + + for target, (ell, m) in enumerate(self.layout.pairs): + value = jnp.zeros_like(f[..., target, :]) + lower = self.layout.index(ell - 1, m) + upper = self.layout.index(ell + 1, m) + if lower >= 0: + value += ex * (ell - m) / (2 * ell - 1) * g[..., lower, :] + if upper >= 0: + value += ex * (ell + m + 1) / (2 * ell + 3) * h[..., upper, :] + + if m > 0: + lm = self.layout.index(ell - 1, m - 1) + lp = self.layout.index(ell - 1, m + 1) + um = self.layout.index(ell + 1, m - 1) + up = self.layout.index(ell + 1, m + 1) + if lm >= 0: + value += 0.5 * ey_minus_iez / (2 * ell - 1) * g[..., lm, :] + if lp >= 0: + value -= 0.5 * ey_plus_iez * (ell - m) * (ell - m - 1) / (2 * ell - 1) * g[..., lp, :] + if um >= 0: + value -= 0.5 * ey_minus_iez / (2 * ell + 3) * h[..., um, :] + if up >= 0: + value += 0.5 * ey_plus_iez * (ell + m + 1) * (ell + m + 2) / (2 * ell + 3) * h[..., up, :] + else: + lower1 = self.layout.index(ell - 1, 1) + upper1 = self.layout.index(ell + 1, 1) + transverse = jnp.zeros_like(value) + if lower1 >= 0: + transverse -= ell * (ell - 1) / (2 * ell - 1) * g[..., lower1, :] + if upper1 >= 0: + transverse += (ell + 1) * (ell + 2) / (2 * ell + 3) * h[..., upper1, :] + value += jnp.real(ey_plus_iez * transverse) + + out = out.at[..., target, :].set(value) + return out + + def magnetic(self, f: Array, magnetic_field: Array) -> Array: + """Magnetic-rotation contribution, Eqs. (23)--(24).""" + + bx = magnetic_field[..., 0, None] + bz_minus_iby = (magnetic_field[..., 2] - 1j * magnetic_field[..., 1])[..., None] + bz_plus_iby = (magnetic_field[..., 2] + 1j * magnetic_field[..., 1])[..., None] + out = jnp.zeros_like(f) + + for target, (ell, m) in enumerate(self.layout.pairs): + value = jnp.zeros_like(f[..., target, :]) + if m > 0: + value -= 1j * bx * m * f[..., target, :] + plus = self.layout.index(ell, m + 1) + minus = self.layout.index(ell, m - 1) + if plus >= 0: + value += 0.5 * (ell - m) * (ell + m + 1) * bz_minus_iby * f[..., plus, :] + if minus >= 0: + value -= 0.5 * bz_plus_iby * f[..., minus, :] + elif ell > 0: + one = self.layout.index(ell, 1) + if one >= 0: + value += ell * (ell + 1) * jnp.real(bz_minus_iby * f[..., one, :]) + out = out.at[..., target, :].set(value) + return out + + def __call__( + self, + f: Array, + electric_field: Array, + magnetic_field: Array, + dfdz: Array | None = None, + ) -> Array: + result = self.streaming(f, dfdz=dfdz) + self.electric(f, electric_field) + self.magnetic(f, magnetic_field) + # m=0 coefficients represent real surface harmonics. Project away + # roundoff-level imaginary parts so the invariant is explicit. + for i, (_ell, m) in enumerate(self.layout.pairs): + if m == 0: + result = result.at[..., i, :].set(jnp.real(result[..., i, :])) + return result + + +def density(f: Array, layout: HarmonicLayout, v: Array, dv: float) -> Array: + """Electron number density, Eq. (10) with ``g=1``.""" + + f00 = f[..., layout.index(0, 0), :] + return 4.0 * jnp.pi * jnp.sum(jnp.real(f00) * v**2, axis=-1) * dv + + +def conservative_f00_positivity(f: Array, layout: HarmonicLayout, v: Array, dv: float) -> Array: + """Clip negative ``f00`` cells while preserving each spatial density. + + Explicit harmonic transport is not positivity preserving. Small high-speed + undershoots can therefore grow until the Coulomb collision coefficients are + undefined. The angular average must be non-negative physically; after + clipping, rescale its positive part so the ``4 pi integral(f00 v^2 dv)`` + density is unchanged wherever the pre-projection density is positive. + """ + + i00 = layout.index(0, 0) + f00 = jnp.real(f[..., i00, :]) + weights = 4.0 * jnp.pi * jnp.asarray(v) ** 2 * float(dv) + target_density = jnp.sum(f00 * weights, axis=-1, keepdims=True) + positive = jnp.maximum(f00, 0.0) + positive_density = jnp.sum(positive * weights, axis=-1, keepdims=True) + tiny = jnp.finfo(f00.dtype).tiny + scale = jnp.where( + (target_density > 0.0) & (positive_density > tiny), + target_density / jnp.maximum(positive_density, tiny), + 0.0, + ) + return f.at[..., i00, :].set((positive * scale).astype(f.dtype)) + + +class HouLiFilter2D: + """Separable Hou--Li filter on the two configuration-space axes. + + The VFP state is complex and periodic in configuration space, so use full + FFTs rather than the real transforms used by the Cartesian Vlasov solvers. + Velocity and harmonic axes are deliberately untouched. + """ + + def __init__( + self, + nx: int, + ny: int, + alpha: float = 36.0, + order: int = 36, + dimensions: tuple[str, ...] = ("x", "y"), + mesh: Mesh | None = None, + ): + if alpha <= 0.0: + raise ValueError("Hou-Li filter alpha must be positive") + if order < 1: + raise ValueError("Hou-Li filter order must be at least one") + + def _kernel(n: int) -> Array: + # abs(2*fftfreq) runs from zero to one in standard FFT ordering. + eta = jnp.abs(2.0 * jnp.fft.fftfreq(n)) + return jnp.exp(-float(alpha) * eta ** (2 * int(order))) + + unknown = set(dimensions) - {"x", "y"} + if unknown: + raise ValueError(f"Unknown Hou-Li filter dimensions: {sorted(unknown)}") + self.filter_x = _kernel(nx) if "x" in dimensions else None + self.filter_y = _kernel(ny) if "y" in dimensions else None + self.mesh = mesh + + @staticmethod + def _apply(a: Array, kernel: Array, axis: int) -> Array: + shape = [1] * a.ndim + shape[axis] = kernel.size + return jnp.fft.ifft(jnp.fft.fft(a, axis=axis) * kernel.reshape(shape), axis=axis) + + def __call__(self, a: Array) -> Array: + """Filter an array whose leading axes are ``(x, y)``.""" + + if self.filter_x is not None: + if self.mesh is None: + a = self._apply(a, self.filter_x, axis=0) + else: + a = partitioned_high_order_filter(a, self.mesh) + if self.filter_y is not None: + if self.mesh is None: + a = self._apply(a, self.filter_y, axis=1) + else: + spec = P("x", *(None for _ in range(a.ndim - 1))) + a = shard_map( + lambda local: self._apply(local, self.filter_y, axis=1), + mesh=self.mesh, + in_specs=spec, + out_specs=spec, + check_vma=False, + )(a) + return a + + +def current( + f: Array, + layout: HarmonicLayout, + v: Array, + dv: float, + charge: float = -1.0, + streaming_speed: Array | None = None, +) -> Array: + """Current vector from the ``l=1`` modes, using Eq. (11).""" + + shape = f.shape[:-2] + i10 = layout.index(1, 0) + i11 = layout.index(1, 1) + if i10 < 0: + return jnp.zeros((*shape, 3), dtype=jnp.real(f).dtype) + speed = v if streaming_speed is None else streaming_speed + weight = v**2 * speed + m10 = jnp.sum(jnp.real(f[..., i10, :]) * weight, axis=-1) * dv + if i11 >= 0: + m11 = jnp.sum(f[..., i11, :] * weight, axis=-1) * dv + else: + m11 = jnp.zeros_like(m10, dtype=f.dtype) + velocity_moment = (4.0 * jnp.pi / 3.0) * jnp.stack((m10, 2.0 * jnp.real(m11), -2.0 * jnp.imag(m11)), axis=-1) + return charge * velocity_moment + + +def scalar_velocity_moment(f: Array, layout: HarmonicLayout, v: Array, dv: float, power: int) -> Array: + """Return ```` using the Joglekar et al. (2014) convention. + + The moment is normalized by the local electron density, so ``power=0`` + returns one (up to velocity-grid truncation). Keeping these definitions + next to the harmonic convention avoids duplicating delicate angular + normalization factors in diagnostics and Ohm-law closures. + """ + + ne = density(f, layout, v, dv) + f00 = jnp.real(f[..., layout.index(0, 0), :]) + numerator = 4.0 * jnp.pi * jnp.sum(f00 * v ** (power + 2), axis=-1) * dv + return numerator / jnp.maximum(ne, jnp.finfo(ne.dtype).tiny) + + +def vector_velocity_moment(f: Array, layout: HarmonicLayout, v: Array, dv: float, power: int) -> Array: + """Return ```` from the packed ``l=1`` harmonics.""" + + ne = density(f, layout, v, dv) + i10 = layout.index(1, 0) + i11 = layout.index(1, 1) + if i10 < 0: + return jnp.zeros((*f.shape[:-2], 3), dtype=jnp.real(f).dtype) + weight = v ** (power + 3) + m10 = jnp.sum(jnp.real(f[..., i10, :]) * weight, axis=-1) * dv + m11 = jnp.sum(f[..., i11, :] * weight, axis=-1) * dv if i11 >= 0 else jnp.zeros_like(m10, dtype=f.dtype) + numerator = (4.0 * jnp.pi / 3.0) * jnp.stack((m10, 2.0 * jnp.real(m11), -2.0 * jnp.imag(m11)), axis=-1) + return numerator / jnp.maximum(ne[..., None], jnp.finfo(ne.dtype).tiny) + + +def cartesian_l2(f: Array, layout: HarmonicLayout) -> Array: + """Convert packed ``l=2`` coefficients to a symmetric traceless tensor. + + Tzoufras uses the x axis as the polar axis. With that convention the + tensor coefficients are ``Fxx=f20``, ``Fxy=3 Re(f21)``, + ``Fxz=-3 Im(f21)``, ``Fyy=-f20/2+6 Re(f22)``, and + ``Fyz=-6 Im(f22)``; ``Fzz`` follows from tracelessness. + """ + + i20, i21, i22 = (layout.index(2, m) for m in range(3)) + if i20 < 0: + return jnp.zeros((*f.shape[:-2], 3, 3, f.shape[-1]), dtype=jnp.real(f).dtype) + f20 = jnp.real(f[..., i20, :]) + f21 = f[..., i21, :] if i21 >= 0 else jnp.zeros_like(f20, dtype=f.dtype) + f22 = f[..., i22, :] if i22 >= 0 else jnp.zeros_like(f20, dtype=f.dtype) + fxx = f20 + fxy = 3.0 * jnp.real(f21) + fxz = -3.0 * jnp.imag(f21) + fyy = -0.5 * f20 + 6.0 * jnp.real(f22) + fyz = -6.0 * jnp.imag(f22) + fzz = -0.5 * f20 - 6.0 * jnp.real(f22) + return jnp.stack( + ( + jnp.stack((fxx, fxy, fxz), axis=-2), + jnp.stack((fxy, fyy, fyz), axis=-2), + jnp.stack((fxz, fyz, fzz), axis=-2), + ), + axis=-3, + ) + + +def tensor_velocity_moment(f: Array, layout: HarmonicLayout, v: Array, dv: float, power: int) -> Array: + """Return the traceless ```` moment used in kinetic Ohm's law.""" + + ne = density(f, layout, v, dv) + tensor = cartesian_l2(f, layout) + numerator = (8.0 * jnp.pi / 15.0) * jnp.sum(tensor * v ** (power + 4), axis=-1) * dv + return numerator / jnp.maximum(ne[..., None, None], jnp.finfo(ne.dtype).tiny) + + +def nernst_velocity( + f: Array, + layout: HarmonicLayout, + v: Array, + dv: float, + plasma_current: Array | None = None, +) -> Array: + """Return the distribution-function Nernst velocity from PRL Eq. (2).""" + + ne = density(f, layout, v, dv) + if plasma_current is None: + plasma_current = current(f, layout, v, dv) + v3 = scalar_velocity_moment(f, layout, v, dv, power=3) + vv3 = vector_velocity_moment(f, layout, v, dv, power=3) + safe_v3 = jnp.maximum(v3, jnp.finfo(v3.dtype).tiny) + return vv3 / (2.0 * safe_v3[..., None]) + plasma_current / jnp.maximum(ne[..., None], jnp.finfo(ne.dtype).tiny) + + +def complex_to_real(f: Array) -> Array: + """Store a complex harmonic array as a final real/imaginary axis.""" + + return jnp.stack((jnp.real(f), jnp.imag(f)), axis=-1) + + +def real_to_complex(f: Array) -> Array: + """Restore a complex harmonic array from a final real/imaginary axis.""" + + if f.shape[-1] != 2: + raise ValueError("real-embedded harmonic arrays must have a final axis of length 2") + return f[..., 0] + 1j * f[..., 1] diff --git a/adept/vfp2d/ohm.py b/adept/vfp2d/ohm.py new file mode 100644 index 00000000..cc4ca220 --- /dev/null +++ b/adept/vfp2d/ohm.py @@ -0,0 +1,151 @@ +"""Long-timescale kinetic Ohm-law utilities for VFP-2D. + +The diagnostic closure follows Joglekar et al., PRL 112, 105004 (2014), +Eq. (2). It intentionally neglects electron inertia. This makes it useful +for collisional transport and the PRL benchmark, but distinct from the future +fully implicit kinetic-current response solve. +""" + +from __future__ import annotations + +import jax.numpy as jnp +from jax import Array +from jax.sharding import Mesh + +from adept.vfp2d.harmonics import ( + HarmonicLayout, + current, + density, + nernst_velocity, + periodic_central_derivative, + scalar_velocity_moment, + tensor_velocity_moment, +) + + +class KineticOhm2D: + """Evaluate the moment-resolved, inertia-free generalized Ohm law.""" + + def __init__( + self, + layout: HarmonicLayout, + v: Array, + dv: float, + kx: Array, + ky: Array, + *, + resistivity_coefficient: float = 0.0, + dx: float | None = None, + dy: float | None = None, + mesh: Mesh | None = None, + ): + self.layout = layout + self.v = jnp.asarray(v) + self.dv = float(dv) + self.kx = jnp.asarray(kx) + self.ky = jnp.asarray(ky) + self.resistivity_coefficient = float(resistivity_coefficient) + self.dx = None if dx is None else float(dx) + self.dy = None if dy is None else float(dy) + self.mesh = mesh + + def ddx(self, value: Array) -> Array: + if self.dx is not None: + return periodic_central_derivative(value, self.dx, axis=0, mesh=self.mesh) + shape = (self.kx.size,) + (1,) * (value.ndim - 1) + return jnp.fft.ifft(1j * self.kx.reshape(shape) * jnp.fft.fft(value, axis=0), axis=0).real + + def ddy(self, value: Array) -> Array: + if self.dy is not None: + return periodic_central_derivative(value, self.dy, axis=1) + shape = (1, self.ky.size) + (1,) * (value.ndim - 2) + return jnp.fft.ifft(1j * self.ky.reshape(shape) * jnp.fft.fft(value, axis=1), axis=1).real + + def __call__( + self, + flm: Array, + magnetic_field: Array, + *, + plasma_current: Array | None = None, + hidden_dndz: Array | float = 0.0, + ) -> tuple[Array, dict[str, Array]]: + """Return ``E`` and its five PRL Eq. (2) contributions. + + ``hidden_dndz`` represents the prescribed density derivative in the + unresolved z direction. The kinetic step supplies the matching + ``df/dz = (dndz/ne) f`` streaming term, so an isothermal Maxwellian is + in pressure balance instead of being spuriously accelerated by Ez. + """ + + ne = density(flm, self.layout, self.v, self.dv) + safe_ne = jnp.maximum(ne, jnp.finfo(ne.dtype).tiny) + if plasma_current is None: + plasma_current = current(flm, self.layout, self.v, self.dv) + + v3 = scalar_velocity_moment(flm, self.layout, self.v, self.dv, power=3) + v5 = scalar_velocity_moment(flm, self.layout, self.v, self.dv, power=5) + tensor_v3 = tensor_velocity_moment(flm, self.layout, self.v, self.dv, power=3) + v_nernst = nernst_velocity( + flm, + self.layout, + self.v, + self.dv, + plasma_current=plasma_current, + ) + safe_v3 = jnp.maximum(v3, jnp.finfo(v3.dtype).tiny) + hidden_dndz = jnp.broadcast_to(jnp.asarray(hidden_dndz), ne.shape) + + eta = self.resistivity_coefficient / safe_v3 + resistive = eta[..., None] * plasma_current + hall = jnp.cross(plasma_current, magnetic_field) / safe_ne[..., None] + nernst = -jnp.cross(v_nernst, magnetic_field) + + scalar_flux = ne * v5 + scalar_gradient = jnp.stack((self.ddx(scalar_flux), self.ddy(scalar_flux), hidden_dndz * v5), axis=-1) + scalar_pressure = -scalar_gradient / (6.0 * safe_ne * safe_v3)[..., None] + + tensor_flux = ne[..., None, None] * tensor_v3 + tensor_divergence = self.ddx(tensor_flux[..., :, 0]) + self.ddy(tensor_flux[..., :, 1]) + tensor_divergence = tensor_divergence + hidden_dndz[..., None] * tensor_v3[..., :, 2] + tensor_pressure = -tensor_divergence / (2.0 * safe_ne * safe_v3)[..., None] + + terms = { + "resistive": resistive, + "hall": hall, + "nernst": nernst, + "scalar_pressure": scalar_pressure, + "tensor_pressure": tensor_pressure, + } + electric_field = sum(terms.values(), start=jnp.zeros_like(magnetic_field)) + return electric_field, terms + + +def project_current_moment( + flm: Array, + layout: HarmonicLayout, + v: Array, + dv: float, + target_current: Array, +) -> Array: + """Project only the bulk-current moment of ``f1`` onto ``target_current``. + + The correction is proportional to the local ``f00`` radial shape. It leaves + density and every ``l != 1`` harmonic unchanged, while retaining the + non-Maxwellian part of ``f1`` that carries the heat flux. + """ + + i10, i11 = layout.index(1, 0), layout.index(1, 1) + if i10 < 0: + raise ValueError("current projection requires l_max >= 1") + measured = current(flm, layout, v, dv) + correction = target_current - measured + f00 = jnp.real(flm[..., layout.index(0, 0), :]) + response = (4.0 * jnp.pi / 3.0) * jnp.sum(f00 * v**3, axis=-1) * dv + response = jnp.maximum(response, jnp.finfo(response.dtype).tiny) + # current = -response * [a_x, 2 Re(a_1), -2 Im(a_1)] + ax = -correction[..., 0] / response + a1 = (-correction[..., 1] + 1j * correction[..., 2]) / (2.0 * response) + result = flm.at[..., i10, :].add(ax[..., None] * f00) + if i11 >= 0: + result = result.at[..., i11, :].add(a1[..., None] * f00) + return result diff --git a/adept/vfp2d/plotting.py b/adept/vfp2d/plotting.py new file mode 100644 index 00000000..50c787e3 --- /dev/null +++ b/adept/vfp2d/plotting.py @@ -0,0 +1,588 @@ +"""Artifact generation and reconnection diagnostics for VFP-2D runs.""" + +from __future__ import annotations + +import os +from collections.abc import Iterable + +import matplotlib.pyplot as plt +import numpy as np +import xarray as xr + +plt.switch_backend("Agg") + +COMPONENTS = ("x", "y", "z") + + +def _selected_indices(nt: int, n_panels: int) -> np.ndarray: + return np.unique(np.linspace(0, nt - 1, min(nt, n_panels)).round().astype(int)) + + +def _physical_axes(ds: xr.Dataset) -> tuple[np.ndarray, np.ndarray, np.ndarray]: + x = np.asarray(ds.x) * float(ds.attrs.get("length_unit_um", 1.0)) + y = np.asarray(ds.y) * float(ds.attrs.get("length_unit_um", 1.0)) + t = np.asarray(ds.t) * float(ds.attrs.get("time_unit_ps", 1.0)) + return x, y, t + + +def save_xy_facet( + field: xr.DataArray, + ds: xr.Dataset, + path: str, + *, + n_panels: int = 9, + diverging: bool = True, + title: str | None = None, + xlim: tuple[float, float] | None = None, + ylim: tuple[float, float] | None = None, + aspect: str = "auto", +) -> None: + """Save evenly spaced x-y panels with one color scale for all times.""" + + x, y, t = _physical_axes(ds) + indices = _selected_indices(field.sizes["t"], n_panels) + x_indices = np.arange(x.size) + if xlim is not None: + x_indices = np.flatnonzero((x >= xlim[0]) & (x <= xlim[1])) + if not x_indices.size: + raise ValueError(f"No x coordinates lie inside requested facet limits {xlim}") + x = x[x_indices] + y_indices = np.arange(y.size) + if ylim is not None: + y_indices = np.flatnonzero((y >= ylim[0]) & (y <= ylim[1])) + if not y_indices.size: + raise ValueError(f"No y coordinates lie inside requested facet limits {ylim}") + y = y[y_indices] + values = np.asarray(field.isel(t=indices, x=x_indices, y=y_indices)) + finite = values[np.isfinite(values)] + if finite.size == 0: + vmin, vmax = -1.0, 1.0 + elif diverging: + vmax = max(float(np.max(np.abs(finite))), np.finfo(float).tiny) + vmin = -vmax + else: + vmin, vmax = float(np.min(finite)), float(np.max(finite)) + if vmin == vmax: + vmax = vmin + 1.0 + + ncols = min(3, indices.size) + nrows = int(np.ceil(indices.size / ncols)) + fig, axes = plt.subplots( + nrows, + ncols, + figsize=(4.0 * ncols, 3.25 * nrows), + constrained_layout=True, + squeeze=False, + ) + image = None + for panel, (ax, index) in enumerate(zip(axes.flat, indices, strict=False)): + image = ax.pcolormesh( + x, + y, + values[panel].T, + shading="auto", + cmap="RdBu_r" if diverging else "viridis", + vmin=vmin, + vmax=vmax, + ) + ax.set_title(f"t = {t[index]:.3g} ps") + ax.set_xlabel("x [μm]") + ax.set_ylabel("y [μm]") + ax.set_aspect(aspect) + for ax in axes.flat[indices.size :]: + ax.set_visible(False) + if image is not None: + fig.colorbar(image, ax=list(axes.flat[: indices.size]), shrink=0.82) + if title: + fig.suptitle(title, y=1.01) + fig.savefig(path, dpi=140, bbox_inches="tight") + plt.close(fig) + + +def _vector_potential(bx: np.ndarray, by: np.ndarray, dx: float, dy: float) -> np.ndarray: + """Return periodic A_z with B_x=∂_y A_z and B_y=-∂_x A_z.""" + + nx, ny = bx.shape[-2:] + kx = 2.0 * np.pi * np.fft.fftfreq(nx, d=dx) + ky = 2.0 * np.pi * np.fft.fftfreq(ny, d=dy) + kx2d, ky2d = np.meshgrid(kx, ky, indexing="ij") + k2 = kx2d**2 + ky2d**2 + bx_k = np.fft.fftn(bx, axes=(-2, -1)) + by_k = np.fft.fftn(by, axes=(-2, -1)) + numerator = -1j * ky2d * bx_k + 1j * kx2d * by_k + az_k = np.divide(numerator, k2, out=np.zeros_like(numerator), where=k2 > 0.0) + return np.fft.ifftn(az_k, axes=(-2, -1)).real + + +def add_reconnection_diagnostics(ds: xr.Dataset) -> xr.Dataset: + """Add topology diagnostics, gated so Biermann rings are not called reconnection.""" + + bx = np.asarray(ds.b.sel(component="x")) + by = np.asarray(ds.b.sel(component="y")) + x = np.asarray(ds.x) + y = np.asarray(ds.y) + dx = float(np.mean(np.diff(x))) + dy = float(np.mean(np.diff(y))) + az = _vector_potential(bx, by, dx, dy) + x_center = 0.0 if x[0] <= 0.0 <= x[-1] else 0.5 * (x[0] + x[-1]) + y_center = 0.0 if y[0] <= 0.0 <= y[-1] else 0.5 * (y[0] + y[-1]) + ix0, iy0 = int(np.argmin(np.abs(x - x_center))), int(np.argmin(np.abs(y - y_center))) + + lower = np.flatnonzero(y < y_center) + upper = np.flatnonzero(y > y_center) + if not lower.size or not upper.size: + raise ValueError("Reconnection diagnostics require y coordinates on both sides of zero") + + nt = bx.shape[0] + bx_line = bx[:, ix0, :] + lower_iy = lower[np.argmax(np.abs(bx_line[:, lower]), axis=-1)] + upper_iy = upper[np.argmax(np.abs(bx_line[:, upper]), axis=-1)] + bx_lower = bx_line[np.arange(nt), lower_iy] + bx_upper = bx_line[np.arange(nt), upper_iy] + abs_lower, abs_upper = np.abs(bx_lower), np.abs(bx_upper) + upstream_bx = 0.5 * (abs_lower + abs_upper) + upstream_balance = np.divide( + 2.0 * np.minimum(abs_lower, abs_upper), + abs_lower + abs_upper, + out=np.zeros(nt), + where=(abs_lower + abs_upper) > 1e-30, + ) + antiparallel = bx_lower * bx_upper < 0.0 + + upstream_az = 0.5 * (az[np.arange(nt), ix0, lower_iy] + az[np.arange(nt), ix0, upper_iy]) + raw_reconnected_flux = az[:, ix0, iy0] - upstream_az + + vn_y = np.asarray(ds.v_nernst.sel(component="y"))[:, ix0, :] + upstream_vn = 0.5 * ( + np.maximum(vn_y[np.arange(nt), lower_iy], 0.0) + np.maximum(-vn_y[np.arange(nt), upper_iy], 0.0) + ) + ez_x = np.asarray(ds.e.sel(component="z"))[:, ix0, iy0] + + inplane_null_ratio = np.divide( + np.hypot(bx[:, ix0, iy0], by[:, ix0, iy0]), + upstream_bx, + out=np.full(nt, np.inf), + where=upstream_bx > 1e-30, + ) + az_xx = np.gradient(np.gradient(az, dx, axis=1), dx, axis=1) + az_yy = np.gradient(np.gradient(az, dy, axis=2), dy, axis=2) + az_xy = np.gradient(np.gradient(az, dx, axis=1), dy, axis=2) + xpoint_hessian_determinant = az_xx[:, ix0, iy0] * az_yy[:, ix0, iy0] - az_xy[:, ix0, iy0] ** 2 + + jz = np.abs(np.asarray(ds.current.sel(component="z"))) + jz_line = jz[:, ix0, :] + y_from_sheet = y - y[iy0] + sheet_width = np.full(nt, np.nan) + sheet_dominance = np.zeros(nt) + quadrupole_purity = np.zeros(nt) + quadrupole_projection = np.zeros(nt) + quadrupole_central_fraction = np.zeros(nt) + bz = np.asarray(ds.b.sel(component="z")) + for it in range(nt): + upstream_distance = min(abs(y[lower_iy[it]]), abs(y[upper_iy[it]])) + sheet_half_width = max(2.0 * dy, 0.4 * upstream_distance) + sheet_mask = np.abs(y_from_sheet) <= sheet_half_width + weights = jz_line[it, sheet_mask] + if np.sum(weights) > 1e-30: + sheet_width[it] = np.sqrt(np.sum(weights * y_from_sheet[sheet_mask] ** 2) / np.sum(weights)) + global_peak = np.max(jz[it]) + if global_peak > 1e-30: + sheet_dominance[it] = np.max(weights, initial=0.0) / global_peak + + x_half_width = max(2.0 * dx, 2.0 * upstream_distance) + spatial_mask = (np.abs(x[:, None]) <= x_half_width) & (np.abs(y[None, :]) <= upstream_distance) + sign_pattern = np.sign(x[:, None] * y[None, :]) + local_bz = bz[it][spatial_mask] + local_sign = sign_pattern[spatial_mask] + local_l1 = np.sum(np.abs(local_bz)) + if local_l1 > 1e-30: + signed_sum = np.sum(local_bz * local_sign) + quadrupole_purity[it] = abs(signed_sum) / local_l1 + quadrupole_projection[it] = signed_sum / np.count_nonzero(spatial_mask) + quadrupole_central_fraction[it] = local_l1 / np.sum(np.abs(bz[it])) + + reconnection_valid = ( + antiparallel + & (upstream_balance >= 0.25) + & (inplane_null_ratio <= 0.5) + & (xpoint_hessian_determinant < 0.0) + & (sheet_dominance >= 0.25) + ) + peak_upstream_vn = np.max(upstream_vn, initial=0.0) + rate_normalization_valid = reconnection_valid & (upstream_vn >= 0.1 * peak_upstream_vn) + scale = upstream_bx * upstream_vn + normalized_rate = np.divide( + ez_x, + scale, + out=np.full_like(ez_x, np.nan), + where=rate_normalization_valid & (scale > 1e-30), + ) + reconnected_flux = np.where(reconnection_valid, raw_reconnected_flux, np.nan) + + result = ds.assign( + az=(("t", "x", "y"), az), + xpoint_ez=("t", ez_x), + upstream_bx=("t", upstream_bx), + upstream_v_nernst_y=("t", upstream_vn), + upstream_field_balance=("t", upstream_balance), + inplane_null_ratio=("t", inplane_null_ratio), + xpoint_hessian_determinant=("t", xpoint_hessian_determinant), + current_sheet_dominance=("t", sheet_dominance), + reconnection_valid=("t", reconnection_valid), + rate_normalization_valid=("t", rate_normalization_valid), + normalized_reconnection_rate=("t", normalized_rate), + reconnected_flux=("t", reconnected_flux), + current_sheet_rms_width=("t", sheet_width), + bz_quadrupole_purity=("t", quadrupole_purity), + bz_quadrupole_projection=("t", quadrupole_projection), + bz_quadrupole_central_fraction=("t", quadrupole_central_fraction), + ) + result.az.attrs["definition"] = "periodic A_z: B_x=d_y A_z, B_y=-d_x A_z" + result.reconnection_valid.attrs["definition"] = ( + "antiparallel balanced upstream Bx, central in-plane null/Az saddle, and central |jz| >= 25% global peak" + ) + result.normalized_reconnection_rate.attrs["definition"] = ( + "E_z(X)/(<|B_x,up|> ), NaN unless reconnection_valid and inward v_N is at " + "least 10% of its run maximum" + ) + result.reconnected_flux.attrs["definition"] = "A_z(X)-mean(A_z at both upstream locations), gated" + result.current_sheet_rms_width.attrs["definition"] = "central-window |j_z|-weighted RMS y width at x=0" + result.bz_quadrupole_purity.attrs["definition"] = "absolute L1 projection of local Bz onto sign(x*y)" + for name in ( + "ohm_resistive", + "ohm_hall", + "ohm_nernst", + "ohm_scalar_pressure", + "ohm_tensor_pressure", + ): + if name in result: + result[f"xpoint_{name}"] = ("t", np.asarray(result[name].sel(component="z"))[:, ix0, iy0]) + return result + + +def _write_binary(ds: xr.Dataset, binary_dir: str) -> None: + moments = ds.drop_vars(["flm_real", "flm_imag"], errors="ignore") + distribution = ds[[name for name in ("flm_real", "flm_imag") if name in ds]] + + def encoding(dataset: xr.Dataset) -> dict: + return { + name: {"compression": "gzip", "compression_opts": 1, "shuffle": True} + for name, value in dataset.data_vars.items() + if np.issubdtype(value.dtype, np.number) and value.ndim > 0 + } + + moments.to_netcdf(os.path.join(binary_dir, "moments.nc"), engine="h5netcdf", encoding=encoding(moments)) + if distribution.data_vars: + distribution.to_netcdf( + os.path.join(binary_dir, "distribution_flm.nc"), + engine="h5netcdf", + encoding=encoding(distribution), + ) + + +def _plot_regular_moments(ds: xr.Dataset, plot_dir: str, n_panels: int) -> None: + scalar_fields = { + "density": (ds.ne, False), + "temperature": (ds.temperature, False), + "magnetic_field_magnitude": (np.sqrt((ds.b**2).sum("component")), False), + "current_magnitude": (np.sqrt((ds.current**2).sum("component")), False), + "nernst_velocity_magnitude": (np.sqrt((ds.v_nernst**2).sum("component")), False), + } + for name, (field, diverging) in scalar_fields.items(): + save_xy_facet( + field, + ds, + os.path.join(plot_dir, f"xy_facet_{name}.png"), + n_panels=n_panels, + diverging=diverging, + title=name.replace("_", " "), + ) + + for variable in ("e", "b", "current", "v_nernst"): + for component in COMPONENTS: + save_xy_facet( + ds[variable].sel(component=component), + ds, + os.path.join(plot_dir, f"xy_facet_{variable}_{component}.png"), + n_panels=n_panels, + title=f"{variable.replace('_', ' ')} {component}", + ) + + pressure_components: Iterable[tuple[str, str]] = ( + ("x", "x"), + ("x", "y"), + ("x", "z"), + ("y", "y"), + ("y", "z"), + ("z", "z"), + ) + for first, second in pressure_components: + field = ds.pressure_anisotropy.sel(component=first, component_2=second) + save_xy_facet( + field, + ds, + os.path.join(plot_dir, f"xy_facet_pressure_{first}{second}.png"), + n_panels=n_panels, + title=f"pressure anisotropy {first}{second}", + ) + + +def _plot_reconnection_region( + ds: xr.Dataset, + plot_dir: str, + n_panels: int, + xlim: tuple[float, float], + ylim: tuple[float, float], +) -> None: + """Plot a compact set of moments around the X-point, away from domain edges.""" + + fields: dict[str, tuple[xr.DataArray, bool]] = { + "density": (ds.ne, False), + "temperature": (ds.temperature, False), + "magnetic_field_magnitude": (np.sqrt((ds.b**2).sum("component")), False), + "current_magnitude": (np.sqrt((ds.current**2).sum("component")), False), + "nernst_velocity_magnitude": (np.sqrt((ds.v_nernst**2).sum("component")), False), + "b_x": (ds.b.sel(component="x"), True), + "b_z": (ds.b.sel(component="z"), True), + "e_z": (ds.e.sel(component="z"), True), + "current_z": (ds.current.sel(component="z"), True), + "v_nernst_y": (ds.v_nernst.sel(component="y"), True), + "vector_potential": (ds.az, True), + } + ohm_terms = ("ohm_resistive", "ohm_hall", "ohm_nernst", "ohm_scalar_pressure", "ohm_tensor_pressure") + for name in ohm_terms: + if name in ds: + fields[f"{name}_z"] = (ds[name].sel(component="z"), True) + + for name, (field, diverging) in fields.items(): + save_xy_facet( + field, + ds, + os.path.join(plot_dir, f"xy_facet_{name}.png"), + n_panels=n_panels, + diverging=diverging, + title=f"reconnection region: {name.replace('_', ' ')}", + xlim=xlim, + ylim=ylim, + ) + + +def _plot_xpoint_history(ds: xr.Dataset, path: str) -> None: + _, _, t = _physical_axes(ds) + fig, axes = plt.subplots(4, 1, figsize=(8, 12), constrained_layout=True, sharex=True) + axes[0].plot(t, ds.xpoint_ez, color="black", linewidth=2, label="total E_z") + for name in ("ohm_resistive", "ohm_hall", "ohm_nernst", "ohm_scalar_pressure", "ohm_tensor_pressure"): + xpoint_name = f"xpoint_{name}" + if xpoint_name in ds: + axes[0].plot(t, ds[xpoint_name], label=name.removeprefix("ohm_").replace("_", " ")) + axes[0].set_ylabel("X-point E_z [norm.] ") + axes[0].legend(ncol=2, fontsize=8) + axes[0].grid(alpha=0.3) + + axes[1].plot(t, ds.normalized_reconnection_rate, label="E_z/(B_up v_N,in)") + axes[1].plot(t, ds.reconnected_flux, label="reconnected flux") + axes[1].set_ylabel("reconnection diagnostics") + axes[1].legend() + axes[1].grid(alpha=0.3) + + length_um = float(ds.attrs.get("length_unit_um", 1.0)) + axes[2].plot(t, ds.current_sheet_rms_width * length_um, label="current-sheet RMS width") + axes[2].plot(t, ds.upstream_bx, label="upstream |B_x|") + axes[2].plot(t, ds.upstream_v_nernst_y, label="inflow |v_N,y|") + axes[2].set_ylabel("width [μm] / normalized amplitude") + axes[2].legend() + axes[2].grid(alpha=0.3) + + axes[3].plot(t, ds.bz_quadrupole_purity, label="Bz quadrupole purity") + axes[3].plot(t, ds.bz_quadrupole_central_fraction, label="central Bz fraction") + axes[3].step(t, ds.reconnection_valid.astype(float), where="mid", label="valid sheet/X-point") + axes[3].step( + t, + ds.rate_normalization_valid.astype(float), + where="mid", + linestyle=":", + label="valid rate normalization", + ) + axes[3].set_xlabel("t [ps]") + axes[3].set_ylabel("structure score") + axes[3].set_ylim(-0.05, 1.05) + axes[3].legend() + axes[3].grid(alpha=0.3) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def _plot_ohm_lineouts(ds: xr.Dataset, path: str, n_panels: int) -> None: + x, y, t = _physical_axes(ds) + del x + ix0 = int(np.argmin(np.abs(np.asarray(ds.x)))) + indices = _selected_indices(ds.sizes["t"], min(n_panels, 4)) + fig, axes = plt.subplots(indices.size, 1, figsize=(8, 2.8 * indices.size), constrained_layout=True, squeeze=False) + for ax, index in zip(axes[:, 0], indices, strict=True): + ax.plot(y, ds.e.isel(t=index, x=ix0).sel(component="z"), color="black", linewidth=2, label="total E_z") + for name in ("ohm_resistive", "ohm_hall", "ohm_nernst", "ohm_scalar_pressure", "ohm_tensor_pressure"): + if name in ds: + ax.plot( + y, + ds[name].isel(t=index, x=ix0).sel(component="z"), + label=name.removeprefix("ohm_").replace("_", " "), + ) + ax.set_title(f"x = 0, t = {t[index]:.3g} ps") + ax.set_xlabel("y [μm]") + ax.set_ylabel("E_z [norm.]") + ax.grid(alpha=0.3) + axes[0, 0].legend(ncol=3, fontsize=8) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def _plot_topology(ds: xr.Dataset, path: str) -> None: + x, y, t = _physical_axes(ds) + rate_valid_indices = np.flatnonzero(np.asarray(ds.rate_normalization_valid)) + index = int(rate_valid_indices[-1]) if rate_valid_indices.size else ds.sizes["t"] - 1 + bmag = np.sqrt(np.asarray((ds.b.isel(t=index) ** 2).sum("component"))) + az = np.asarray(ds.az.isel(t=index)) + vx = np.asarray(ds.v_nernst.isel(t=index).sel(component="x")) + vy = np.asarray(ds.v_nernst.isel(t=index).sel(component="y")) + central_half_width = 3.0 * max(abs(float(y[0])), abs(float(y[-1]))) + central_x_indices = np.flatnonzero(np.abs(x) <= central_half_width) + fig, ax = plt.subplots(figsize=(8, 6), constrained_layout=True) + image = ax.pcolormesh(x, y, bmag.T, shading="auto", cmap="magma") + if np.ptp(az) > 0.0: + ax.contour(x, y, az.T, colors="white", linewidths=0.7, levels=14, alpha=0.8) + stride_x, stride_y = max(1, central_x_indices.size // 20), max(1, len(y) // 12) + quiver_x_indices = central_x_indices[::stride_x] + quiver_y_indices = np.arange(0, len(y), stride_y) + central_speed = np.hypot(vx[central_x_indices], vy[central_x_indices]) + velocity_scale = max(float(np.percentile(central_speed, 95.0)), np.finfo(float).tiny) + ax.quiver( + x[quiver_x_indices], + y[quiver_y_indices], + (vx[np.ix_(quiver_x_indices, quiver_y_indices)] / velocity_scale).T, + (vy[np.ix_(quiver_x_indices, quiver_y_indices)] / velocity_scale).T, + color="cyan", + pivot="mid", + scale=22.0, + width=0.0025, + ) + ax.set_title(f"|B|, A_z contours, and Nernst velocity at t={t[index]:.3g} ps") + ax.set_xlabel("x [μm]") + ax.set_ylabel("y [μm]") + ax.set_xlim(-central_half_width, central_half_width) + ax.set_aspect("equal") + fig.colorbar(image, ax=ax, label="|B| [norm.]") + fig.savefig(path, dpi=160, bbox_inches="tight") + plt.close(fig) + + +def _plot_sheet_lineouts(ds: xr.Dataset, path: str, n_panels: int) -> None: + """Plot the opposing Bx ribbons and central Jz sheet on x=0.""" + + _, y, t = _physical_axes(ds) + ix0 = int(np.argmin(np.abs(np.asarray(ds.x)))) + indices = _selected_indices(ds.sizes["t"], min(n_panels, 4)) + fig, axes = plt.subplots(indices.size, 1, figsize=(8, 2.8 * indices.size), constrained_layout=True, squeeze=False) + for ax, index in zip(axes[:, 0], indices, strict=True): + ax.plot(y, ds.b.isel(t=index, x=ix0).sel(component="x"), color="tab:blue", label="B_x") + twin = ax.twinx() + twin.plot(y, ds.current.isel(t=index, x=ix0).sel(component="z"), color="tab:red", label="j_z") + ax.axvline(0.0, color="black", linewidth=0.7, alpha=0.4) + ax.set_title(f"x = 0, t = {t[index]:.3g} ps") + ax.set_xlabel("y [μm]") + ax.set_ylabel("B_x", color="tab:blue") + twin.set_ylabel("j_z", color="tab:red") + ax.grid(alpha=0.3) + fig.savefig(path, dpi=150, bbox_inches="tight") + plt.close(fig) + + +def save_artifacts(ds: xr.Dataset, td: str, *, n_panels: int = 9) -> None: + """Write compact binaries, standard facets, and reconnection plots.""" + + binary_dir = os.path.join(td, "binary") + moments_dir = os.path.join(td, "plots", "moments") + reconnection_dir = os.path.join(td, "plots", "reconnection") + reconnection_region_dir = os.path.join(td, "plots", "reconnection_region") + for directory in (binary_dir, moments_dir, reconnection_dir, reconnection_region_dir): + os.makedirs(directory, exist_ok=True) + _write_binary(ds, binary_dir) + _plot_regular_moments(ds, moments_dir, n_panels) + + for name in ("az", "xpoint_ez"): + if name == "az": + save_xy_facet( + ds.az, + ds, + os.path.join(reconnection_dir, "xy_facet_vector_potential.png"), + n_panels=n_panels, + title="reconnection flux A_z", + ) + for variable, component in ( + ("b", "x"), + ("b", "y"), + ("b", "z"), + ("e", "z"), + ("current", "z"), + ("v_nernst", "y"), + ): + save_xy_facet( + ds[variable].sel(component=component), + ds, + os.path.join(reconnection_dir, f"xy_facet_{variable}_{component}.png"), + n_panels=n_panels, + title=f"reconnection: {variable} {component}", + ) + y_um = _physical_axes(ds)[1] + central_half_width = 3.0 * max(abs(float(y_um[0])), abs(float(y_um[-1]))) + reconnection_xlim = (-central_half_width, central_half_width) + reconnection_y_half_width = 0.75 * max(abs(float(y_um[0])), abs(float(y_um[-1]))) + reconnection_ylim = (-reconnection_y_half_width, reconnection_y_half_width) + _plot_reconnection_region(ds, reconnection_region_dir, n_panels, reconnection_xlim, reconnection_ylim) + save_xy_facet( + ds.b.sel(component="z"), + ds, + os.path.join(reconnection_dir, "xy_facet_b_z_reconnection_region.png"), + n_panels=n_panels, + title="reconnection-region B_z quadrupole", + xlim=reconnection_xlim, + ) + for name in ("ohm_resistive", "ohm_hall", "ohm_nernst", "ohm_scalar_pressure", "ohm_tensor_pressure"): + if name in ds: + save_xy_facet( + ds[name].sel(component="z"), + ds, + os.path.join(reconnection_dir, f"xy_facet_{name}_z.png"), + n_panels=n_panels, + title=f"{name.replace('_', ' ')} z", + ) + _plot_xpoint_history(ds, os.path.join(reconnection_dir, "xpoint_history.png")) + _plot_ohm_lineouts(ds, os.path.join(reconnection_dir, "ohm_z_lineouts_x0.png"), n_panels) + _plot_sheet_lineouts(ds, os.path.join(reconnection_dir, "bx_jz_sheet_lineouts_x0.png"), n_panels) + _plot_topology(ds, os.path.join(reconnection_dir, "topology_nernst_final.png")) + + +def reconnection_metrics(ds: xr.Dataset) -> dict[str, float]: + """Small scalar summary suitable for MLflow comparisons.""" + + rate = np.asarray(ds.normalized_reconnection_rate) + flux = np.asarray(ds.reconnected_flux) + finite_rate = rate[np.isfinite(rate)] + finite_flux = flux[np.isfinite(flux)] + final_valid = bool(ds.reconnection_valid[-1]) + final_rate_valid = bool(ds.rate_normalization_valid[-1]) + return { + "vfp2d_peak_abs_b": float(np.max(np.abs(ds.b))), + "vfp2d_peak_temperature": float(np.max(ds.temperature)), + "vfp2d_peak_abs_reconnection_rate": float(np.max(np.abs(finite_rate))) if finite_rate.size else 0.0, + "vfp2d_final_reconnection_rate": float(rate[-1]) if np.isfinite(rate[-1]) else 0.0, + "vfp2d_final_reconnected_flux": float(flux[-1]) if np.isfinite(flux[-1]) else 0.0, + "vfp2d_final_current_sheet_rms_width": ( + float(ds.current_sheet_rms_width[-1]) if np.isfinite(ds.current_sheet_rms_width[-1]) else 0.0 + ), + "vfp2d_reconnection_valid_fraction": float(np.mean(ds.reconnection_valid)), + "vfp2d_final_reconnection_valid": float(final_valid), + "vfp2d_rate_normalization_valid_fraction": float(np.mean(ds.rate_normalization_valid)), + "vfp2d_final_rate_normalization_valid": float(final_rate_valid), + "vfp2d_final_bz_quadrupole_purity": float(ds.bz_quadrupole_purity[-1]), + "vfp2d_peak_abs_reconnected_flux": float(np.max(np.abs(finite_flux))) if finite_flux.size else 0.0, + } diff --git a/adept/vfp2d/vector_field.py b/adept/vfp2d/vector_field.py new file mode 100644 index 00000000..ffceb174 --- /dev/null +++ b/adept/vfp2d/vector_field.py @@ -0,0 +1,337 @@ +"""Coupled 2D3P spherical-harmonic Vlasov--Maxwell vector field.""" + +from __future__ import annotations + +import jax.numpy as jnp +import jax.tree_util as jtu +from jax import Array +from jax.sharding import Mesh + +from adept.vfp2d.collisions import CollisionStep +from adept.vfp2d.harmonics import ( + HarmonicLayout, + HouLiFilter2D, + TzoufrasVlasov, + complex_to_real, + conservative_f00_positivity, + current, + density, + periodic_central_derivative, + real_to_complex, +) +from adept.vfp2d.ohm import KineticOhm2D, project_current_moment + + +def _ib_gate(t: float, args: dict | None) -> Array: + """Return the optional smooth temporal envelope for IB heating.""" + + if not args: + return jnp.asarray(1.0) + width = jnp.asarray(args.get("ib_switch_width", 0.0)) + gate = jnp.asarray(1.0) + if "ib_t_on" in args: + t_on = jnp.asarray(args["ib_t_on"]) + sharp_on = jnp.where(t >= t_on, 1.0, 0.0) + smooth_on = 0.5 * (1.0 + jnp.tanh((t - t_on) / jnp.maximum(width, 1e-30))) + gate = gate * jnp.where(width > 0.0, smooth_on, sharp_on) + if "ib_t_off" in args: + t_off = jnp.asarray(args["ib_t_off"]) + sharp_off = jnp.where(t < t_off, 1.0, 0.0) + smooth_off = 0.5 * (1.0 - jnp.tanh((t - t_off) / jnp.maximum(width, 1e-30))) + gate = gate * jnp.where(width > 0.0, smooth_off, sharp_off) + return gate + + +class Maxwell2D: + """Full three-component Maxwell curl operator with ``d/dz = 0``.""" + + def __init__( + self, + kx: Array, + ky: Array, + c: float, + *, + dx: float | None = None, + dy: float | None = None, + mesh: Mesh | None = None, + ): + self.kx = jnp.asarray(kx) + self.ky = jnp.asarray(ky) + self.c2 = float(c) ** 2 + self.dx = None if dx is None else float(dx) + self.dy = None if dy is None else float(dy) + self.mesh = mesh + + def ddx(self, a: Array) -> Array: + if self.dx is not None: + return periodic_central_derivative(a, self.dx, axis=0, mesh=self.mesh) + return jnp.fft.ifft(1j * self.kx[:, None] * jnp.fft.fft(a, axis=0), axis=0).real + + def ddy(self, a: Array) -> Array: + if self.dy is not None: + return periodic_central_derivative(a, self.dy, axis=1) + return jnp.fft.ifft(1j * self.ky[None, :] * jnp.fft.fft(a, axis=1), axis=1).real + + def curl(self, a: Array) -> Array: + ax, ay, az = a[..., 0], a[..., 1], a[..., 2] + return jnp.stack((self.ddy(az), -self.ddx(az), self.ddx(ay) - self.ddy(ax)), axis=-1) + + def __call__(self, electric_field: Array, magnetic_field: Array, plasma_current: Array) -> tuple[Array, Array]: + dedt = self.c2 * self.curl(magnetic_field) - plasma_current + dbdt = -self.curl(electric_field) + return dedt, dbdt + + +class SpectralPoisson2D: + """Periodic initial Gauss-law solve for ``div(E)=rho``.""" + + def __init__(self, kx: Array, ky: Array): + self.kx = jnp.asarray(kx)[:, None] + self.ky = jnp.asarray(ky)[None, :] + k2 = self.kx**2 + self.ky**2 + self.inv_k2 = jnp.where(k2 > 0, 1.0 / k2, 0.0) + + def __call__(self, charge_density: Array) -> Array: + rho_k = jnp.fft.fft2(charge_density) + phi_k = self.inv_k2 * rho_k + ex = jnp.fft.ifft2(-1j * self.kx * phi_k).real + ey = jnp.fft.ifft2(-1j * self.ky * phi_k).real + return jnp.stack((ex, ey, jnp.zeros_like(ex)), axis=-1) + + +class VlasovMaxwell: + """Explicit collisionless RHS for packed arbitrary-``f_lm`` state arrays. + + The state is ``{"flm": ..., "e": ..., "b": ...}``. Optional external + fields and current drivers may be supplied through ``args`` with keys + ``external_e``, ``external_b``, and ``driver_current``. Each value may be + either an array or a callable of time. + """ + + def __init__( + self, + vlasov: TzoufrasVlasov, + maxwell: Maxwell2D, + layout: HarmonicLayout, + v: Array, + dv: float, + charge: float = -1.0, + real_storage: bool = False, + streaming_speed: Array | None = None, + ): + self.vlasov = vlasov + self.maxwell = maxwell + self.layout = layout + self.v = jnp.asarray(v) + self.dv = float(dv) + self.charge = float(charge) + self.real_storage = bool(real_storage) + self.streaming_speed = self.v if streaming_speed is None else jnp.asarray(streaming_speed) + + @staticmethod + def _arg(args: dict | None, key: str, t: float, template: Array) -> Array: + if not args or key not in args: + return jnp.zeros_like(template) + value = args[key] + return value(t) if callable(value) else value + + def __call__(self, t: float, state: dict[str, Array], args: dict | None = None) -> dict[str, Array]: + e = state["e"] + b = state["b"] + total_e = e + self._arg(args, "external_e", t, e) + total_b = b + self._arg(args, "external_b", t, b) + flm = real_to_complex(state["flm"]) if self.real_storage else state["flm"] + dfdt = self.vlasov(flm, total_e, total_b) + plasma_current = current( + flm, + self.layout, + self.v, + self.dv, + charge=self.charge, + streaming_speed=self.streaming_speed, + ) + driver_current = self._arg(args, "driver_current", t, e) + dedt, dbdt = self.maxwell(e, b, plasma_current + driver_current) + return {"flm": complex_to_real(dfdt) if self.real_storage else dfdt, "e": dedt, "b": dbdt} + + +class SplitStepVFP2D: + """One second-order explicit Vlasov--Maxwell step with collision splitting. + + The return value is the advanced state (rather than a derivative), matching + ADEPT's map-style ``Stepper`` interface. Collisions are applied in two half + steps around an explicit midpoint Vlasov--Maxwell update. + """ + + def __init__(self, rhs: VlasovMaxwell, dt: float, collisions: CollisionStep | None = None): + self.rhs = rhs + self.dt = float(dt) + self.collisions = collisions + + def _collide(self, t: float, state: dict[str, Array], args: dict | None, dt: float) -> dict[str, Array]: + if self.collisions is None: + return state + z = 1.0 if not args else args.get("Z", 1.0) + ni = 1.0 if not args else args.get("ni", 1.0) + heating = {} + if args: + for key in ("D0_heating", "ib_vosc2", "ib_Z2ni_w0"): + if key in args: + heating[key] = args[key] + if "ib_vosc2" in heating: + heating["ib_vosc2"] = heating["ib_vosc2"] * _ib_gate(t, args) + flm = real_to_complex(state["flm"]) if self.rhs.real_storage else state["flm"] + flm = self.collisions(flm, Z=z, ni=ni, dt=dt, **heating) + if self.rhs.real_storage: + flm = complex_to_real(flm) + return {**state, "flm": flm} + + def __call__(self, t: float, state: dict[str, Array], args: dict | None = None) -> dict[str, Array]: + state = self._collide(t, state, args, 0.5 * self.dt) + k1 = self.rhs(t, state, args) + midpoint = jtu.tree_map(lambda value, slope: value + 0.5 * self.dt * slope, state, k1) + k2 = self.rhs(t + 0.5 * self.dt, midpoint, args) + result = jtu.tree_map(lambda value, slope: value + self.dt * slope, state, k2) + return self._collide(t + self.dt, result, args, 0.5 * self.dt) + + +class KineticOhmStep: + """Long-timescale RK4 step using the inertia-free kinetic Ohm law. + + The quasistatic Ampere current is enforced by a minimal projection of the + bulk-current moment of ``f1``. This removes light and plasma oscillations + while preserving the velocity-dependent ``f1`` structure responsible for + nonlocal heat flow. It is deliberately separate from the future fully + implicit kinetic-current-response algorithm. + """ + + def __init__( + self, + vlasov: TzoufrasVlasov, + maxwell: Maxwell2D, + ohm: KineticOhm2D, + layout: HarmonicLayout, + v: Array, + dv: float, + dt: float, + collisions: CollisionStep | None = None, + real_storage: bool = False, + enforce_f00_positivity: bool = False, + spatial_filter: HouLiFilter2D | None = None, + ): + self.vlasov = vlasov + self.maxwell = maxwell + self.ohm = ohm + self.layout = layout + self.v = jnp.asarray(v) + self.dv = float(dv) + self.dt = float(dt) + self.collisions = collisions + self.real_storage = bool(real_storage) + self.enforce_f00_positivity = bool(enforce_f00_positivity) + self.spatial_filter = spatial_filter + + def _positive_f00(self, flm: Array) -> Array: + if not self.enforce_f00_positivity: + return flm + return conservative_f00_positivity(flm, self.layout, self.v, self.dv) + + def _filter(self, value: Array) -> Array: + return value if self.spatial_filter is None else self.spatial_filter(value) + + def _collide(self, t: float, flm: Array, args: dict | None, dt: float) -> Array: + if self.collisions is None: + return flm + z = 1.0 if not args else args.get("Z", 1.0) + ni = 1.0 if not args else args.get("ni", 1.0) + heating = {} + if args: + for key in ("D0_heating", "ib_vosc2", "ib_Z2ni_w0"): + if key in args: + heating[key] = args[key] + if "ib_vosc2" in heating: + heating["ib_vosc2"] = heating["ib_vosc2"] * _ib_gate(t, args) + return self.collisions(flm, Z=z, ni=ni, dt=dt, **heating) + + def _target_current(self, magnetic_field: Array) -> Array: + return self.maxwell.c2 * self.maxwell.curl(magnetic_field) + + def _project(self, flm: Array, magnetic_field: Array) -> Array: + return project_current_moment( + flm, + self.layout, + self.v, + self.dv, + self._target_current(magnetic_field), + ) + + @staticmethod + def _hidden_dndz(t: float, args: dict | None, template: Array) -> Array: + if not args or "hidden_dndz" not in args: + return jnp.zeros_like(template) + source = jnp.broadcast_to(jnp.asarray(args["hidden_dndz"]), template.shape) + if "hidden_gradient_t_off" not in args: + return source + t_off = jnp.asarray(args["hidden_gradient_t_off"]) + width = jnp.asarray(args.get("hidden_gradient_switch_width", 0.0)) + sharp_gate = jnp.where(t < t_off, 1.0, 0.0) + smooth_gate = 0.5 * (1.0 - jnp.tanh((t - t_off) / jnp.maximum(width, 1e-30))) + return source * jnp.where(width > 0.0, smooth_gate, sharp_gate) + + def _rates(self, t: float, flm: Array, magnetic_field: Array, args: dict | None) -> tuple[Array, Array, Array]: + flm = self._project(flm, magnetic_field) + hidden_dndz = self._hidden_dndz(t, args, magnetic_field[..., 0]) + electric_field, _terms = self.ohm( + flm, + magnetic_field, + plasma_current=self._target_current(magnetic_field), + hidden_dndz=hidden_dndz, + ) + ne = density(flm, self.layout, self.v, self.dv) + safe_ne = jnp.maximum(ne, jnp.finfo(ne.dtype).tiny) + dfdz = hidden_dndz[..., None, None] * flm / safe_ne[..., None, None] + return ( + self.vlasov(flm, electric_field, magnetic_field, dfdz=dfdz), + -self.maxwell.curl(electric_field), + electric_field, + ) + + def __call__(self, t: float, state: dict[str, Array], args: dict | None = None) -> dict[str, Array]: + flm = real_to_complex(state["flm"]) if self.real_storage else state["flm"] + magnetic_field = state["b"] + flm = self._positive_f00(self._collide(t, flm, args, 0.5 * self.dt)) + flm = self._project(flm, magnetic_field) + + df1, db1, _electric1 = self._rates(t, flm, magnetic_field, args) + stage2_b = magnetic_field + 0.5 * self.dt * db1 + stage2_f = self._positive_f00(self._project(flm + 0.5 * self.dt * df1, stage2_b)) + df2, db2, _electric2 = self._rates(t + 0.5 * self.dt, stage2_f, stage2_b, args) + + stage3_b = magnetic_field + 0.5 * self.dt * db2 + stage3_f = self._positive_f00(self._project(flm + 0.5 * self.dt * df2, stage3_b)) + df3, db3, _electric3 = self._rates(t + 0.5 * self.dt, stage3_f, stage3_b, args) + + stage4_b = magnetic_field + self.dt * db3 + stage4_f = self._positive_f00(self._project(flm + self.dt * df3, stage4_b)) + df4, db4, _electric4 = self._rates(t + self.dt, stage4_f, stage4_b, args) + + result_b = magnetic_field + (self.dt / 6.0) * (db1 + 2.0 * db2 + 2.0 * db3 + db4) + result_f = flm + (self.dt / 6.0) * (df1 + 2.0 * df2 + 2.0 * df3 + df4) + result_f = self._positive_f00(self._project(result_f, result_b)) + result_f = self._positive_f00(self._collide(t + self.dt, result_f, args, 0.5 * self.dt)) + # Pseudospectral field products feed unresolved power into the grid + # scale during long heated runs. Filter only configuration space, once + # per full step, then restore the f00 and Ampere-moment invariants. + result_b = jnp.real(self._filter(result_b)) + result_f = self._positive_f00(self._filter(result_f)) + result_f = self._project(result_f, result_b) + hidden_dndz = self._hidden_dndz(t + self.dt, args, result_b[..., 0]) + result_e, _terms = self.ohm( + result_f, + result_b, + plasma_current=self._target_current(result_b), + hidden_dndz=hidden_dndz, + ) + if self.real_storage: + result_f = complex_to_real(result_f) + return {"flm": result_f, "e": result_e, "b": result_b} diff --git a/configs/vfp-2d/joglekar-2014-prl-nersc-ny128.yaml b/configs/vfp-2d/joglekar-2014-prl-nersc-ny128.yaml new file mode 100644 index 00000000..5757c37d --- /dev/null +++ b/configs/vfp-2d/joglekar-2014-prl-nersc-ny128.yaml @@ -0,0 +1,83 @@ +# Y-refinement check for the sharded Joglekar et al. reconstruction. +# This changes only ny, dt, tmax, and the save cadence relative to the wide-box +# production config so the comparison isolates the choppy y structure in the +# late-time magnetic field and kinetic Ohm-law terms. +solver: vfp-2d + +mlflow: + experiment: vfp2d-joglekar-2014 + run: wide-x-sharded-4gpu-ny128-dt1fs-30ps-houli-p12 + +units: + laser_wavelength: 351nm + reference electron temperature: 1600eV + reference ion temperature: 160eV + reference electron density: 2.5e22/cm^3 + Z: 1 + Ion: Au+ + logLambda: nrl + +density: + quasineutrality: true + species-electron: + m: 2.0 + n: {basis: uniform, baseline: 1.0} + T: {basis: uniform, baseline: 1.0} + +grid: + xmin: -510um + xmax: 510um + nx: 224 + ymin: -34um + ymax: 34um + # dy = 0.53125 um, twice the published central y resolution. + ny: 128 + tmin: 0ps + tmax: 30ps + # Halving dt protects the dispersive kinetic-Ohm update when dy is halved. + dt: 1fs + nv: 96 + vmax: 8.0 + lmax: 2 + mmax: 2 + sharding: {enabled: true, axis: x} + +terms: + # The p=36 taper left a grid-locked 2.5--4-cell y mode in Bz and the kinetic + # Ohm terms. A p=12 convergence scan removes the upper-k shelf while keeping + # the 30 ps sheet width and reconnected flux within 1.5% of p=36. + hou_li_filter: {is_on: true, alpha: 36.0, order: 12, dimensions: [x, y]} + field_solver: + mode: kinetic-ohm + hidden_density_gradient: + active: true + scale_length: 17um + switch_off: 17.78ps + switch_width: 0.5ps + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + y_centers: [-34um, 34um] + y_radius: 17um + fokker_planck: + active: true + flm: {ee: true} + f00: {model: CoulombianKernel, scheme: chang_cooper, positivity: conservative} + +drivers: + ib: + intensity_1e15_Wcm2: 0.25 + polarisation: linear + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + y_centers: [-34um, 34um] + y_radius: 17um + +save: + t: {tmin: 0ps, tmax: 30ps, nt: 13} + +output: + n_panels: 9 diff --git a/configs/vfp-2d/joglekar-2014-prl-nersc.yaml b/configs/vfp-2d/joglekar-2014-prl-nersc.yaml new file mode 100644 index 00000000..0eccb18d --- /dev/null +++ b/configs/vfp-2d/joglekar-2014-prl-nersc.yaml @@ -0,0 +1,95 @@ +# Short NERSC benchmark for the published Joglekar et al. geometry. +# The Letter uses a central uniform x mesh out to +/-136 um and a stretched +# mesh to +/-510 um. VFP-2D is still Fourier-periodic, so this first benchmark +# uses the full wide x extent at the published central dx: periodic images are +# 1.02 mm apart and act as effectively open boundaries on this timescale. +solver: vfp-2d + +mlflow: + experiment: vfp2d-joglekar-2014 + run: published-geometry-wide-x-60ps-sharded + +units: + laser_wavelength: 351nm + reference electron temperature: 1600eV + reference ion temperature: 160eV + reference electron density: 2.5e22/cm^3 + # The Letter does not report the mean ionization in its five pages. + Z: 1 + Ion: Au+ + logLambda: nrl + +density: + quasineutrality: true + species-electron: + m: 2.0 + n: {basis: uniform, baseline: 1.0} + T: {basis: uniform, baseline: 1.0} + +grid: + # r0=17 um and lambda_mfp=0.34 um in Joglekar et al. PRL 112, 105004. + # dx=13.4 lambda_mfp and dy=3.125 lambda_mfp match the published central mesh. + xmin: -510um + xmax: 510um + nx: 224 + ymin: -34um + ymax: 34um + ny: 64 + tmin: 0ps + tmax: 60ps + # The linear RK4 bound permits a larger step, but the coupled kinetic-Ohm + # system needs this margin for velocity-space acceleration and positivity. + # max(k v dt)=0.526, versus the RK4 imaginary-axis limit 2 sqrt(2). + dt: 2fs + nv: 96 + vmax: 8.0 + lmax: 2 + mmax: 2 + sharding: {enabled: true, axis: x} + +terms: + # Suppress the grid-scale pseudospectral cascade without touching velocity + # space. This is the same high-order Hou-Li regularization used elsewhere in + # ADEPT; resolved spot and transport scales are effectively unchanged. + # The sharded path maps x filtering to an eighth-difference halo stencil. + hou_li_filter: {is_on: true, alpha: 36.0, order: 36, dimensions: [x, y]} + field_solver: + mode: kinetic-ohm + hidden_density_gradient: + active: true + scale_length: 17um + switch_off: 17.78ps # 800 tau_n in the Letter + switch_width: 0.5ps + # Joglekar et al. prescribe this same two-spot envelope for d(n)/dz. + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + # The published source is centered at the two y boundaries. + y_centers: [-34um, 34um] + y_radius: 17um + fokker_planck: + active: true + flm: {ee: true} + # Chang-Cooper is positivity preserving and is the safer choice once IB + # heating drives f00 far from Maxwellian. The 2 fs Strang step applies + # collisions in 1 fs half-steps, about 0.045 tau_n for this benchmark. + f00: {model: CoulombianKernel, scheme: chang_cooper, positivity: conservative} + +drivers: + ib: + intensity_1e15_Wcm2: 0.25 + polarisation: linear + # Heating remains on after the imposed hidden density gradient is removed. + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + y_centers: [-34um, 34um] + y_radius: 17um + +save: + t: {tmin: 0ps, tmax: 60ps, nt: 13} + +output: + n_panels: 9 diff --git a/configs/vfp-2d/joglekar-2014-prl.yaml b/configs/vfp-2d/joglekar-2014-prl.yaml new file mode 100644 index 00000000..e72a1b46 --- /dev/null +++ b/configs/vfp-2d/joglekar-2014-prl.yaml @@ -0,0 +1,83 @@ +# Reduced, runnable heating/moment benchmark derived from Joglekar et al. PRL 112, +# 105004 (2014). This is not yet the fully implicit PRL reproduction: +# `kinetic-ohm` removes light/plasma waves using Eq. (2) plus a current-moment +# projection, while the original IMPACTA calculation solved the kinetic current +# response fully implicitly. +solver: vfp-2d + +mlflow: + experiment: vfp2d-joglekar-2014 + run: two-spot-ib-smoke + +units: + laser_wavelength: 351nm + reference electron temperature: 1600eV + reference ion temperature: 160eV + reference electron density: 2.5e22/cm^3 + # The five-page Letter does not report ionization. Keep this explicit and + # replace it from the original input deck before a reference comparison. + Z: 1 + Ion: Au+ + logLambda: nrl + +density: + quasineutrality: true + species-electron: + m: 2.0 + n: {basis: uniform, baseline: 1.0} + T: {basis: uniform, baseline: 1.0} + +grid: + # Central reduced box. Published lambda_mfp=0.34 um and r0=50 lambda_mfp=17 um. + xmin: -34um + xmax: 34um + nx: 32 + ymin: -68um + ymax: 68um + ny: 32 + tmin: 0fs + tmax: 0.02fs + dt: 0.01fs + nv: 96 + vmax: 8.0 + lmax: 2 + mmax: 2 + +terms: + hou_li_filter: {is_on: true, alpha: 36.0, order: 36, dimensions: [x, y]} + field_solver: + mode: kinetic-ohm + hidden_density_gradient: + active: true + scale_length: 17um # Ln=50 lambda_mfp + switch_off: 17.78ps # 800 tau_n in the Letter + switch_width: 0.5ps + # Joglekar et al. prescribe this same two-spot envelope for d(n)/dz. + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + y_centers: [-25.5um, 25.5um] + y_radius: 17um + fokker_planck: + active: true + flm: {ee: true} + # Chang-Cooper remains robust as IB heating drives f00 away from Maxwellian. + f00: {model: CoulombianKernel, scheme: chang_cooper, positivity: conservative} + +drivers: + ib: + intensity_1e15_Wcm2: 0.25 + polarisation: linear + switch_off: 17.78ps # 800 tau_n in the Letter + switch_width: 0.5ps + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + # Midplane heating is 2 exp[-(1.5)^2] = 0.21 of one spot peak. + y_centers: [-25.5um, 25.5um] + y_radius: 17um + +save: + t: {tmin: 0fs, tmax: 0.02fs, nt: 3} diff --git a/configs/vfp-2d/landau-damping.yaml b/configs/vfp-2d/landau-damping.yaml new file mode 100644 index 00000000..7d7a462b --- /dev/null +++ b/configs/vfp-2d/landau-damping.yaml @@ -0,0 +1,67 @@ +solver: vfp-2d + +mlflow: + experiment: vfp2d + run: landau-damping + +units: + laser_wavelength: 351nm + reference electron temperature: 3000eV + reference ion temperature: 300eV + reference electron density: 2.275e21/cm^3 + Z: 6 + Ion: Au+ + logLambda: nrl + +density: + # A fixed uniform ion background leaves the electron perturbation charged, + # so the initial spectral Poisson solve seeds the electrostatic field. + quasineutrality: false + species-electron: + m: 2.0 + n: + x: + basis: cosine + baseline: 1.0 + amplitude: 1.0e-4 + wavelength: 20um + y: + basis: cosine + baseline: 1.0 + amplitude: 1.0e-4 + wavelength: 20um + T: + basis: uniform + baseline: 1.0 + +grid: + xmin: 0um + xmax: 20um + nx: 32 + ymin: 0um + ymax: 20um + ny: 32 + tmin: 0fs + tmax: 100fs + dt: 0.05fs + nv: 96 + vmax: 8.0 + lmax: 5 + mmax: 5 + +terms: + fokker_planck: + active: false + flm: + ee: true + f00: + model: CoulombianKernel + scheme: central + +drivers: {} + +save: + t: + tmin: 0fs + tmax: 100fs + nt: 101 diff --git a/docs/RUNNING_A_SIM.md b/docs/RUNNING_A_SIM.md index 51974501..1e6724a1 100644 --- a/docs/RUNNING_A_SIM.md +++ b/docs/RUNNING_A_SIM.md @@ -12,6 +12,7 @@ uv run run.py --cfg path_to_my_config - [Vlasov-1D2V](source/solvers/vlasov1d2v/config.md) - 1D2V Vlasov-Poisson-Fokker-Planck in cylindrical velocity space, with a full-geometry Coulomb collision operator - [Vlasov-2D](source/solvers/vlasov2d/config.md) - 2D2V Vlasov-Maxwell solver - [VFP-1D](source/solvers/vfp1d/config.md) - Vlasov-Fokker-Planck electron transport solver +- [VFP-2D](source/solvers/vfp2d/config.md) - 2D3P arbitrary-spherical-harmonic Vlasov-Maxwell-Fokker-Planck solver - [LPSE-2D (Envelope-2D)](source/solvers/lpse2d/config.md) - 2D laser-plasma envelope solver for TPD/SRS - [Spectrax-1D](source/solvers/spectrax1d/config.md) - 1D Hermite-Fourier Vlasov-Maxwell solver - [Hermite-Legendre-1D](source/solvers/hermite_legendre_1d/config.md) - 1D-1V mixed Hermite-Legendre electrostatic Vlasov-Poisson solver diff --git a/docs/source/index.rst b/docs/source/index.rst index b65d10f3..47ed35a7 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -32,6 +32,8 @@ Documentation solvers/vlasov1d2v/overview solvers/vlasov2d/overview solvers/vfp1d/overview + solvers/vfp2d/overview + solvers/vfp2d/joglekar2014 solvers/spectrax1d/overview solvers/hermite_legendre_1d/overview solvers/pic1d/overview @@ -51,6 +53,7 @@ Documentation solvers/hermite_legendre_1d/config solvers/pic1d/config solvers/lpse2d/config + solvers/vfp2d/config solvers/tf1d/config solvers/osiris/config diff --git a/docs/source/solvers.md b/docs/source/solvers.md index 1cd573dd..d28dcff0 100644 --- a/docs/source/solvers.md +++ b/docs/source/solvers.md @@ -10,6 +10,7 @@ at the top of a configuration file selects which one runs. | `vlasov-1d2v` | [Vlasov-1D2V](solvers/vlasov1d2v/overview.md) | 1D2V in cylindrical velocity space with a full-geometry Coulomb collision operator | | `vlasov-2d` | [Vlasov-2D](solvers/vlasov2d/overview.md) | 2D2V Vlasov-Maxwell | | `vfp-1d` | [VFP-1D](solvers/vfp1d/overview.md) | Vlasov-Fokker-Planck electron transport | +| `vfp-2d` | [VFP-2D](solvers/vfp2d/overview.md) | 2D3P arbitrary-harmonic Vlasov-Maxwell-Fokker-Planck | | `envelope-2d` | [LPSE-2D](solvers/lpse2d/overview.md) | 2D laser-plasma envelope solver | | `spectrax-1d` | [Spectrax-1D](solvers/spectrax1d/overview.md) | Hermite-Fourier Vlasov-Maxwell | | `hermite-epw-1d` | [Spectrax-1D](solvers/spectrax1d/overview.md) | Spectrax-1D with electron plasma wave diagnostics | @@ -57,6 +58,13 @@ Vlasov-Fokker-Planck solver for electron transport over collisional time-scales. - [Overview & Equations](solvers/vfp1d/overview.md) - [Configuration Reference](solvers/vfp1d/config.md) +### [VFP-2D](solvers/vfp2d/overview.md) + +2D3P Vlasov-Maxwell-Fokker-Planck solver with arbitrary complex spherical harmonics, packed JAX-native storage, and full Tzoufras/KALOS angular couplings. + +- [Overview & Equations](solvers/vfp2d/overview.md) +- [Configuration Reference](solvers/vfp2d/config.md) + ## Spectral Solvers These represent velocity space with a spectral basis rather than a grid, so the state is a set of diff --git a/docs/source/solvers/vfp1d/config.md b/docs/source/solvers/vfp1d/config.md index 35b68cb6..8f44fdac 100644 --- a/docs/source/solvers/vfp1d/config.md +++ b/docs/source/solvers/vfp1d/config.md @@ -272,7 +272,7 @@ drivers: ### drivers.ib -Inverse Bremsstrahlung (IB) laser heating. Augments the Fokker-Planck diffusion coefficient $D$ by $v_\text{osc}^2 g(v) / (6v)$ where $g(v) = [1 + (Z^2 n_i / (\omega_0 v^3))^2]^{-1}$ (Ridgers eq 4.39). Drives the distribution toward a Langdon/super-Gaussian shape. +Inverse Bremsstrahlung (IB) laser heating. Augments the Fokker-Planck diffusion coefficient $D$ by $v_\text{osc}^2 g(v) / (6v)$ where $g(v) = [1 + (\nu_{ei}(v) / \omega_0)^2]^{-1}$ and $\nu_{ei}(v)=\nu_{ee,0}(\ln\Lambda_{ei}/\ln\Lambda_{ee})Z^2n_i/v^3$ in normalized units (Ridgers eq 4.39). Drives the distribution toward a Langdon/super-Gaussian shape. | Field | Type | Default | Description | |-------|------|---------|-------------| diff --git a/docs/source/solvers/vfp2d/config.md b/docs/source/solvers/vfp2d/config.md new file mode 100644 index 00000000..9bc75564 --- /dev/null +++ b/docs/source/solvers/vfp2d/config.md @@ -0,0 +1,154 @@ +# VFP-2D Configuration + +Set `solver: vfp-2d`. + +## Grid + +```yaml +grid: + xmin: 0um + xmax: 20um + nx: 32 + ymin: 0um + ymax: 20um + ny: 32 + tmin: 0fs + tmax: 100fs + dt: 0.05fs + nv: 96 + vmax: 8.0 + lmax: 5 + mmax: 3 + relativistic: false + sharding: {enabled: false, axis: x} +``` + +`lmax` is the highest retained $\ell$. `mmax` defaults to `lmax`; lowering it provides a controlled transverse-angular truncation. For compatibility, `nl` is accepted as an alias for `lmax`. + +By default, `vmax` is expressed in the same number-of-thermal-speeds convention as VFP-1D. Set `vmax_is_normalized: true` to provide the radial coordinate directly in code units. In relativistic mode this direct coordinate is $p/(m_ec)$. + +Set `sharding.enabled: true` to partition the state, spatial drivers, and collision batches +along $x$ over every visible JAX device. `nx` must be divisible by the device count. The +sharded path uses fourth-order periodic finite differences (with two-cell halo exchange) +for spatial derivatives, because a global Fourier transform along a partitioned axis would +replicate the dominant distribution array. Saved snapshots are replicated only when they +are written. On this path, an requested $x$ Hou--Li filter is implemented as a shard-local +eighth-difference Nyquist filter with four-cell halo exchange; $y$ retains the spectral +Hou--Li filter. + +## Initial distribution + +VFP-1D `species-*` components are accepted. A profile without an axis is applied along $x$. Separable 2D profiles use `x` and `y` children: + +```yaml +density: + quasineutrality: false + species-electron: + m: 2.0 + n: + x: {basis: cosine, baseline: 1.0, amplitude: 1.0e-4, wavelength: 20um} + y: {basis: cosine, baseline: 1.0, amplitude: 1.0e-4, wavelength: 20um} + T: {basis: uniform, baseline: 1.0} +``` + +Supported analytic bases are `uniform`, `sine`, `cosine`, and `tanh`; file profiles retain the VFP-1D loader behavior. With `quasineutrality: true`, the stationary ion charge follows the initial electron density. With `false`, it is spatially uniform at the mean density and the initial Poisson solve produces the field associated with electron-density perturbations. + +## Laser heating + +VFP2D shares the conservative inverse-bremsstrahlung and Maxwellian heating operators with +VFP1D. Heating amplitudes may be spatially uniform or multiplied by a two-dimensional +profile. The two-spot profile used by the Joglekar benchmark is: + +```yaml +drivers: + ib: + intensity_1e15_Wcm2: 0.25 + polarisation: linear + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + y_centers: [-8.5um, 8.5um] + y_radius: 17um +``` + +`gaussian_spots` evaluates +$A\exp[-((x-x_0)/r_x)^2]\sum_i\exp[-((y-y_i)/r_y)^2]$. +`maxwellian_heating` accepts the same optional `profile` child with a scalar `D0`. + +## Collisions + +```yaml +terms: + fokker_planck: + active: true + flm: + ee: true + f00: + model: CoulombianKernel + scheme: chang_cooper +``` + +`flm.ee: true` uses the full linearized anisotropic electron-electron terms. `false` uses the Epperlein-Haines $Z_*$ approximation. The `f00` model and differencing choices are shared with VFP-1D. +`chang_cooper` is recommended when heating produces a strongly non-Maxwellian distribution +because it is positivity preserving. `log_mean` has a zero semidiscrete spherical-energy +derivative for the kernel model at the frozen distribution, but a finite implicit update is +not exactly energy conserving away from a Maxwellian. Collision-step convergence is required; +keep the collision half-step well below the shortest relevant collision time. + +## Long-timescale kinetic Ohm mode + +`maxwell` is the default field solver. For collisional transport times, `kinetic-ohm` suppresses +displacement current and electron plasma oscillations, evaluates the full Joglekar Eq. (2), +and projects the current moment onto quasistatic Ampere's law: + +```yaml +terms: + field_solver: + mode: kinetic-ohm + hidden_density_gradient: + active: true + scale_length: 17um + switch_off: 17.78ps + profile: + basis: gaussian_spots + x_radius: 17um + y_centers: [-8.5um, 8.5um] + y_radius: 17um +``` + +The optional hidden gradient is the unresolved $\partial_z n$ used by the 2.5D PRL geometry. +It enters the pressure-gradient Ohm residual and can be switched sharply (`switch_width` +omitted) or with a differentiable tanh gate (`switch_width` set). Output variables prefixed +with `ohm_` contain the resistive, Hall, Nernst, scalar-pressure, and $f_2$ tensor-pressure +contributions. + +The reconnection diagnostics report a normalized rate and flux only when the upstream +$B_x$ fields are antiparallel and balanced and the origin contains both an in-plane null/ +$A_z$ saddle and a central current sheet. Invalid samples are stored as NaN rather than +turning Biermann-ring motion into a false reconnection rate. The normalized rate is also +suppressed when the signed inward Nernst speed is below 10% of its maximum over the saved +history, where division by a vanishing inflow would otherwise create a spurious spike. +`bz_quadrupole_purity` is the local L1 projection of $B_z$ onto the expected four-lobe +`sign(x*y)` pattern. Full-domain facets use a free display aspect ratio so wide boxes remain +readable. The `plots/reconnection_region` artifact directory also contains a curated set of +moments and Ohm-law terms cropped to the central three y-half-widths in x, excluding the +distant x boundaries and the outermost 25% of the y domain. This isolates the X-point and +current-sheet neighborhood from the outflow extent and laser-source boundary cells. + +## Saving + +```yaml +save: + t: + tmin: 0fs + tmax: 100fs + nt: 101 +``` + +Post-processing returns an xarray dataset with `flm_real`, `flm_imag`, `e`, `b`, density, +temperature, current, Nernst velocity, and the traceless pressure-anisotropy moment. Harmonics +are labeled by the `ell` and `m` coordinates. + +See the [Joglekar 2014 reconstruction and hydro-coupling design](joglekar2014.md) for the +distinction between the runnable reduced benchmark and the planned long-time implicit solve. diff --git a/docs/source/solvers/vfp2d/joglekar2014.md b/docs/source/solvers/vfp2d/joglekar2014.md new file mode 100644 index 00000000..f5ddb3e3 --- /dev/null +++ b/docs/source/solvers/vfp2d/joglekar2014.md @@ -0,0 +1,155 @@ +# Joglekar 2014 reconstruction and hydro coupling + +This page is both a benchmark specification and an implementation boundary. The target is +A. S. Joglekar *et al.*, *Physical Review Letters* **112**, 105004 (2014), +[doi:10.1103/PhysRevLett.112.105004](https://doi.org/10.1103/PhysRevLett.112.105004). +The reduced heating configuration in `configs/vfp-2d/joglekar-2014-prl.yaml` exercises the +parts that are implemented today. It is not yet the fully implicit PRL reproduction. + +## What the benchmark requires + +The paper retained the Cartesian-tensor equivalent of all spherical harmonics through +$\ell=2$. Its generalized kinetic Ohm law was + +$$ +\mathbf E=\bar\eta\mathbf j+\frac{\mathbf j\times\mathbf B}{e n_e} +-\mathbf v_T\times\mathbf B +-\frac{\nabla(n_em_e\langle v^5\rangle)}{6en_e\langle v^3\rangle} +-\frac{\nabla\cdot(n_em_e\langle\mathbf{vv}v^3\rangle)}{2en_e\langle v^3\rangle}, +$$ + +with + +$$ +\mathbf v_T=\frac{\langle\mathbf v v^3\rangle}{2\langle v^3\rangle} ++\frac{\mathbf j}{e n_e}. +$$ + +The last term is an $f_2$ pressure-anisotropy contribution and is essential at the X point; +an $f_0+f_1$ model is therefore insufficient. VFP2D now emits the exact scalar, vector, and +traceless tensor moments in these equations, together with $\mathbf j$, $T_e$, and +$\mathbf v_T$. + +The published setup used: + +- $T_{e0}=1.6$ keV and $n_{e0}=2.5\times10^{22}\,\mathrm{cm}^{-3}$; +- $v_{th}/c=0.08$, $\omega_{pe}\tau_n=125$, and $\lambda_{mfp}=0.34\,\mu$m; +- $-1500