drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/sol_flux_tube.py

@@ -0,0 +1,133 @@
+"""Reduced isothermal scrape-off-layer (SOL) flux-tube parallel transport.
+
+An open-field-line model: on the open slab geometry
+(:func:`drbx.geometry.build_open_slab_geometry`), plasma is transported along
+the field to the two target plates, where a Bohm sheath drains it at the sound
+speed. The model evolves the parallel density ``n`` and parallel momentum
+``m = n v`` as an isothermal Euler system in the field-parallel coordinate ``z``,
+
+    d n / dt + d (n v) / dz = S_n
+    d m / dt + d (n v^2 + n c_s^2) / dz = 0
+
+with an upstream particle source ``S_n`` and Bohm sheath outflow (``|v| >= c_s``)
+at ``z = 0`` and ``z = L``. Faces use a Rusanov (local Lax--Friedrichs) flux; the
+whole update is pure JAX (``jit``/``grad``/``vmap`` transparent). The steady
+state is the classic two-point SOL solution: the flow accelerates from a
+stagnation point to the sound speed at each target, and the target density is
+half the upstream density.
+
+The transport acts along the parallel (``z``) axis of the ``(nx, ny, nz)`` field
+arrays, so each ``(x, y)`` column is an independent flux tube sharing the open
+target plates.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+
+from ..geometry import FciGeometry3D
+
+__all__ = [
+    "SolFluxTubeParameters",
+    "sol_flux_tube_source",
+    "sol_flux_tube_rhs",
+    "sol_flux_tube_step",
+    "sol_flux_tube_run",
+]
+
+
+@dataclass(frozen=True)
+class SolFluxTubeParameters:
+    """Parameters of the reduced isothermal SOL flux tube."""
+
+    sound_speed: float = 1.0            # normalized isothermal sound speed c_s
+    source_amplitude: float = 0.02      # peak upstream particle source
+    source_width: float = 4.0           # parallel width of the upstream source
+    density_floor: float = 1.0e-6
+
+
+def sol_flux_tube_source(geometry: FciGeometry3D, params: SolFluxTubeParameters) -> jnp.ndarray:
+    """Upstream particle source: a Gaussian centered at the parallel midplane."""
+
+    z = jnp.asarray(geometry.grid.z.centers, dtype=jnp.float64)
+    midplane = 0.5 * (float(z[0]) + float(z[-1]))
+    profile = params.source_amplitude * jnp.exp(-((z - midplane) ** 2) / params.source_width**2)
+    return jnp.broadcast_to(profile[None, None, :], geometry.shape)
+
+
+def _flux(density: jnp.ndarray, momentum: jnp.ndarray, c_s: float) -> tuple[jnp.ndarray, jnp.ndarray]:
+    velocity = momentum / density
+    return momentum, momentum * velocity + density * c_s**2
+
+
+def _rusanov(nL, mL, nR, mR, c_s):
+    fL_n, fL_m = _flux(nL, mL, c_s)
+    fR_n, fR_m = _flux(nR, mR, c_s)
+    speed = jnp.maximum(jnp.abs(mL / nL), jnp.abs(mR / nR)) + c_s
+    flux_n = 0.5 * (fL_n + fR_n) - 0.5 * speed * (nR - nL)
+    flux_m = 0.5 * (fL_m + fR_m) - 0.5 * speed * (mR - mL)
+    return flux_n, flux_m
+
+
+def sol_flux_tube_rhs(
+    density: jnp.ndarray,
+    momentum: jnp.ndarray,
+    geometry: FciGeometry3D,
+    params: SolFluxTubeParameters,
+    source: jnp.ndarray,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    """Parallel-transport RHS ``(dn/dt, dm/dt)`` with Bohm sheath targets."""
+
+    c_s = float(params.sound_speed)
+    dz = jnp.asarray(geometry.spacing.dz, dtype=jnp.float64)
+    n = jnp.maximum(density, float(params.density_floor))
+    v = momentum / n
+
+    # Interior faces along z (between adjacent cells).
+    interior_n, interior_m = _rusanov(n[:, :, :-1], momentum[:, :, :-1], n[:, :, 1:], momentum[:, :, 1:], c_s)
+
+    # Bohm sheath outflow at the two targets: |v| >= c_s pointing into the plate.
+    v_left = jnp.minimum(v[:, :, 0], -c_s)
+    v_right = jnp.maximum(v[:, :, -1], c_s)
+    left_n, left_m = _flux(n[:, :, 0], n[:, :, 0] * v_left, c_s)
+    right_n, right_m = _flux(n[:, :, -1], n[:, :, -1] * v_right, c_s)
+
+    face_n = jnp.concatenate([left_n[:, :, None], interior_n, right_n[:, :, None]], axis=2)
+    face_m = jnp.concatenate([left_m[:, :, None], interior_m, right_m[:, :, None]], axis=2)
+
+    d_density = -(face_n[:, :, 1:] - face_n[:, :, :-1]) / dz + source
+    d_momentum = -(face_m[:, :, 1:] - face_m[:, :, :-1]) / dz
+    return d_density, d_momentum
+
+
+def sol_flux_tube_step(density, momentum, geometry, params, source, dt):
+    """One RK4 step of the reduced SOL flux tube (density floored to stay positive)."""
+
+    def rhs(n, m):
+        return sol_flux_tube_rhs(n, m, geometry, params, source)
+
+    k1n, k1m = rhs(density, momentum)
+    k2n, k2m = rhs(density + 0.5 * dt * k1n, momentum + 0.5 * dt * k1m)
+    k3n, k3m = rhs(density + 0.5 * dt * k2n, momentum + 0.5 * dt * k2m)
+    k4n, k4m = rhs(density + dt * k3n, momentum + dt * k3m)
+    density = density + dt / 6.0 * (k1n + 2 * k2n + 2 * k3n + k4n)
+    momentum = momentum + dt / 6.0 * (k1m + 2 * k2m + 2 * k3m + k4m)
+    return jnp.maximum(density, float(params.density_floor)), momentum
+
+
+def sol_flux_tube_run(density, momentum, geometry, params, source, *, dt, steps):
+    """Advance ``steps`` RK4 steps with a jitted ``lax.scan``; return the final state."""
+
+    @jax.jit
+    def _run(n0, m0):
+        def body(carry, _):
+            n, m = carry
+            return sol_flux_tube_step(n, m, geometry, params, source, dt), None
+
+        (n_final, m_final), _ = jax.lax.scan(body, (n0, m0), None, length=steps)
+        return n_final, m_final
+
+    return _run(density, momentum)

drbx/native/stellarator_turbulence.py

@@ -0,0 +1,343 @@
+"""Reusable four-field stellarator turbulence driver (closed and open field lines).
+
+A thin, reusable harness around the validated four-field interchange RK4 step:
+seed a multi-mode density perturbation on a stellarator FCI geometry (rotating
+ellipse, island divertor, ...) and let the curvature drive develop
+turbulence-like interchange dynamics. On an open geometry (endpoint masks set,
+e.g. via ``limiter_radius`` or traced island-divertor masks) a Bohm sheath
+density sink drains the open-endpoint cells each step, so the scrape-off layer
+empties onto the target while the core stays closed.
+
+Public pieces, each reusable on its own:
+
+- :func:`build_free_decay_boundary_conditions` -- Dirichlet phi / Neumann
+  fields at the radial walls, bundled as :class:`FourFieldBoundaryConditions`;
+- :func:`build_four_field_phi_solver` -- the GMRES perpendicular-Laplacian
+  inverter used by the phi inversion;
+- :func:`four_field_rk4_step` -- one classic RK4 step of the four-field blob
+  RHS (``compute_4field_blob_rhs``), carrying the phi guess between stages;
+- :func:`multi_mode_state` -- the seeded multi-mode initial state;
+- :func:`run_stellarator_turbulence` -- the full loop, returning a
+  :class:`TurbulenceRun` with frames and particle/sheath-flux traces.
+
+Used by ``tests/test_stellarator_turbulence.py`` and
+``tests/test_island_divertor.py`` (fast physics gates). The example scripts in
+``examples/stellarator/`` show the same anatomy written out step by step.
+"""
+
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from ..geometry import (
+    ConservativeStencilBuilder,
+    FciGeometry3D,
+    LocalStencilBuilder,
+    RegularFaceGeometry3D,
+    build_conservative_stencil_from_field,
+    build_curvature_coefficients,
+    build_local_stencil_from_field,
+)
+from .fci_4_field_rhs import (
+    Fci4FieldBlobParameters,
+    Fci4FieldState,
+    compute_4field_blob_rhs,
+)
+from .fci_boundaries import (
+    BC_DIRICHLET,
+    BC_NEUMANN,
+    BoundaryFaceBC3D,
+    CutWallBC3D,
+    CutWallGeometry3D,
+    FourFieldBoundaryConditions,
+)
+from .fci_operators import PerpLaplacianInverseSolver, build_perp_laplacian_face_projectors
+from .fci_sheath_recycling import compute_fci_sheath_recycling
+
+TEMPERATURE = 1.0  # normalized Te = Ti for the sheath sound speed
+SHEATH_DENSITY_FLOOR = 1.0e-4  # positivity floor after the sheath sink
+
+__all__ = [
+    "TEMPERATURE",
+    "SHEATH_DENSITY_FLOOR",
+    "FourFieldBoundaryConditions",
+    "TurbulenceRun",
+    "apply_sheath_sink",
+    "build_free_decay_boundary_conditions",
+    "build_four_field_phi_solver",
+    "four_field_rk4_step",
+    "multi_mode_state",
+    "radial_dirichlet_face_bc",
+    "radial_neumann_face_bc",
+    "run_stellarator_turbulence",
+]
+
+
+def radial_dirichlet_face_bc(geometry: FciGeometry3D) -> BoundaryFaceBC3D:
+    """Zero-Dirichlet at the two radial (x) walls; periodic elsewhere."""
+
+    face = RegularFaceGeometry3D.unit(geometry)
+    return BoundaryFaceBC3D(
+        kind_x=jnp.zeros_like(face.x_area, dtype=jnp.int32).at[0].set(BC_DIRICHLET).at[-1].set(BC_DIRICHLET),
+        kind_y=jnp.zeros_like(face.y_area, dtype=jnp.int32),
+        kind_z=jnp.zeros_like(face.z_area, dtype=jnp.int32),
+        value_x=jnp.zeros_like(face.x_area, dtype=jnp.float64),
+        value_y=jnp.zeros_like(face.y_area, dtype=jnp.float64),
+        value_z=jnp.zeros_like(face.z_area, dtype=jnp.float64),
+        mask_x=jnp.zeros_like(face.x_open_mask, dtype=bool).at[0].set(True).at[-1].set(True),
+        mask_y=jnp.zeros_like(face.y_open_mask, dtype=bool),
+        mask_z=jnp.zeros_like(face.z_open_mask, dtype=bool),
+    )
+
+
+def radial_neumann_face_bc(geometry: FciGeometry3D) -> BoundaryFaceBC3D:
+    """Zero-Neumann at the two radial (x) walls; periodic elsewhere."""
+
+    face = RegularFaceGeometry3D.unit(geometry)
+    return BoundaryFaceBC3D(
+        kind_x=jnp.zeros_like(face.x_area, dtype=jnp.int32).at[0].set(BC_NEUMANN).at[-1].set(BC_NEUMANN),
+        kind_y=jnp.zeros_like(face.y_area, dtype=jnp.int32),
+        kind_z=jnp.zeros_like(face.z_area, dtype=jnp.int32),
+        value_x=jnp.zeros_like(face.x_area, dtype=jnp.float64),
+        value_y=jnp.zeros_like(face.y_area, dtype=jnp.float64),
+        value_z=jnp.zeros_like(face.z_area, dtype=jnp.float64),
+        mask_x=jnp.zeros_like(face.x_open_mask, dtype=bool).at[0].set(True).at[-1].set(True),
+        mask_y=jnp.zeros_like(face.y_open_mask, dtype=bool),
+        mask_z=jnp.zeros_like(face.z_open_mask, dtype=bool),
+    )
+
+
+def build_free_decay_boundary_conditions(geometry: FciGeometry3D) -> FourFieldBoundaryConditions:
+    """Free-decay walls: Dirichlet phi = 0, zero-Neumann for the evolved fields."""
+
+    neumann = radial_neumann_face_bc(geometry)
+    return FourFieldBoundaryConditions(
+        phi_face_bc=radial_dirichlet_face_bc(geometry),
+        density_face_bc=neumann,
+        omega_face_bc=neumann,
+        v_ion_parallel_face_bc=neumann,
+        v_electron_parallel_face_bc=neumann,
+    )
+
+
+def build_four_field_phi_solver(
+    geometry: FciGeometry3D,
+    parameters: Fci4FieldBlobParameters,
+    *,
+    conservative_stencil_builder: ConservativeStencilBuilder,
+    face_projectors: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] | None,
+) -> PerpLaplacianInverseSolver:
+    """GMRES inverter for the perpendicular Laplacian used by the phi inversion."""
+
+    return PerpLaplacianInverseSolver(
+        geometry,
+        conservative_stencil_builder,
+        tol=float(parameters.phi_inversion_tol),
+        maxiter=int(parameters.phi_inversion_maxiter),
+        restart=int(parameters.phi_inversion_restart),
+        face_projectors=face_projectors,
+        regular_face_geometry=RegularFaceGeometry3D.unit(geometry),
+        cut_wall_geometry=CutWallGeometry3D.empty(),
+        cut_wall_bc=CutWallBC3D.empty(),
+        periodic_axes=(False, True, True),
+        pin_point=None,
+        pin_value=0.0,
+        project_mean_zero=False,
+        target_mean_phi=None,
+        regularization_epsilon=0.0,
+        gmres_debug=False,
+        # Hot-loop configuration: skip the per-solve host-synced residual
+        # check so the phi inversion (and any step built on it) can run fully
+        # inside jit. The validation harnesses build their own solvers with
+        # check_residual=True.
+        check_residual=False,
+    )
+
+
+def four_field_rk4_step(
+    state: Fci4FieldState,
+    *,
+    geometry: FciGeometry3D,
+    timestep: float,
+    parameters: Fci4FieldBlobParameters,
+    curvature_coefficients: jnp.ndarray,
+    stencil_builder: LocalStencilBuilder,
+    conservative_stencil_builder: ConservativeStencilBuilder,
+    boundary_conditions: FourFieldBoundaryConditions,
+    phi_face_projectors: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray] | None,
+    phi_inverse_solver: PerpLaplacianInverseSolver,
+    phi_guess: jnp.ndarray | None = None,
+) -> tuple[Fci4FieldState, jnp.ndarray]:
+    """One classic RK4 step of the four-field blob RHS.
+
+    The electrostatic potential from each stage seeds the GMRES phi inversion of
+    the next, so passing the returned phi back in as ``phi_guess`` keeps the
+    inversion warm across steps. Returns ``(next_state, phi_of_last_stage)``.
+    """
+
+    common_kwargs = dict(
+        geometry=geometry,
+        stencil_builder=stencil_builder,
+        conservative_stencil_builder=conservative_stencil_builder,
+        parameters=parameters,
+        curvature_coefficients=curvature_coefficients,
+        boundary_conditions=boundary_conditions,
+        phi_face_projectors=phi_face_projectors,
+        phi_inverse_solver=phi_inverse_solver,
+        gmres_debug=False,
+        return_phi=True,
+    )
+
+    rhs_1, _timings, phi_1 = compute_4field_blob_rhs(state, phi_guess=phi_guess, **common_kwargs)
+    k1 = rhs_1.rhs
+    stage_1 = state.axpy(k1, scale=0.5 * timestep)
+
+    rhs_2, _timings, phi_2 = compute_4field_blob_rhs(stage_1, phi_guess=phi_1, **common_kwargs)
+    k2 = rhs_2.rhs
+    stage_2 = state.axpy(k2, scale=0.5 * timestep)
+
+    rhs_3, _timings, phi_3 = compute_4field_blob_rhs(stage_2, phi_guess=phi_2, **common_kwargs)
+    k3 = rhs_3.rhs
+    stage_3 = state.axpy(k3, scale=timestep)
+
+    rhs_4, _timings, phi_4 = compute_4field_blob_rhs(stage_3, phi_guess=phi_3, **common_kwargs)
+    k4 = rhs_4.rhs
+
+    increment = Fci4FieldState(
+        density=(k1.density + 2.0 * k2.density + 2.0 * k3.density + k4.density) / 6.0,
+        omega=(k1.omega + 2.0 * k2.omega + 2.0 * k3.omega + k4.omega) / 6.0,
+        v_ion_parallel=(
+            k1.v_ion_parallel + 2.0 * k2.v_ion_parallel + 2.0 * k3.v_ion_parallel + k4.v_ion_parallel
+        )
+        / 6.0,
+        v_electron_parallel=(
+            k1.v_electron_parallel
+            + 2.0 * k2.v_electron_parallel
+            + 2.0 * k3.v_electron_parallel
+            + k4.v_electron_parallel
+        )
+        / 6.0,
+    )
+    return state.axpy(increment, scale=timestep), phi_4
+
+
+def apply_sheath_sink(
+    state: Fci4FieldState,
+    geometry: FciGeometry3D,
+    dt: float,
+    *,
+    temperature: float = TEMPERATURE,
+) -> tuple[Fci4FieldState, float]:
+    """Bohm sheath density sink on the open-endpoint cells (explicit Euler).
+
+    Returns the drained state and the total ion particle loss rate of the
+    *pre-sink* state (the instantaneous sheath flux to the target).
+    """
+
+    temperature_field = jnp.full(geometry.shape, float(temperature))
+    sheath = compute_fci_sheath_recycling(state.density, temperature_field, temperature_field, geometry.maps)
+    drained = Fci4FieldState(
+        density=jnp.maximum(state.density - dt * sheath.ion_particle_loss, SHEATH_DENSITY_FLOOR),
+        omega=state.omega,
+        v_ion_parallel=state.v_ion_parallel,
+        v_electron_parallel=state.v_electron_parallel,
+    )
+    return drained, float(sheath.total_ion_particle_loss)
+
+
+def multi_mode_state(
+    geometry: FciGeometry3D,
+    *,
+    amplitude: float = 0.08,
+    seed: int = 0,
+    modes: tuple[tuple[int, int], ...] = ((2, 1), (3, 2), (4, 1), (5, 3)),
+) -> Fci4FieldState:
+    """Seed a random-phase multi-mode density perturbation with a radial envelope."""
+
+    rng = np.random.default_rng(seed)
+    x = np.asarray(geometry.grid.x.centers)[:, None, None]
+    theta = np.asarray(geometry.grid.y.centers)[None, :, None]
+    zeta = np.asarray(geometry.grid.z.centers)[None, None, :]
+    x_norm = (x - x.min()) / (x.max() - x.min())
+    envelope = np.sin(np.pi * x_norm)
+    perturbation = np.zeros(geometry.shape)
+    for m, n in modes:
+        perturbation += rng.uniform(0.5, 1.0) * np.cos(m * theta + n * zeta + rng.uniform(0, 2 * np.pi))
+    density = 1.0 + amplitude * envelope * perturbation
+    zeros = jnp.zeros(geometry.shape, dtype=jnp.float64)
+    return Fci4FieldState(density=jnp.asarray(density), omega=zeros,
+                          v_ion_parallel=zeros, v_electron_parallel=zeros)
+
+
+class TurbulenceRun(NamedTuple):
+    """Density/vorticity frames plus particle-content and sheath-flux traces."""
+
+    density_frames: np.ndarray      # (n_frames, nx, ny, nz), float32
+    omega_frames: np.ndarray
+    times: np.ndarray
+    particle_content: np.ndarray    # Jacobian-weighted total density per frame
+    target_flux: np.ndarray         # total sheath ion loss per frame (0 when closed)
+
+
+def run_stellarator_turbulence(geometry: FciGeometry3D, *, steps: int, dt: float,
+                               amplitude: float = 0.08, seed: int = 0,
+                               sheath_sink: bool = False, frame_stride: int = 1) -> TurbulenceRun:
+    """Advance the seeded four-field state; optionally drain open endpoints."""
+
+    parameters = Fci4FieldBlobParameters(rho_star=1.0, phi_inversion_tol=5.0e-5,
+                                         phi_inversion_maxiter=100, phi_inversion_restart=200)
+    stencil_builder = LocalStencilBuilder(build_local_stencil_from_field.build_fn)
+    conservative_builder = ConservativeStencilBuilder(build_conservative_stencil_from_field.build_fn)
+    boundary_conditions = build_free_decay_boundary_conditions(geometry)
+    curvature = build_curvature_coefficients(geometry, periodic_axes=(False, True, True))
+    projectors = build_perp_laplacian_face_projectors(geometry)
+    phi_solver = build_four_field_phi_solver(
+        geometry, parameters,
+        conservative_stencil_builder=conservative_builder, face_projectors=projectors,
+    )
+
+    temperature = jnp.full(geometry.shape, TEMPERATURE)
+    jacobian = np.asarray(geometry.cell_metric.J)
+    state = multi_mode_state(geometry, amplitude=amplitude, seed=seed)
+    phi_guess = jnp.zeros(geometry.shape, dtype=jnp.float64)
+    density_frames, omega_frames, times, content, flux = [], [], [], [], []
+
+    # The whole RK4 step (four RHS evaluations, each with a GMRES phi
+    # inversion) is sync-free, so compile it once and run it as a single
+    # XLA program.
+    jitted_step = jax.jit(
+        lambda current, guess: four_field_rk4_step(
+            current, geometry=geometry, timestep=dt, parameters=parameters,
+            curvature_coefficients=curvature, stencil_builder=stencil_builder,
+            conservative_stencil_builder=conservative_builder,
+            boundary_conditions=boundary_conditions, phi_face_projectors=projectors,
+            phi_inverse_solver=phi_solver, phi_guess=guess,
+        )
+    )
+
+    def record(step_index, current):
+        density = np.asarray(current.density, dtype=np.float32)
+        density_frames.append(density)
+        omega_frames.append(np.asarray(current.omega, dtype=np.float32))
+        times.append(step_index * dt)
+        content.append(float(np.sum(np.asarray(current.density, dtype=np.float64) * jacobian)))
+        if sheath_sink:
+            sheath = compute_fci_sheath_recycling(current.density, temperature, temperature, geometry.maps)
+            flux.append(float(sheath.total_ion_particle_loss))
+        else:
+            flux.append(0.0)
+
+    record(0, state)
+    for step_index in range(1, steps + 1):
+        state, phi_guess = jitted_step(state, phi_guess)
+        if sheath_sink:
+            state, _loss = apply_sheath_sink(state, geometry, dt)
+        if step_index % frame_stride == 0 or step_index == steps:
+            record(step_index, state)
+
+    return TurbulenceRun(np.stack(density_frames), np.stack(omega_frames),
+                         np.asarray(times), np.asarray(content), np.asarray(flux))

drbx/native/transport.py

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+from jax.scipy.linalg import expm
+
+from ..runtime import runtime_jax_dtype
+from .mesh import StructuredMesh, apply_field_boundaries
+from .metrics import StructuredMetrics
+
+
+@dataclass(frozen=True)
+class OneStepDiffusionResult:
+    density: jnp.ndarray
+    pressure: jnp.ndarray
+
+
+@dataclass(frozen=True)
+class DiffusionHistoryResult:
+    density_history: jnp.ndarray
+    pressure_history: jnp.ndarray
+
+
+def advance_anomalous_diffusion_one_step(
+    density: jnp.ndarray,
+    pressure: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    anomalous_D: float,
+    density_boundary: str,
+    pressure_boundary: str,
+    timestep: float,
+) -> OneStepDiffusionResult:
+    dtype = runtime_jax_dtype()
+    if mesh.nz != 1:
+        raise NotImplementedError("Native one-step anomalous diffusion currently supports nz = 1 only.")
+    if density_boundary.strip().lower() != "neumann" or pressure_boundary.strip().lower() != "neumann":
+        raise NotImplementedError("Native one-step anomalous diffusion currently supports Neumann X boundaries only.")
+    tolerance = 1e-6 if dtype == jnp.float32 else 1e-12
+    if not jnp.allclose(density, pressure, rtol=tolerance, atol=tolerance):
+        raise NotImplementedError(
+            "Native one-step anomalous diffusion currently requires identical density and pressure initial states."
+        )
+    if not jnp.allclose(metrics.g23, 0.0, rtol=tolerance, atol=tolerance):
+        raise NotImplementedError("Native one-step anomalous diffusion currently supports g23 = 0 structured metrics only.")
+
+    operator = _build_radial_diffusion_operator(mesh, metrics, anomalous_D)
+    propagator = expm(operator * timestep)
+    density_next = _advance_field_with_operator(density, propagator, mesh, boundary_kind=density_boundary)
+    pressure_next = _advance_field_with_operator(pressure, propagator, mesh, boundary_kind=pressure_boundary)
+    return OneStepDiffusionResult(density=density_next, pressure=pressure_next)
+
+
+def advance_anomalous_diffusion_history(
+    density: jnp.ndarray,
+    pressure: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    anomalous_D: float,
+    density_boundary: str,
+    pressure_boundary: str,
+    timestep: float,
+    steps: int,
+) -> DiffusionHistoryResult:
+    if steps < 0:
+        raise ValueError("steps must be non-negative")
+    if mesh.nz != 1:
+        raise NotImplementedError("Native anomalous diffusion history currently supports nz = 1 only.")
+    if density_boundary.strip().lower() != "neumann" or pressure_boundary.strip().lower() != "neumann":
+        raise NotImplementedError("Native anomalous diffusion history currently supports Neumann X boundaries only.")
+
+    dtype = runtime_jax_dtype()
+    operator = _build_radial_diffusion_operator(mesh, metrics, anomalous_D)
+    propagator = expm(operator * timestep)
+    density_history = [jnp.asarray(density, dtype=dtype)]
+    pressure_history = [jnp.asarray(pressure, dtype=dtype)]
+    current_density = density_history[0]
+    current_pressure = pressure_history[0]
+    for _ in range(steps):
+        current_density = _advance_field_with_operator(current_density, propagator, mesh, boundary_kind=density_boundary)
+        current_pressure = _advance_field_with_operator(current_pressure, propagator, mesh, boundary_kind=pressure_boundary)
+        density_history.append(current_density)
+        pressure_history.append(current_pressure)
+    return DiffusionHistoryResult(
+        density_history=jnp.stack(density_history, axis=0),
+        pressure_history=jnp.stack(pressure_history, axis=0),
+    )
+
+
+def _build_radial_diffusion_operator(
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    anomalous_D: float,
+) -> jnp.ndarray:
+    y_index = mesh.ystart
+    z_index = 0
+    dx = metrics.dx[:, y_index, z_index]
+    J = metrics.J[:, y_index, z_index]
+    g11 = metrics.g11[:, y_index, z_index]
+    dtype = runtime_jax_dtype()
+
+    interior_nx = mesh.xend - mesh.xstart + 1
+    matrix = jnp.zeros((interior_nx, interior_nx), dtype=dtype)
+
+    for global_index in range(mesh.xstart, mesh.xend):
+        left = global_index - mesh.xstart
+        right = left + 1
+        face_coef = anomalous_D * (J[global_index] * g11[global_index] + J[global_index + 1] * g11[global_index + 1])
+        face_coef /= dx[global_index] + dx[global_index + 1]
+
+        matrix = matrix.at[left, left].add(-face_coef / (dx[global_index] * J[global_index]))
+        matrix = matrix.at[left, right].add(face_coef / (dx[global_index] * J[global_index]))
+        matrix = matrix.at[right, left].add(face_coef / (dx[global_index + 1] * J[global_index + 1]))
+        matrix = matrix.at[right, right].add(-face_coef / (dx[global_index + 1] * J[global_index + 1]))
+
+    return matrix
+
+
+def _advance_field_with_operator(
+    field: jnp.ndarray,
+    propagator: jnp.ndarray,
+    mesh: StructuredMesh,
+    *,
+    boundary_kind: str,
+) -> jnp.ndarray:
+    result = jnp.asarray(field, dtype=runtime_jax_dtype())
+    interior = result[mesh.xstart : mesh.xend + 1, mesh.ystart : mesh.yend + 1, 0]
+    updated = propagator @ interior
+    result = result.at[mesh.xstart : mesh.xend + 1, mesh.ystart : mesh.yend + 1, 0].set(updated)
+    result = apply_field_boundaries(result, mesh, x_boundary=boundary_kind)
+    return result

drbx/native/units.py

@@ -0,0 +1,32 @@
+from __future__ import annotations
+
+from ..config.normalization import ELEMENTARY_CHARGE, PROTON_MASS
+from ..runtime.run_config import RunConfiguration
+
+
+def resolved_dataset_scalars(run_config: RunConfiguration) -> dict[str, float]:
+    if run_config.normalization is not None:
+        normalization = run_config.normalization
+        return {
+            "Nnorm": normalization.Nnorm,
+            "Tnorm": normalization.Tnorm,
+            "Bnorm": normalization.Bnorm,
+            "Cs0": normalization.Cs0,
+            "Omega_ci": normalization.Omega_ci,
+            "rho_s0": normalization.rho_s0,
+        }
+
+    Nnorm = float(run_config.model_scalars.get("Nnorm", 1.0e19))
+    Tnorm = float(run_config.model_scalars.get("Tnorm", 100.0))
+    Bnorm = float(run_config.model_scalars.get("Bnorm", 1.0))
+    Cs0 = float((ELEMENTARY_CHARGE * Tnorm / PROTON_MASS) ** 0.5)
+    Omega_ci = float(ELEMENTARY_CHARGE * Bnorm / PROTON_MASS)
+    rho_s0 = float(Cs0 / Omega_ci)
+    return {
+        "Nnorm": Nnorm,
+        "Tnorm": Tnorm,
+        "Bnorm": Bnorm,
+        "Cs0": Cs0,
+        "Omega_ci": Omega_ci,
+        "rho_s0": rho_s0,
+    }