From 825197090a9c0c3d8c840e23739b0750cc7dac96 Mon Sep 17 00:00:00 2001
From: Rogerio Jorge <rogerio.jorge@wisc.edu>
Date: Thu, 15 Oct 2026 15:57:37 -0500
Subject: [PATCH 1/5] Add source-weighted helium residence diagnostics

---
 README.md                                     |   6 +
 docs/helium_residence.rst                     |  65 +++++
 docs/index.rst                                |   3 +-
 essos/ash.py                                  | 226 ++++++++++++++++++
 .../particle_tracing/helium_ash_removal.py    | 110 +++++++++
 pyproject.toml                                |   2 +
 tests/test_ash.py                             | 196 +++++++++++++++
 7 files changed, 607 insertions(+), 1 deletion(-)
 create mode 100644 docs/helium_residence.rst
 create mode 100644 essos/ash.py
 create mode 100644 examples/particle_tracing/helium_ash_removal.py
 create mode 100644 tests/test_ash.py

diff --git a/README.md b/README.md
index e29d9545..85c27726 100644
--- a/README.md
+++ b/README.md
@@ -32,6 +32,12 @@ 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.** Optional `essos[ash]` coefficient-driven,
+  source-weighted residence/occupation with conservative fluxes, checked SOLVAX
+  solves and native JAX derivatives. See the [explicit-input example](examples/particle_tracing/helium_ash_removal.py)
+  and [scope/validation notes](docs/helium_residence.rst). Core residence requires
+  qualified transport coefficients, profiles and boundary conditions; it does
+  not predict pumped exhaust from a WOUT.
 - **Field-line tracing.** Adaptive, arclength and toroidal-angle models, with
   Poincare sections.
 - **Fields.** Biot-Savart from coils, VMEC equilibria (stellarator-symmetric or
diff --git a/docs/helium_residence.rst b/docs/helium_residence.rst
new file mode 100644
index 00000000..b5d1edd2
--- /dev/null
+++ b/docs/helium_residence.rst
@@ -0,0 +1,65 @@
+Thermal-helium core residence
+============================
+
+Install ``pip install 'essos[ash]'`` for the optional SOLVAX dependency.
+``essos.ash`` accepts explicit local trace-helium transport coefficients,
+physical cell volumes and a source. It solves the conservative steady equation
+``K n = M S`` and its transpose ``K.T T = M c``. The source-weighted occupation
+is ``(M c).T n / sum(M S) = (M S).T T / sum(M S)``. For ``c=1`` this is mean
+core residence; weighted occupation can describe a frozen core-dilution cost.
+This standard forward/backward identity is exact for the stated discrete model.
+
+Use ``jax.config.update('jax_enable_x64', True)`` before preparing arrays for the
+strict residual and balance tolerances. Face arrays D [m^2/s], V [m/s] and
+W=dV/dr [m^2] have n+1 entries; cell volumes [m^3], source [m^-3/s] and optional
+dimensionless cost have n. The first radial edge is zero and reflects.
+``boundary_kappa_m_s`` is mandatory: positive finite values impose
+``Gamma=kappa*n`` and positive infinity absorbs. The half-cell transport and
+boundary resistances are both included. A closed boundary, nonpositive D,
+invalid values, unresolved inward barriers or unqualified solves yield
+``valid=False`` and NaN physical outputs. Shape errors raise ``ValueError``.
+``check_helium_residence(result)`` raises on an invalid result outside JIT.
+
+The functions retain native JAX arrays for JIT and differentiation through all
+provided coefficients, geometry measures, source, cost and boundary inputs.
+File loading is host-only. Differentiation cannot recover dependence omitted
+by a coefficient provider, including a frozen electric field or ambipolar root.
+Pivot and residual checks do not provide a condition-number error bound;
+refine the radial mesh and kinetic response and test the local transport closure.
+No artificial diffusivity floor is applied.
+
+Run the manufactured numerical check::
+
+    python examples/particle_tracing/helium_ash_removal.py --manufactured
+
+For an actual equilibrium use an explicit exported response::
+
+    python examples/particle_tracing/helium_ash_removal.py \
+      --wout path/to/wout.nc --coefficients path/to/helium_response.json
+
+The JSON has ``coefficients`` matching the evaluator's argument names and
+``metadata`` describing ``equilibrium_sha256``, ``profiles_id``, ``source_id``,
+``collision_model``, ``electric_field`` (prescribed profile or root branch),
+``turbulent_closure``, ``boundary_model``, ``volume_origin``,
+``coefficient_status`` and ``trace_validity``. JSON null kappa means absorbing.
+The example checks the WOUT hash and radius, reports assumptions and numerical
+diagnostics, and does not infer missing transport physics. Volumes and W must
+be obtained consistently from that equilibrium; a hash alone cannot verify
+quadrature. Do not extrapolate an unmodeled outer region. A direct optional
+DKX calculation can export this same bundle; DKX is never imported by this
+module and full kinetic preparation/qualification remains the provider's work.
+
+``trace_response`` tests at least three distinct density-gradient probes of
+``Gamma/n=V-D*d(log n)/dr``. It flags non-affine or nonpositive-D responses.
+``convection_from_marker_drift`` implements ``V=A-D'-D*d(log W)/dr`` away from
+the axis. A Maxwellian-bath orbit response is not automatically a full-FP
+trace-impurity response; match collision physics before comparing coefficients.
+
+``restricted_residence`` uses a common positive horizon and +infinity for
+physical survivors. It returns a weighted empirical survival integral,
+escaped weight, effective sample size and an independent-marker standard error.
+The restricted mean does not estimate an unobserved survival tail. Any numerical
+failure invalidates the estimate rather than becoming a confined/censored
+particle. Hard exit events are validation observables, not smooth optimization
+objectives. Alpha energy retention, wall interception and pumped exhaust require
+separate models and diagnostics; core exit does not establish pumping.
diff --git a/docs/index.rst b/docs/index.rst
index 878336f4..06621929 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -19,4 +19,5 @@ stellarator coils to improve plasma confinement.
    :maxdepth: 2
    :caption: Contents:
 
-   getting_started
\ No newline at end of file
+   getting_started
+   helium_residence
\ No newline at end of file
diff --git a/essos/ash.py b/essos/ash.py
new file mode 100644
index 00000000..ac88e2ac
--- /dev/null
+++ b/essos/ash.py
@@ -0,0 +1,226 @@
+"""Coefficient-driven thermal-helium core residence, not pumped exhaust.
+
+The trace, local closure is Gamma=-D*dn/dr+V*n. Cell volumes and W=dV/dr
+must describe the same physical radial coordinate. Profiles, electric field,
+collision model and boundary assumptions belong to the coefficient provider.
+This module neither infers them from a WOUT nor imports a kinetic solver.
+All numerical functions are JIT/AD compatible; invalid values return a false
+``valid`` flag and NaN physical outputs. Static shape mistakes raise ValueError.
+Use :func:`check_helium_residence` at a host/reporting boundary.
+"""
+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 with checked SOLVAX tridiagonal solves.
+
+    Edges, W, D and V have shape (n+1,), volumes, S and dimensionless c (n,).
+    At least two cells start at r=0; the axis reflects. ``boundary_kappa_m_s``
+    is explicit: +inf absorbs; finite positive kappa gives Gamma=kappa*n at
+    the outer face, including half-cell transport resistance. Zero kappa has
+    no finite steady residence. Positive diffusivity is required without floors.
+
+    ``occupation_s`` and ``mean_occupation_s`` use c; ``mean_residence_s`` is
+    unweighted inventory/source. Validity qualifies this discrete solve only:
+    pivot, RHS residual, positivity, duality and balance checks do not bound
+    continuum error or certify a kinetic closure or its geometry derivatives.
+    Gradients include every supplied coefficient/volume/source/cost dependence;
+    host loading and an unprovided electric-field/root response are outside it.
+    """
+    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
+    tolerance = 2.e-8
+    valid = (input_ok & 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(cost/rate), mask(inventory/rate), 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)
+    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")
+    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.sum(w) > 0) & ~jnp.any(f))
+    w = w/jnp.sum(w)
+    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))
diff --git a/examples/particle_tracing/helium_ash_removal.py b/examples/particle_tracing/helium_ash_removal.py
new file mode 100644
index 00000000..39b55093
--- /dev/null
+++ b/examples/particle_tracing/helium_ash_removal.py
@@ -0,0 +1,110 @@
+"""Explicit exported-coefficient/WOUT helium residence example (no DKX import).
+
+    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
+
+An exporter supplies ``coefficients`` matching evaluate_helium_residence kwargs
+and ``metadata`` with the required keys below. JSON null kappa means absorbing.
+Coefficients must cover the declared physical radial domain; this example does
+not extrapolate a sparse kinetic calculation to the LCFS. WOUT identifies and
+checks the geometry; it supplies no missing profiles, electric field or D/V.
+The manufactured mode checks the numerical pipeline only, without an equilibrium.
+"""
+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()
diff --git a/pyproject.toml b/pyproject.toml
index 798c4924..bbd3beda 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -70,6 +70,8 @@ Homepage = "https://github.com/uwplasma/ESSOS"
 essos = "essos.__main__:main"
 
 [project.optional-dependencies]
+# Optional coefficient-driven helium residence diagnostics; orbit tracing stays independent.
+ash = ["solvax>=0.28.2"]
 # Examples that build VMEX equilibria, mgrid fields or free-boundary
 # (virtual-casing) fields. 0.11.3 is the first VMEX with MgridField(order=...).
 vmex = ["vmex[freeb]>=0.11.3"]
diff --git a/tests/test_ash.py b/tests/test_ash.py
new file mode 100644
index 00000000..5a353969
--- /dev/null
+++ b/tests/test_ash.py
@@ -0,0 +1,196 @@
+"""Independent manufactured/continuum checks; no stellarator physics certificate."""
+import jax
+jax.config.update("jax_enable_x64", True)
+import jax.numpy as jnp
+import numpy as np
+import pytest
+from scipy.integrate import quad
+
+pytest.importorskip("solvax")
+
+from essos.ash import (
+    _bernoulli, evaluate_helium_residence, check_helium_residence,
+    restricted_residence, trace_response, convection_from_marker_drift,
+)
+
+
+def cylinder(n=64, *, d=1., v=0., kappa=np.inf, power=1., source=None, cost=None):
+    e = jnp.linspace(0, 1, n+1)**power
+    return evaluate_helium_residence(
+        e, jnp.diff(e**2), 2*e, jnp.ones_like(e)*d, v*e,
+        jnp.ones(n) if source is None else source,
+        boundary_kappa_m_s=kappa, core_cost=cost)
+
+
+def test_bernoulli_limits_and_native_derivatives():
+    x = jnp.array([-1000., -50., -.001, 0., .001, 50., 1000.])
+    np.testing.assert_allclose(_bernoulli(-x)-_bernoulli(x), x, atol=1.e-13)
+    assert float(_bernoulli(0.)) == 1.
+    assert float(jax.grad(_bernoulli)(0.)) == -.5
+    assert np.isfinite(jax.vmap(jax.grad(_bernoulli))(x)).all()
+
+
+def test_absorbing_robin_and_second_order_convergence():
+    errors = []
+    for n in (16, 32, 64):
+        r = check_helium_residence(cylinder(n))
+        errors.append(float(r.mean_residence_s)-.125)
+        robin = check_helium_residence(cylinder(n, kappa=.2))
+        np.testing.assert_allclose(robin.mean_residence_s-r.mean_residence_s, 2.5, rtol=1.e-11)
+    np.testing.assert_allclose(np.array(errors[:-1])/errors[1:], 4., rtol=1.e-9)
+    nonuniform = check_helium_residence(cylinder(256, power=1.5))
+    np.testing.assert_allclose(nonuniform.mean_residence_s, .125, rtol=5.e-5)
+
+
+def test_variable_geometry_against_independent_continuum_integral():
+    # Spherical W=3r^2, D=1+r, V=2r(1+r), hence integrating factor exp(r^2).
+    n = 256
+    e = jnp.linspace(0, 1, n+1)
+    r = evaluate_helium_residence(e, jnp.diff(e**3), 3*e**2, 1+e,
+                                 2*e*(1+e), jnp.ones(n), boundary_kappa_m_s=.8)
+    assert bool(r.valid)
+    def inner(t):
+        return quad(lambda y: 3*y*y*np.exp(y*y), 0, t, epsabs=1.e-12)[0]
+    def exact(x):
+        return (quad(lambda t: np.exp(-t*t)*inner(t)/(3*t*t*(1+t)), x, 1)[0]
+                + np.exp(-1)*inner(1)/(3*.8))
+    centers = np.asarray((e[:-1]+e[1:])/2)
+    np.testing.assert_allclose(r.occupation_s[::32], [exact(x) for x in centers[::32]], rtol=8.e-5)
+
+
+def test_weighted_duality_balance_source_scaling_and_location():
+    n = 64
+    source = 1+jnp.exp(-jnp.linspace(0, 1, n)**2/.1)
+    cost = 1+3*jnp.linspace(1, 0, n)
+    a = cylinder(n, v=-2., kappa=.4, source=source, cost=cost)
+    b = cylinder(n, v=-2., kappa=.4, source=7*source, cost=cost)
+    assert bool(a.valid) and bool(b.valid)
+    assert float(a.duality_relative_error) < 1.e-11
+    assert float(a.balance_relative_error) < 1.e-11
+    np.testing.assert_allclose(b.inventory/a.inventory, 7., rtol=1.e-11)
+    np.testing.assert_allclose(b.mean_occupation_s, a.mean_occupation_s, rtol=1.e-11)
+    assert a.cost > a.inventory
+    core, edge = np.zeros(n), np.zeros(n)
+    core[4], edge[-5] = 1., 1.
+    assert cylinder(n, source=core).mean_residence_s > cylinder(n, source=edge).mean_residence_s
+    assert cylinder(n, d=.03, v=.09).mean_residence_s > 10*cylinder(n).mean_residence_s
+
+
+def test_native_jit_grad_includes_geometry_source_cost_and_boundary():
+    n = 18
+    def objective(p):
+        # Every input category depends on p; finite differences exercise all paths.
+        e = jnp.linspace(0, 1, n+1)*p[0]
+        s = 1+p[3]*jnp.linspace(0, 1, n)
+        return evaluate_helium_residence(
+            e, jnp.diff(e**2), 2*e, jnp.ones_like(e)*p[1], p[2]*e,
+            s, boundary_kappa_m_s=p[4], core_cost=1+p[5]*jnp.linspace(1, 0, n)).mean_occupation_s
+    p = jnp.array([1.1, .8, .3, .2, .6, .4])
+    native = jax.jit(jax.grad(objective))(p)
+    steps = np.array([1.e-3, 1.e-4, 1.e-5])
+    errors = []
+    for step in steps:
+        fd = np.array([(objective(p.at[i].add(step))-objective(p.at[i].add(-step)))/(2*step)
+                       for i in range(len(p))])
+        errors.append(float(np.max(np.abs(native-fd)/np.maximum(np.abs(native), 1.e-8))))
+        np.testing.assert_allclose(native, fd, rtol=5.e-6, atol=2.e-8)
+    assert errors[-1] < .02*errors[0]
+    # Absorbing boundary must not introduce inf*0 gradients through W/geometry.
+    def absorbing(radius):
+        e = jnp.linspace(0, 1, n+1)*radius
+        return evaluate_helium_residence(e, jnp.diff(e**2), 2*e, jnp.ones_like(e),
+            jnp.zeros_like(e), jnp.ones(n), boundary_kappa_m_s=jnp.inf).mean_residence_s
+    np.testing.assert_allclose(jax.grad(absorbing)(1.), 2*absorbing(1.), rtol=1.e-11)
+    np.testing.assert_allclose(jax.grad(lambda d: cylinder(16, d=d).mean_residence_s)(1.),
+                               -cylinder(16).mean_residence_s, rtol=1.e-11)
+
+
+@pytest.mark.parametrize('kwargs', [dict(d=0.), dict(d=-1.), dict(d=np.nan),
+    dict(kappa=0.), dict(kappa=-1.), dict(kappa=np.nan), dict(source=-np.ones(8)),
+    dict(source=np.zeros(8)), dict(cost=-np.ones(8))])
+def test_invalid_values_are_explicit_under_jit(kwargs):
+    r = jax.jit(lambda: cylinder(8, **kwargs))()
+    assert not bool(r.valid)
+    assert np.isnan(r.mean_occupation_s)
+    with pytest.raises(ArithmeticError):
+        check_helium_residence(r)
+
+
+def test_bad_mesh_shape_and_numerically_closed_escape_rejected():
+    with pytest.raises(ValueError):
+        cylinder(8, source=np.ones(9))
+    e = np.linspace(0, 1, 9)
+    r = evaluate_helium_residence(e, np.diff(e**2), 2*e, np.ones(9),
+                                 -1.e6*np.ones(9), np.ones(8), boundary_kappa_m_s=np.inf)
+    assert not bool(r.valid)
+    bad = e.copy(); bad[3] = bad[2]
+    r = evaluate_helium_residence(bad, np.diff(e**2), 2*e, np.ones(9),
+                                 np.zeros(9), np.ones(8), boundary_kappa_m_s=1.)
+    assert not bool(r.valid)
+
+
+def test_common_horizon_survival_integral_and_weighted_uncertainty():
+    r = restricted_residence([1., 2., np.inf], 4., weights=[1., 2., 1.])
+    assert bool(r.valid) and not bool(r.no_censoring_in_sample)
+    assert float(r.restricted_mean_s) == 2.25
+    assert float(r.escaped_fraction) == .75
+    np.testing.assert_allclose(r.standard_error_s, np.sqrt(.375*1.1875/.625))
+    np.testing.assert_allclose(r.effective_sample_size, 8/3)
+    # Integrating empirical survival independently agrees with the restricted mean.
+    t = np.array([0., 1., 2., 4.])
+    survival = np.array([1., .75, .25])
+    assert float(r.restricted_mean_s) == np.diff(t) @ survival
+    for times, failed in (([1., 2., np.inf], [False, True, False]),
+                           ([1., np.nan, np.inf], None), ([1., 5., np.inf], None),
+                           ([1., -np.inf, np.inf], None)):
+        bad = restricted_residence(times, 4., failed=failed)
+        assert not bool(bad.valid) and np.isnan(bad.restricted_mean_s)
+    zero_weight_failure = restricted_residence([1., np.inf], 4., weights=[1., 0.], failed=[False, True])
+    assert not bool(zero_weight_failure.valid)
+
+
+def test_trace_response_and_marker_drift_correction():
+    x = jnp.array([-2., 0., 3.])
+    r = jax.jit(trace_response)(x, .4-.7*x)
+    assert bool(r.valid)
+    np.testing.assert_allclose([r.diffusion_m2_s, r.convection_m_s], [.7, .4], atol=1.e-14)
+    assert not bool(trace_response(x, .4-.7*x+.1*x*x).valid)
+    assert not bool(trace_response(x, .4+.7*x).valid)
+    assert not bool(trace_response(jnp.zeros(3), jnp.ones(3)).valid)
+    # W=3r^2, D=1+r, V=.4r: A=V+1+(1+r)*2/r.
+    radius = jnp.array([.1, .3, .8])
+    d = 1+radius
+    drift = .4*radius+1+d*2/radius
+    np.testing.assert_allclose(convection_from_marker_drift(drift, d, 1., 2/radius), .4*radius, atol=2.e-15)
+
+
+def test_explicit_coefficient_wout_example_and_provenance(tmp_path):
+    import hashlib
+    import importlib.util
+    import json
+    from pathlib import Path
+    from netCDF4 import Dataset
+    example_path = Path(__file__).parents[1]/'examples/particle_tracing/helium_ash_removal.py'
+    spec = importlib.util.spec_from_file_location('helium_example', example_path)
+    example = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(example)
+    # Minimal WOUT metadata is enough here: no claim of a kinetic/equilibrium run.
+    wout = tmp_path/'wout.nc'
+    with Dataset(wout, 'w') as nc:
+        nc.createVariable('Aminor_p', 'f8').assignValue(1.)
+    coefficients, metadata = example.manufactured_response(16)
+    coefficients = {k: v.tolist() if isinstance(v, np.ndarray) else v for k, v in coefficients.items()}
+    coefficients['boundary_kappa_m_s'] = None
+    metadata['equilibrium_sha256'] = hashlib.sha256(wout.read_bytes()).hexdigest()
+    bundle = tmp_path/'response.json'
+    bundle.write_text(json.dumps(dict(coefficients=coefficients, metadata=metadata)))
+    output = tmp_path/'report.json'
+    report = example.main(['--coefficients', str(bundle), '--wout', str(wout), '--output', str(output)])
+    assert report['pumped_fraction'] == 'not modeled'
+    assert report['thermal_orbit_verification'] == 'not evaluated'
+    assert report['metadata']['modeled_boundary_s'] == 1.
+    assert json.loads(output.read_text()) == report
+    metadata['equilibrium_sha256'] = 'wrong'
+    bundle.write_text(json.dumps(dict(coefficients=coefficients, metadata=metadata)))
+    with pytest.raises(ValueError, match='hash'):
+        example.load_response(bundle, wout)

From 7507c4ce99018ed035fa6157b5025234c04c1aea Mon Sep 17 00:00:00 2001
From: Rogerio Jorge <rogerio.jorge@wisc.edu>
Date: Thu, 15 Oct 2026 16:05:55 -0500
Subject: [PATCH 2/5] Reject overflowing residence outputs and normalize large
 marker weights

---
 essos/ash.py      | 22 +++++++++++++++++-----
 tests/test_ash.py | 33 +++++++++++++++++++++++++++++++++
 2 files changed, 50 insertions(+), 5 deletions(-)

diff --git a/essos/ash.py b/essos/ash.py
index ac88e2ac..1eeba91f 100644
--- a/essos/ash.py
+++ b/essos/ash.py
@@ -128,15 +128,18 @@ def residual(lower, upper, x, rhs):
     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(density)) & jnp.all(jnp.isfinite(occupation)) &
+    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(cost/rate), mask(inventory/rate), mask(flux), rf, rb, duality, balance,
+        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)
@@ -171,14 +174,23 @@ def restricted_residence(exit_times_s, horizon_s, *, weights=None, failed=None):
     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)
+    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.sum(w) > 0) & ~jnp.any(f))
-    w = w/jnp.sum(w)
+             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
diff --git a/tests/test_ash.py b/tests/test_ash.py
index 5a353969..834666a3 100644
--- a/tests/test_ash.py
+++ b/tests/test_ash.py
@@ -194,3 +194,36 @@ def test_explicit_coefficient_wout_example_and_provenance(tmp_path):
     bundle.write_text(json.dumps(dict(coefficients=coefficients, metadata=metadata)))
     with pytest.raises(ValueError, match='hash'):
         example.load_response(bundle, wout)
+
+
+def test_huge_finite_marker_weights_are_normalized_without_overflow():
+    huge = jnp.asarray([1.e308, 1.e308])
+    r = jax.jit(lambda w: restricted_residence(jnp.array([1., 2.]), 3., weights=w))(huge)
+    assert bool(r.valid)
+    np.testing.assert_allclose(r.restricted_mean_s, 1.5)
+    np.testing.assert_allclose(r.escaped_fraction, 1.)
+    np.testing.assert_allclose(r.standard_error_s, .5)
+    np.testing.assert_allclose(r.effective_sample_size, 2.)
+    # Large integer counts are converted before reduction, avoiding wraparound.
+    counts = jnp.asarray([9_000_000_000_000_000_000]*3, dtype=jnp.int64)
+    integer_result = restricted_residence([1., 2., 3.], 4., weights=counts)
+    assert bool(integer_result.valid)
+    np.testing.assert_allclose(integer_result.restricted_mean_s, 2.)
+    np.testing.assert_allclose(integer_result.escaped_fraction, 1.)
+    np.testing.assert_allclose(integer_result.effective_sample_size, 3.)
+
+
+def test_overflowing_inventory_invalidates_otherwise_accurate_weighted_solve():
+    # Tiny cost weights keep occupation/cost finite while ordinary inventory
+    # overflows. Residual/duality/balance checks alone cannot detect this.
+    r = jax.jit(lambda: evaluate_helium_residence(
+        jnp.array([0., .5, 1.]), jnp.array([1.e200, 1.e200]),
+        jnp.ones(3), jnp.ones(3), jnp.zeros(3), jnp.ones(2),
+        boundary_kappa_m_s=jnp.inf, core_cost=jnp.array([1.e-200, 1.e-200])))()
+    assert not bool(r.valid)
+    assert np.isnan(r.inventory) and np.isnan(r.mean_residence_s)
+    assert np.isnan(r.mean_occupation_s)
+    assert float(r.forward_residual) == 0.
+    assert float(r.duality_relative_error) == 0.
+    with pytest.raises(ArithmeticError):
+        check_helium_residence(r)

From d35fa728b6b34d3e3acd4c79592c9a59aa982bc5 Mon Sep 17 00:00:00 2001
From: Rogerio Jorge <rogerio.jorge@wisc.edu>
Date: Thu, 15 Oct 2026 16:32:17 -0500
Subject: [PATCH 3/5] Run optional helium ash tests in CI

---
 .github/workflows/build_test.yml | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/.github/workflows/build_test.yml b/.github/workflows/build_test.yml
index 064db38d..0823fa27 100644
--- a/.github/workflows/build_test.yml
+++ b/.github/workflows/build_test.yml
@@ -30,7 +30,8 @@ jobs:
     - name: Install ESSOS
       run: |
         python -m pip install --upgrade pip
-        pip install .
+        # Exercise the optional residence solver in both supported Python lanes.
+        pip install ".[ash]"
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip

From 830de32bd6679e66ab2445479ba468a83ec803d7 Mon Sep 17 00:00:00 2001
From: Rogerio Jorge <rogerio.jorge@wisc.edu>
Date: Tue, 20 Oct 2026 13:09:59 -0500
Subject: [PATCH 4/5] Shorten residence API documentation

---
 docs/helium_residence.rst | 103 +++++++++++++++++---------------------
 essos/ash.py              |  39 ++++++---------
 2 files changed, 62 insertions(+), 80 deletions(-)

diff --git a/docs/helium_residence.rst b/docs/helium_residence.rst
index b5d1edd2..e167fd61 100644
--- a/docs/helium_residence.rst
+++ b/docs/helium_residence.rst
@@ -1,65 +1,54 @@
 Thermal-helium core residence
 ============================
 
-Install ``pip install 'essos[ash]'`` for the optional SOLVAX dependency.
-``essos.ash`` accepts explicit local trace-helium transport coefficients,
-physical cell volumes and a source. It solves the conservative steady equation
-``K n = M S`` and its transpose ``K.T T = M c``. The source-weighted occupation
-is ``(M c).T n / sum(M S) = (M S).T T / sum(M S)``. For ``c=1`` this is mean
-core residence; weighted occupation can describe a frozen core-dilution cost.
-This standard forward/backward identity is exact for the stated discrete model.
-
-Use ``jax.config.update('jax_enable_x64', True)`` before preparing arrays for the
-strict residual and balance tolerances. Face arrays D [m^2/s], V [m/s] and
-W=dV/dr [m^2] have n+1 entries; cell volumes [m^3], source [m^-3/s] and optional
-dimensionless cost have n. The first radial edge is zero and reflects.
-``boundary_kappa_m_s`` is mandatory: positive finite values impose
-``Gamma=kappa*n`` and positive infinity absorbs. The half-cell transport and
-boundary resistances are both included. A closed boundary, nonpositive D,
-invalid values, unresolved inward barriers or unqualified solves yield
-``valid=False`` and NaN physical outputs. Shape errors raise ``ValueError``.
-``check_helium_residence(result)`` raises on an invalid result outside JIT.
-
-The functions retain native JAX arrays for JIT and differentiation through all
-provided coefficients, geometry measures, source, cost and boundary inputs.
-File loading is host-only. Differentiation cannot recover dependence omitted
-by a coefficient provider, including a frozen electric field or ambipolar root.
-Pivot and residual checks do not provide a condition-number error bound;
-refine the radial mesh and kinetic response and test the local transport closure.
-No artificial diffusivity floor is applied.
-
-Run the manufactured numerical check::
+Install ``pip install 'essos[ash]'`` and enable JAX float64. ``essos.ash``
+solves the local trace-helium closure ``Gamma=-D*dn/dr+V*n`` using supplied
+transport coefficients, geometry measures, a source and a boundary condition.
+The forward and backward equations are ``K n = M S`` and ``K.T T = M c``.
+Source-weighted occupation is ``(M c).T n / sum(M S) = (M S).T T / sum(M S)``;
+``c=1`` gives mean core residence. Recycling and pumping require separate models.
+
+Inputs and checks
+-----------------
+
+Face arrays D [m^2/s], V [m/s] and W=dV/dr [m^2] have n+1 entries. Cell volumes
+[m^3], source [m^-3/s] and optional dimensionless cost have n entries. The radial
+mesh starts at zero with a reflecting axis. Positive finite ``boundary_kappa_m_s``
+imposes ``Gamma=kappa*n``, including half-cell resistance; +infinity absorbs.
+D must be positive. Volumes and W must use the same physical radial coordinate.
+
+Invalid inputs or failed discrete residual, positivity, duality or balance checks
+return ``valid=False`` and NaN physical outputs. Shape errors raise ``ValueError``;
+``check_helium_residence`` raises on an invalid result outside JIT. These checks
+require float64 and do not establish radial convergence or kinetic validity.
+Gradients include supplied coefficient, geometry, source, cost and boundary
+dependence; omitted responses, such as an ambipolar root, remain omitted.
+
+Example
+-------
+
+Run the manufactured cylinder or supply an exported kinetic response::
 
     python examples/particle_tracing/helium_ash_removal.py --manufactured
-
-For an actual equilibrium use an explicit exported response::
-
     python examples/particle_tracing/helium_ash_removal.py \
       --wout path/to/wout.nc --coefficients path/to/helium_response.json
 
-The JSON has ``coefficients`` matching the evaluator's argument names and
-``metadata`` describing ``equilibrium_sha256``, ``profiles_id``, ``source_id``,
-``collision_model``, ``electric_field`` (prescribed profile or root branch),
-``turbulent_closure``, ``boundary_model``, ``volume_origin``,
-``coefficient_status`` and ``trace_validity``. JSON null kappa means absorbing.
-The example checks the WOUT hash and radius, reports assumptions and numerical
-diagnostics, and does not infer missing transport physics. Volumes and W must
-be obtained consistently from that equilibrium; a hash alone cannot verify
-quadrature. Do not extrapolate an unmodeled outer region. A direct optional
-DKX calculation can export this same bundle; DKX is never imported by this
-module and full kinetic preparation/qualification remains the provider's work.
-
-``trace_response`` tests at least three distinct density-gradient probes of
-``Gamma/n=V-D*d(log n)/dr``. It flags non-affine or nonpositive-D responses.
-``convection_from_marker_drift`` implements ``V=A-D'-D*d(log W)/dr`` away from
-the axis. A Maxwellian-bath orbit response is not automatically a full-FP
-trace-impurity response; match collision physics before comparing coefficients.
-
-``restricted_residence`` uses a common positive horizon and +infinity for
-physical survivors. It returns a weighted empirical survival integral,
-escaped weight, effective sample size and an independent-marker standard error.
-The restricted mean does not estimate an unobserved survival tail. Any numerical
-failure invalidates the estimate rather than becoming a confined/censored
-particle. Hard exit events are validation observables, not smooth optimization
-objectives. Alpha energy retention, wall interception and pumped exhaust require
-separate models and diagnostics; core exit does not establish pumping.
+The JSON contains ``coefficients`` matching the evaluator arguments and
+``metadata`` with ``equilibrium_sha256``, ``profiles_id``, ``source_id``,
+``collision_model``, ``electric_field``, ``turbulent_closure``, ``boundary_model``,
+``volume_origin``, ``coefficient_status`` and ``trace_validity``. Null kappa means
+absorbing. The example checks the WOUT hash and modeled radius. The provider
+must supply consistent volumes, profiles and coefficients over that domain;
+a hash does not verify quadrature, and WOUT alone cannot supply transport.
+
+Other diagnostics
+-----------------
+
+``trace_response`` fits ``Gamma/n=V-D*d(log n)/dr`` using at least three distinct
+probes with matched backgrounds and electric field. ``convection_from_marker_drift``
+uses ``V=A-D'-D*d(log W)/dr`` away from the singular axis.
+
+``restricted_residence`` estimates ``E[min(T,H)]`` at a common positive horizon,
+using +infinity for survivors. Any numerical failure invalidates the estimate.
+Its standard error assumes independent markers with fixed weights; the
+unobserved tail and correlated-marker uncertainty are not estimated.
diff --git a/essos/ash.py b/essos/ash.py
index 1eeba91f..50687141 100644
--- a/essos/ash.py
+++ b/essos/ash.py
@@ -1,12 +1,9 @@
-"""Coefficient-driven thermal-helium core residence, not pumped exhaust.
-
-The trace, local closure is Gamma=-D*dn/dr+V*n. Cell volumes and W=dV/dr
-must describe the same physical radial coordinate. Profiles, electric field,
-collision model and boundary assumptions belong to the coefficient provider.
-This module neither infers them from a WOUT nor imports a kinetic solver.
-All numerical functions are JIT/AD compatible; invalid values return a false
-``valid`` flag and NaN physical outputs. Static shape mistakes raise ValueError.
-Use :func:`check_helium_residence` at a host/reporting boundary.
+"""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 docs/helium_residence.rst for the model and input contract.
 """
 from typing import NamedTuple
 
@@ -52,20 +49,16 @@ def evaluate_helium_residence(
     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 with checked SOLVAX tridiagonal solves.
-
-    Edges, W, D and V have shape (n+1,), volumes, S and dimensionless c (n,).
-    At least two cells start at r=0; the axis reflects. ``boundary_kappa_m_s``
-    is explicit: +inf absorbs; finite positive kappa gives Gamma=kappa*n at
-    the outer face, including half-cell transport resistance. Zero kappa has
-    no finite steady residence. Positive diffusivity is required without floors.
-
-    ``occupation_s`` and ``mean_occupation_s`` use c; ``mean_residence_s`` is
-    unweighted inventory/source. Validity qualifies this discrete solve only:
-    pivot, RHS residual, positivity, duality and balance checks do not bound
-    continuum error or certify a kinetic closure or its geometry derivatives.
-    Gradients include every supplied coefficient/volume/source/cost dependence;
-    host loading and an unprovided electric-field/root response are outside it.
+    """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

From a1037b386e8f1ff8ff7bec391a66f451f808b517 Mon Sep 17 00:00:00 2001
From: Rogerio Jorge <rogerio.jorge@wisc.edu>
Date: Fri, 23 Oct 2026 15:03:50 -0500
Subject: [PATCH 5/5] Consolidate helium residence documentation in README

---
 README.md                                     | 37 ++++++++++---
 docs/helium_residence.rst                     | 54 -------------------
 docs/index.rst                                |  3 +-
 essos/ash.py                                  |  2 +-
 .../particle_tracing/helium_ash_removal.py    | 12 ++---
 5 files changed, 38 insertions(+), 70 deletions(-)
 delete mode 100644 docs/helium_residence.rst

diff --git a/README.md b/README.md
index 7d62918e..4f24c6f0 100644
--- a/README.md
+++ b/README.md
@@ -32,12 +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.** Optional `essos[ash]` coefficient-driven,
-  source-weighted residence/occupation with conservative fluxes, checked SOLVAX
-  solves and native JAX derivatives. See the [explicit-input example](examples/particle_tracing/helium_ash_removal.py)
-  and [scope/validation notes](docs/helium_residence.rst). Core residence requires
-  qualified transport coefficients, profiles and boundary conditions; it does
-  not predict pumped exhaust from a WOUT.
+- **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
@@ -147,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
diff --git a/docs/helium_residence.rst b/docs/helium_residence.rst
deleted file mode 100644
index e167fd61..00000000
--- a/docs/helium_residence.rst
+++ /dev/null
@@ -1,54 +0,0 @@
-Thermal-helium core residence
-============================
-
-Install ``pip install 'essos[ash]'`` and enable JAX float64. ``essos.ash``
-solves the local trace-helium closure ``Gamma=-D*dn/dr+V*n`` using supplied
-transport coefficients, geometry measures, a source and a boundary condition.
-The forward and backward equations are ``K n = M S`` and ``K.T T = M c``.
-Source-weighted occupation is ``(M c).T n / sum(M S) = (M S).T T / sum(M S)``;
-``c=1`` gives mean core residence. Recycling and pumping require separate models.
-
-Inputs and checks
------------------
-
-Face arrays D [m^2/s], V [m/s] and W=dV/dr [m^2] have n+1 entries. Cell volumes
-[m^3], source [m^-3/s] and optional dimensionless cost have n entries. The radial
-mesh starts at zero with a reflecting axis. Positive finite ``boundary_kappa_m_s``
-imposes ``Gamma=kappa*n``, including half-cell resistance; +infinity absorbs.
-D must be positive. Volumes and W must use the same physical radial coordinate.
-
-Invalid inputs or failed discrete residual, positivity, duality or balance checks
-return ``valid=False`` and NaN physical outputs. Shape errors raise ``ValueError``;
-``check_helium_residence`` raises on an invalid result outside JIT. These checks
-require float64 and do not establish radial convergence or kinetic validity.
-Gradients include supplied coefficient, geometry, source, cost and boundary
-dependence; omitted responses, such as an ambipolar root, remain omitted.
-
-Example
--------
-
-Run the manufactured cylinder or supply an exported kinetic response::
-
-    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
-
-The JSON contains ``coefficients`` matching the evaluator arguments and
-``metadata`` with ``equilibrium_sha256``, ``profiles_id``, ``source_id``,
-``collision_model``, ``electric_field``, ``turbulent_closure``, ``boundary_model``,
-``volume_origin``, ``coefficient_status`` and ``trace_validity``. Null kappa means
-absorbing. The example checks the WOUT hash and modeled radius. The provider
-must supply consistent volumes, profiles and coefficients over that domain;
-a hash does not verify quadrature, and WOUT alone cannot supply transport.
-
-Other diagnostics
------------------
-
-``trace_response`` fits ``Gamma/n=V-D*d(log n)/dr`` using at least three distinct
-probes with matched backgrounds and electric field. ``convection_from_marker_drift``
-uses ``V=A-D'-D*d(log W)/dr`` away from the singular axis.
-
-``restricted_residence`` estimates ``E[min(T,H)]`` at a common positive horizon,
-using +infinity for survivors. Any numerical failure invalidates the estimate.
-Its standard error assumes independent markers with fixed weights; the
-unobserved tail and correlated-marker uncertainty are not estimated.
diff --git a/docs/index.rst b/docs/index.rst
index 06621929..878336f4 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -19,5 +19,4 @@ stellarator coils to improve plasma confinement.
    :maxdepth: 2
    :caption: Contents:
 
-   getting_started
-   helium_residence
\ No newline at end of file
+   getting_started
\ No newline at end of file
diff --git a/essos/ash.py b/essos/ash.py
index 50687141..d5e6e7f8 100644
--- a/essos/ash.py
+++ b/essos/ash.py
@@ -3,7 +3,7 @@
 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 docs/helium_residence.rst for the model and input contract.
+and pumping. See the README for the model and input contract.
 """
 from typing import NamedTuple
 
diff --git a/examples/particle_tracing/helium_ash_removal.py b/examples/particle_tracing/helium_ash_removal.py
index 39b55093..8455a8dc 100644
--- a/examples/particle_tracing/helium_ash_removal.py
+++ b/examples/particle_tracing/helium_ash_removal.py
@@ -1,15 +1,13 @@
-"""Explicit exported-coefficient/WOUT helium residence example (no DKX import).
+"""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
 
-An exporter supplies ``coefficients`` matching evaluate_helium_residence kwargs
-and ``metadata`` with the required keys below. JSON null kappa means absorbing.
-Coefficients must cover the declared physical radial domain; this example does
-not extrapolate a sparse kinetic calculation to the LCFS. WOUT identifies and
-checks the geometry; it supplies no missing profiles, electric field or D/V.
-The manufactured mode checks the numerical pipeline only, without an equilibrium.
+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