drbx 2.0.0.dev0__py3-none-any.whl

drbx/native/fci_sheath_recycling.py

@@ -0,0 +1,206 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax.numpy as jnp
+
+from ..geometry import FciMaps
+
+
+@dataclass(frozen=True)
+class FciTargetMasks:
+    """Open-field endpoint masks derived from traced FCI plane intersections."""
+
+    forward: jnp.ndarray
+    backward: jnp.ndarray
+    endpoint_count: jnp.ndarray
+    active: jnp.ndarray
+
+
+@dataclass(frozen=True)
+class FciSheathRecyclingResult:
+    """JAX-native sheath-loss and target-recycling diagnostic fields."""
+
+    masks: FciTargetMasks
+    sound_speed: jnp.ndarray
+    forward_ion_particle_flux: jnp.ndarray
+    backward_ion_particle_flux: jnp.ndarray
+    ion_particle_loss: jnp.ndarray
+    electron_particle_loss: jnp.ndarray
+    electron_heat_loss: jnp.ndarray
+    ion_heat_loss: jnp.ndarray
+    target_heat_load: jnp.ndarray
+    recycled_particle_source: jnp.ndarray
+    recycled_neutral_energy_source: jnp.ndarray
+    plasma_particle_sink: jnp.ndarray
+    plasma_energy_sink: jnp.ndarray
+    current_residual: jnp.ndarray
+    total_ion_particle_loss: jnp.ndarray
+    total_electron_particle_loss: jnp.ndarray
+    total_recycled_particle_source: jnp.ndarray
+    total_target_heat_load: jnp.ndarray
+    total_recycled_neutral_energy: jnp.ndarray
+    particle_recycling_residual: jnp.ndarray
+    neutral_energy_recycling_residual: jnp.ndarray
+    current_balance_residual: jnp.ndarray
+
+
+def build_fci_target_masks(maps: FciMaps) -> FciTargetMasks:
+    """Build endpoint masks from forward/backward field-line map exits."""
+
+    forward = jnp.asarray(maps.forward_boundary, dtype=jnp.float64)
+    backward = jnp.asarray(maps.backward_boundary, dtype=jnp.float64)
+    endpoint_count = forward + backward
+    active = endpoint_count > 0.0
+    return FciTargetMasks(
+        forward=forward,
+        backward=backward,
+        endpoint_count=endpoint_count,
+        active=active,
+    )
+
+
+def fci_bohm_sound_speed(
+    electron_temperature: jnp.ndarray,
+    ion_temperature: jnp.ndarray,
+    *,
+    ion_mass: float = 1.0,
+    temperature_floor: float = 1.0e-12,
+) -> jnp.ndarray:
+    """Return the normalized sheath-entry sound speed."""
+
+    te = jnp.maximum(jnp.asarray(electron_temperature, dtype=jnp.float64), float(temperature_floor))
+    ti = jnp.maximum(jnp.asarray(ion_temperature, dtype=jnp.float64), float(temperature_floor))
+    return jnp.sqrt((te + ti) / float(ion_mass))
+
+
+def compute_fci_sheath_recycling(
+    density: jnp.ndarray,
+    electron_temperature: jnp.ndarray,
+    ion_temperature: jnp.ndarray,
+    maps: FciMaps,
+    *,
+    recycling_fraction: float = 0.98,
+    electron_sheath_transmission: float = 5.0,
+    ion_sheath_transmission: float = 3.5,
+    recycled_neutral_energy: float = 0.03,
+    ion_mass: float = 1.0,
+    density_floor: float = 1.0e-12,
+    temperature_floor: float = 1.0e-12,
+) -> FciSheathRecyclingResult:
+    """Evaluate target loss and recycled neutral sources on FCI endpoints.
+
+    The closure is the first non-axisymmetric target gate. It applies a
+    normalized Bohm flux on every traced endpoint and closes exact accounting
+    identities for recycled particle and neutral-energy sources. The electron
+    particle flux is reconstructed from zero-current balance for singly charged
+    ions, so the current residual is an explicit regression field.
+    """
+
+    masks = build_fci_target_masks(maps)
+    n = jnp.maximum(jnp.asarray(density, dtype=jnp.float64), float(density_floor))
+    te = jnp.maximum(jnp.asarray(electron_temperature, dtype=jnp.float64), float(temperature_floor))
+    ti = jnp.maximum(jnp.asarray(ion_temperature, dtype=jnp.float64), float(temperature_floor))
+    cs = fci_bohm_sound_speed(
+        te,
+        ti,
+        ion_mass=ion_mass,
+        temperature_floor=temperature_floor,
+    )
+
+    forward_flux = masks.forward * n * cs
+    backward_flux = masks.backward * n * cs
+    ion_loss = forward_flux + backward_flux
+    electron_loss = ion_loss
+    electron_heat_loss = float(electron_sheath_transmission) * ion_loss * te
+    ion_heat_loss = float(ion_sheath_transmission) * ion_loss * ti
+    target_heat_load = electron_heat_loss + ion_heat_loss
+    recycled_particle_source = float(recycling_fraction) * ion_loss
+    recycled_neutral_energy_source = float(recycled_neutral_energy) * recycled_particle_source
+    plasma_particle_sink = -ion_loss
+    plasma_energy_sink = -target_heat_load
+    current_residual = electron_loss - ion_loss
+
+    total_ion_loss = jnp.sum(ion_loss)
+    total_electron_loss = jnp.sum(electron_loss)
+    total_recycled_source = jnp.sum(recycled_particle_source)
+    total_heat_load = jnp.sum(target_heat_load)
+    total_neutral_energy = jnp.sum(recycled_neutral_energy_source)
+    particle_recycling_residual = total_recycled_source - float(recycling_fraction) * total_ion_loss
+    neutral_energy_recycling_residual = total_neutral_energy - float(recycled_neutral_energy) * total_recycled_source
+    current_balance_residual = total_electron_loss - total_ion_loss
+
+    return FciSheathRecyclingResult(
+        masks=masks,
+        sound_speed=cs,
+        forward_ion_particle_flux=forward_flux,
+        backward_ion_particle_flux=backward_flux,
+        ion_particle_loss=ion_loss,
+        electron_particle_loss=electron_loss,
+        electron_heat_loss=electron_heat_loss,
+        ion_heat_loss=ion_heat_loss,
+        target_heat_load=target_heat_load,
+        recycled_particle_source=recycled_particle_source,
+        recycled_neutral_energy_source=recycled_neutral_energy_source,
+        plasma_particle_sink=plasma_particle_sink,
+        plasma_energy_sink=plasma_energy_sink,
+        current_residual=current_residual,
+        total_ion_particle_loss=total_ion_loss,
+        total_electron_particle_loss=total_electron_loss,
+        total_recycled_particle_source=total_recycled_source,
+        total_target_heat_load=total_heat_load,
+        total_recycled_neutral_energy=total_neutral_energy,
+        particle_recycling_residual=particle_recycling_residual,
+        neutral_energy_recycling_residual=neutral_energy_recycling_residual,
+        current_balance_residual=current_balance_residual,
+    )
+
+
+def fci_sheath_recycling_field_rhs(
+    fields: dict[str, jnp.ndarray],
+    maps: FciMaps,
+    *,
+    ion_density_name: str = "Ni",
+    electron_density_name: str = "Ne",
+    neutral_density_name: str = "Nn",
+    ion_pressure_name: str = "Pi",
+    electron_pressure_name: str = "Pe",
+    neutral_pressure_name: str = "Pn",
+    recycling_fraction: float = 0.98,
+    electron_sheath_transmission: float = 5.0,
+    ion_sheath_transmission: float = 3.5,
+    recycled_neutral_energy: float = 0.03,
+    ion_mass: float = 1.0,
+    density_floor: float = 1.0e-12,
+    temperature_floor: float = 1.0e-12,
+) -> dict[str, jnp.ndarray]:
+    """Return fixed-field RHS arrays for FCI sheath and recycling sources."""
+
+    ion_density = jnp.asarray(fields[ion_density_name], dtype=jnp.float64)
+    electron_density = jnp.asarray(fields.get(electron_density_name, ion_density), dtype=jnp.float64)
+    ion_pressure = jnp.asarray(fields[ion_pressure_name], dtype=jnp.float64)
+    electron_pressure = jnp.asarray(fields[electron_pressure_name], dtype=jnp.float64)
+    ion_temperature = ion_pressure / jnp.maximum(ion_density, float(density_floor))
+    electron_temperature = electron_pressure / jnp.maximum(electron_density, float(density_floor))
+    result = compute_fci_sheath_recycling(
+        ion_density,
+        electron_temperature,
+        ion_temperature,
+        maps,
+        recycling_fraction=recycling_fraction,
+        electron_sheath_transmission=electron_sheath_transmission,
+        ion_sheath_transmission=ion_sheath_transmission,
+        recycled_neutral_energy=recycled_neutral_energy,
+        ion_mass=ion_mass,
+        density_floor=density_floor,
+        temperature_floor=temperature_floor,
+    )
+    rhs = {
+        ion_density_name: -result.ion_particle_loss,
+        electron_density_name: -result.electron_particle_loss,
+        neutral_density_name: result.recycled_particle_source,
+        ion_pressure_name: -result.ion_heat_loss,
+        electron_pressure_name: -result.electron_heat_loss,
+        neutral_pressure_name: result.recycled_neutral_energy_source,
+    }
+    return {name: value for name, value in rhs.items() if name in fields}

drbx/native/fci_time_integrator.py

@@ -0,0 +1,96 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+from functools import reduce
+from typing import Callable, Generic, TypeVar
+
+import jax
+
+from .fci_model import FciModelState
+
+
+StateT = TypeVar("StateT", bound=FciModelState)
+CarryT = TypeVar("CarryT")
+AuxT = TypeVar("AuxT")
+
+
+@jax.tree_util.register_pytree_node_class
+@dataclass(frozen=True)
+class Rk4StepResult(Generic[StateT, CarryT, AuxT]):
+    """Container for a single RK4 advance.
+
+    ``stage_aux`` stores the four auxiliary payloads returned by the stage RHS
+    calls. The auxiliary payload is intentionally opaque so that callers can
+    thread timings, solver diagnostics, warm-start carries, or any other
+    model-specific stage information without the RK4 core knowing about it.
+    """
+
+    state: StateT
+    carry: CarryT
+    stage_aux: tuple[AuxT, AuxT, AuxT, AuxT]
+
+    def tree_flatten(self):
+        aux_1, aux_2, aux_3, aux_4 = self.stage_aux
+        return (self.state, self.carry, aux_1, aux_2, aux_3, aux_4), None
+
+    @classmethod
+    def tree_unflatten(cls, _aux_data, children):
+        state, carry, aux_1, aux_2, aux_3, aux_4 = children
+        return cls(state=state, carry=carry, stage_aux=(aux_1, aux_2, aux_3, aux_4))
+
+
+def _rk4_weighted_rhs(k1: StateT, k2: StateT, k3: StateT, k4: StateT) -> StateT:
+    """Return ``k1 + 2*k2 + 2*k3 + k4`` using the shared state algebra."""
+
+    return k1.axpy(k2, scale=2.0).axpy(k3, scale=2.0).axpy(k4, scale=1.0)
+
+
+def rk4_step(
+    state: StateT,
+    *,
+    time: float | jax.Array,
+    timestep: float | jax.Array,
+    rhs_fn: Callable[[StateT, float | jax.Array, CarryT], tuple[StateT, CarryT, AuxT]],
+    carry: CarryT,
+) -> Rk4StepResult[StateT, CarryT, AuxT]:
+    """Advance ``state`` by one classical RK4 step.
+
+    The RK4 algebra is model-agnostic:
+
+    - ``rhs_fn`` computes the stage RHS for the current stage state and returns
+      the RHS, the next carry value, and an arbitrary auxiliary payload.
+    - The carry is threaded from stage to stage so models can warm-start local
+      solves, keep stage caches, or propagate other stage-local context.
+    - The final carry returned by the step is the carry produced by the fourth
+      stage evaluation.
+    """
+
+    k1, carry_1, aux_1 = rhs_fn(state, time, carry)
+    stage_1 = state.axpy(k1, scale=0.5 * timestep)
+
+    k2, carry_2, aux_2 = rhs_fn(stage_1, time + 0.5 * timestep, carry_1)
+    stage_2 = state.axpy(k2, scale=0.5 * timestep)
+
+    k3, carry_3, aux_3 = rhs_fn(stage_2, time + 0.5 * timestep, carry_2)
+    stage_3 = state.axpy(k3, scale=timestep)
+
+    k4, carry_4, aux_4 = rhs_fn(stage_3, time + timestep, carry_3)
+    next_state = state.axpy(_rk4_weighted_rhs(k1, k2, k3, k4), scale=timestep / 6.0)
+    return Rk4StepResult(state=next_state, carry=carry_4, stage_aux=(aux_1, aux_2, aux_3, aux_4))
+
+
+def sum_stage_outputs(stage_outputs: tuple[AuxT, AuxT, AuxT, AuxT]) -> AuxT:
+    """Reduce four stage payloads by addition.
+
+    This is handy when each stage returns a timing vector or another additive
+    PyTree. Models with non-additive diagnostics can ignore this helper and
+    reduce their stage payloads manually.
+    """
+
+    def _add(left: AuxT, right: AuxT) -> AuxT:
+        return jax.tree_util.tree_map(lambda lhs, rhs: lhs + rhs, left, right)
+
+    return reduce(_add, stage_outputs[1:], stage_outputs[0])
+
+
+__all__ = ["Rk4StepResult", "rk4_step", "sum_stage_outputs"]

drbx/native/fci_vorticity.py

@@ -0,0 +1,198 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+import jax
+import jax.numpy as jnp
+
+from ..geometry import MetricTensor3D
+from .fci import conservative_perp_diffusion_xz
+
+
+@dataclass(frozen=True)
+class FciVorticitySolveResult:
+    potential: jnp.ndarray
+    residual: jnp.ndarray
+    residual_l2: jnp.ndarray
+    iterations: int
+    preconditioner: str | None = None
+
+
+def apply_fci_vorticity_operator(
+    potential: jnp.ndarray,
+    density: jnp.ndarray,
+    metric: MetricTensor3D,
+    *,
+    boussinesq: bool = True,
+    regularization: float = 1.0e-9,
+) -> jnp.ndarray:
+    """Apply the positive perpendicular vorticity operator to ``phi``."""
+
+    phi = _remove_mean(jnp.asarray(potential, dtype=jnp.float64), metric)
+    n = jnp.asarray(density, dtype=jnp.float64)
+    if boussinesq:
+        coefficient = jnp.ones_like(phi) * jnp.mean(n / jnp.square(metric.Bxy))
+    else:
+        coefficient = n / jnp.maximum(jnp.square(metric.Bxy), 1.0e-30)
+    operator = -conservative_perp_diffusion_xz(phi, coefficient, metric)
+    return _remove_mean(operator, metric) + float(regularization) * phi
+
+
+def solve_fci_vorticity_potential_cg(
+    vorticity: jnp.ndarray,
+    density: jnp.ndarray,
+    metric: MetricTensor3D,
+    *,
+    iterations: int = 80,
+    boussinesq: bool = True,
+    regularization: float = 1.0e-9,
+    preconditioner: str | None = None,
+) -> FciVorticitySolveResult:
+    """Solve the metric-weighted perpendicular vorticity inversion with CG."""
+
+    rhs = _remove_mean(jnp.asarray(vorticity, dtype=jnp.float64), metric)
+    x0 = jnp.zeros_like(rhs, dtype=jnp.float64)
+    normalized_preconditioner = _normalize_preconditioner_name(preconditioner)
+
+    def apply_operator(value: jnp.ndarray) -> jnp.ndarray:
+        return apply_fci_vorticity_operator(
+            value,
+            density,
+            metric,
+            boussinesq=boussinesq,
+            regularization=regularization,
+        )
+
+    r0 = rhs - apply_operator(x0)
+    if normalized_preconditioner == "jacobi":
+        inverse_diagonal = _fci_vorticity_jacobi_inverse_diagonal(
+            density,
+            metric,
+            boussinesq=boussinesq,
+            regularization=regularization,
+        )
+
+        def apply_preconditioner(residual: jnp.ndarray) -> jnp.ndarray:
+            return _remove_mean(inverse_diagonal * residual, metric)
+
+        z0 = apply_preconditioner(r0)
+        p0 = z0
+        rz0 = _inner(r0, z0, metric)
+
+        def body(
+            _index: int,
+            carry: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray],
+        ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+            x, r, p, rz = carry
+            ap = apply_operator(p)
+            alpha = rz / jnp.maximum(_inner(p, ap, metric), 1.0e-30)
+            x_next = _remove_mean(x + alpha * p, metric)
+            r_next = r - alpha * ap
+            z_next = apply_preconditioner(r_next)
+            rz_next = _inner(r_next, z_next, metric)
+            beta = rz_next / jnp.maximum(rz, 1.0e-30)
+            p_next = z_next + beta * p
+            return x_next, r_next, p_next, rz_next
+
+        potential, residual, _direction, _rz = jax.lax.fori_loop(
+            0,
+            int(iterations),
+            body,
+            (x0, r0, p0, rz0),
+        )
+    else:
+        p0 = r0
+
+        def body(
+            _index: int,
+            carry: tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray],
+        ) -> tuple[jnp.ndarray, jnp.ndarray, jnp.ndarray]:
+            x, r, p = carry
+            ap = apply_operator(p)
+            rs = _inner(r, r, metric)
+            alpha = rs / jnp.maximum(_inner(p, ap, metric), 1.0e-30)
+            x_next = _remove_mean(x + alpha * p, metric)
+            r_next = r - alpha * ap
+            rs_next = _inner(r_next, r_next, metric)
+            beta = rs_next / jnp.maximum(rs, 1.0e-30)
+            p_next = r_next + beta * p
+            return x_next, r_next, p_next
+
+        potential, residual, _direction = jax.lax.fori_loop(
+            0,
+            int(iterations),
+            body,
+            (x0, r0, p0),
+        )
+    residual = rhs - apply_operator(potential)
+    residual_l2 = jnp.sqrt(
+        _inner(residual, residual, metric)
+        / jnp.maximum(_inner(rhs, rhs, metric), 1.0e-30)
+    )
+    return FciVorticitySolveResult(
+        potential=potential,
+        residual=residual,
+        residual_l2=residual_l2,
+        iterations=int(iterations),
+        preconditioner=normalized_preconditioner,
+    )
+
+
+def _normalize_preconditioner_name(preconditioner: str | None) -> str | None:
+    if preconditioner is None:
+        return None
+    normalized = str(preconditioner).strip().lower().replace("-", "_")
+    if normalized in {"", "none", "unpreconditioned"}:
+        return None
+    if normalized in {"jacobi", "diagonal", "diag"}:
+        return "jacobi"
+    raise ValueError(f"Unsupported FCI vorticity preconditioner {preconditioner!r}.")
+
+
+def _fci_vorticity_jacobi_inverse_diagonal(
+    density: jnp.ndarray,
+    metric: MetricTensor3D,
+    *,
+    boussinesq: bool,
+    regularization: float,
+) -> jnp.ndarray:
+    n = jnp.asarray(density, dtype=jnp.float64)
+    if boussinesq:
+        coefficient = jnp.ones_like(n) * jnp.mean(n / jnp.square(metric.Bxy))
+    else:
+        coefficient = n / jnp.maximum(jnp.square(metric.Bxy), 1.0e-30)
+    jac = jnp.asarray(metric.J, dtype=jnp.float64)
+    dx = jnp.asarray(metric.dx, dtype=jnp.float64)
+    dz = jnp.asarray(metric.dz, dtype=jnp.float64)
+    kx = jac * coefficient * jnp.asarray(metric.g11, dtype=jnp.float64)
+    kz = jac * coefficient * jnp.asarray(metric.g33, dtype=jnp.float64)
+
+    dx_face = 0.5 * (dx[1:, :, :] + dx[:-1, :, :])
+    kx_face = 0.5 * (kx[1:, :, :] + kx[:-1, :, :])
+    x_plus = jnp.zeros_like(n).at[:-1, :, :].set(
+        kx_face / jnp.maximum(dx_face, 1.0e-30)
+    )
+    x_minus = jnp.zeros_like(n).at[1:, :, :].set(
+        kx_face / jnp.maximum(dx_face, 1.0e-30)
+    )
+
+    dz_plus = 0.5 * (dz + jnp.roll(dz, -1, axis=2))
+    kz_face = 0.5 * (kz + jnp.roll(kz, -1, axis=2))
+    z_plus = kz_face / jnp.maximum(dz_plus, 1.0e-30)
+    z_minus = jnp.roll(z_plus, 1, axis=2)
+    diagonal = (
+        (x_plus + x_minus) / jnp.maximum(dx, 1.0e-30)
+        + (z_plus + z_minus) / jnp.maximum(dz, 1.0e-30)
+    ) / jnp.maximum(jac, 1.0e-30)
+    diagonal = diagonal + float(regularization)
+    return 1.0 / jnp.maximum(diagonal, 1.0e-30)
+
+
+def _inner(left: jnp.ndarray, right: jnp.ndarray, metric: MetricTensor3D) -> jnp.ndarray:
+    return jnp.sum(jnp.asarray(metric.J, dtype=jnp.float64) * left * right)
+
+
+def _remove_mean(value: jnp.ndarray, metric: MetricTensor3D) -> jnp.ndarray:
+    weights = jnp.asarray(metric.J, dtype=jnp.float64)
+    mean = jnp.sum(weights * value) / jnp.maximum(jnp.sum(weights), 1.0e-30)
+    return value - mean