drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/neutrals/detachment_sol_model.py

@@ -0,0 +1,221 @@
+"""Self-consistent detaching 1D SOL model (evolved temperature).
+
+Extends the recycling SOL to an *evolved* plasma pressure so the target
+temperature responds to the plasma conditions -- the ingredient a detachment
+study needs. Along the field it evolves ion density, parallel momentum, plasma
+pressure, and a diffusive neutral density, with:
+
+- **Spitzer parallel conduction** ``kappa ~ T^{5/2}`` solved implicitly (a solvax
+  tridiagonal), so the stiff parabolic heat transport is unconditionally stable;
+- **radiative / ionization energy loss** from the AMJUEL fits, applied as a
+  *self-limiting* semi-implicit sink (``P <- P / (1 + dt * loss_rate)``) so the
+  loss cannot drive the pressure negative and switches off as the plasma cools;
+- a **Bohm sheath heat sink** at the target and an **upstream power source**; and
+- the neutral **recycling**, parallel diffusion, and ionization/recombination /
+  charge-exchange coupling of the recycling model.
+
+As the upstream density rises at fixed upstream power the target cools, crosses
+into the recombining regime below ~1 eV, and the target ion flux rolls over --
+the classic SD1D detachment signature. Everything is pure ``jax.numpy`` (with the
+solvax tridiagonal solves) and differentiable. Fields are hermes-3 normalized;
+time is normalized to the parallel transit time.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+
+from solvax import tridiagonal_solve
+
+from .atomic_rates import (
+    charge_exchange_rate_coefficient,
+    energy_loss_coefficient,
+    rate_coefficient,
+)
+from .reactions import PlasmaNormalization
+
+__all__ = [
+    "DetachmentSolParameters",
+    "DetachmentSolState",
+    "detachment_sol_step",
+    "detachment_sol_run",
+    "detachment_diagnostics",
+]
+
+
+@dataclass(frozen=True)
+class DetachmentSolParameters:
+    """Parameters of the self-consistent detaching SOL model."""
+
+    parallel_length: float = 30.0
+    upstream_density: float = 4.0
+    upstream_power: float = 6.0              # normalized upstream power source
+    power_width: float = 0.2                 # parallel width of the power source
+    conduction_coefficient: float = 2.0      # Spitzer kappa0 (normalized)
+    sheath_transmission: float = 7.0         # gamma (Te + Ti)
+    neutral_diffusion: float = 8.0
+    recycling_fraction: float = 0.95
+    ion_mass: float = 2.0
+    density_floor: float = 1.0e-4
+    pressure_floor: float = 1.0e-5
+    temperature_floor: float = 1.0e-3
+    # A scrape-off-layer reference temperature (upstream ~ tens of eV); the
+    # detachment transition and target-flux rollover live in this regime.
+    normalization: PlasmaNormalization = PlasmaNormalization(Tnorm=50.0)
+
+
+class DetachmentSolState(NamedTuple):
+    """Evolved fields: ion density, ion momentum, plasma pressure, neutral density."""
+
+    ion_density: jnp.ndarray
+    ion_momentum: jnp.ndarray
+    plasma_pressure: jnp.ndarray
+    neutral_density: jnp.ndarray
+
+
+def _temperature(density, pressure, params):
+    return jnp.maximum(pressure / (2.0 * jnp.maximum(density, params.density_floor)), params.temperature_floor)
+
+
+def _sound_speed(density, pressure, params):
+    return jnp.sqrt(jnp.maximum(pressure / (params.ion_mass * jnp.maximum(density, params.density_floor)), 0.0))
+
+
+def _hyperbolic_rhs(density, momentum, pressure, params):
+    mass = params.ion_mass
+    dz = 1.0 / density.shape[0]
+    velocity = momentum / (mass * jnp.maximum(density, params.density_floor))
+    speed = jnp.abs(velocity) + _sound_speed(density, pressure, params)
+    face_speed = jnp.maximum(speed[:-1], speed[1:])
+    flux_n = 0.5 * (density[:-1] * velocity[:-1] + density[1:] * velocity[1:]) - 0.5 * face_speed * (density[1:] - density[:-1])
+    flux_m = 0.5 * (momentum[:-1] * velocity[:-1] + pressure[:-1] + momentum[1:] * velocity[1:] + pressure[1:]) - 0.5 * face_speed * (momentum[1:] - momentum[:-1])
+    flux_p = 0.5 * ((5.0 / 3.0) * pressure[:-1] * velocity[:-1] + (5.0 / 3.0) * pressure[1:] * velocity[1:]) - 0.5 * face_speed * (pressure[1:] - pressure[:-1])
+    target_velocity = jnp.maximum(velocity[-1], _sound_speed(density, pressure, params)[-1])
+    flux_n = jnp.concatenate([jnp.zeros(1), flux_n, (density[-1] * target_velocity)[None]])
+    flux_m = jnp.concatenate([pressure[:1], flux_m, (mass * density[-1] * target_velocity**2 + pressure[-1])[None]])
+    flux_p = jnp.concatenate([jnp.zeros(1), flux_p, ((5.0 / 3.0) * pressure[-1] * target_velocity)[None]])
+    pressure_gradient = jnp.gradient(pressure, dz)
+    d_density = -(flux_n[1:] - flux_n[:-1]) / dz
+    d_momentum = -(flux_m[1:] - flux_m[:-1]) / dz
+    d_pressure = -(flux_p[1:] - flux_p[:-1]) / dz + (2.0 / 3.0) * velocity * pressure_gradient
+    return d_density, d_momentum, d_pressure
+
+
+def _implicit_diffusion(field, face_diffusivity, dt):
+    n = field.shape[0]
+    dz = 1.0 / n
+    coefficient = dt / dz**2
+    lower = jnp.zeros(n).at[1:].set(-coefficient * face_diffusivity)
+    upper = jnp.zeros(n).at[:-1].set(-coefficient * face_diffusivity)
+    diagonal = jnp.ones(n).at[1:-1].set(1.0 + coefficient * (face_diffusivity[:-1] + face_diffusivity[1:]))
+    diagonal = diagonal.at[0].set(1.0 + coefficient * face_diffusivity[0])
+    diagonal = diagonal.at[-1].set(1.0 + coefficient * face_diffusivity[-1])
+    return tridiagonal_solve(lower, diagonal, upper, field, method="thomas")
+
+
+def detachment_sol_step(state, params, dt):
+    """Advance the self-consistent detaching SOL one operator-split step."""
+
+    norm = params.normalization
+    nz = state.ion_density.shape[0]
+    dz = 1.0 / nz
+    z = (jnp.arange(nz, dtype=jnp.float64) + 0.5) / nz
+    density_floor = params.density_floor
+    pressure_floor = params.pressure_floor
+    rate_scale = norm.Nnorm * (params.parallel_length / norm.sound_speed)
+    energy_scale = rate_scale / norm.Tnorm
+
+    density = jnp.maximum(state.ion_density, density_floor)
+    momentum = state.ion_momentum
+    pressure = jnp.maximum(state.plasma_pressure, pressure_floor)
+    neutral_density = state.neutral_density
+
+    # Hyperbolic transport + upstream power source.
+    d_density, d_momentum, d_pressure = _hyperbolic_rhs(density, momentum, pressure, params)
+    power_source = params.upstream_power * jnp.exp(-(z**2) / params.power_width**2)
+    density = jnp.maximum(density + dt * d_density, density_floor)
+    momentum = momentum + dt * d_momentum
+    pressure = jnp.maximum(pressure + dt * d_pressure + dt * (2.0 / 3.0) * power_source, pressure_floor)
+
+    # Atomic reactions: particle exchange (implicit) + momentum friction.
+    temperature = _temperature(density, pressure, params)
+    physical_temperature = temperature * norm.Tnorm
+    electron_density_m3 = jnp.maximum(density, density_floor) * norm.Nnorm
+    electron = jnp.maximum(density, density_floor)
+    ionization = rate_coefficient("d", "iz", physical_temperature, electron_density_m3) * rate_scale
+    recombination = rate_coefficient("d", "rec", physical_temperature, electron_density_m3) * rate_scale
+    ionization_frequency = ionization * electron
+    recombination_frequency = recombination * electron
+    total = density + neutral_density
+    density = jnp.maximum((density + dt * ionization_frequency * total) / (1.0 + dt * (ionization_frequency + recombination_frequency)), density_floor)
+    neutral_density = jnp.maximum(total - density, 0.0)
+    effective_temperature = jnp.clip(2.0 * temperature * norm.Tnorm / params.ion_mass, 0.01, 1.0e4)
+    charge_exchange_frequency = charge_exchange_rate_coefficient(effective_temperature) * electron * rate_scale
+    momentum = momentum / (1.0 + dt * (charge_exchange_frequency * neutral_density + recombination_frequency))
+
+    # Self-limiting radiative / ionization energy loss.
+    ionization_energy = energy_loss_coefficient("d", "iz", physical_temperature, electron_density_m3)
+    recombination_energy = energy_loss_coefficient("d", "rec", physical_temperature, electron_density_m3)
+    energy_loss = (neutral_density * electron * ionization_energy + electron * electron * recombination_energy) * energy_scale
+    loss_rate = (2.0 / 3.0) * jnp.maximum(energy_loss, 0.0) / jnp.maximum(pressure, pressure_floor)
+    pressure = pressure / (1.0 + dt * loss_rate)
+
+    # Bohm sheath heat sink at the target (semi-implicit).
+    sound_speed = _sound_speed(density, pressure, params)
+    sheath_rate = jnp.zeros(nz).at[-1].set(
+        (2.0 / 3.0) * params.sheath_transmission * density[-1] * sound_speed[-1] * _temperature(density, pressure, params)[-1] / dz / jnp.maximum(pressure[-1], pressure_floor)
+    )
+    pressure = pressure / (1.0 + dt * sheath_rate)
+
+    # Implicit Spitzer conduction.
+    temperature = _temperature(density, pressure, params)
+    conductivity = params.conduction_coefficient * temperature**2.5
+    diffusivity = conductivity / (3.0 * jnp.maximum(density, density_floor))
+    pressure = jnp.maximum(_implicit_diffusion(pressure, 0.5 * (diffusivity[:-1] + diffusivity[1:]), dt), pressure_floor)
+
+    density = density.at[0].set(params.upstream_density)
+
+    # Neutral parallel diffusion + target recycling.
+    neutral_diffusivity = jnp.full(nz - 1, params.neutral_diffusion * 0.04 / (params.ion_mass * 0.05))
+    sound_speed = _sound_speed(density, pressure, params)
+    recycling_source = jnp.zeros(nz).at[-1].set(params.recycling_fraction * density[-1] * sound_speed[-1] / dz)
+    neutral_density = jnp.maximum(_implicit_diffusion(neutral_density + dt * recycling_source, neutral_diffusivity, dt), 0.0)
+
+    return DetachmentSolState(density, momentum, pressure, neutral_density)
+
+
+def detachment_sol_run(state, params, *, dt, steps):
+    """Advance ``steps`` operator-split steps with a jitted ``lax.scan``."""
+
+    @jax.jit
+    def _run(initial):
+        def body(carry, _):
+            return detachment_sol_step(carry, params, dt), None
+
+        final, _ = jax.lax.scan(body, initial, None, length=steps)
+        return final
+
+    return _run(state)
+
+
+class DetachmentDiagnostics(NamedTuple):
+    target_ion_flux: jnp.ndarray
+    target_temperature_ev: jnp.ndarray
+    target_density: jnp.ndarray
+
+
+def detachment_diagnostics(state, params) -> DetachmentDiagnostics:
+    """Target ion flux (``n c_s``), target temperature [eV], and target density."""
+
+    density = jnp.maximum(state.ion_density, params.density_floor)
+    temperature = _temperature(density, jnp.maximum(state.plasma_pressure, params.pressure_floor), params)
+    sound_speed = _sound_speed(density, jnp.maximum(state.plasma_pressure, params.pressure_floor), params)
+    return DetachmentDiagnostics(
+        target_ion_flux=density[-1] * sound_speed[-1],
+        target_temperature_ev=temperature[-1] * params.normalization.Tnorm,
+        target_density=density[-1],
+    )

drbx/native/neutrals/reactions.py

@@ -0,0 +1,164 @@
+"""Plasma <-> neutral atomic reaction sources for the 1D SOL model.
+
+Assembles the ionization / recombination / charge-exchange source channels that
+couple a hydrogenic plasma fluid (ion density ``Ni``, parallel momentum, ion
+temperature ``Ti``, electron temperature ``Te``) to a neutral fluid (``Nn``,
+momentum, ``Tn``), following the hermes-3 hydrogen reaction closure. Particle,
+momentum, and (thermal) energy transfers use the Galilean-invariant form: each
+particle transfer carries the source species' momentum ``m V`` and thermal
+energy ``1.5 T``, and charge exchange adds a frictional heating ``0.5 m R dV^2``
+from the ion-atom velocity difference. The electron channel is the ionization
+cost / recombination radiation from the AMJUEL energy-loss fits.
+
+All fields are in hermes-3 normalized units (density / ``Nnorm``, temperature /
+``Tnorm`` [eV], velocity / ``Cs0``); the returned source rates are normalized to
+``Nnorm * Omega_ci`` (particle / momentum) and ``Nnorm * Tnorm * Omega_ci``
+(energy). Everything is pure ``jax.numpy`` and differentiable.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from .atomic_rates import (
+    charge_exchange_rate_coefficient,
+    energy_loss_coefficient,
+    rate_coefficient,
+)
+
+__all__ = ["PlasmaNormalization", "HydrogenReactionSources", "compute_hydrogen_reaction_sources"]
+
+_ELEMENTARY_CHARGE = 1.602176634e-19
+_PROTON_MASS = 1.67262192369e-27
+
+
+@dataclass(frozen=True)
+class PlasmaNormalization:
+    """hermes-3 normalization constants (deck defaults: 1e19 m^-3, 100 eV, 1 T)."""
+
+    Nnorm: float = 1.0e19
+    Tnorm: float = 100.0
+    Bnorm: float = 1.0
+    ion_mass: float = 2.0  # deuterium, in proton masses
+
+    @property
+    def sound_speed(self) -> float:
+        return float((_ELEMENTARY_CHARGE * self.Tnorm / _PROTON_MASS) ** 0.5)
+
+    @property
+    def cyclotron_frequency(self) -> float:
+        return float(_ELEMENTARY_CHARGE * self.Bnorm / _PROTON_MASS)
+
+
+@dataclass(frozen=True)
+class HydrogenReactionSources:
+    """Normalized reaction source channels for the ion, neutral, and electron fluids."""
+
+    ion_density: jnp.ndarray
+    neutral_density: jnp.ndarray
+    ion_momentum: jnp.ndarray
+    neutral_momentum: jnp.ndarray
+    ion_energy: jnp.ndarray
+    neutral_energy: jnp.ndarray
+    electron_energy: jnp.ndarray
+    ionization_rate: jnp.ndarray
+    recombination_rate: jnp.ndarray
+    charge_exchange_rate: jnp.ndarray
+
+
+def compute_hydrogen_reaction_sources(
+    ion_density: jnp.ndarray,
+    ion_velocity: jnp.ndarray,
+    ion_temperature: jnp.ndarray,
+    electron_temperature: jnp.ndarray,
+    neutral_density: jnp.ndarray,
+    neutral_velocity: jnp.ndarray,
+    neutral_temperature: jnp.ndarray,
+    *,
+    normalization: PlasmaNormalization = PlasmaNormalization(),
+    charge_exchange_multiplier: float = 1.0,
+) -> HydrogenReactionSources:
+    """Return the normalized ion/neutral/electron reaction source channels."""
+
+    mass = float(normalization.ion_mass)
+    nnorm = float(normalization.Nnorm)
+    tnorm = float(normalization.Tnorm)
+    omega = normalization.cyclotron_frequency
+    rate_scale = nnorm / omega
+    energy_scale = nnorm / (tnorm * omega)
+
+    # Physical inputs for the rate fits.
+    electron_temperature_ev = electron_temperature * tnorm
+    electron_density_m3 = ion_density * nnorm  # quasineutral, Z = 1
+
+    ionization_sigmav = rate_coefficient("d", "iz", electron_temperature_ev, electron_density_m3)
+    recombination_sigmav = rate_coefficient("d", "rec", electron_temperature_ev, electron_density_m3)
+    ionization_energy = energy_loss_coefficient("d", "iz", electron_temperature_ev, electron_density_m3)
+    recombination_energy = energy_loss_coefficient("d", "rec", electron_temperature_ev, electron_density_m3)
+
+    effective_temperature = jnp.clip(
+        (neutral_temperature / mass + ion_temperature / mass) * tnorm, 0.01, 1.0e4
+    )
+    cx_sigmav = charge_exchange_rate_coefficient(effective_temperature)
+
+    # Normalized volumetric reaction rates.
+    ionization_rate = neutral_density * ion_density * ionization_sigmav * rate_scale
+    recombination_rate = ion_density * ion_density * recombination_sigmav * rate_scale
+    charge_exchange_rate = (
+        ion_density * neutral_density * cx_sigmav * rate_scale * float(charge_exchange_multiplier)
+    )
+    ionization_radiation = neutral_density * ion_density * ionization_energy * energy_scale
+    recombination_radiation = ion_density * ion_density * recombination_energy * energy_scale
+
+    # Momentum carried by each particle transfer (source species' m V).
+    ionization_momentum = ionization_rate * mass * neutral_velocity
+    recombination_momentum = recombination_rate * mass * ion_velocity
+    cx_neutral_momentum = charge_exchange_rate * mass * neutral_velocity
+    cx_ion_momentum = charge_exchange_rate * mass * ion_velocity
+
+    # Thermal energy carried by each particle transfer, plus CX frictional heating.
+    ionization_energy_transfer = 1.5 * ionization_rate * neutral_temperature
+    recombination_energy_transfer = 1.5 * recombination_rate * ion_temperature
+    cx_neutral_thermal = 1.5 * charge_exchange_rate * neutral_temperature
+    cx_ion_thermal = 1.5 * charge_exchange_rate * ion_temperature
+    velocity_delta = ion_velocity - neutral_velocity
+    cx_frictional_heat = 0.5 * mass * charge_exchange_rate * velocity_delta**2
+
+    ion_density_source = ionization_rate - recombination_rate
+    neutral_density_source = -ionization_rate + recombination_rate
+    ion_momentum_source = (
+        ionization_momentum - recombination_momentum + cx_neutral_momentum - cx_ion_momentum
+    )
+    neutral_momentum_source = (
+        -ionization_momentum + recombination_momentum - cx_neutral_momentum + cx_ion_momentum
+    )
+    ion_energy_source = (
+        ionization_energy_transfer
+        - recombination_energy_transfer
+        + cx_neutral_thermal
+        - cx_ion_thermal
+        + cx_frictional_heat
+    )
+    neutral_energy_source = (
+        -ionization_energy_transfer
+        + recombination_energy_transfer
+        - cx_neutral_thermal
+        + cx_ion_thermal
+        + cx_frictional_heat
+    )
+    electron_energy_source = -ionization_radiation - recombination_radiation
+
+    return HydrogenReactionSources(
+        ion_density=ion_density_source,
+        neutral_density=neutral_density_source,
+        ion_momentum=ion_momentum_source,
+        neutral_momentum=neutral_momentum_source,
+        ion_energy=ion_energy_source,
+        neutral_energy=neutral_energy_source,
+        electron_energy=electron_energy_source,
+        ionization_rate=ionization_rate,
+        recombination_rate=recombination_rate,
+        charge_exchange_rate=charge_exchange_rate,
+    )

drbx/native/neutrals/recycling_sol_model.py

@@ -0,0 +1,197 @@
+"""Coupled 1D scrape-off-layer recycling model: plasma + diffusive neutral.
+
+A reduced SOL model that couples a hydrogenic plasma fluid to a recycled neutral
+fluid through the packaged atomic reactions. Along the field coordinate ``z``
+(stagnation midplane at ``z = 0``, target plate at ``z = L``) it evolves the ion
+density and parallel momentum on a *prescribed* parallel temperature profile
+(hot upstream, cold target -- the imposed-temperature closure common in reduced
+SOL studies, which sidesteps the stiff self-consistent conduction/radiation
+energy balance), together with a neutral density that is:
+
+- **recycled** from the Bohm ion flux at the target (fraction ``R``),
+- **transported** by parallel diffusion (solved implicitly with a tridiagonal
+  solve, so the parabolic term is unconditionally stable), and
+- **ionized / recombined** back into the plasma via the AMJUEL rate coefficients
+  (an operator-split, per-cell implicit update that is stable against the stiff
+  ionization source).
+
+The result is an attached recycling SOL: neutrals born at the target penetrate
+upstream, ionize where the plasma is hot, and the plasma flows back to the target
+at the Bohm speed. Self-consistent detachment (an *evolved* temperature with
+conduction and radiative rollover) is a further extension.
+
+Everything is pure ``jax.numpy`` (with the solvax tridiagonal solve) and therefore
+``jit``/``grad``/``vmap`` transparent. Fields are hermes-3 normalized
+(density / ``Nnorm``, temperature / ``Tnorm``, velocity / the reference sound
+speed); the reaction rates use physical temperature and density internally.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import NamedTuple
+
+import jax
+import jax.numpy as jnp
+
+from solvax import tridiagonal_solve
+
+from .atomic_rates import charge_exchange_rate_coefficient, rate_coefficient
+from .reactions import PlasmaNormalization
+
+__all__ = [
+    "SolRecyclingParameters",
+    "SolRecyclingState",
+    "linear_target_temperature_profile",
+    "sol_recycling_step",
+    "sol_recycling_run",
+    "target_ion_flux",
+]
+
+
+@dataclass(frozen=True)
+class SolRecyclingParameters:
+    """Parameters of the reduced recycling SOL model."""
+
+    parallel_length: float = 30.0            # m, stagnation -> target
+    upstream_density: float = 1.0            # fixed upstream ion density (normalized)
+    recycling_fraction: float = 1.0          # target recycling fraction R
+    neutral_diffusion: float = 8.0           # neutral parallel diffusion enhancement
+    neutral_temperature: float = 0.04        # normalized (cold, ~2 eV) neutral temperature
+    ion_mass: float = 2.0                    # deuterium
+    density_floor: float = 1.0e-4
+    normalization: PlasmaNormalization = PlasmaNormalization()
+
+
+class SolRecyclingState(NamedTuple):
+    """Evolved fields: ion density, ion momentum (A n v), neutral density."""
+
+    ion_density: jnp.ndarray
+    ion_momentum: jnp.ndarray
+    neutral_density: jnp.ndarray
+
+
+def linear_target_temperature_profile(nz: int, *, upstream_ev: float = 30.0, target_ev: float = 2.0,
+                                      normalization: PlasmaNormalization = PlasmaNormalization()) -> jnp.ndarray:
+    """Prescribed parallel temperature (normalized): quadratic hot->cold to the target."""
+
+    z = (jnp.arange(nz, dtype=jnp.float64) + 0.5) / nz
+    profile_ev = upstream_ev - (upstream_ev - target_ev) * z**2
+    return profile_ev / normalization.Tnorm
+
+
+def _sound_speed(density, temperature, params):
+    return jnp.sqrt(jnp.maximum(2.0 * temperature / params.ion_mass, 0.0))
+
+
+def _velocity(momentum, density, params):
+    return momentum / (params.ion_mass * jnp.maximum(density, params.density_floor))
+
+
+def _hyperbolic_rhs(density, momentum, temperature, params):
+    mass = params.ion_mass
+    dz = 1.0 / density.shape[0]
+    velocity = _velocity(momentum, density, params)
+    pressure = 2.0 * density * temperature
+    speed = jnp.abs(velocity) + _sound_speed(density, temperature, params)
+    face_speed = jnp.maximum(speed[:-1], speed[1:])
+    flux_n = 0.5 * (density[:-1] * velocity[:-1] + density[1:] * velocity[1:]) - 0.5 * face_speed * (density[1:] - density[:-1])
+    flux_m = 0.5 * (momentum[:-1] * velocity[:-1] + pressure[:-1] + momentum[1:] * velocity[1:] + pressure[1:]) - 0.5 * face_speed * (momentum[1:] - momentum[:-1])
+    # noflow stagnation at z=0 (pressure force only); Bohm outflow at the target.
+    target_velocity = jnp.maximum(velocity[-1], _sound_speed(density, temperature, params)[-1])
+    flux_n = jnp.concatenate([jnp.zeros(1), flux_n, (density[-1] * target_velocity)[None]])
+    flux_m = jnp.concatenate([pressure[:1], flux_m, (mass * density[-1] * target_velocity**2 + pressure[-1])[None]])
+    d_density = -(flux_n[1:] - flux_n[:-1]) / dz
+    d_momentum = -(flux_m[1:] - flux_m[:-1]) / dz
+    return d_density, d_momentum
+
+
+def _implicit_diffusion(field, face_diffusivity, source, dt):
+    """Backward-Euler parallel diffusion with no-flux ends (tridiagonal solve)."""
+
+    n = field.shape[0]
+    dz = 1.0 / n
+    coefficient = dt / dz**2
+    lower = jnp.zeros(n).at[1:].set(-coefficient * face_diffusivity)
+    upper = jnp.zeros(n).at[:-1].set(-coefficient * face_diffusivity)
+    diagonal = jnp.ones(n).at[1:-1].set(1.0 + coefficient * (face_diffusivity[:-1] + face_diffusivity[1:]))
+    diagonal = diagonal.at[0].set(1.0 + coefficient * face_diffusivity[0])
+    diagonal = diagonal.at[-1].set(1.0 + coefficient * face_diffusivity[-1])
+    return tridiagonal_solve(lower, diagonal, upper, field + dt * source, method="thomas")
+
+
+def _reaction_update(density, momentum, neutral_density, temperature, params, dt):
+    """Operator-split implicit ionization/recombination particle exchange, plus the
+    charge-exchange + recombination momentum friction that drags the plasma flow
+    toward the (stationary) neutrals -- the drag that keeps the target sonic."""
+
+    norm = params.normalization
+    physical_temperature = temperature * norm.Tnorm
+    density_floor = params.density_floor
+    electron_density_m3 = jnp.maximum(density, density_floor) * norm.Nnorm
+    # Time is normalized to the parallel transit time L / c_s, so a physical
+    # volumetric rate n^2 <sigma v> maps to n^2 <sigma v> * Nnorm * (L / c_s).
+    rate_scale = norm.Nnorm * (params.parallel_length / norm.sound_speed)
+    ionization = rate_coefficient("d", "iz", physical_temperature, electron_density_m3) * rate_scale
+    recombination = rate_coefficient("d", "rec", physical_temperature, electron_density_m3) * rate_scale
+    effective_temperature = jnp.clip(2.0 * temperature * norm.Tnorm / params.ion_mass, 0.01, 1.0e4)
+    charge_exchange = charge_exchange_rate_coefficient(effective_temperature) * rate_scale
+    electron = jnp.maximum(density, density_floor)
+    ionization_frequency = ionization * electron
+    recombination_frequency = recombination * electron
+    total = density + neutral_density
+    new_density = (density + dt * ionization_frequency * total) / (1.0 + dt * (ionization_frequency + recombination_frequency))
+    # Momentum friction: charge exchange and recombination remove ion parallel
+    # momentum (neutrals are ~stationary); implicit drag for stability.
+    friction_frequency = charge_exchange * neutral_density + recombination_frequency
+    new_momentum = momentum / (1.0 + dt * friction_frequency)
+    return jnp.maximum(new_density, density_floor), new_momentum, jnp.maximum(total - new_density, 0.0)
+
+
+def sol_recycling_step(state, temperature, params, dt):
+    """Advance the recycling SOL one operator-split step."""
+
+    density = jnp.maximum(state.ion_density, params.density_floor)
+    momentum = state.ion_momentum
+    neutral_density = state.neutral_density
+
+    d_density, d_momentum = _hyperbolic_rhs(density, momentum, temperature, params)
+    density = jnp.maximum(density + dt * d_density, params.density_floor)
+    momentum = momentum + dt * d_momentum
+
+    density, momentum, neutral_density = _reaction_update(density, momentum, neutral_density, temperature, params, dt)
+
+    diffusivity = params.neutral_diffusion * params.neutral_temperature / (params.ion_mass * 0.05)
+    face_diffusivity = jnp.full(density.shape[0] - 1, diffusivity)
+    sound_speed = _sound_speed(density, temperature, params)
+    dz = 1.0 / density.shape[0]
+    recycling_source = jnp.zeros(density.shape[0]).at[-1].set(
+        params.recycling_fraction * density[-1] * sound_speed[-1] / dz
+    )
+    neutral_density = jnp.maximum(_implicit_diffusion(neutral_density, face_diffusivity, recycling_source, dt), 0.0)
+
+    # Pin the upstream ion density (Dirichlet) at the end of the step.
+    density = density.at[0].set(params.upstream_density)
+
+    return SolRecyclingState(ion_density=density, ion_momentum=momentum, neutral_density=neutral_density)
+
+
+def sol_recycling_run(state, temperature, params, *, dt, steps):
+    """Advance ``steps`` operator-split steps with a jitted ``lax.scan``."""
+
+    @jax.jit
+    def _run(initial):
+        def body(carry, _):
+            return sol_recycling_step(carry, temperature, params, dt), None
+
+        final, _ = jax.lax.scan(body, initial, None, length=steps)
+        return final
+
+    return _run(state)
+
+
+def target_ion_flux(state, temperature, params) -> jnp.ndarray:
+    """Bohm ion particle flux to the target plate ``n c_s``."""
+
+    density = jnp.maximum(state.ion_density, params.density_floor)
+    return density[-1] * _sound_speed(density, temperature, params)[-1]