From f2449d2af43331109bf0ccc824695b6a8fc0b11a Mon Sep 17 00:00:00 2001 From: archis Date: Tue, 18 Aug 2026 20:54:21 -0700 Subject: [PATCH 01/27] Add arbitrary-harmonic 2D3P VFP solver --- adept/__init__.py | 2 +- adept/_base_.py | 3 + adept/vfp1d/fokker_planck.py | 72 ++- adept/vfp2d/__init__.py | 51 +++ adept/vfp2d/base.py | 505 ++++++++++++++++++++++ adept/vfp2d/collisions.py | 95 ++++ adept/vfp2d/grid.py | 79 ++++ adept/vfp2d/harmonics.py | 427 ++++++++++++++++++ adept/vfp2d/ohm.py | 146 +++++++ adept/vfp2d/vector_field.py | 260 +++++++++++ configs/vfp-2d/joglekar-2014-prl.yaml | 78 ++++ configs/vfp-2d/landau-damping.yaml | 67 +++ docs/RUNNING_A_SIM.md | 1 + docs/source/index.rst | 3 + docs/source/solvers.md | 8 + docs/source/solvers/vfp2d/config.md | 126 ++++++ docs/source/solvers/vfp2d/joglekar2014.md | 138 ++++++ docs/source/solvers/vfp2d/overview.md | 56 +++ tests/test_vfp2d/test_base.py | 164 +++++++ tests/test_vfp2d/test_harmonics.py | 318 ++++++++++++++ 20 files changed, 2575 insertions(+), 24 deletions(-) create mode 100644 adept/vfp2d/__init__.py create mode 100644 adept/vfp2d/base.py create mode 100644 adept/vfp2d/collisions.py create mode 100644 adept/vfp2d/grid.py create mode 100644 adept/vfp2d/harmonics.py create mode 100644 adept/vfp2d/ohm.py create mode 100644 adept/vfp2d/vector_field.py create mode 100644 configs/vfp-2d/joglekar-2014-prl.yaml create mode 100644 configs/vfp-2d/landau-damping.yaml create mode 100644 docs/source/solvers/vfp2d/config.md create mode 100644 docs/source/solvers/vfp2d/joglekar2014.md create mode 100644 docs/source/solvers/vfp2d/overview.md create mode 100644 tests/test_vfp2d/test_base.py create mode 100644 tests/test_vfp2d/test_harmonics.py 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/vfp1d/fokker_planck.py b/adept/vfp1d/fokker_planck.py index 910a0975..7ec8e7bf 100644 --- a/adept/vfp1d/fokker_planck.py +++ b/adept/vfp1d/fokker_planck.py @@ -352,7 +352,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 +489,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) + 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 = lx.linear_solve(op, jnp.real(f10), solver=lx.AutoLinearSolver(well_posed=True)).value + imag = lx.linear_solve(op, jnp.imag(f10), solver=lx.AutoLinearSolver(well_posed=True)).value + return real + 1j * imag 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): + 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 +515,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/vfp2d/__init__.py b/adept/vfp2d/__init__.py new file mode 100644 index 00000000..0e9ec143 --- /dev/null +++ b/adept/vfp2d/__init__.py @@ -0,0 +1,51 @@ +"""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.ohm import KineticOhm2D, project_current_moment +from adept.vfp2d.harmonics import ( + HarmonicLayout, + TzoufrasVlasov, + complex_to_real, + cartesian_l2, + current, + density, + nernst_velocity, + real_to_complex, + scalar_velocity_moment, + tensor_velocity_moment, + vector_velocity_moment, +) +from adept.vfp2d.vector_field import ( + KineticOhmStep, + Maxwell2D, + SpectralPoisson2D, + SplitStepVFP2D, + VlasovMaxwell, +) + +__all__ = [ + "Grid", + "BaseVFP2D", + "HarmonicLayout", + "AnisotropicCollisions", + "CollisionStep", + "Maxwell2D", + "KineticOhm2D", + "KineticOhmStep", + "SpectralPoisson2D", + "SplitStepVFP2D", + "TzoufrasVlasov", + "VlasovMaxwell", + "current", + "cartesian_l2", + "complex_to_real", + "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..3cab2a07 --- /dev/null +++ b/adept/vfp2d/base.py @@ -0,0 +1,505 @@ +"""ADEPTModule entry point for the arbitrary-harmonic VFP-2D solver.""" + +from __future__ import annotations + +from dataclasses import asdict + +import jax.numpy as jnp +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, +) +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.grid import Grid +from adept.vfp2d.harmonics import ( + HarmonicLayout, + 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.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 + + 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 + ) + self.args["ib_Z2ni_w0"] = ( + self.args["Z"] ** 2 * self.args["ni"] / self.cfg["units"]["derived"]["w0_norm"] + ) + + 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)) + + def init_diffeqsolve(self): + 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 + vlasov = TzoufrasVlasov( + self.layout, + self.grid.v, + self.grid.dv, + self.grid.kx, + self.grid.ky, + streaming_speed=streaming_speed, + ) + maxwell = Maxwell2D(self.grid.kx, self.grid.ky, c=self.plasma_norm.speed_of_light_norm()) + 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, + ) + step = KineticOhmStep( + vlasov, + maxwell, + self._kinetic_ohm, + self.layout, + self.grid.v, + self.grid.dv, + self.grid.dt, + collisions=collisions, + real_storage=True, + ) + 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} + self.diffeqsolve_quants = {"terms": ODETerm(step), "solver": Stepper(), "saveat": SaveAt(ts=self.save_times)} + + def __call__(self, trainable_modules: dict | None, args: dict | None): + return { + "solver result": 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"], + ) + } + + 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)): + target_current = self._maxwell.c2 * self._maxwell.curl(result.ys["b"][index]) + hidden_dndz = KineticOhmStep._hidden_dndz( + float(time), self.args, result.ys["b"][index, ..., 0] + ) + _electric, terms = self._kinetic_ohm( + flm_jax[index], + result.ys["b"][index], + 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)"}, + ) + return {"vfp2d": ds} diff --git a/adept/vfp2d/collisions.py b/adept/vfp2d/collisions.py new file mode 100644 index 00000000..e9bf2178 --- /dev/null +++ b/adept/vfp2d/collisions.py @@ -0,0 +1,95 @@ +"""Collision adapters for packed arbitrary-``f_lm`` VFP-2D states.""" + +from __future__ import annotations + +import jax.numpy as jnp +from jax import Array + +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, + ): + self.layout = layout + self.isotropic = isotropic + self.anisotropic = anisotropic + + 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: + 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 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..42387d79 --- /dev/null +++ b/adept/vfp2d/harmonics.py @@ -0,0 +1,427 @@ +"""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 + + +@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) + + +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, + ): + 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) + + 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) -> Array: + """Spatial-advection contribution, Eqs. (17)--(19), with ``d/dz = 0``.""" + + dfdx = _spectral_derivative(f, self.kx, axis=0) + dfdy = _spectral_derivative(f, self.ky, axis=1) + 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) * dfdy[..., lm, :] + if lp >= 0: + value += ( + 0.5 + * v + * (ell - m) + * (ell - m - 1) + / (2 * ell - 1) + * dfdy[..., lp, :] + ) + if um >= 0: + value += 0.5 * v / (2 * ell + 3) * dfdy[..., um, :] + if up >= 0: + value -= ( + 0.5 + * v + * (ell + m + 1) + * (ell + m + 2) + / (2 * ell + 3) + * dfdy[..., 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) * dfdy[..., lower1, :] + if upper1 >= 0: + transverse += (ell + 1) * (ell + 2) / (2 * ell + 3) * dfdy[..., 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) -> Array: + result = self.streaming(f) + 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 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. + """ + + shape = (*f.shape[:-2], 3, 3, f.shape[-1]) + result = jnp.zeros(shape, dtype=jnp.real(f).dtype) + i20, i21, i22 = (layout.index(2, m) for m in range(3)) + if i20 < 0: + return result + 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) + result = result.at[..., 0, 0, :].set(fxx) + result = result.at[..., 0, 1, :].set(fxy) + result = result.at[..., 1, 0, :].set(fxy) + result = result.at[..., 0, 2, :].set(fxz) + result = result.at[..., 2, 0, :].set(fxz) + result = result.at[..., 1, 1, :].set(fyy) + result = result.at[..., 1, 2, :].set(fyz) + result = result.at[..., 2, 1, :].set(fyz) + return result.at[..., 2, 2, :].set(fzz) + + +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..1ba90830 --- /dev/null +++ b/adept/vfp2d/ohm.py @@ -0,0 +1,146 @@ +"""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 adept.vfp2d.harmonics import ( + HarmonicLayout, + current, + density, + nernst_velocity, + 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, + ): + 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) + + def ddx(self, value: Array) -> Array: + 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: + 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. Moments other than density are assumed not to + vary in z, exactly matching the 2.5D source construction in the Letter. + """ + + 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/vector_field.py b/adept/vfp2d/vector_field.py new file mode 100644 index 00000000..65c6c9ea --- /dev/null +++ b/adept/vfp2d/vector_field.py @@ -0,0 +1,260 @@ +"""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 adept.vfp2d.collisions import CollisionStep +from adept.vfp2d.harmonics import HarmonicLayout, TzoufrasVlasov, complex_to_real, current, real_to_complex +from adept.vfp2d.ohm import KineticOhm2D, project_current_moment + + +class Maxwell2D: + """Full three-component Maxwell curl operator with ``d/dz = 0``.""" + + def __init__(self, kx: Array, ky: Array, c: float): + self.kx = jnp.asarray(kx) + self.ky = jnp.asarray(ky) + self.c2 = float(c) ** 2 + + def ddx(self, a: Array) -> Array: + return jnp.fft.ifft(1j * self.kx[:, None] * jnp.fft.fft(a, axis=0), axis=0).real + + def ddy(self, a: Array) -> Array: + 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, 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] + 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(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(result, args, 0.5 * self.dt) + + +class KineticOhmStep: + """Long-timescale midpoint 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, + ): + 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) + + def _collide(self, 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] + 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, + ) + return ( + self.vlasov(flm, electric_field, magnetic_field), + -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._collide(flm, args, 0.5 * self.dt) + flm = self._project(flm, magnetic_field) + + df1, db1, _electric1 = self._rates(t, flm, magnetic_field, args) + midpoint_f = flm + 0.5 * self.dt * df1 + midpoint_b = magnetic_field + 0.5 * self.dt * db1 + midpoint_f = self._project(midpoint_f, midpoint_b) + df2, db2, _electric2 = self._rates(t + 0.5 * self.dt, midpoint_f, midpoint_b, args) + + result_f = self._project(flm + self.dt * df2, magnetic_field + self.dt * db2) + result_f = self._collide(result_f, args, 0.5 * self.dt) + result_b = magnetic_field + self.dt * db2 + 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.yaml b/configs/vfp-2d/joglekar-2014-prl.yaml new file mode 100644 index 00000000..5bf5036b --- /dev/null +++ b/configs/vfp-2d/joglekar-2014-prl.yaml @@ -0,0 +1,78 @@ +# 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: -34um + ymax: 34um + ny: 32 + tmin: 0fs + tmax: 0.02fs + dt: 0.01fs + nv: 96 + vmax: 8.0 + lmax: 2 + mmax: 2 + +terms: + field_solver: + mode: kinetic-ohm + hidden_density_gradient: + active: true + scale_length: 17um # Ln=50 lambda_mfp + switch_off: 17.78ps # 800 tau_n, using 27000 tau_n = 0.6 ns + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + # Placeholder separation; replace from the original input deck. + y_centers: [-8.5um, 8.5um] + y_radius: 17um + fokker_planck: + active: true + flm: {ee: true} + f00: {model: CoulombianKernel, scheme: central} + +drivers: + ib: + intensity_1e15_Wcm2: 0.25 + polarisation: linear + profile: + basis: gaussian_spots + x_center: 0um + x_radius: 17um + # Placeholder separation: the Letter prints y_max but not its value. + y_centers: [-8.5um, 8.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/vfp2d/config.md b/docs/source/solvers/vfp2d/config.md new file mode 100644 index 00000000..9a488f52 --- /dev/null +++ b/docs/source/solvers/vfp2d/config.md @@ -0,0 +1,126 @@ +# 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 +``` + +`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)$. + +## 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: central +``` + +`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. + +## 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. + +## 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..17bafac4 --- /dev/null +++ b/docs/source/solvers/vfp2d/joglekar2014.md @@ -0,0 +1,138 @@ +# 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