drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/mesh.py

@@ -0,0 +1,238 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from ..config.boutinp import BoutConfig
+from ..runtime import runtime_jax_dtype
+from ..runtime.run_config import RunConfiguration
+
+
+@dataclass(frozen=True)
+class StructuredMesh:
+    nx: int
+    ny: int
+    nz: int
+    mxg: int
+    myg: int
+    symmetric_global_x: bool
+    symmetric_global_y: bool
+    jyseps1_1: int
+    jyseps2_1: int
+    jyseps1_2: int
+    jyseps2_2: int
+    ny_inner: int
+    has_lower_y_target: bool
+    has_upper_y_target: bool
+    x: jnp.ndarray
+    y: jnp.ndarray
+    z: jnp.ndarray
+
+    @property
+    def xstart(self) -> int:
+        return self.mxg
+
+    @property
+    def xend(self) -> int:
+        return self.nx - self.mxg - 1
+
+    @property
+    def ystart(self) -> int:
+        return self.myg
+
+    @property
+    def yend(self) -> int:
+        return self.myg + self.ny - 1
+
+    @property
+    def local_ny(self) -> int:
+        return self.ny + 2 * self.myg
+
+    def expression_context(self, *, time: float = 0.0) -> dict[str, jnp.ndarray]:
+        dtype = runtime_jax_dtype()
+        x3 = self.x[:, None, None]
+        y3 = self.y[None, :, None]
+        z3 = self.z[None, None, :]
+        return {
+            "x": x3,
+            "y": 2.0 * jnp.pi * y3,
+            "z": 2.0 * jnp.pi * z3,
+            "t": jnp.array(time, dtype=dtype),
+        }
+
+
+def build_structured_mesh(config: BoutConfig, run_config: RunConfiguration) -> StructuredMesh:
+    if run_config.mesh.nx is None or run_config.mesh.ny is None or run_config.mesh.nz is None:
+        raise ValueError("Structured native execution requires explicit mesh nx, ny, and nz.")
+
+    nx = run_config.mesh.nx
+    ny = run_config.mesh.ny
+    nz = run_config.mesh.nz
+    mxg = run_config.mesh.mxg
+    myg = run_config.mesh.myg
+    symmetric_global_x = _mesh_bool(config, "symmetricGlobalX", default=True)
+    symmetric_global_y = _mesh_bool(config, "symmetricGlobalY", default=True)
+    jyseps1_1 = _mesh_int(config, "jyseps1_1", default=-1)
+    jyseps2_1 = _mesh_int(config, "jyseps2_1", default=ny // 2)
+    jyseps1_2 = _mesh_int(config, "jyseps1_2", default=ny // 2)
+    jyseps2_2 = _mesh_int(config, "jyseps2_2", default=ny - 1)
+    ny_inner = _mesh_int(config, "ny_inner", default=jyseps2_1)
+
+    x = _global_x_coordinates(nx=nx, mxg=mxg, symmetric=symmetric_global_x)
+    y = _global_y_coordinates(
+        ny=ny,
+        myg=myg,
+        symmetric=symmetric_global_y,
+        jyseps1_1=jyseps1_1,
+        jyseps2_1=jyseps2_1,
+        jyseps1_2=jyseps1_2,
+        jyseps2_2=jyseps2_2,
+        ny_inner=ny_inner,
+    )
+    dtype = runtime_jax_dtype()
+    z = jnp.arange(nz, dtype=dtype) / float(nz)
+    has_explicit_y_boundaries = config.has_option("mesh", "ixseps1") or config.has_option("mesh", "ixseps2")
+    has_lower_y_target = False
+    has_upper_y_target = bool(has_explicit_y_boundaries)
+    return StructuredMesh(
+        nx=nx,
+        ny=ny,
+        nz=nz,
+        mxg=mxg,
+        myg=myg,
+        symmetric_global_x=symmetric_global_x,
+        symmetric_global_y=symmetric_global_y,
+        jyseps1_1=jyseps1_1,
+        jyseps2_1=jyseps2_1,
+        jyseps1_2=jyseps1_2,
+        jyseps2_2=jyseps2_2,
+        ny_inner=ny_inner,
+        has_lower_y_target=has_lower_y_target,
+        has_upper_y_target=has_upper_y_target,
+        x=x,
+        y=y,
+        z=z,
+    )
+
+
+def broadcast_to_field_shape(value: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    array = jnp.asarray(value, dtype=dtype)
+    target_shape = (mesh.nx, mesh.local_ny, mesh.nz)
+    if array.shape == target_shape:
+        return array
+    return jnp.broadcast_to(array, target_shape)
+
+
+def apply_zero_dirichlet_x_guards(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    result = jnp.asarray(field, dtype=dtype)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    if mesh.mxg <= 0:
+        return result
+
+    result = result.at[mesh.xstart - 1, y_slice, :].set(-result[mesh.xstart, y_slice, :])
+    result = result.at[mesh.xend + 1, y_slice, :].set(-result[mesh.xend, y_slice, :])
+    for offset in range(2, mesh.mxg + 1):
+        result = result.at[mesh.xstart - offset, y_slice, :].set(0.0)
+        result = result.at[mesh.xend + offset, y_slice, :].set(0.0)
+    return result
+
+
+def apply_neumann_x_guards(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    result = jnp.asarray(field, dtype=dtype)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    if mesh.mxg <= 0:
+        return result
+
+    for offset in range(1, mesh.mxg + 1):
+        result = result.at[mesh.xstart - offset, y_slice, :].set(result[mesh.xstart - 1 + offset, y_slice, :])
+        result = result.at[mesh.xend + offset, y_slice, :].set(result[mesh.xend + 1 - offset, y_slice, :])
+    return result
+
+
+def communicate_y_guards(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    result = jnp.asarray(field, dtype=dtype)
+    for offset in range(mesh.myg):
+        result = result.at[:, mesh.ystart - 1 - offset, :].set(result[:, mesh.ystart + offset, :])
+        result = result.at[:, mesh.yend + 1 + offset, :].set(result[:, mesh.yend - offset, :])
+    return result
+
+
+def project_nonnegative_x_boundaries(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    result = jnp.asarray(field, dtype=dtype)
+    y_slice = slice(mesh.ystart, mesh.yend + 1)
+    if mesh.mxg > 0:
+        left = result[:mesh.xstart, y_slice, :]
+        right = result[mesh.xend + 1 :, y_slice, :]
+        result = result.at[:mesh.xstart, y_slice, :].set(jnp.maximum(left, 0.0))
+        result = result.at[mesh.xend + 1 :, y_slice, :].set(jnp.maximum(right, 0.0))
+    return result
+
+
+def apply_x_boundary(field: jnp.ndarray, mesh: StructuredMesh, boundary_kind: str) -> jnp.ndarray:
+    normalized = boundary_kind.strip().lower()
+    if normalized == "neumann":
+        return apply_neumann_x_guards(field, mesh)
+    if normalized in {"dirichlet", "dirichlet_zero", "zero"}:
+        return apply_zero_dirichlet_x_guards(field, mesh)
+    raise NotImplementedError(f"Unsupported X boundary kind {boundary_kind!r}")
+
+
+def apply_field_boundaries(field: jnp.ndarray, mesh: StructuredMesh, *, x_boundary: str) -> jnp.ndarray:
+    result = apply_x_boundary(field, mesh, x_boundary)
+    result = communicate_y_guards(result, mesh)
+    if x_boundary.strip().lower() in {"dirichlet", "dirichlet_zero", "zero"}:
+        result = project_nonnegative_x_boundaries(result, mesh)
+    return result
+
+
+def _global_x_coordinates(*, nx: int, mxg: int, symmetric: bool) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    mx = nx - 2 * mxg
+    global_indices = jnp.arange(nx, dtype=dtype)
+    if symmetric:
+        return (0.5 + global_indices - (nx - mx) * 0.5) / float(mx)
+    return global_indices / float(mx)
+
+
+def _global_y_coordinates(
+    *,
+    ny: int,
+    myg: int,
+    symmetric: bool,
+    jyseps1_1: int,
+    jyseps2_1: int,
+    jyseps1_2: int,
+    jyseps2_2: int,
+    ny_inner: int,
+) -> jnp.ndarray:
+    dtype = runtime_jax_dtype()
+    local_indices = jnp.arange(ny + 2 * myg, dtype=dtype) - float(myg)
+    nycore = float((jyseps2_1 - jyseps1_1) + (jyseps2_2 - jyseps1_2))
+    if symmetric:
+        before = (local_indices - (jyseps1_1 + 0.5)) / nycore
+        after = (local_indices - (jyseps1_1 + 0.5 + (jyseps1_2 - jyseps2_1))) / nycore
+        return jnp.where(local_indices < float(ny_inner), before, after)
+
+    core_indices = local_indices
+    lower = core_indices - (jyseps1_1 + 1.0)
+    upper = core_indices - (jyseps1_1 + 1.0 + (jyseps1_2 - jyseps2_1))
+    return jnp.where(core_indices <= float(jyseps2_1), lower / nycore, upper / nycore)
+
+
+def _mesh_bool(config: BoutConfig, key: str, *, default: bool) -> bool:
+    if not config.has_option("mesh", key):
+        return default
+    return bool(config.parsed("mesh", key))
+
+
+def _mesh_int(config: BoutConfig, key: str, *, default: int) -> int:
+    if not config.has_option("mesh", key):
+        return default
+    value = config.parsed("mesh", key)
+    return int(round(float(value)))

drbx/native/metrics.py

@@ -0,0 +1,234 @@
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import jax.numpy as jnp
+
+from ..config.boutinp import BoutConfig
+from ..config.model import has_model_section, locate_model_section
+from ..runtime import runtime_jax_dtype
+from ..runtime.run_config import RunConfiguration
+from .expression import ArrayExpressionEvaluator
+from .mesh import StructuredMesh, broadcast_to_field_shape
+from .units import resolved_dataset_scalars
+
+
+@dataclass(frozen=True)
+class StructuredMetrics:
+    dx: jnp.ndarray
+    dy: jnp.ndarray
+    dz: jnp.ndarray
+    J: jnp.ndarray
+    g11: jnp.ndarray
+    g33: jnp.ndarray
+    g22: jnp.ndarray
+    g_22: jnp.ndarray
+    g23: jnp.ndarray
+    Bxy: jnp.ndarray
+    g_23: jnp.ndarray | None = None
+
+
+def build_structured_metrics(
+    config: BoutConfig,
+    run_config: RunConfiguration,
+    mesh: StructuredMesh,
+) -> StructuredMetrics:
+    dtype = runtime_jax_dtype()
+    evaluator = ArrayExpressionEvaluator(config, local_values=mesh.expression_context())
+    scalars = resolved_dataset_scalars(run_config)
+    rho_s0 = scalars["rho_s0"]
+    Bnorm = scalars["Bnorm"]
+    normalize_metric = _normalize_metric_enabled(config)
+
+    raw_dx = _metric_value(config, evaluator, "dx", default=1.0)
+    raw_dy = _metric_value(config, evaluator, "dy", default=1.0)
+    raw_dz = _metric_value(config, evaluator, "dz", default=(2.0 * math.pi) / float(mesh.nz))
+    raw_J = _metric_value(config, evaluator, "J", default=1.0)
+    raw_g11 = _metric_value(config, evaluator, "g11", default=1.0)
+    raw_g33 = _metric_value(config, evaluator, "g33", default=1.0)
+    raw_g22 = _metric_value(config, evaluator, "g22", default=1.0)
+    if config.has_section("mesh") and config.has_option("mesh", "g_22"):
+        raw_g_22 = _metric_value(config, evaluator, "g_22", default=1.0)
+    else:
+        raw_g_22 = 1.0 / jnp.asarray(raw_g22, dtype=dtype)
+    raw_g23 = _metric_value(config, evaluator, "g23", default=0.0)
+    raw_g_23 = _metric_value(config, evaluator, "g_23", default=0.0)
+    raw_Bxy = _metric_value(config, evaluator, "Bxy", default=1.0)
+
+    if _recalculate_metric_enabled(run_config):
+        recalculated = _recalculate_orthogonal_metrics(
+            config,
+            evaluator,
+            mesh,
+            rho_s0=rho_s0,
+            Bnorm=Bnorm,
+            raw_dx=raw_dx,
+            raw_dy=raw_dy,
+            raw_dz=raw_dz,
+            raw_Bxy=raw_Bxy,
+            raw_g_23=raw_g_23,
+            dtype=dtype,
+        )
+        raw_dx = recalculated.dx
+        raw_dy = recalculated.dy
+        raw_dz = recalculated.dz
+        raw_J = recalculated.J
+        raw_g11 = recalculated.g11
+        raw_g33 = recalculated.g33
+        raw_g22 = recalculated.g22
+        raw_g_22 = recalculated.g_22
+        raw_g23 = recalculated.g23
+        raw_g_23 = recalculated.g_23
+        raw_Bxy = recalculated.Bxy
+
+    normalized_g22 = _normalize_g22(
+        broadcast_to_field_shape(raw_g22, mesh),
+        normalize_metric=normalize_metric,
+        rho_s0=rho_s0,
+    )
+    normalized_g_22 = _normalize_g_22(
+        broadcast_to_field_shape(raw_g_22, mesh),
+        normalize_metric=normalize_metric,
+        rho_s0=rho_s0,
+    )
+    return StructuredMetrics(
+        dx=_normalize_dx(broadcast_to_field_shape(raw_dx, mesh), normalize_metric=normalize_metric, rho_s0=rho_s0, Bnorm=Bnorm),
+        dy=broadcast_to_field_shape(raw_dy, mesh),
+        dz=broadcast_to_field_shape(raw_dz, mesh),
+        J=_normalize_J(broadcast_to_field_shape(raw_J, mesh), normalize_metric=normalize_metric, rho_s0=rho_s0),
+        g11=_normalize_g11(broadcast_to_field_shape(raw_g11, mesh), normalize_metric=normalize_metric, rho_s0=rho_s0),
+        g33=_normalize_g33(broadcast_to_field_shape(raw_g33, mesh), normalize_metric=normalize_metric, rho_s0=rho_s0),
+        g22=normalized_g22,
+        g_22=normalized_g_22,
+        g23=broadcast_to_field_shape(raw_g23, mesh),
+        g_23=broadcast_to_field_shape(raw_g_23, mesh),
+        Bxy=_normalize_Bxy(broadcast_to_field_shape(raw_Bxy, mesh), normalize_metric=normalize_metric, Bnorm=Bnorm),
+    )
+
+
+def _metric_value(
+    config: BoutConfig,
+    evaluator: ArrayExpressionEvaluator,
+    key: str,
+    *,
+    default: float,
+) -> Any:
+    if not config.has_section("mesh") or not config.has_option("mesh", key):
+        return default
+    return evaluator.resolve_option("mesh", key)
+
+
+def _normalize_metric_enabled(config: BoutConfig) -> bool:
+    if not has_model_section(config):
+        return True
+    model_section = locate_model_section(config)
+    if not config.has_option(model_section, "normalise_metric"):
+        return True
+    return bool(config.parsed(model_section, "normalise_metric"))
+
+
+def _recalculate_metric_enabled(run_config: RunConfiguration) -> bool:
+    return bool(run_config.normalization and run_config.normalization.metric_policy.recalculate_metric)
+
+
+def _recalculate_orthogonal_metrics(
+    config: BoutConfig,
+    evaluator: ArrayExpressionEvaluator,
+    mesh: StructuredMesh,
+    *,
+    rho_s0: float,
+    Bnorm: float,
+    raw_dx: Any,
+    raw_dy: Any,
+    raw_dz: Any,
+    raw_Bxy: Any,
+    raw_g_23: Any,
+    dtype: Any,
+) -> StructuredMetrics:
+    Rxy = _metric_value(config, evaluator, "Rxy", default=1.0)
+    Bpxy = _metric_value(config, evaluator, "Bpxy", default=1.0)
+    Btxy = _metric_value(config, evaluator, "Btxy", default=0.0)
+    hthe = _metric_value(config, evaluator, "hthe", default=1.0)
+    sinty = _metric_value(config, evaluator, "sinty", default=0.0)
+
+    normalized_Rxy = broadcast_to_field_shape(Rxy, mesh) / rho_s0
+    normalized_hthe = broadcast_to_field_shape(hthe, mesh) / rho_s0
+    normalized_sinty = broadcast_to_field_shape(sinty, mesh) * (rho_s0 * rho_s0 * Bnorm)
+    normalized_Bpxy = broadcast_to_field_shape(Bpxy, mesh) / Bnorm
+    normalized_Btxy = broadcast_to_field_shape(Btxy, mesh) / Bnorm
+    normalized_Bxy = broadcast_to_field_shape(raw_Bxy, mesh) / Bnorm
+
+    if run_parallel_transform_type(config) == "shifted":
+        normalized_sinty = jnp.zeros_like(normalized_sinty, dtype=dtype)
+
+    sign_Bp = jnp.where(jnp.min(normalized_Bpxy) < 0.0, -1.0, 1.0)
+    g11_normalized = jnp.square(normalized_Rxy * normalized_Bpxy)
+    g22 = 1.0 / jnp.square(broadcast_to_field_shape(hthe, mesh))
+    g_22 = jnp.square(broadcast_to_field_shape(hthe, mesh))
+    g33_normalized = jnp.square(normalized_sinty) * g11_normalized + jnp.square(normalized_Bxy) / g11_normalized
+    g23 = -sign_Bp * normalized_Btxy / (normalized_hthe * normalized_Bpxy * normalized_Rxy)
+    J_normalized = normalized_hthe / normalized_Bpxy
+
+    return StructuredMetrics(
+        dx=broadcast_to_field_shape(raw_dx, mesh),
+        dy=broadcast_to_field_shape(raw_dy, mesh),
+        dz=broadcast_to_field_shape(raw_dz, mesh),
+        J=J_normalized * rho_s0,
+        g11=g11_normalized * (rho_s0 * rho_s0),
+        g33=g33_normalized / (rho_s0 * rho_s0),
+        g22=g22,
+        g_22=g_22,
+        g23=g23,
+        g_23=broadcast_to_field_shape(raw_g_23, mesh),
+        Bxy=broadcast_to_field_shape(raw_Bxy, mesh),
+    )
+
+
+def run_parallel_transform_type(config: BoutConfig) -> str:
+    if not config.has_section("mesh:paralleltransform") or not config.has_option("mesh:paralleltransform", "type"):
+        return "identity"
+    return str(config.parsed("mesh:paralleltransform", "type"))
+
+
+def _normalize_dx(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float, Bnorm: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value / (rho_s0 * rho_s0 * Bnorm)
+
+
+def _normalize_J(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value / rho_s0
+
+
+def _normalize_g11(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value / (rho_s0 * rho_s0)
+
+
+def _normalize_g33(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value * (rho_s0 * rho_s0)
+
+
+def _normalize_g22(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value * (rho_s0 * rho_s0)
+
+
+def _normalize_g_22(value: jnp.ndarray, *, normalize_metric: bool, rho_s0: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value / (rho_s0 * rho_s0)
+
+
+def _normalize_Bxy(value: jnp.ndarray, *, normalize_metric: bool, Bnorm: float) -> jnp.ndarray:
+    if not normalize_metric:
+        return value
+    return value / Bnorm

drbx/native/neutrals/__init__.py

@@ -0,0 +1,58 @@
+"""Neutral / recycling physics for the 1D scrape-off-layer model.
+
+Self-contained JAX implementations of the hydrogenic atomic reaction rates
+(AMJUEL ionization / recombination fits, the Janev/AMJUEL charge-exchange
+polynomial) used to couple a plasma fluid to a neutral fluid in the
+scrape-off layer. All coefficients ship with the package
+(``drbx.data.atomic_rates``); there is no external-database dependency.
+"""
+
+from .atomic_rates import (
+    charge_exchange_rate_coefficient,
+    energy_loss_coefficient,
+    eval_amjuel_fit,
+    load_amjuel_coefficients,
+    rate_coefficient,
+)
+from .reactions import (
+    HydrogenReactionSources,
+    PlasmaNormalization,
+    compute_hydrogen_reaction_sources,
+)
+from .detachment_sol_model import (
+    DetachmentSolParameters,
+    DetachmentSolState,
+    detachment_diagnostics,
+    detachment_sol_run,
+    detachment_sol_step,
+)
+from .recycling_sol_model import (
+    SolRecyclingParameters,
+    SolRecyclingState,
+    linear_target_temperature_profile,
+    sol_recycling_run,
+    sol_recycling_step,
+    target_ion_flux,
+)
+
+__all__ = [
+    "rate_coefficient",
+    "energy_loss_coefficient",
+    "charge_exchange_rate_coefficient",
+    "eval_amjuel_fit",
+    "load_amjuel_coefficients",
+    "PlasmaNormalization",
+    "HydrogenReactionSources",
+    "compute_hydrogen_reaction_sources",
+    "SolRecyclingParameters",
+    "SolRecyclingState",
+    "linear_target_temperature_profile",
+    "sol_recycling_run",
+    "sol_recycling_step",
+    "target_ion_flux",
+    "DetachmentSolParameters",
+    "DetachmentSolState",
+    "detachment_diagnostics",
+    "detachment_sol_run",
+    "detachment_sol_step",
+]

drbx/native/neutrals/atomic_rates.py

@@ -0,0 +1,134 @@
+"""Hydrogenic atomic reaction-rate coefficients (self-contained, pure JAX).
+
+Ionization and recombination rate coefficients ``<sigma v>(T_e, n_e)`` come from
+the packaged AMJUEL double-polynomial fits; the charge-exchange coefficient
+``<sigma v>(T_eff)`` comes from the AMJUEL H.2 3.1.8 one-dimensional polynomial
+(inlined here, as hermes-3 does, since it carries no tabulated file). Every
+routine is pure ``jax.numpy`` and therefore ``jit``/``grad``/``vmap``
+transparent; the coefficient tables ship with the package under
+``drbx.data.atomic_rates`` (real published AMJUEL fits, verified against the
+hermes-3 database provenance), so there is no runtime external dependency.
+
+Temperatures are in eV and densities in m^-3; rate coefficients are returned in
+m^3 / s. The AMJUEL fits are clamped to their fitted range
+(``T in [0.1, 1e4] eV``, ``n in [1e14, 1e22] m^-3``).
+"""
+
+from __future__ import annotations
+
+import json
+from functools import lru_cache
+from importlib import resources
+
+import jax.numpy as jnp
+import numpy as np
+
+__all__ = [
+    "rate_coefficient",
+    "energy_loss_coefficient",
+    "charge_exchange_rate_coefficient",
+    "eval_amjuel_fit",
+    "load_amjuel_coefficients",
+]
+
+# AMJUEL fit files, keyed by (species, reaction). Deuterium and tritium share the
+# hydrogen fits; helium has its own.
+_AMJUEL_FILENAMES = {
+    ("d", "iz"): "iz_AMJUEL_H.x_2.1.5.json",
+    ("d", "rec"): "rec_AMJUEL_H.x_2.1.8.json",
+    ("t", "iz"): "iz_AMJUEL_H.x_2.1.5.json",
+    ("t", "rec"): "rec_AMJUEL_H.x_2.1.8.json",
+    ("he", "iz"): "iz_AMJUEL_H.x_2.3.9a.json",
+    ("he", "rec"): "rec_AMJUEL_H.x_2.3.13a.json",
+}
+
+# AMJUEL H.2 3.1.8 charge-exchange polynomial: ln(sigma v [cm^3/s]) as a degree-8
+# polynomial in ln(T_eff [eV]) (density independent), as inlined in hermes-3's
+# hydrogen_charge_exchange.cxx.
+_CX_CONSTANT = -18.5028
+_CX_COEFFICIENTS = (
+    0.3708409,
+    7.949876e-3,
+    -6.143769e-4,
+    -4.698969e-4,
+    -4.096807e-4,
+    1.440382e-4,
+    -1.514243e-5,
+    5.122435e-7,
+)
+
+
+@lru_cache(maxsize=None)
+def load_amjuel_coefficients(species: str, reaction: str) -> tuple[np.ndarray, np.ndarray, float]:
+    """Load one AMJUEL fit: ``(sigma_v_coeffs, sigma_v_E_coeffs, electron_heating)``.
+
+    ``sigma_v_coeffs`` and ``sigma_v_E_coeffs`` are 9x9 log-log polynomial
+    coefficient matrices; ``electron_heating`` is the potential energy [eV]
+    returned to the electrons (e.g. 13.6 eV on recombination).
+    """
+
+    filename = _AMJUEL_FILENAMES[(species, reaction)]
+    payload = json.loads(
+        resources.files("drbx.data.atomic_rates").joinpath(filename).read_text(encoding="utf-8")
+    )
+    return (
+        np.asarray(payload["sigma_v_coeffs"], dtype=np.float64),
+        np.asarray(payload["sigma_v_E_coeffs"], dtype=np.float64),
+        float(payload["electron_heating"]),
+    )
+
+
+def eval_amjuel_fit(temperature_ev, density_m3, coefficients) -> jnp.ndarray:
+    """Evaluate an AMJUEL double-polynomial fit, returning ``<sigma v>`` in m^3/s.
+
+    ``<sigma v> = 1e-6 * exp( sum_ij c_ij (ln T)^i (ln(n/1e14))^j )`` with ``T``
+    and ``n`` clamped to the fit's validity range.
+    """
+
+    temperature = jnp.clip(jnp.asarray(temperature_ev, dtype=jnp.float64), 0.1, 1.0e4)
+    density = jnp.clip(jnp.asarray(density_m3, dtype=jnp.float64), 1.0e14, 1.0e22)
+    log_t = jnp.log(temperature)
+    log_n = jnp.log(density / 1.0e14)
+
+    coeffs = jnp.asarray(coefficients, dtype=jnp.float64)
+    result = jnp.zeros_like(log_t)
+    for row in coeffs[::-1]:
+        row_result = jnp.zeros_like(log_n)
+        for coefficient in row[::-1]:
+            row_result = row_result * log_n + coefficient
+        result = result * log_t + row_result
+    return jnp.exp(result) * 1.0e-6
+
+
+def rate_coefficient(species: str, reaction: str, electron_temperature, electron_density) -> jnp.ndarray:
+    """Ionization (``reaction="iz"``) or recombination (``"rec"``) ``<sigma v>`` [m^3/s]."""
+
+    sigma_v_coeffs, _, _ = load_amjuel_coefficients(species, reaction)
+    return eval_amjuel_fit(electron_temperature, electron_density, sigma_v_coeffs)
+
+
+def energy_loss_coefficient(species: str, reaction: str, electron_temperature, electron_density) -> jnp.ndarray:
+    """Electron energy-loss coefficient ``<sigma v E>`` [eV m^3/s] for the reaction.
+
+    The recombination channel returns the ``electron_heating`` potential energy
+    (13.6 eV for hydrogen) as a heating term, so the net electron energy sink is
+    ``<sigma v E> - electron_heating * <sigma v>``; ionization has no potential
+    add-back (``electron_heating = 0``).
+    """
+
+    sigma_v_coeffs, sigma_v_E_coeffs, electron_heating = load_amjuel_coefficients(species, reaction)
+    energy_loss = eval_amjuel_fit(electron_temperature, electron_density, sigma_v_E_coeffs)
+    rate = eval_amjuel_fit(electron_temperature, electron_density, sigma_v_coeffs)
+    return energy_loss - float(electron_heating) * rate
+
+
+def charge_exchange_rate_coefficient(effective_temperature_ev) -> jnp.ndarray:
+    """Hydrogenic charge-exchange ``<sigma v>(T_eff)`` [m^3/s] (AMJUEL H.2 3.1.8)."""
+
+    log_t = jnp.log(jnp.clip(jnp.asarray(effective_temperature_ev, dtype=jnp.float64), 0.1, 1.0e4))
+    log_sigma_v = jnp.full_like(log_t, _CX_CONSTANT)
+    power = jnp.ones_like(log_t)
+    for coefficient in _CX_COEFFICIENTS:
+        power = power * log_t
+        log_sigma_v = log_sigma_v + coefficient * power
+    return jnp.exp(log_sigma_v) * 1.0e-6