Add source-weighted helium residence diagnostics and example

.github/workflows/build_test.yml

@@ -34,7 +34,7 @@ jobs:
     - name: Install ESSOS and test tools
       run: |
         python -m pip install --upgrade pip
-        pip install . flake8 pytest pytest-cov pytest-xdist
+        pip install ".[ash]" flake8 pytest pytest-cov pytest-xdist
         if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
         echo "JAX_VERSION=$(python -c 'import jax; print(jax.__version__)')" >> $GITHUB_ENV
     - name: Cache JAX compilations

README.md

@@ -32,6 +32,9 @@ pip install essos
 - **Boozer-coordinate tracing.** Guiding-centre diagnostics from the Boozer
   `|B|` spectrum and flux functions `iota`, `G` and `I`, with fixed-step
   (RK4, Dopri5, Dopri8) or adaptive Dormand-Prince stepping and optional collisions.
+- **Thermal-helium residence.** Source-weighted core residence from supplied
+  transport coefficients, with checked SOLVAX solves and JAX derivatives.
+  See [inputs and example](#thermal-helium-residence).
 - **Field-line tracing.** Adaptive, arclength and toroidal-angle models, with
   Poincare sections.
 - **Fields.** Biot-Savart from coils, VMEC equilibria (stellarator-symmetric or
@@ -141,6 +144,34 @@ print(tracing.loss_fractions)
 More in [`examples/particle_tracing`](examples/particle_tracing), including
 full-orbit, collisional and electric-field variants.
 
+## Thermal-helium residence
+
+Install `pip install 'essos[ash]'` and enable JAX float64. `essos.ash` solves
+`Gamma=-D*dn/dr+V*n` with supplied transport, geometry, source and boundary data.
+The forward/backward equations are `K n = M S` and `K.T T = M c`;
+`(M c).T n / sum(M S) = (M S).T T / sum(M S)` gives source-weighted occupation.
+Cost `c=1` gives mean core residence; recycling and pumping need separate models.
+
+Face arrays D [m²/s], V [m/s] and W=dV/dr [m²] have n+1 entries; cell volumes
+[m³], source [m⁻³/s] and dimensionless cost have n entries. Use consistent
+physical radial coordinates, positive D and a reflecting axis at r=0.
+Positive finite `boundary_kappa_m_s` imposes `Gamma=kappa*n`, including half-cell
+resistance; +infinity absorbs. Shape errors raise `ValueError`; invalid values
+or failed discrete checks give `valid=False` and NaN outputs.
+`check_helium_residence` raises outside JIT. These checks do not establish mesh
+convergence or kinetic validity; gradients include only supplied dependencies.
+
+Run the [example](examples/particle_tracing/helium_ash_removal.py) with
+`--manufactured`, or `--wout path/to/wout.nc --coefficients path/to/response.json`.
+Its loader documents the required coefficient/provenance bundle and checks the
+WOUT hash and modeled radius; the provider must qualify volumes and transport.
+
+`trace_response` fits matched gradient/flux probes; `convection_from_marker_drift`
+converts marker drift away from the axis. `restricted_residence` estimates
+`E[min(T,H)]`, using +infinity for survivors and rejecting numerical failures.
+Its uncertainty assumes independent markers with fixed weights; it omits the
+unobserved tail and correlated-marker uncertainty.
+
 ## Tracing to the wall
 
 Inside the last closed flux surface (LCFS), guiding centres move in VMEC

essos/ash.py

@@ -0,0 +1,231 @@
+"""Source-weighted thermal-helium residence for Gamma = -D*dn/dr + V*n.
+
+Supply transport coefficients, consistent radial volumes, a source and a boundary
+condition. JIT/AD-compatible functions return valid=False and NaN outputs for
+invalid values; shape errors raise ValueError. Core residence excludes recycling
+and pumping. See the README for the model and input contract.
+"""
+from typing import NamedTuple
+
+import jax.numpy as jnp
+
+
+class HeliumResidence(NamedTuple):
+    density_m3: object
+    occupation_s: object
+    source_rate_s: object
+    inventory: object
+    cost: object
+    mean_occupation_s: object
+    mean_residence_s: object
+    boundary_particle_rate_s: object
+    forward_residual: object
+    backward_residual: object
+    duality_relative_error: object
+    balance_relative_error: object
+    minimum_relative_pivot: object
+    maximum_peclet: object
+    valid: object
+
+
+def _bernoulli(x):
+    """Stable x/expm1(x); safe inactive branches also have finite derivatives."""
+    small = jnp.abs(x) < 1.e-4
+    # Evaluate only a bounded nonzero argument in the ordinary branch.
+    z = jnp.where(small, 1., jnp.clip(x, -50., 50.))
+    ordinary = z / jnp.expm1(z)
+    positive = jnp.maximum(x, 1.)
+    negative = jnp.minimum(x, -1.)
+    high = positive * jnp.exp(-positive) / (-jnp.expm1(-positive))
+    low = -negative / (-jnp.expm1(negative))
+    zsmall = jnp.where(small, x, 0.)
+    series = 1 - zsmall/2 + zsmall*zsmall/12 - zsmall**4/720 + zsmall**6/30240
+    return jnp.where(small, series,
+                     jnp.where(x > 50., high, jnp.where(x < -50., low, ordinary)))
+
+
+def evaluate_helium_residence(
+    radial_edges_m, cell_volume_m3, volume_derivative_m2,
+    diffusion_m2_s, convection_m_s, source_m3_s, *,
+    boundary_kappa_m_s, core_cost=None,
+):
+    """Solve K*n=M*S and K.T*T=M*c using checked SOLVAX tridiagonal solves.
+
+    Edges, W=dV/dr, D and V have shape (n+1,); volumes, S and dimensionless
+    cost c have shape (n,). At least two cells start at r=0 with a reflecting
+    axis. Positive finite kappa imposes Gamma=kappa*n, including half-cell
+    resistance; +inf absorbs. D and kappa must be positive; no floors are used.
+
+    occupation_s and mean_occupation_s use c; mean_residence_s is inventory
+    divided by source rate. Validity checks the discrete solve, not continuum
+    convergence or the kinetic closure. Gradients include supplied inputs only.
+    """
+    try:
+        from solvax import tridiagonal_factor, tridiagonal_solve_factored_checked
+    except ImportError as error:
+        raise ImportError("Helium residence requires ESSOS's optional ash extra: pip install 'essos[ash]'") from error
+
+    e, m, w, d, v, s = (jnp.asarray(x, dtype=jnp.result_type(float, jnp.asarray(x))) for x in
+                        (radial_edges_m, cell_volume_m3, volume_derivative_m2,
+                         diffusion_m2_s, convection_m_s, source_m3_s))
+    if e.ndim != 1 or e.size < 3:
+        raise ValueError("Require at least two radial cells")
+    n = e.size - 1
+    c = jnp.ones_like(m) if core_cost is None else jnp.asarray(core_cost)
+    kappa = jnp.asarray(boundary_kappa_m_s)
+    if any(x.shape != (n+1,) for x in (w, d, v)) or any(x.shape != (n,) for x in (m, s, c)) or kappa.ndim != 0:
+        raise ValueError("Edges/faces must have n+1 entries, cells n, and kappa be scalar")
+    input_ok = ((e[0] == 0) & jnp.all(jnp.diff(e) > 0) &
+                jnp.all(m > 0) & (w[0] >= 0) & jnp.all(w[1:] > 0) &
+                jnp.all(d > 0) & jnp.all(s >= 0) & jnp.all(c >= 0) &
+                (kappa > 0) & ~jnp.isnan(kappa))
+    for x in (e, m, w, d, v, s, c):
+        input_ok &= jnp.all(jnp.isfinite(x))
+    centers = (e[1:] + e[:-1])/2
+    h = jnp.diff(centers)
+    pe = v[1:-1]*h/d[1:-1]
+    g = w[1:-1]*d[1:-1]/h
+    out, inc = g*_bernoulli(-pe), g*_bernoulli(pe)
+    lo = jnp.concatenate((jnp.zeros(1), -out))
+    up = jnp.concatenate((-inc, jnp.zeros(1)))
+    diag = jnp.concatenate((out, jnp.zeros(1))) + jnp.concatenate((jnp.zeros(1), inc))
+    hb = e[-1] - centers[-1]
+    pb = v[-1]*hb/d[-1]
+    gb = w[-1]*d[-1]/hb
+    finite_kappa = jnp.where(jnp.isposinf(kappa), 1., kappa)
+    inverse_kappa = jnp.where(jnp.isposinf(kappa), 0., 1/finite_kappa)
+    escape = gb*_bernoulli(-pb)/(1 + gb*_bernoulli(pb)*inverse_kappa/w[-1])
+    diag = diag.at[-1].add(escape)
+    b, weight = m*s, m*c
+    rate = jnp.sum(b)
+    input_ok &= (rate > 0) & jnp.isfinite(escape) & (escape > 0)
+    f = tridiagonal_factor(lo, diag, up)
+    forward = tridiagonal_solve_factored_checked(f, b, fallback="nan")
+    lot = jnp.concatenate((jnp.zeros(1), up[:-1]))
+    upt = jnp.concatenate((lo[1:], jnp.zeros(1)))
+    ft = tridiagonal_factor(lot, diag, upt)
+    backward = tridiagonal_solve_factored_checked(ft, weight, fallback="nan")
+    density, occupation = forward.solution, backward.solution
+
+    def residual(lower, upper, x, rhs):
+        applied = diag*x
+        applied = applied.at[1:].add(lower[1:]*x[:-1])
+        applied = applied.at[:-1].add(upper[:-1]*x[1:])
+        scale = jnp.max(jnp.abs(rhs))
+        return jnp.max(jnp.abs(applied-rhs))/jnp.where(scale > 0, scale, 1.)
+
+    rf = residual(lo, up, density, b)
+    rb = residual(lot, upt, occupation, weight)
+    inventory = m @ density
+    cost, dual = weight @ density, b @ occupation
+    flux = escape*density[-1]
+    duality = jnp.abs(cost-dual)/jnp.maximum(jnp.abs(cost), jnp.finfo(diag.dtype).tiny)
+    balance = jnp.abs(flux-rate)/rate
+    mean_occupation, mean_residence = cost/rate, inventory/rate
+    physical_scalars = jnp.stack((rate, inventory, cost, mean_occupation, mean_residence, flux))
+    tolerance = 2.e-8
+    valid = (input_ok & jnp.all(jnp.isfinite(physical_scalars)) &
+             jnp.all(jnp.isfinite(density)) & jnp.all(jnp.isfinite(occupation)) &
+             jnp.all(density >= 0) & jnp.all(occupation >= 0) &
+             (rf <= tolerance) & (rb <= tolerance) &
+             (duality <= tolerance) & (balance <= tolerance))
+    mask = lambda x: jnp.where(valid, x, jnp.nan)
+    return HeliumResidence(
+        mask(density), mask(occupation), mask(rate), mask(inventory), mask(cost),
+        mask(mean_occupation), mask(mean_residence), mask(flux), rf, rb, duality, balance,
+        jnp.minimum(forward.diagnostics.minimum_relative_pivot,
+                    backward.diagnostics.minimum_relative_pivot),
+        jnp.maximum(jnp.max(jnp.abs(pe)), jnp.abs(pb)), valid)
+
+
+def check_helium_residence(result):
+    """Host-only failure boundary; return a valid result or raise ArithmeticError."""
+    if not bool(result.valid):
+        raise ArithmeticError("Invalid helium inputs or unqualified radial solve; inspect residuals, balance and pivots")
+    return result
+
+
+class RestrictedResidence(NamedTuple):
+    restricted_mean_s: object
+    escaped_fraction: object
+    failed_fraction: object
+    standard_error_s: object
+    effective_sample_size: object
+    no_censoring_in_sample: object
+    valid: object
+
+
+def restricted_residence(exit_times_s, horizon_s, *, weights=None, failed=None):
+    """Common-horizon E[min(T,H)], with +inf for physical survivors.
+
+    Numerical failures invalidate the estimate, even at zero marker weight.
+    This is not a full mean or a fitted tail. Standard error assumes independent
+    identically distributed markers with fixed weights, using the weighted sample variance; it does not
+    account for uncertain coefficients or correlated/common-noise cohorts.
+    """
+    t = jnp.asarray(exit_times_s, dtype=float)
+    if t.ndim != 1 or t.size == 0:
+        raise ValueError("Require a nonempty vector of first-exit times")
+    h = jnp.asarray(horizon_s)
+    w = jnp.ones_like(t) if weights is None else jnp.asarray(weights, dtype=float)
+    f = jnp.zeros_like(t, dtype=bool) if failed is None else jnp.asarray(failed, dtype=bool)
+    if w.shape != t.shape or f.shape != t.shape or h.ndim != 0:
+        raise ValueError("Times, weights and failures must match; horizon must be scalar")
+    # Scaling before summation preserves normalized weights even when their
+    # finite raw sum overflows. Invalid weights are still rejected below.
+    largest_weight = jnp.max(w)
+    safe_largest = jnp.where(jnp.isfinite(largest_weight) & (largest_weight > 0), largest_weight, 1.)
+    # Use a power-of-two scale: direct division by a huge maximum may be
+    # lowered to a reciprocal that underflows/flushes to zero under XLA.
+    exponent = jnp.frexp(safe_largest)[1]
+    scaled_weights = jnp.ldexp(w, -exponent)
+    total = jnp.sum(scaled_weights)
+    valid = (jnp.isfinite(h) & (h > 0) & jnp.all(t >= 0) &
+             jnp.all(~jnp.isnan(t)) & jnp.all((t <= h) | jnp.isposinf(t)) &
+             jnp.all(jnp.isfinite(w)) & jnp.all(w >= 0) & jnp.isfinite(total) & (total > 0) & ~jnp.any(f))
+    w = scaled_weights/jnp.where(jnp.isfinite(total) & (total > 0), total, 1.)
+    x = jnp.minimum(t, h)
+    mean = w @ x
+    sumsq = w @ w
+    se = jnp.sqrt(sumsq*jnp.sum(w*(x-mean)**2)/jnp.where(sumsq < 1, 1-sumsq, jnp.nan))
+    mask = lambda x: jnp.where(valid, x, jnp.nan)
+    return RestrictedResidence(mask(mean), mask(w @ jnp.isfinite(t)), w @ f,
+                               mask(se), 1/sumsq, valid & jnp.all(jnp.isfinite(t)), valid)
+
+
+class TraceResponse(NamedTuple):
+    diffusion_m2_s: object
+    convection_m_s: object
+    fit_defect_m_s: object
+    valid: object
+
+
+def trace_response(log_gradients_m_inv, flux_over_density_m_s, *, relative_tolerance=1.e-7):
+    """Fit Gamma/n=V-D*d(log n)/dr; >=3 distinct probes test affine response.
+
+    Providers must match all other profiles, kinetic models and Er across probes.
+    This diagnostic alone cannot qualify trace abundance or full-FP physics.
+    """
+    x, y = map(jnp.asarray, (log_gradients_m_inv, flux_over_density_m_s))
+    if x.ndim != 1 or x.size < 3 or y.shape != x.shape:
+        raise ValueError("Require at least three paired gradient/flux probes")
+    xc, yc = x-jnp.mean(x), y-jnp.mean(y)
+    slope = (xc @ yc)/(xc @ xc)
+    drift = jnp.mean(y)-slope*jnp.mean(x)
+    defect = jnp.max(jnp.abs(y-(drift+slope*x)))
+    distinct = 1 + jnp.sum(jnp.diff(jnp.sort(x)) > 0)
+    valid = (jnp.all(jnp.isfinite(x)) & jnp.all(jnp.isfinite(y)) & (distinct >= 3) &
+             (jnp.isfinite(relative_tolerance)) & (relative_tolerance >= 0) & (slope < 0) & jnp.isfinite(slope) &
+             (defect <= relative_tolerance*jnp.maximum(jnp.max(jnp.abs(y)), 1.e-12)))
+    return TraceResponse(jnp.where(valid, -slope, jnp.nan),
+                         jnp.where(valid, drift, jnp.nan), defect, valid)
+
+
+def convection_from_marker_drift(drift_m_s, diffusion_m2_s,
+                                 diffusion_gradient_m_s, log_volume_gradient_m_inv):
+    """V=A-D'-D*d(log W)/dr away from the axis, in the same physical r.
+
+    W'/W is singular at a regular axis; use conservative fluxes there.
+    """
+    return (jnp.asarray(drift_m_s) - jnp.asarray(diffusion_gradient_m_s) -
+            jnp.asarray(diffusion_m2_s)*jnp.asarray(log_volume_gradient_m_inv))

examples/particle_tracing/helium_ash_removal.py

@@ -0,0 +1,108 @@
+"""Helium residence from exported coefficients or a manufactured cylinder.
+
+    python examples/particle_tracing/helium_ash_removal.py --manufactured
+    python examples/particle_tracing/helium_ash_removal.py \
+        --wout path/to/wout.nc --coefficients path/to/helium_response.json
+
+Supply JSON ``coefficients`` matching evaluate_helium_residence kwargs and
+``metadata`` with _REQUIRED_METADATA keys. Null kappa means absorbing.
+Coefficients must cover the modeled radius; WOUT supplies geometry identity,
+not missing transport or profiles. Manufactured mode tests the pipeline only.
+"""
+import argparse
+import hashlib
+import json
+from pathlib import Path
+
+import jax
+jax.config.update("jax_enable_x64", True)
+import numpy as np
+
+from essos.ash import evaluate_helium_residence, check_helium_residence
+
+_REQUIRED_METADATA = (
+    "equilibrium_sha256", "profiles_id", "source_id", "collision_model",
+    "electric_field", "turbulent_closure", "boundary_model", "volume_origin",
+    "coefficient_status", "trace_validity",
+)
+
+
+def _hash(path):
+    return hashlib.sha256(path.read_bytes()).hexdigest()
+
+
+def load_response(path, wout):
+    """Host-only input audit; no file operations occur in the native evaluator."""
+    from netCDF4 import Dataset
+    bundle = json.loads(path.read_text())
+    coefficients, metadata = bundle["coefficients"], bundle["metadata"]
+    missing = set(_REQUIRED_METADATA)-metadata.keys()
+    if missing:
+        raise ValueError(f"Missing physical/provenance metadata: {sorted(missing)}")
+    if any(metadata[key] is None or metadata[key] == "" for key in _REQUIRED_METADATA):
+        raise ValueError("Metadata must explicitly describe assumptions or say 'not evaluated'")
+    if metadata["equilibrium_sha256"] != _hash(wout):
+        raise ValueError("Coefficient equilibrium hash does not match the supplied WOUT")
+    if coefficients["boundary_kappa_m_s"] is None:
+        coefficients["boundary_kappa_m_s"] = np.inf
+    with Dataset(wout) as data:
+        minor_radius = float(np.asarray(data.variables["Aminor_p"][:]).squeeze())
+    outer_radius = float(coefficients["radial_edges_m"][-1])
+    if not np.isfinite(minor_radius) or minor_radius <= 0 or outer_radius > minor_radius*(1+1.e-10):
+        raise ValueError("Require a positive WOUT minor radius and modeled r <= Aminor_p")
+    metadata = dict(metadata, wout=str(wout.resolve()), coefficient_sha256=_hash(path),
+                    radial_coordinate="r=Aminor_p*sqrt(s)",
+                    modeled_boundary_radius_m=outer_radius,
+                    modeled_boundary_s=(outer_radius/minor_radius)**2)
+    return coefficients, metadata
+
+
+def manufactured_response(cells=128):
+    e = np.linspace(0, 1, cells+1)
+    return dict(radial_edges_m=e, cell_volume_m3=np.diff(e**2),
+                volume_derivative_m2=2*e, diffusion_m2_s=np.ones_like(e),
+                convection_m_s=np.zeros_like(e), source_m3_s=np.ones(cells),
+                boundary_kappa_m_s=np.inf), dict(
+                    scope="Manufactured cylinder only; no stellarator evaluated",
+                    profiles_id="constant unit source", source_id="unit-volume cylinder",
+                    collision_model="not evaluated", electric_field="not evaluated",
+                    turbulent_closure="not evaluated", boundary_model="absorbing at r=1 m",
+                    volume_origin="analytic W=2r, total volume 1 m^3",
+                    coefficient_status="manufactured D=1 m^2/s, V=0",
+                    trace_validity="not evaluated", modeled_boundary_radius_m=1.)
+
+
+def main(argv=None):
+    parser = argparse.ArgumentParser(description=__doc__)
+    choice = parser.add_mutually_exclusive_group(required=True)
+    choice.add_argument("--coefficients", type=Path)
+    choice.add_argument("--manufactured", action="store_true")
+    parser.add_argument("--wout", type=Path)
+    parser.add_argument("--output", type=Path, help="Write the same compact report as JSON")
+    args = parser.parse_args(argv)
+    if args.coefficients:
+        if args.wout is None:
+            parser.error("--coefficients requires --wout")
+        coefficients, metadata = load_response(args.coefficients, args.wout)
+    else:
+        if args.wout is not None:
+            parser.error("--manufactured has no WOUT; use explicit exported coefficients")
+        coefficients, metadata = manufactured_response()
+    result = check_helium_residence(evaluate_helium_residence(**coefficients))
+    fields = ("source_rate_s", "inventory", "cost", "mean_residence_s", "mean_occupation_s",
+              "boundary_particle_rate_s", "forward_residual", "backward_residual",
+              "duality_relative_error", "balance_relative_error", "minimum_relative_pivot", "maximum_peclet")
+    report = dict(metadata=metadata, result={name: float(getattr(result, name)) for name in fields},
+                  numerical_status="discrete checks passed; radial convergence not evaluated",
+                  alpha_energy_retention="not evaluated", thermal_core_passage="coefficient-model estimate",
+                  wall_interception="not evaluated", pumped_fraction="not modeled",
+                  thermal_orbit_verification="not evaluated")
+    text = json.dumps(report, indent=2, allow_nan=False)+"\n"
+    print(text, end="")
+    if args.output:
+        args.output.write_text(text)
+    return report
+
+
+if __name__ == "__main__":
+    main()