co2-eos 0.2.0__tar.gz → 0.3.1__tar.gz

{co2_eos-0.2.0 → co2_eos-0.3.1}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: co2-eos
-Version: 0.2.0
+Version: 0.3.1
 Summary: Differentiable CO₂ thermodynamic properties in JAX
 Author-email: John Patrick Therrien <john@flux-tech.de>
 License-Expression: Apache-2.0
@@ -83,7 +83,7 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
 CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
 
 - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. The inversions carry hand-written `custom_jvp` rules (implicit function theorem), so forward- and reverse-mode gradients are exact and cheap; validated against central finite differences.
-- **Fast.** JIT-compiled and vectorisable, with hand-coded analytic α-derivatives instead of a nested-`jax.grad` chain. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery — it solves `T(ρ, u)` once (table-seeded, fixed-iteration analytic-Cv Newton) and reuses the derivatives for every property. On an Apple M-class CPU at a 4096-point float64 batch this is **4.8× faster** for the `T(ρ, u)` inversion and **3.6× faster** for the full `T + P + μ + k` recovery than the v0.1 autodiff path. On an OVH Tesla V100S the same comparison is 1.3–2.2× at these (latency-bound) batch sizes and grows to **7.7× / 4.9×** once the GPU is compute-bound at large batches (`bench/compare.py`, `bench/PROFILING.md`).
+- **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows — the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches — 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
 - **Also fast for `(P, T)`.** The `state_from_PT` workflow that mirrors `PropsSI` includes an iterative density solve and flattens to a few tens of microseconds per point at large batches, roughly 2.7× faster than `CoolProp.PropsSI` in a Python loop on the same inputs (10,000-point batch, Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU it runs at 27,800 states/sec on a 10⁶-point batch (36 μs/pt) versus 1,600 states/sec on Colab CPU (615 μs/pt), a 17× speedup. Its ceiling is the density solve's data-dependent `while_loop` (lockstep across GPU warps); the feed-forward `(ρ, u)` hot path above does not have that bottleneck.
 - **Composable.** `jit`, `vmap`, `grad`, `custom_vjp`: the full JAX transformation stack works. Embed property evaluations inside your own JIT-compiled simulation and differentiate end-to-end.
 - **Phase-aware.** Robust inversions near the critical point using Halley's method with step damping and bisection fallback. Two-phase dome detection that avoids convergence to thermodynamically unstable spinodal states.

{co2_eos-0.2.0 → co2_eos-0.3.1}/README.md

@@ -49,7 +49,7 @@ But CoolProp is a C++ library with Python bindings. You can't `jax.grad` through
 CO2-EOS solves this for CO₂ by implementing the same reference EOS (Span-Wagner 1996) directly in JAX:
 
 - **Differentiable.** `jax.grad` through any property, any inversion, any combination. No finite differences. The inversions carry hand-written `custom_jvp` rules (implicit function theorem), so forward- and reverse-mode gradients are exact and cheap; validated against central finite differences.
-- **Fast.** JIT-compiled and vectorisable, with hand-coded analytic α-derivatives instead of a nested-`jax.grad` chain. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery — it solves `T(ρ, u)` once (table-seeded, fixed-iteration analytic-Cv Newton) and reuses the derivatives for every property. On an Apple M-class CPU at a 4096-point float64 batch this is **4.8× faster** for the `T(ρ, u)` inversion and **3.6× faster** for the full `T + P + μ + k` recovery than the v0.1 autodiff path. On an OVH Tesla V100S the same comparison is 1.3–2.2× at these (latency-bound) batch sizes and grows to **7.7× / 4.9×** once the GPU is compute-bound at large batches (`bench/compare.py`, `bench/PROFILING.md`).
+- **Fast.** JIT-compiled and vectorisable, with hand-coded, transcendental-economized analytic α-derivatives: the Span-Wagner exponent structure (¼-integer τ-powers, integer δ-powers, shared exponential envelopes) is evaluated by sqrt/multiply ladders, so a full six-derivative bundle costs ~12 exps + 3 pows per point instead of ~150 array-exponent pows — the same arithmetic, re-associated, exact to round-off. For simulation codes that carry `(ρ, u)` as conserved variables, `properties_from_rho_u` fuses the whole primitive recovery: a dome-safe table seed plus three unrolled analytic-Cv Newton steps recover `T(ρ, u)` to 1e-12 K, and one shared derivative bundle feeds every property including the transport critical enhancement. On an Apple M2 Pro (float64) the full recovery runs **123 µs for a 64-point batch and 0.53 µs/point at large batches — 7.6–19.9× faster than the v0.1 autodiff path** and 2.6–6.3× faster than v0.2. On an OVH Tesla V100S the inversion reaches **63.7×** vs v0.1 at a 2²⁰-point batch (2 ns/point; full recovery 16 ns/point), and the launch-latency floor at small batches dropped ~1.7× (`bench/compare.py`, `bench/PROFILING.md`, validation in `docs/validation-fastpath-2026-07.md`).
 - **Also fast for `(P, T)`.** The `state_from_PT` workflow that mirrors `PropsSI` includes an iterative density solve and flattens to a few tens of microseconds per point at large batches, roughly 2.7× faster than `CoolProp.PropsSI` in a Python loop on the same inputs (10,000-point batch, Apple M2 Pro; see `examples/launch_demo.ipynb`, section 2). On a Colab T4 GPU it runs at 27,800 states/sec on a 10⁶-point batch (36 μs/pt) versus 1,600 states/sec on Colab CPU (615 μs/pt), a 17× speedup. Its ceiling is the density solve's data-dependent `while_loop` (lockstep across GPU warps); the feed-forward `(ρ, u)` hot path above does not have that bottleneck.
 - **Composable.** `jit`, `vmap`, `grad`, `custom_vjp`: the full JAX transformation stack works. Embed property evaluations inside your own JIT-compiled simulation and differentiate end-to-end.
 - **Phase-aware.** Robust inversions near the critical point using Halley's method with step damping and bisection fallback. Two-phase dome detection that avoids convergence to thermodynamically unstable spinodal states.

{co2_eos-0.2.0 → co2_eos-0.3.1}/co2_eos/__init__.py

@@ -52,7 +52,7 @@ import jax.numpy as jnp
 # Enable float64 for the entire library (also done by submodules).
 jax.config.update("jax_enable_x64", True)
 
-__version__ = "0.2.0"
+__version__ = "0.3.1"
 
 # ── Phase hint constants ────────────────────────────────────────────────────
 LIQUID = 0

{co2_eos-0.2.0 → co2_eos-0.3.1}/co2_eos/core.py

@@ -14,12 +14,13 @@ Design:
   * The Newton inner loop touches only τ-derivatives (`helmholtz.residual_tau_*`)
     with the δ-dependent envelopes precomputed once (they are loop-invariant at
     fixed ρ) — all u and Cv need at fixed density.
-  * A precomputed (ρ, u) → T₀ table seeds Newton within a fraction of a kelvin,
-    so a short FIXED, unrolled iteration (no `lax.while_loop`) reaches float64
-    round-off.  Fixed iteration keeps the GPU kernel uniform and branchless; the
-    table seed is purely a convergence accelerator (the polish sets accuracy and
-    the IFT JVP sets gradients), so it carries no accuracy or differentiability
-    risk.
+  * A precomputed (ρ, u) → T₀ table (dome-safe, uniform grid) seeds Newton so
+    that a short FIXED, unrolled iteration (no `lax.while_loop`) reaches
+    float64 round-off — measured 8.5e-13 K max over the single-phase envelope
+    at 3 steps.  Fixed iteration keeps the GPU kernel uniform and branchless;
+    the table seed is purely a convergence accelerator (the polish sets
+    accuracy and the IFT JVP sets gradients), so it carries no accuracy or
+    differentiability risk.
   * Gradients use the implicit function theorem via `custom_jvp`, with the
     Jacobian entries (Cv, ∂u/∂ρ) taken analytically — correct jvp and vjp.
 
@@ -39,10 +40,13 @@ from co2_eos.span_wagner import TC, RHOC, R, T_TRIPLE
 from co2_eos import helmholtz as hz
 from co2_eos import transport as _tr
 
-# Fixed Newton iterations for the (ρ, u) → T inversion.  The table seed lands
-# < ~0.5 K from the root across its range; 4 analytic-Newton steps then reach
-# float64 round-off, 5 leaves margin.  (Validated in tests/test_inversion_*.)
-_NEWTON_ITERS = 5
+# Fixed Newton iterations for the (ρ, u) → T inversion.  With the dome-filled
+# 256×512 table seed, 3 analytic-Newton steps reach float64 round-off over the
+# single-phase envelope (measured max 8.5e-13 K on 40k samples incl. the
+# critical neighbourhood; see bench/proto_seed_v2.py and tests).  The affine
+# fallback seed is coarse (±tens of K), so it gets a longer fixed iteration.
+_NEWTON_ITERS_TABLE = 3
+_NEWTON_ITERS_AFFINE = 8
 _T_MIN = T_TRIPLE
 _T_MAX = 800.0
 
@@ -62,30 +66,51 @@ _AFFINE_A, _AFFINE_B = 255.6, 1.656e-4
 _HAVE_SEED = False
 try:
     _seed = np.load(_SEED_PATH)
-    _SEED_RHO = jnp.asarray(_seed["rho_grid"])
-    _SEED_U = jnp.asarray(_seed["u_grid"])
-    _SEED_T0 = jnp.asarray(_seed["T0_table"])
-    _SEED_NRHO = int(_SEED_RHO.shape[0])
-    _SEED_NU = int(_SEED_U.shape[0])
+    # Stored float32 (seed-only precision; halves the embedded-constant
+    # footprint, which XLA:CPU gathers care about).  Kept flat so the four
+    # bilinear corners come from ONE gather of (base + static offsets) —
+    # four separate gather ops cost ~1.3 µs each per call.
+    _SEED_T0_FLAT = jnp.asarray(_seed["T0_table"], dtype=jnp.float32).ravel()
+    _np_rho = np.asarray(_seed["rho_grid"])
+    _np_u = np.asarray(_seed["u_grid"])
+    _SEED_NRHO = int(_np_rho.shape[0])
+    _SEED_NU = int(_np_u.shape[0])
+    # The grids are uniform (linspace) — index by arithmetic, not searchsorted.
+    _SEED_R_LO = float(_np_rho[0])
+    _SEED_R_STEP = float(_np_rho[1] - _np_rho[0])
+    _SEED_U_LO = float(_np_u[0])
+    _SEED_U_STEP = float(_np_u[1] - _np_u[0])
+    for _g, _step in ((_np_rho, _SEED_R_STEP), (_np_u, _SEED_U_STEP)):
+        if not np.allclose(np.diff(_g), _step, rtol=1e-12):
+            raise ValueError("seed table grids must be uniform")
     _HAVE_SEED = True
 except FileNotFoundError:
     pass
 
 
 def _seed_T(rho, u):
-    """Initial T guess for the (ρ, u) inversion via the bilinear table."""
+    """Initial T guess for the (ρ, u) inversion via the bilinear table.
+
+    Uniform-grid direct indexing (the generator writes linspace grids), with
+    the query clamped to the table box — outside it the Newton polish still
+    owns the accuracy, the seed is only a starting point.
+    """
     if not _HAVE_SEED:
         return jnp.clip(_AFFINE_A + _AFFINE_B * u, _T_MIN, _T_MAX)
-    ri = jnp.clip(jnp.searchsorted(_SEED_RHO, rho) - 1, 0, _SEED_NRHO - 2)
-    uj = jnp.clip(jnp.searchsorted(_SEED_U, u) - 1, 0, _SEED_NU - 2)
-    r0, r1 = _SEED_RHO[ri], _SEED_RHO[ri + 1]
-    u0, u1 = _SEED_U[uj], _SEED_U[uj + 1]
-    fr = jnp.clip((rho - r0) / (r1 - r0), 0.0, 1.0)
-    fu = jnp.clip((u - u0) / (u1 - u0), 0.0, 1.0)
-    c00 = _SEED_T0[ri, uj]
-    c01 = _SEED_T0[ri, uj + 1]
-    c10 = _SEED_T0[ri + 1, uj]
-    c11 = _SEED_T0[ri + 1, uj + 1]
+    fi = jnp.clip((rho - _SEED_R_LO) / _SEED_R_STEP, 0.0,
+                  _SEED_NRHO - 1 - 1e-9)
+    fj = jnp.clip((u - _SEED_U_LO) / _SEED_U_STEP, 0.0, _SEED_NU - 1 - 1e-9)
+    ri = fi.astype(jnp.int32)
+    uj = fj.astype(jnp.int32)
+    fr = fi - ri
+    fu = fj - uj
+    base = ri * _SEED_NU + uj
+    corners = _SEED_T0_FLAT[base + jnp.array([0, 1, _SEED_NU, _SEED_NU + 1],
+                                             dtype=jnp.int32)]
+    c00, c01, c10, c11 = (corners[0].astype(jnp.float64),
+                          corners[1].astype(jnp.float64),
+                          corners[2].astype(jnp.float64),
+                          corners[3].astype(jnp.float64))
     return ((c00 * (1.0 - fr) + c10 * fr) * (1.0 - fu)
             + (c01 * (1.0 - fr) + c11 * fr) * fu)
 
@@ -148,22 +173,23 @@ def _solve_T(rho, u):
     """Fixed-iteration Newton solve for T given (ρ, u). Scalar.
 
     δ-invariant envelopes are precomputed once; each step evaluates only the
-    τ-dependent transcendentals.
+    τ-dependent transcendentals.  The short fixed count is unrolled (no loop
+    construct) so XLA fuses the steps and the kernel stays branchless.
     """
     delta = rho / RHOC
     dstate = hz.residual_tau_prep(delta)
     T = _seed_T(rho, u)
+    iters = _NEWTON_ITERS_TABLE if _HAVE_SEED else _NEWTON_ITERS_AFFINE
 
-    def step(_, T):
+    for _ in range(iters):
         tau = TC / T
         a0_t, a0_tt = hz.ideal_tau_only(tau)
         ar_t, ar_tt = hz.residual_tau_fast(tau, dstate)
         f = R * TC * (a0_t + ar_t) - u
         cv = -R * tau ** 2 * (a0_tt + ar_tt)
         cv = jnp.where(jnp.abs(cv) > 1e-30, cv, 1e-30)
-        return jnp.clip(T - f / cv, _T_MIN, _T_MAX)
-
-    return jax.lax.fori_loop(0, _NEWTON_ITERS, step, T)
+        T = jnp.clip(T - f / cv, _T_MIN, _T_MAX)
+    return T
 
 
 # ── Implicit differentiation via custom_jvp ─────────────────────────────────

co2_eos-0.3.1/co2_eos/helmholtz.py

@@ -0,0 +1,449 @@
+"""Span-Wagner (1996) reduced Helmholtz energy with hand-coded analytic
+derivatives — pure JAX, transcendental-economized.
+
+This is the numerical core of the EOS hot path.  It computes α and all the
+derivatives a property/inversion pass needs — α_δ, α_τ, α_δδ, α_ττ, α_δτ —
+analytically, and exploits the *structure* of the Span-Wagner exponents so the
+per-point transcendental count collapses:
+
+  * every τ-exponent in the polynomial/exponential/Gaussian terms is a
+    multiple of ¼ → all τ-powers come from ``sqrt(sqrt(τ))`` plus a multiply
+    ladder (zero ``pow`` calls);
+  * every δ-exponent is a small integer (d ≤ 10, l ≤ 6) → multiply ladders;
+  * the 27 ``exp(-δ^l)`` envelopes take only six distinct values (l = 1..6);
+  * the five Gaussian exponentials share (η, β, γ) pairwise → four distinct;
+  * the non-analytic terms share θ (A, β identical across terms) and terms
+    40/41 share Δ; Δ^(7/8) is a sqrt ladder; only Δ^0.925, s^(5/3) and
+    s^(2/3) remain as real ``pow`` calls.
+
+The structural facts are asserted at import against the coefficient tables in
+``span_wagner`` (the single source of truth) — if the tables ever change shape,
+import fails loudly rather than computing a stale factorization.
+
+The result is the same Span-Wagner sum re-associated: agreement with the
+autodiff derivatives of ``span_wagner.alphar`` / ``alpha0`` is round-off level
+(~1e-12; see ``tests/test_analytic_derivs.py``), and every function remains
+JIT-compilable, vmappable and differentiable to higher order.
+
+Convention: τ = Tc/T (reduced inverse temperature), δ = ρ/ρc (reduced
+density).  Derivatives are plain partials w.r.t. τ and δ (not the δ·∂/∂δ
+reduced form), matching the formulas in ``span_wagner``.
+"""
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+jax.config.update("jax_enable_x64", True)
+
+# Pull the coefficient tables straight from span_wagner so there is a single
+# source of truth for the constants.
+from co2_eos.span_wagner import (
+    _A0_A1, _A0_A2, _A0_LOGTAU, _A0_PE_N, _A0_PE_T,
+    _AR_N, _AR_D, _AR_T, _AR_L,
+    _AR_GAUSS_N, _AR_GAUSS_D, _AR_GAUSS_T,
+    _AR_GAUSS_ETA, _AR_GAUSS_BETA, _AR_GAUSS_GAMMA, _AR_GAUSS_EPS,
+    _NA_N, _NA_A, _NA_B, _NA_BETA, _NA_BIG_A, _NA_BIG_B, _NA_BIG_C, _NA_BIG_D,
+)
+
+# ── Static (Python-scalar) views of the tables — folded into the trace ─────
+_PN = np.asarray(_AR_N)
+_PD = np.asarray(_AR_D).astype(int)
+_PT = np.asarray(_AR_T)
+_PL = np.asarray(_AR_L).astype(int)
+
+_GN = np.asarray(_AR_GAUSS_N)
+_GD = np.asarray(_AR_GAUSS_D).astype(int)
+_GT = np.asarray(_AR_GAUSS_T)
+_GETA = np.asarray(_AR_GAUSS_ETA)
+_GBETA = np.asarray(_AR_GAUSS_BETA)
+_GGAMMA = np.asarray(_AR_GAUSS_GAMMA)
+
+_NAN = np.asarray(_NA_N)
+_NAA = np.asarray(_NA_A)
+_NAB = np.asarray(_NA_B)
+_NABIGB = np.asarray(_NA_BIG_B)
+_NABIGC = np.asarray(_NA_BIG_C)
+
+# Structural facts the economization relies on — asserted, not assumed.
+assert np.all(_PD == np.asarray(_AR_D)) and _PD.max() <= 10, "δ-exponent not a small integer"
+assert np.all(_PL == np.asarray(_AR_L)) and _PL.max() <= 6, "l-exponent not a small integer"
+assert np.all(np.abs(_PT * 4 - np.round(_PT * 4)) < 1e-12), "τ-exponent not a multiple of ¼"
+assert np.all(np.abs(_GT * 4 - np.round(_GT * 4)) < 1e-12)
+assert np.all(np.asarray(_AR_GAUSS_EPS) == 1.0), "Gaussian ε ≠ 1"
+_NA_P_SCALAR = float(1.0 / (2.0 * np.asarray(_NA_BETA)[0]))
+assert np.all(np.asarray(_NA_BETA) == np.asarray(_NA_BETA)[0]), "NA β not shared"
+_NA_BIGA_SCALAR = float(np.asarray(_NA_BIG_A)[0])
+assert np.all(np.asarray(_NA_BIG_A) == _NA_BIGA_SCALAR), "NA A not shared"
+_NA_BIGD_SCALAR = float(np.asarray(_NA_BIG_D)[0])
+assert np.all(np.asarray(_NA_BIG_D) == _NA_BIGD_SCALAR), "NA D not shared"
+assert set(np.round(_NAA * 2).astype(int)) <= {6, 7}, "NA a must be 3 or 3.5"
+assert _NAB[0] == 0.875 and _NAB[2] == 0.875, "Δ^(7/8) sqrt ladder assumption"
+
+# Reciprocal of (2β) for the non-analytic θ exponent, precomputed (array view
+# kept for external callers/tests; the kernels use the shared scalar).
+_NA_P = 1.0 / (2.0 * _NA_BETA)
+
+# Distinct τ-exponents of the polynomial+exponential terms, and per-exponent
+# static weights for the τ-derivative dot products.
+_TPOLY = sorted(set(_PT.tolist()))
+_TP_IDX = {t: j for j, t in enumerate(_TPOLY)}
+_TALL = sorted(set(_TPOLY) | set(_GT.tolist()))
+_W_T = np.array(_TPOLY)                       # t_j
+_W_TT = np.array([t * (t - 1.0) for t in _TPOLY])
+
+
+def _tau_ladder(tau):
+    """All distinct τ^t needed by the terms, via 2 sqrts + multiplies.
+
+    Every exponent is a multiple of ¼, so τ^t = τ^int · (τ^¼)^frac with the
+    integer parts built by a multiply chain over increasing exponents.
+    """
+    st = jnp.sqrt(tau)
+    qt = jnp.sqrt(st)
+    p = {0.0: jnp.ones_like(tau), 0.25: qt, 0.5: st, 0.75: st * qt, 1.0: tau}
+    for t in _TALL:
+        if t in p:
+            continue
+        if t == int(t):
+            t_int = int(t)
+            half = float(t_int // 2)
+            other = float(t_int - t_int // 2)
+            if half in p and other in p:
+                p[t] = p[half] * p[other]
+            else:
+                avail = max(k for k in p if k == int(k) and 0 < k <= t)
+                acc, rem = p[avail], t_int - int(avail)
+                while rem > 0:
+                    step = max(k for k in p if k == int(k) and 0 < k <= rem)
+                    acc = acc * p[step]
+                    rem -= int(step)
+                p[t] = acc
+        else:
+            base, frac = float(int(t)), t - int(t)
+            p[t] = p[base] * p[frac]        # base already built (sorted order)
+    return p
+
+
+def _delta_ladder(delta, dmax=10):
+    """δ^0..δ^dmax by successive multiplication."""
+    p = {0: jnp.ones_like(delta), 1: delta}
+    for k in range(2, dmax + 1):
+        p[k] = p[k - 1] * delta
+    return p
+
+
+def _exp_envelopes(dp):
+    """The six distinct exponential envelopes E_l = exp(-δ^l), l = 0..6."""
+    E = {0: jnp.ones_like(dp[1])}
+    for l in range(1, 7):
+        E[l] = jnp.exp(-dp[l])
+    return E
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Ideal-gas part α⁰ — τ derivatives (δ part is only ln δ, not needed here)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def ideal_derivs(tau, delta):
+    """Return (α⁰, α⁰_τ, α⁰_ττ) at scalar (τ, δ).
+
+    α⁰ = a1 + a2·τ + L·ln τ + ln δ + Σ nk·ln(1 - exp(-θk·τ))
+    Only the τ-derivatives and the value are used by the property/inversion
+    formulas (the ideal δ-dependence cancels out of every measurable property
+    except through the additive ln δ in the value, which entropy/Gibbs need).
+    """
+    L = _A0_LOGTAU
+    e = jnp.exp(-_A0_PE_T * tau)            # exp(-θk τ)
+    one_minus_e = 1.0 - e
+
+    a0 = (_A0_A1 + _A0_A2 * tau + L * jnp.log(tau) + jnp.log(delta)
+          + jnp.sum(_A0_PE_N * jnp.log(one_minus_e)))
+
+    # d/dτ ln(1-e) = θ·e/(1-e)
+    a0_t = (_A0_A2 + L / tau
+            + jnp.sum(_A0_PE_N * _A0_PE_T * e / one_minus_e))
+
+    # d²/dτ² ln(1-e) = -θ²·e/(1-e)²
+    a0_tt = (-L / tau ** 2
+             - jnp.sum(_A0_PE_N * _A0_PE_T ** 2 * e / one_minus_e ** 2))
+
+    return a0, a0_t, a0_tt
+
+
+def ideal_tau_only(tau):
+    """Return (α⁰_τ, α⁰_ττ) — the only ideal quantities the Newton needs.
+
+    Independent of δ, so cheaper than ``ideal_derivs`` for the inner loop.
+    """
+    L = _A0_LOGTAU
+    e = jnp.exp(-_A0_PE_T * tau)
+    one_minus_e = 1.0 - e
+    a0_t = (_A0_A2 + L / tau
+            + jnp.sum(_A0_PE_N * _A0_PE_T * e / one_minus_e))
+    a0_tt = (-L / tau ** 2
+             - jnp.sum(_A0_PE_N * _A0_PE_T ** 2 * e / one_minus_e ** 2))
+    return a0_t, a0_tt
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Non-analytic terms — shared helpers
+# ═══════════════════════════════════════════════════════════════════════════
+
+def _pow78(Db):
+    """Db^(7/8) by three sqrts and three multiplies (Db ≥ 1e-300 guarded)."""
+    r8 = jnp.sqrt(jnp.sqrt(jnp.sqrt(Db)))
+    q2 = r8 * r8
+    q4 = q2 * q2
+    return q4 * q2 * r8
+
+
+def _na_delta_state(delta):
+    """δ-side quantities of the non-analytic terms (loop-invariant at fixed ρ).
+
+    Returns (s, dme, s_p, Bsa, e_cs) where Bsa[(a, B)] = B·s^a and
+    e_cs[C] = exp(-C·s) for the distinct (a, B) / C among the three terms.
+    """
+    dme = delta - 1.0
+    s = dme * dme
+    s_p = s ** _NA_P_SCALAR                  # s^(5/3)
+    s3 = s * s * s
+    abs_dm1 = jnp.abs(dme)
+    sa = {3.5: s3 * abs_dm1, 3.0: s3}        # s^a by multiplies
+    Bsa = {}
+    for k in range(3):
+        key = (float(_NAA[k]), float(_NABIGB[k]))
+        if key not in Bsa:
+            Bsa[key] = key[1] * sa[key[0]]
+    e_cs = {float(C): jnp.exp(-float(C) * s) for C in sorted(set(_NABIGC.tolist()))}
+    return s, dme, s_p, Bsa, e_cs
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# Residual part αʳ — full analytic derivative bundle (one economized pass)
+# ═══════════════════════════════════════════════════════════════════════════
+
+def residual_derivs(tau, delta):
+    """Return (αʳ, αʳ_δ, αʳ_τ, αʳ_δδ, αʳ_ττ, αʳ_δτ) at scalar (τ, δ).
+
+    All six are accumulated from the same per-term values; the transcendental
+    budget per point is ~12 exps + ~8 sqrts + 3 pows (see module docstring).
+    """
+    inv_d = 1.0 / delta
+    inv_d2 = inv_d * inv_d
+    inv_t = 1.0 / tau
+    inv_t2 = inv_t * inv_t
+
+    tp = _tau_ladder(tau)
+    dp = _delta_ladder(delta)
+    E = _exp_envelopes(dp)
+
+    ar = ar_d = ar_t = ar_dd = ar_tt = ar_dt = 0.0
+
+    # ── Polynomial + exponential terms (static unroll; constants folded) ──
+    for k in range(len(_PN)):
+        n, d, t, l = float(_PN[k]), int(_PD[k]), float(_PT[k]), int(_PL[k])
+        Tk = n * dp[d] * tp[t] * E[l]
+        ld = l * dp[l] if l else 0.0         # l·δ^l (0 when l = 0)
+        D1 = (d - ld) * inv_d
+        D1p = (-d - (l * (l - 1) * dp[l] if l else 0.0)) * inv_d2
+        Tt = t * inv_t
+        ar += Tk
+        ar_d += Tk * D1
+        ar_t += Tk * Tt
+        ar_dd += Tk * (D1 * D1 + D1p)
+        ar_tt += Tk * (t * (t - 1.0) * inv_t2)
+        ar_dt += Tk * D1 * Tt
+
+    # ── Gaussian bell-shaped terms (four distinct exponentials) ──
+    dme = delta - 1.0                        # ε = 1 for all five terms
+    dme2 = dme * dme
+    e_g = {}
+    for k in range(5):
+        key = (float(_GETA[k]), float(_GBETA[k]), float(_GGAMMA[k]))
+        if key not in e_g:
+            tmg = tau - key[2]
+            e_g[key] = (tmg, jnp.exp(-key[0] * dme2 - key[1] * tmg * tmg))
+    for k in range(5):
+        n, d, t = float(_GN[k]), int(_GD[k]), float(_GT[k])
+        eta, beta = float(_GETA[k]), float(_GBETA[k])
+        tmg, eg = e_g[(eta, beta, float(_GGAMMA[k]))]
+        G = n * dp[d] * tp[t] * eg
+        Gd = d * inv_d - 2.0 * eta * dme
+        Gdp = -d * inv_d2 - 2.0 * eta
+        Gt = t * inv_t - 2.0 * beta * tmg
+        Gtp = -t * inv_t2 - 2.0 * beta
+        ar += G
+        ar_d += G * Gd
+        ar_t += G * Gt
+        ar_dd += G * (Gd * Gd + Gdp)
+        ar_tt += G * (Gt * Gt + Gtp)
+        ar_dt += G * Gd * Gt
+
+    # ── Non-analytic terms (shared θ; terms 40/41 share Δ) ──
+    s, dme, s_p, Bsa, e_cs = _na_delta_state(delta)
+    s_pm1 = s ** (_NA_P_SCALAR - 1.0)        # s^(2/3)
+    s2 = s * s
+    sam1 = {3.5: s2 * jnp.abs(dme), 3.0: s2}  # s^(a-1) by multiplies
+    tm1 = tau - 1.0
+
+    theta = -tm1 + _NA_BIGA_SCALAR * s_p
+    th_d = 2.0 * _NA_BIGA_SCALAR * _NA_P_SCALAR * dme * s_pm1
+    th_dd = 2.0 * _NA_BIGA_SCALAR * _NA_P_SCALAR * (2.0 * _NA_P_SCALAR - 1.0) * s_pm1
+    De_t = -2.0 * theta
+    De_tt = 2.0
+    exp_tm1 = jnp.exp(-_NA_BIGD_SCALAR * tm1 * tm1)
+
+    Delta_cache = {}
+    for k in range(3):
+        a, b = float(_NAA[k]), float(_NAB[k])
+        B, Cc = float(_NABIGB[k]), float(_NABIGC[k])
+        dkey = (a, B)
+        if dkey not in Delta_cache:
+            Delta = theta * theta + Bsa[dkey]
+            Db = jnp.maximum(Delta, 1e-300)
+            De_d = 2.0 * theta * th_d + 2.0 * B * a * dme * sam1[a]
+            De_dd = (2.0 * th_d * th_d + 2.0 * theta * th_dd
+                     + 2.0 * B * a * (2.0 * a - 1.0) * sam1[a])
+            De_dt = -2.0 * th_d
+            Delta_cache[dkey] = (Db, De_d, De_dd, De_dt)
+        Db, De_d, De_dd, De_dt = Delta_cache[dkey]
+
+        F = _pow78(Db) if b == 0.875 else Db ** b
+        Fm1 = F / Db
+        Fm2 = Fm1 / Db
+        F_d = b * Fm1 * De_d
+        F_t = b * Fm1 * De_t
+        F_dd = b * (b - 1.0) * Fm2 * De_d * De_d + b * Fm1 * De_dd
+        F_tt = b * (b - 1.0) * Fm2 * De_t * De_t + b * Fm1 * De_tt
+        F_dt = b * (b - 1.0) * Fm2 * De_d * De_t + b * Fm1 * De_dt
+
+        Psi = e_cs[Cc] * exp_tm1
+        Ps_d = -2.0 * Cc * dme * Psi
+        Ps_t = -2.0 * _NA_BIGD_SCALAR * tm1 * Psi
+        Ps_dd = (4.0 * Cc * Cc * s - 2.0 * Cc) * Psi
+        Ps_tt = (4.0 * _NA_BIGD_SCALAR * _NA_BIGD_SCALAR * tm1 * tm1
+                 - 2.0 * _NA_BIGD_SCALAR) * Psi
+        Ps_dt = 4.0 * Cc * _NA_BIGD_SCALAR * dme * tm1 * Psi
+
+        n = float(_NAN[k])
+        W = F * delta * Psi
+        W_d = F_d * delta * Psi + F * Psi + F * delta * Ps_d
+        W_t = F_t * delta * Psi + F * delta * Ps_t
+        W_dd = (F_dd * delta * Psi + 2.0 * F_d * Psi + 2.0 * F_d * delta * Ps_d
+                + 2.0 * F * Ps_d + F * delta * Ps_dd)
+        W_tt = F_tt * delta * Psi + 2.0 * F_t * delta * Ps_t + F * delta * Ps_tt
+        W_dt = (F_dt * delta * Psi + F_d * delta * Ps_t + F_t * Psi + F * Ps_t
+                + F_t * delta * Ps_d + F * delta * Ps_dt)
+
+        ar += n * W
+        ar_d += n * W_d
+        ar_t += n * W_t
+        ar_dd += n * W_dd
+        ar_tt += n * W_tt
+        ar_dt += n * W_dt
+
+    return ar, ar_d, ar_t, ar_dd, ar_tt, ar_dt
+
+
+# ═══════════════════════════════════════════════════════════════════════════
+# τ-only inner-loop pair: δ-invariant prep + per-iteration fast eval
+# ═══════════════════════════════════════════════════════════════════════════
+
+def residual_tau_prep(delta):
+    """Precompute the δ-invariant state for the τ-derivative inner loop.
+
+    At fixed ρ (the Newton solves T at fixed density) every δ-dependent power
+    and exponential is loop-invariant.  The polynomial+exponential terms
+    collapse to one coefficient per distinct τ-exponent:
+    αʳ_poly(τ; δ) = Σ_j Ct[j]·τ^t_j with Ct[j] = Σ_{terms with t_j} n·δ^d·E_l.
+    """
+    dp = _delta_ladder(delta)
+    E = _exp_envelopes(dp)
+    Ct = [0.0] * len(_TPOLY)
+    for k in range(len(_PN)):
+        Ct[_TP_IDX[float(_PT[k])]] += float(_PN[k]) * dp[int(_PD[k])] * E[int(_PL[k])]
+    Ct = jnp.stack([jnp.asarray(c) for c in Ct])
+
+    dme = delta - 1.0
+    dme2 = dme * dme
+    Cg = jnp.stack([float(_GN[k]) * dp[int(_GD[k])]
+                    * jnp.exp(-float(_GETA[k]) * dme2) for k in range(5)])
+
+    _, _, s_p, Bsa, e_cs = _na_delta_state(delta)
+    return Ct, Cg, s_p, Bsa, e_cs, delta
+
+
+def residual_tau_fast(tau, dstate):
+    """Return (αʳ_τ, αʳ_ττ) from precomputed δ-state ``dstate``.
+
+    Identical result to ``residual_tau_derivs(tau, delta)`` but per call only
+    the τ-side transcendentals are evaluated: ~6 exps + 8 sqrts + 1 pow.
+    """
+    Ct, Cg, s_p, Bsa, e_cs, delta = dstate
+    inv_t = 1.0 / tau
+    inv_t2 = inv_t * inv_t
+
+    tp = _tau_ladder(tau)
+    taupows = jnp.stack([tp[t] for t in _TPOLY])
+    base = Ct * taupows
+    p_t = jnp.sum(base * _W_T) * inv_t
+    p_tt = jnp.sum(base * _W_TT) * inv_t2
+
+    g_t = 0.0
+    g_tt = 0.0
+    e_g = {}
+    for k in range(5):
+        key = (float(_GBETA[k]), float(_GGAMMA[k]))
+        if key not in e_g:
+            tmg = tau - key[1]
+            e_g[key] = (tmg, jnp.exp(-key[0] * tmg * tmg))
+    for k in range(5):
+        t, beta = float(_GT[k]), float(_GBETA[k])
+        tmg, eg = e_g[(beta, float(_GGAMMA[k]))]
+        gk = Cg[k] * tp[t] * eg
+        Gt = t * inv_t - 2.0 * beta * tmg
+        Gtp = -t * inv_t2 - 2.0 * beta
+        g_t += gk * Gt
+        g_tt += gk * (Gt * Gt + Gtp)
+
+    tm1 = tau - 1.0
+    theta = -tm1 + _NA_BIGA_SCALAR * s_p
+    De_t = -2.0 * theta
+    De_tt = 2.0
+    exp_tm1 = jnp.exp(-_NA_BIGD_SCALAR * tm1 * tm1)
+
+    na_t = 0.0
+    na_tt = 0.0
+    Db_cache = {}
+    for k in range(3):
+        a, b = float(_NAA[k]), float(_NAB[k])
+        B, Cc = float(_NABIGB[k]), float(_NABIGC[k])
+        dkey = (a, B)
+        if dkey not in Db_cache:
+            Db_cache[dkey] = jnp.maximum(theta * theta + Bsa[dkey], 1e-300)
+        Db = Db_cache[dkey]
+        F = _pow78(Db) if b == 0.875 else Db ** b
+        Fm1 = F / Db
+        Fm2 = Fm1 / Db
+        F_t = b * Fm1 * De_t
+        F_tt = b * (b - 1.0) * Fm2 * De_t * De_t + b * Fm1 * De_tt
+        Psi = e_cs[Cc] * exp_tm1
+        Ps_t = -2.0 * _NA_BIGD_SCALAR * tm1 * Psi
+        Ps_tt = (4.0 * _NA_BIGD_SCALAR * _NA_BIGD_SCALAR * tm1 * tm1
+                 - 2.0 * _NA_BIGD_SCALAR) * Psi
+        W_t = F_t * delta * Psi + F * delta * Ps_t
+        W_tt = F_tt * delta * Psi + 2.0 * F_t * delta * Ps_t + F * delta * Ps_tt
+        na_t += float(_NAN[k]) * W_t
+        na_tt += float(_NAN[k]) * W_tt
+
+    return p_t + g_t + na_t, p_tt + g_tt + na_tt
+
+
+def residual_tau_derivs(tau, delta):
+    """Return (αʳ_τ, αʳ_ττ) — the residual quantities the Newton needs.
+
+    Computes only the τ-derivatives; the δ-dependent envelopes are
+    loop-invariant at fixed ρ and XLA hoists them across an unrolled Newton.
+    """
+    return residual_tau_fast(tau, residual_tau_prep(delta))