drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/fluid_1d.py

@@ -0,0 +1,330 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from ..config.boutinp import BoutConfig
+from .expression import ArrayExpressionEvaluator
+from .limiters import periodic_monotonic_centered_edges_jax as _periodic_monotonic_centered_edges
+from .mesh import StructuredMesh, broadcast_to_field_shape
+from .metrics import StructuredMetrics
+
+
+@dataclass(frozen=True)
+class Fluid1DState:
+    density: jnp.ndarray
+    pressure: jnp.ndarray
+    momentum: jnp.ndarray
+
+
+@dataclass(frozen=True)
+class Fluid1DRhsResult:
+    density: jnp.ndarray
+    pressure: jnp.ndarray
+    momentum: jnp.ndarray
+
+
+@dataclass(frozen=True)
+class Fluid1DHistoryResult:
+    density_history: jnp.ndarray
+    pressure_history: jnp.ndarray
+    momentum_history: jnp.ndarray
+
+
+def initialize_mms_state(config: BoutConfig, *, section: str, mesh: StructuredMesh) -> Fluid1DState:
+    return Fluid1DState(
+        density=evaluate_field_option(config, f"N{section}", "solution", mesh=mesh, time=0.0),
+        pressure=evaluate_field_option(config, f"P{section}", "solution", mesh=mesh, time=0.0),
+        momentum=evaluate_field_option(config, f"NV{section}", "solution", mesh=mesh, time=0.0),
+    )
+
+
+def evaluate_field_option(
+    config: BoutConfig,
+    variable_name: str,
+    option: str,
+    *,
+    mesh: StructuredMesh,
+    time: float,
+) -> jnp.ndarray:
+    evaluator = ArrayExpressionEvaluator(config, local_values=mesh.expression_context(time=time))
+    field = broadcast_to_field_shape(evaluator.resolve_option(variable_name, option), mesh)
+    return apply_periodic_y_guards(field, mesh)
+
+
+def compute_mms_rhs(
+    config: BoutConfig,
+    state: Fluid1DState,
+    *,
+    section: str,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    atomic_mass: float,
+    time: float,
+) -> Fluid1DRhsResult:
+    guarded = Fluid1DState(
+        density=apply_periodic_y_guards(state.density, mesh),
+        pressure=apply_periodic_y_guards(state.pressure, mesh),
+        momentum=apply_periodic_y_guards(state.momentum, mesh),
+    )
+
+    density = guarded.density
+    pressure = guarded.pressure
+    momentum = guarded.momentum
+    velocity = momentum / (atomic_mass * density)
+    temperature = pressure / density
+    fastest_wave = jnp.sqrt(temperature / atomic_mass)
+
+    density_rhs = -div_par_mod_periodic(density, velocity, fastest_wave, mesh=mesh, metrics=metrics)
+    density_rhs = density_rhs + evaluate_field_option(config, f"N{section}", "source", mesh=mesh, time=time)
+
+    pressure_rhs = -(5.0 / 3.0) * div_par_mod_periodic(pressure, velocity, fastest_wave, mesh=mesh, metrics=metrics)
+    pressure_rhs = pressure_rhs + (2.0 / 3.0) * velocity * grad_par_periodic(pressure, mesh=mesh, metrics=metrics)
+    pressure_rhs = pressure_rhs + evaluate_field_option(config, f"P{section}", "source", mesh=mesh, time=time)
+
+    momentum_rhs = -atomic_mass * div_par_fvv_periodic(density, velocity, fastest_wave, mesh=mesh, metrics=metrics)
+    momentum_rhs = momentum_rhs - grad_par_periodic(pressure, mesh=mesh, metrics=metrics)
+    momentum_rhs = momentum_rhs + evaluate_field_option(config, f"NV{section}", "source", mesh=mesh, time=time)
+
+    return Fluid1DRhsResult(
+        density=density_rhs,
+        pressure=pressure_rhs,
+        momentum=momentum_rhs,
+    )
+
+
+def advance_mms_history(
+    config: BoutConfig,
+    *,
+    section: str,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    atomic_mass: float,
+    timestep: float,
+    steps: int,
+    substeps: int,
+    initial_state: Fluid1DState | None = None,
+    start_time: float = 0.0,
+) -> Fluid1DHistoryResult:
+    if steps < 0:
+        raise ValueError("steps must be non-negative")
+    if substeps <= 0:
+        raise ValueError("substeps must be positive")
+
+    state = initial_state if initial_state is not None else initialize_mms_state(config, section=section, mesh=mesh)
+    density_history = [state.density]
+    pressure_history = [state.pressure]
+    momentum_history = [state.momentum]
+    current_time = float(start_time)
+    sub_timestep = timestep / float(substeps)
+
+    for _ in range(steps):
+        for _ in range(substeps):
+            state = _rk4_step(
+                config,
+                state,
+                section=section,
+                mesh=mesh,
+                metrics=metrics,
+                atomic_mass=atomic_mass,
+                time=current_time,
+                timestep=sub_timestep,
+            )
+            current_time += sub_timestep
+        density_history.append(state.density)
+        pressure_history.append(state.pressure)
+        momentum_history.append(state.momentum)
+
+    return Fluid1DHistoryResult(
+        density_history=jnp.stack(density_history, axis=0),
+        pressure_history=jnp.stack(pressure_history, axis=0),
+        momentum_history=jnp.stack(momentum_history, axis=0),
+    )
+
+
+def apply_periodic_y_guards(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    result = jnp.asarray(field, dtype=jnp.float64)
+    if mesh.myg <= 0:
+        return result
+
+    interior = result[:, mesh.ystart : mesh.yend + 1, :]
+    for offset in range(mesh.myg):
+        result = result.at[:, mesh.ystart - 1 - offset, :].set(interior[:, -(offset + 1), :])
+        result = result.at[:, mesh.yend + 1 + offset, :].set(interior[:, offset, :])
+    return result
+
+
+def div_par_mod_periodic(
+    field: jnp.ndarray,
+    velocity: jnp.ndarray,
+    wave_speed: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+) -> jnp.ndarray:
+    field_interior = _interior_y(field, mesh)
+    velocity_interior = _interior_y(velocity, mesh)
+    wave_interior = _interior_y(wave_speed, mesh)
+    dy = _interior_y(metrics.dy, mesh)
+    J = _interior_y(metrics.J, mesh)
+    common = _face_common_factor(metrics, mesh)
+
+    field_left_cell, field_right_cell = _mc_cell_edges(field_interior)
+    velocity_left_cell, velocity_right_cell = _mc_cell_edges(velocity_interior)
+
+    field_left = field_right_cell
+    field_right = jnp.roll(field_left_cell, shift=-1, axis=1)
+    velocity_left = velocity_right_cell
+    velocity_right = jnp.roll(velocity_left_cell, shift=-1, axis=1)
+    amax = _face_wave_speed(velocity_interior, wave_interior)
+
+    flux = 0.5 * (
+        field_left * velocity_left
+        + field_right * velocity_right
+        + amax * (field_left - field_right)
+    )
+    return _scatter_face_divergence(flux, dy=dy, J=J, common=common, mesh=mesh)
+
+
+def div_par_fvv_periodic(
+    density: jnp.ndarray,
+    velocity: jnp.ndarray,
+    wave_speed: jnp.ndarray,
+    *,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+) -> jnp.ndarray:
+    density_interior = _interior_y(density, mesh)
+    velocity_interior = _interior_y(velocity, mesh)
+    wave_interior = _interior_y(wave_speed, mesh)
+    dy = _interior_y(metrics.dy, mesh)
+    J = _interior_y(metrics.J, mesh)
+    common = _face_common_factor(metrics, mesh)
+
+    density_left_cell, density_right_cell = _mc_cell_edges(density_interior)
+    velocity_left_cell, velocity_right_cell = _mc_cell_edges(velocity_interior)
+
+    density_left = density_right_cell
+    density_right = jnp.roll(density_left_cell, shift=-1, axis=1)
+    velocity_left = velocity_right_cell
+    velocity_right = jnp.roll(velocity_left_cell, shift=-1, axis=1)
+    amax = _face_wave_speed(velocity_interior, wave_interior)
+
+    flux = 0.5 * (
+        density_left * velocity_left * velocity_left
+        + density_right * velocity_right * velocity_right
+        + amax * (density_left * velocity_left - density_right * velocity_right)
+    )
+    return _scatter_face_divergence(flux, dy=dy, J=J, common=common, mesh=mesh)
+
+
+def grad_par_periodic(field: jnp.ndarray, *, mesh: StructuredMesh, metrics: StructuredMetrics) -> jnp.ndarray:
+    dy = _interior_y(metrics.dy, mesh)
+    field_interior = _interior_y(field, mesh)
+    if not jnp.allclose(dy, dy[:, :1, :], rtol=1e-12, atol=1e-12):
+        raise NotImplementedError("Native periodic fluid gradient currently requires uniform dy.")
+
+    spacing = dy[:, :1, :]
+    gradient = (jnp.roll(field_interior, shift=-1, axis=1) - jnp.roll(field_interior, shift=1, axis=1)) / (2.0 * spacing)
+    result = jnp.zeros_like(field, dtype=jnp.float64)
+    result = result.at[:, mesh.ystart : mesh.yend + 1, :].set(gradient)
+    return apply_periodic_y_guards(result, mesh)
+
+
+def _rk4_step(
+    config: BoutConfig,
+    state: Fluid1DState,
+    *,
+    section: str,
+    mesh: StructuredMesh,
+    metrics: StructuredMetrics,
+    atomic_mass: float,
+    time: float,
+    timestep: float,
+) -> Fluid1DState:
+    k1 = compute_mms_rhs(config, state, section=section, mesh=mesh, metrics=metrics, atomic_mass=atomic_mass, time=time)
+    k2 = compute_mms_rhs(
+        config,
+        _add_state(state, k1, scale=0.5 * timestep),
+        section=section,
+        mesh=mesh,
+        metrics=metrics,
+        atomic_mass=atomic_mass,
+        time=time + 0.5 * timestep,
+    )
+    k3 = compute_mms_rhs(
+        config,
+        _add_state(state, k2, scale=0.5 * timestep),
+        section=section,
+        mesh=mesh,
+        metrics=metrics,
+        atomic_mass=atomic_mass,
+        time=time + 0.5 * timestep,
+    )
+    k4 = compute_mms_rhs(
+        config,
+        _add_state(state, k3, scale=timestep),
+        section=section,
+        mesh=mesh,
+        metrics=metrics,
+        atomic_mass=atomic_mass,
+        time=time + timestep,
+    )
+
+    next_state = Fluid1DState(
+        density=state.density + (timestep / 6.0) * (k1.density + 2.0 * k2.density + 2.0 * k3.density + k4.density),
+        pressure=state.pressure
+        + (timestep / 6.0) * (k1.pressure + 2.0 * k2.pressure + 2.0 * k3.pressure + k4.pressure),
+        momentum=state.momentum
+        + (timestep / 6.0) * (k1.momentum + 2.0 * k2.momentum + 2.0 * k3.momentum + k4.momentum),
+    )
+    return Fluid1DState(
+        density=apply_periodic_y_guards(next_state.density, mesh),
+        pressure=apply_periodic_y_guards(next_state.pressure, mesh),
+        momentum=apply_periodic_y_guards(next_state.momentum, mesh),
+    )
+
+
+def _add_state(state: Fluid1DState, rhs: Fluid1DRhsResult, *, scale: float) -> Fluid1DState:
+    return Fluid1DState(
+        density=state.density + scale * rhs.density,
+        pressure=state.pressure + scale * rhs.pressure,
+        momentum=state.momentum + scale * rhs.momentum,
+    )
+
+
+def _interior_y(field: jnp.ndarray, mesh: StructuredMesh) -> jnp.ndarray:
+    return jnp.asarray(field, dtype=jnp.float64)[:, mesh.ystart : mesh.yend + 1, :]
+
+
+def _mc_cell_edges(field: jnp.ndarray) -> tuple[jnp.ndarray, jnp.ndarray]:
+    return _periodic_monotonic_centered_edges(field, axis=1)
+
+
+def _face_common_factor(metrics: StructuredMetrics, mesh: StructuredMesh) -> jnp.ndarray:
+    J = _interior_y(metrics.J, mesh)
+    g22 = _interior_y(metrics.g_22, mesh)
+    return (J + jnp.roll(J, shift=-1, axis=1)) / (jnp.sqrt(g22) + jnp.sqrt(jnp.roll(g22, shift=-1, axis=1)))
+
+
+def _face_wave_speed(velocity: jnp.ndarray, wave_speed: jnp.ndarray) -> jnp.ndarray:
+    return jnp.maximum(
+        jnp.maximum(wave_speed, jnp.roll(wave_speed, shift=-1, axis=1)),
+        jnp.maximum(jnp.abs(velocity), jnp.abs(jnp.roll(velocity, shift=-1, axis=1))),
+    )
+
+
+def _scatter_face_divergence(
+    flux: jnp.ndarray,
+    *,
+    dy: jnp.ndarray,
+    J: jnp.ndarray,
+    common: jnp.ndarray,
+    mesh: StructuredMesh,
+) -> jnp.ndarray:
+    flux_right = flux * common / (dy * J)
+    flux_left = jnp.roll(flux, shift=1, axis=1) * jnp.roll(common, shift=1, axis=1) / (dy * J)
+    result = jnp.zeros((mesh.nx, mesh.local_ny, mesh.nz), dtype=jnp.float64)
+    result = result.at[:, mesh.ystart : mesh.yend + 1, :].set(flux_right - flux_left)
+    return apply_periodic_y_guards(result, mesh)

drbx/native/hasegawa_wakatani.py

@@ -0,0 +1,196 @@
+"""JAX-native 2-D Hasegawa-Wakatani drift-wave turbulence model.
+
+This is the closed-field-line (periodic flux-tube) drift-wave turbulence
+flagship: a pseudo-spectral solver for the two-field Hasegawa-Wakatani system
+in the perpendicular plane,
+
+    d/dt zeta = -{phi, zeta} + alpha (phi - n) - nu * lap^2 zeta - mu * zeta
+    d/dt n    = -{phi, n} - kappa d/dy phi + alpha (phi - n) - nu * lap^2 n - mu * n
+
+with vorticity ``zeta = lap phi`` (so ``phi_k = -zeta_k / k^2``), adiabaticity
+``alpha``, background density gradient ``kappa``, hyperviscosity ``nu``, and an
+optional scale-independent friction ``mu`` (default 0). The friction models
+large-scale drag (e.g. sheath or neutral damping in the tokamak edge) and, by
+absorbing the 2-D inverse cascade at the box scale, lets a fixed-step run reach
+a statistically steady saturated state.
+``{a, b} = da/dx db/dy - da/dy db/dx`` is the E x B Poisson bracket, evaluated
+pseudo-spectrally with 2/3-rule dealiasing.
+
+The whole right-hand side is written in JAX, so a run is ``jit``-compiled and
+differentiable. Its single-mode linear growth rate reproduces the eigenvalue of
+:func:`drbx.linear.resistive_drift_wave_operator` (benchmark B2); at finite
+amplitude it develops nonlinear E x B transport with an outward particle flux
+that saturates statistically when ``mu > 0`` absorbs the inverse cascade.
+
+Initial spectra must be Hermitian (e.g. the FFT of a real field). The solver
+evolves the complex spectral state as given; a non-Hermitian state corresponds
+to complex-valued fields, i.e. an unphysical complexified system that does not
+obey the real system's energy balance and can blow up nonlinearly.
+
+References: Hasegawa & Wakatani, Phys. Rev. Lett. 50, 682 (1983); Numata
+et al., Phys. Plasmas 14, 102312 (2007).
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+
+__all__ = [
+    "HasegawaWakataniParameters",
+    "HasegawaWakataniGrid",
+    "hw_grid",
+    "potential_from_vorticity",
+    "hw_rhs",
+    "hw_step",
+    "hw_run",
+    "hw_run_flux_history",
+    "particle_flux",
+]
+
+
+@dataclass(frozen=True)
+class HasegawaWakataniParameters:
+    """Physical parameters of the Hasegawa-Wakatani system."""
+
+    adiabaticity: float = 1.0        # alpha: parallel coupling
+    gradient: float = 1.0            # kappa: background density gradient drive
+    hyperviscosity: float = 1.0e-3   # nu: lap^2 damping of the grid scale
+    friction: float = 0.0            # mu: linear drag absorbing the inverse cascade
+
+
+@dataclass(frozen=True)
+class HasegawaWakataniGrid:
+    """Periodic perpendicular grid and its spectral wavenumbers."""
+
+    n: int
+    length: float
+    kx: jnp.ndarray
+    ky: jnp.ndarray
+    k2: jnp.ndarray
+    inv_k2: jnp.ndarray
+    dealias: jnp.ndarray
+
+
+def hw_grid(n: int, length: float) -> HasegawaWakataniGrid:
+    """Build an ``n x n`` periodic grid of side ``length`` with spectral operators."""
+
+    modes = jnp.fft.fftfreq(n, d=length / n) * 2.0 * jnp.pi
+    kx = modes[:, None] * jnp.ones((1, n))
+    ky = jnp.ones((n, 1)) * modes[None, :]
+    k2 = kx**2 + ky**2
+    inv_k2 = jnp.where(k2 > 0, 1.0 / jnp.where(k2 > 0, k2, 1.0), 0.0)
+    # 2/3-rule dealiasing mask.
+    kmax = jnp.max(jnp.abs(modes)) * (2.0 / 3.0)
+    dealias = (jnp.abs(kx) <= kmax) & (jnp.abs(ky) <= kmax)
+    return HasegawaWakataniGrid(
+        n=n, length=length, kx=kx, ky=ky, k2=k2, inv_k2=inv_k2,
+        dealias=dealias.astype(jnp.float64),
+    )
+
+
+def potential_from_vorticity(zeta_hat: jnp.ndarray, grid: HasegawaWakataniGrid) -> jnp.ndarray:
+    """Invert ``lap phi = zeta`` in Fourier space: ``phi_k = -zeta_k / k^2``."""
+
+    return -zeta_hat * grid.inv_k2
+
+
+def _bracket_hat(phi_hat, field_hat, grid: HasegawaWakataniGrid) -> jnp.ndarray:
+    """Poisson bracket ``{phi, field}`` returned in Fourier space (dealiased)."""
+
+    dphi_dx = jnp.fft.ifft2(1j * grid.kx * phi_hat)
+    dphi_dy = jnp.fft.ifft2(1j * grid.ky * phi_hat)
+    df_dx = jnp.fft.ifft2(1j * grid.kx * field_hat)
+    df_dy = jnp.fft.ifft2(1j * grid.ky * field_hat)
+    bracket = dphi_dx * df_dy - dphi_dy * df_dx
+    return jnp.fft.fft2(bracket) * grid.dealias
+
+
+def hw_rhs(zeta_hat, n_hat, grid: HasegawaWakataniGrid, params: HasegawaWakataniParameters):
+    """Spectral right-hand side ``(d zeta_hat/dt, d n_hat/dt)``."""
+
+    phi_hat = potential_from_vorticity(zeta_hat, grid)
+    coupling = params.adiabaticity * (phi_hat - n_hat)
+    damping_zeta = params.hyperviscosity * grid.k2**2 * zeta_hat + params.friction * zeta_hat
+    damping_n = params.hyperviscosity * grid.k2**2 * n_hat + params.friction * n_hat
+    d_zeta = -_bracket_hat(phi_hat, zeta_hat, grid) + coupling - damping_zeta
+    d_n = (
+        -_bracket_hat(phi_hat, n_hat, grid)
+        - params.gradient * (1j * grid.ky) * phi_hat
+        + coupling
+        - damping_n
+    )
+    return d_zeta, d_n
+
+
+def hw_step(zeta_hat, n_hat, grid, params, dt):
+    """One classical RK4 step of the spectral system."""
+
+    def rhs(z, m):
+        return hw_rhs(z, m, grid, params)
+
+    k1z, k1n = rhs(zeta_hat, n_hat)
+    k2z, k2n = rhs(zeta_hat + 0.5 * dt * k1z, n_hat + 0.5 * dt * k1n)
+    k3z, k3n = rhs(zeta_hat + 0.5 * dt * k2z, n_hat + 0.5 * dt * k2n)
+    k4z, k4n = rhs(zeta_hat + dt * k3z, n_hat + dt * k3n)
+    zeta_next = zeta_hat + (dt / 6.0) * (k1z + 2 * k2z + 2 * k3z + k4z)
+    n_next = n_hat + (dt / 6.0) * (k1n + 2 * k2n + 2 * k3n + k4n)
+    return zeta_next, n_next
+
+
+def hw_run(zeta_hat, n_hat, grid, params, *, dt, steps):
+    """Advance ``steps`` RK4 steps with a jitted ``lax.scan``; return final spectra.
+
+    ``grid``, ``params``, ``dt``, and ``steps`` are captured as constants, so the
+    compiled function depends only on the two spectral state arrays.
+    """
+
+    @jax.jit
+    def _run(zeta0, n0):
+        def body(carry, _):
+            z, m = carry
+            return hw_step(z, m, grid, params, dt), None
+
+        (zeta_final, n_final), _ = jax.lax.scan(body, (zeta0, n0), None, length=steps)
+        return zeta_final, n_final
+
+    return _run(zeta_hat, n_hat)
+
+
+def hw_run_flux_history(zeta_hat, n_hat, grid, params, *, dt, steps, sample_every=1):
+    """Advance ``steps`` RK4 steps; return final spectra and the sampled flux history.
+
+    Like :func:`hw_run`, but the jitted ``lax.scan`` also records the
+    domain-averaged radial particle flux every ``sample_every`` steps (an array
+    of length ``steps // sample_every``). Because the whole scan is JAX, any
+    reduction of the history (e.g. its time average over a saturated window) is
+    differentiable with respect to the physical parameters.
+    """
+
+    outer = steps // sample_every
+
+    @jax.jit
+    def _run(zeta0, n0):
+        def inner(carry, _):
+            z, m = carry
+            return hw_step(z, m, grid, params, dt), None
+
+        def body(carry, _):
+            carry, _ = jax.lax.scan(inner, carry, None, length=sample_every)
+            return carry, particle_flux(carry[0], carry[1], grid)
+
+        (zeta_final, n_final), fluxes = jax.lax.scan(body, (zeta0, n0), None, length=outer)
+        return zeta_final, n_final, fluxes
+
+    return _run(zeta_hat, n_hat)
+
+
+def particle_flux(zeta_hat, n_hat, grid: HasegawaWakataniGrid) -> jnp.ndarray:
+    """Domain-averaged radial E x B particle flux ``<n v_x> = <n * (-d phi/dy)>``."""
+
+    phi_hat = potential_from_vorticity(zeta_hat, grid)
+    v_x = jnp.real(jnp.fft.ifft2(-1j * grid.ky * phi_hat))
+    density = jnp.real(jnp.fft.ifft2(n_hat))
+    return jnp.mean(density * v_x)

drbx/native/limiters.py

@@ -0,0 +1,57 @@
+from __future__ import annotations
+
+import jax.numpy as jnp
+import numpy as np
+
+
+def minmod3_jax(a: jnp.ndarray, b: jnp.ndarray, c: jnp.ndarray) -> jnp.ndarray:
+    same_sign = (a * b > 0.0) & (a * c > 0.0)
+    magnitude = jnp.minimum(jnp.abs(a), jnp.minimum(jnp.abs(b), jnp.abs(c)))
+    return jnp.where(same_sign, jnp.sign(a) * magnitude, 0.0)
+
+
+def monotonic_centered_edges_jax(
+    center: jnp.ndarray,
+    minus: jnp.ndarray,
+    plus: jnp.ndarray,
+) -> tuple[jnp.ndarray, jnp.ndarray]:
+    slope = minmod3_jax(2.0 * (plus - center), 0.5 * (plus - minus), 2.0 * (center - minus))
+    return center - 0.5 * slope, center + 0.5 * slope
+
+
+def periodic_monotonic_centered_edges_jax(field: jnp.ndarray, *, axis: int) -> tuple[jnp.ndarray, jnp.ndarray]:
+    center = jnp.asarray(field, dtype=jnp.float64)
+    minus = jnp.roll(center, shift=1, axis=axis)
+    plus = jnp.roll(center, shift=-1, axis=axis)
+    return monotonic_centered_edges_jax(center, minus, plus)
+
+
+def minmod3_numpy(a: np.ndarray, b: np.ndarray, c: np.ndarray) -> np.ndarray:
+    same_sign = (a * b > 0.0) & (a * c > 0.0)
+    magnitude = np.minimum(np.abs(a), np.minimum(np.abs(b), np.abs(c)))
+    return np.where(same_sign, np.sign(a) * magnitude, 0.0)
+
+
+def monotonic_centered_edges_numpy(
+    center: np.ndarray,
+    minus: np.ndarray,
+    plus: np.ndarray,
+) -> tuple[np.ndarray, np.ndarray]:
+    slope = minmod3_numpy(2.0 * (plus - center), 0.5 * (plus - minus), 2.0 * (center - minus))
+    return center - 0.5 * slope, center + 0.5 * slope
+
+
+def minmod3_scalar(a: float, b: float, c: float) -> float:
+    if (a * b <= 0.0) or (a * c <= 0.0):
+        return 0.0
+    magnitude = min(abs(a), abs(b), abs(c))
+    return float(np.sign(a) * magnitude)
+
+
+def monotonic_centered_edges_scalar(center: float, minus: float, plus: float) -> tuple[float, float]:
+    slope = minmod3_scalar(
+        2.0 * (plus - center),
+        0.5 * (plus - minus),
+        2.0 * (center - minus),
+    )
+    return center - 0.5 * slope, center + 0.5 * slope