smsharma · MinhMPA · May 3, 2026 · May 3, 2026 · May 3, 2026 · May 3, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -1,9 +1,97 @@
 # clax Development Progress
 
-## Status: Speed-optimized fit_cl preset (34s V100) + full accuracy pipeline
+## Status: Speed-optimized fit_cl preset (34s V100) + full accuracy pipeline + forward-mode AD
 
 **End-to-end differentiable pipeline from cosmological parameters to P(k),
-C_l^TT/EE/TE/BB, and lensed C_l^TT/EE/TE/BB. AD gradients verified to 0.03%.**
+C_l^TT/EE/TE/BB, and lensed C_l^TT/EE/TE/BB. AD gradients verified to 0.03%.
+`jax.jvp` (forward-mode AD) now works through the full pipeline.**
+
+### May 10, 2026: Forward-mode AD + stable thermodynamics gradients (PR-B + PR-C)
+
+**Three AD fixes making `jax.jvp` work end-to-end:**
+
+**PR-B (`feat/forward-mode-ad`)** — Forward-mode through `z_reio` and `theta_s`:
+- Converted `_find_z_reio` from `custom_vjp` to `custom_jvp` with IFT rule:
+  `z_reio_dot = -F_dot / dF_dz` where `F = tau_reio_model(z_reio) - tau_reio_target`.
+- Fixed inf-inf NaN in He C_He tangent at z~15 (`He_Boltz=exp(~500)`):
+  reformulated `C_He = (inv_A + Λ) / (inv_A + Λ + R_up)`.
+- Converted `shoot_fn` to `custom_jvp` with IFT rule `dh/dθ_s = 1/(dθ_s/dh)`.
+  Also restores reverse-mode via JAX transposition — existing tests unchanged.
+
+**PR-C (`fix/thermo-remaining-gradients`)** — Stable `omega_b` gradients for splines:
+- n_H_0 rescaling for `kappa_dot_of_loga`, `exp_m_kappa_of_loga`, `g_of_loga`.
+  Stops the ~10^8× FD blowup from the Friedmann-scan eigenvalue accumulation.
+  Exact where x_e~const (loga<-8); 10-30% residual near recombination (still finite).
+
+**Test results:** 10/10 thermo, 7/7 shooting, all 5 new gradient tests GREEN.
+
+### May 3, 2026: C_l^phiphi API consolidation + z-aware Halofit injection (BREAKING)
+
+**Collapses six redundant `compute_cl_pp_*` implementations into a single
+public `compute_cl_pp(... nonlinear="none"|"halofit")` backed by the
+source-Limber kernel. Halofit nonlinear corrections are z-aware on-the-fly,
+matching the CLASS approach (`fourier.c:1706-1716`).**
+
+This supersedes the earlier consolidation merged as PR #11 and immediately
+reverted (PR #13). The earlier version had three test failures emerge
+post-merge — vmap-unsafe Halofit modulator (k_eval validation +
+`if z == 0.0`), spurious R values at high z due to the Python-level
+`sigma_convergence_check` being bypassed under vmap, and a missing local
+test fixture. All three are fixed in this PR.
+
+**Public API:**
+
+```python
+clax.compute_cl_pp(pt, params, bg, th, l_max, *, nonlinear="none")
+# nonlinear ∈ {"none", "halofit"}; anything else raises ValueError.
+# nonlinear="ept" is reserved for a follow-up clax-pt PR.
+```
+
+The function is now re-exported at the package root (`clax.compute_cl_pp`)
+along with `clax.lens_cls`.
+
+**Removed** (BREAKING):
+
+- `compute_cl_pp` (Siddharth's original) — superseded by source-Limber.
+- `compute_cl_pp_fast` — inaccurate at l ≥ 300 per its own docstring.
+- `compute_cl_pp_vmap` — Hermite Bessel-table vmap; superseded.
+- `compute_cl_pp_limber` — Poisson-reconstruction Limber (~20% overestimate
+  at l = 2500 vs CLASS).
+
+**Renamed:**
+
+- `compute_cl_pp_source_limber` → `compute_cl_pp` (sole public entry).
+- `compute_cl_pp_transfer` → `_compute_cl_pp_full_bessel` (private oracle
+  retained for cross-impl tests).
+
+**Halofit modulator** (new, private — `_halofit_modulator`):
+
+- Builds R(k, z) = P_NL(k, z) / P_lin(k, z) on a 100-point z-grid via
+  `vmap(compute_pk_nonlinear)` (CLASS-aligned density).
+- Inline `sigma_R(R_min, ...) >= 1` check using `jnp.where` (replaces the
+  Python try/except in `compute_pk_nonlinear` that gets bypassed under
+  vmap). Forces R = 1 where Halofit isn't applicable, matching CLASS
+  `fourier.c:1706-1716`.
+- `k_max_extend = 0` default (no power-law extension; uses pt.k_grid as-is,
+  matching CLASS's no-extrapolation behavior). Pass a positive value to
+  override with log-log extrapolation for narrow k-grids.
+- 2D-interpolates R onto every (pt.k_grid, pt.tau_grid) lattice point via
+  `bg.loga_of_tau` for the τ→z mapping.
+- Multiplies sqrt(R) into S_transfer (CLASS source-multiplication recipe).
+
+**Tests:**
+
+- `tests/test_cl_pp.py` (renamed from `test_cl_pp_source_limber.py`):
+  contract + linear accuracy + cross-impl-vs-`_compute_cl_pp_full_bessel`
+  + Halofit smoke + JIT/AD compatibility.
+- `tests/test_clpp_halofit_ratio.py` (rewritten): NL/linear ratio vs CLASS
+  Halofit reference (≤ 7% at l ≤ 500, ≤ 10% at l ≥ 1000).
+- `tests/test_lensing.py`: callers updated to the new signature, with a
+  local `pipeline` fixture matching upstream-main convention.
+- Deleted: `tests/test_cl_pp_implementations.py`, `tests/test_cl_pp_limber.py`.
+
+**Migration:** pre-1.0, hard break — no deprecation shims. Replace all
+calls to the removed implementations with `compute_cl_pp(... nonlinear=...)`.
 
 ### Apr 20, 2026: Rodas5 Rosenbrock solver + dark energy perturbations + accuracy fixes
 

diff --git a/README.md b/README.md
@@ -53,6 +53,40 @@ Full spin-2 correlation function lensing with Cgl2 corrections, 12 Wigner d-func
 
 Lensing algorithm sub-0.2% at all l=10-2000 for TT and EE (tested with CLASS unlensed+pp as input).
 
+### Lensing potential C_l^phiphi
+
+`clax.compute_cl_pp(pt, params, bg, th, l_max, *, nonlinear="none")` uses
+the source-based Limber kernel (CLASS `transfer.c` + `harmonic.c`):
+
+With `pt_k_max_cl >= 5.0 Mpc^-1` (required for Halofit's σ(R) bisection),
+measured residuals at the default cosmology:
+
+Absolute |C_l^phiphi_clax / C_l^phiphi_CLASS - 1| at the default cosmology:
+
+| Multipole l | nonlinear="none" | nonlinear="halofit" | NL/lin ratio agreement |
+|-------------|-----------------:|--------------------:|-----------------------:|
+| 100         | 0.18%            | 0.16%               | 0.02%                  |
+| 200         | 0.14%            | 0.10%               | 0.05%                  |
+| 500         | 0.23%            | 0.15%               | 0.09%                  |
+| 1000        | 0.12%            | 0.08%               | 0.04%                  |
+| 1500        | 0.49%            | 0.61%               | 0.11%                  |
+| 2000        | 0.71%            | 0.76%               | 0.05%                  |
+| 2500        | 0.39%            | 1.16%               | 0.76%                  |
+
+Both nonlinear modes share the same source-Limber kernel, so they track
+each other closely — the absolute Halofit residual is dominated by the
+common linear-kernel systematic, with a small extra contribution from the
+R(k, z) accuracy that the third column isolates. Halofit injection
+multiplies the lensing source by sqrt(R(k, z(τ))) where R = P_NL/P_lin
+is computed on a 100-point z-grid via `vmap(compute_pk_nonlinear)` with
+log-log k-extension to k_max=20 Mpc^-1 (CLASS itself uses a dedicated
+nonlinear k-grid that extends past the perturbation k-range; the
+extension provides the equivalent coverage). Sub-percent absolute
+residual across all ℓ ≤ 2000; ~1% at ℓ=2500 where linear-kernel and
+R-residual systematics add constructively. Narrower k-grids
+(`pt_k_max_cl < 5`) gracefully degrade to no NL correction (R = 1, per
+CLASS `fourier.c:1706-1716`).
+
 ### Multi-cosmology validation
 
 Validated at 10 LCDM parameter variations (omega_b, omega_cdm, h, n_s, tau_reio at +/-20%):
@@ -61,13 +95,22 @@ Validated at 10 LCDM parameter variations (omega_b, omega_cdm, h, n_s, tau_reio
 
 ### Matter power spectrum P(k)
 
-| k [Mpc^-1] | clax / CLASS | Error |
-|-------------|------------------|-------|
-| 0.001       | 0.970            | 3.0%  |
-| 0.010       | 0.986            | 1.4%  |
-| 0.050       | 0.984            | 1.6%  |
-| 0.100       | 1.013            | 1.3%  |
-| 0.300       | 0.966            | 3.5%  |
+Single-mode `compute_pk(params, prec, k)` at `PrecisionParams.planck_fast()`
+(60 k/decade, l_max=50, ncdm_q_size=5, rtol=1e-6) vs CLASS v3.3.4
+linear `pk_lin(k, z=0)` at default Planck 2018 LCDM:
+
+| k [Mpc^-1] | clax / CLASS | Error    |
+|------------|--------------|----------|
+| 0.003      | 0.9983       | -0.17%   |
+| 0.010      | 1.0025       | +0.25%   |
+| 0.050      | 1.0004       | +0.04%   |
+| 0.100      | 1.0007       | +0.07%   |
+| 0.300      | 0.9997       | -0.03%   |
+| 1.000      | 0.9980       | -0.20%   |
+
+Median |error| ≈ 0.07%, worst-case 0.25%. Underlying perturbation evolution
+(δ_cdm and δ_b at synchronous gauge) matches CLASS to <0.05% — most of any
+remaining residual is from photon-hierarchy and IC corrections at next order.
 
 ### Pipeline accuracy
 
@@ -131,6 +174,16 @@ th = clax.thermodynamics_solve(params, prec, bg)
 pt = perturbations_solve(params, prec, bg, th)
 cl_tt = compute_cl_tt_interp(pt, params, bg, [30, 100, 200])
 cl_ee = compute_cl_ee_interp(pt, params, bg, [30, 100, 200])
+
+# CMB lensing potential C_l^phiphi (linear or Halofit-corrected)
+cl_pp_lin = clax.compute_cl_pp(pt, params, bg, th, l_max=2500)
+cl_pp_nl  = clax.compute_cl_pp(pt, params, bg, th, l_max=2500,
+                                nonlinear="halofit")
+
+# Lens the unlensed CMB spectra
+cl_tt_lensed, cl_ee_lensed, cl_te_lensed, cl_bb_lensed = clax.lens_cls(
+    cl_tt_unlensed, cl_ee_unlensed, cl_te_unlensed, cl_bb_unlensed,
+    cl_pp_nl, l_max=2500)
 ```
 
 ## Installation
@@ -241,7 +294,7 @@ pk_direct = compute_pk(params, prec_direct, k=0.05)
 | `bessel.py`         | Spherical Bessel functions j_l(x)                |
 | `transfer.py`       | Linear matter P(k) from perturbation solve       |
 | `harmonic.py`       | C_l^TT/EE/TE/BB from line-of-sight integration  |
-| `lensing.py`        | Correlation-function lensing method              |
+| `lensing.py`        | C_l^phiphi (source-Limber, optional Halofit NL) and lensed C_l (correlation-function method) |
 | `nonlinear.py`      | HaloFit (Takahashi 2012)                         |
 | `shooting.py`       | theta_s -> H0 via Newton + `custom_vjp`          |
 

diff --git a/clax/__init__.py b/clax/__init__.py
@@ -25,6 +25,7 @@
 from clax.primordial import primordial_scalar_pk, primordial_tensor_pk  # noqa: F401
 from clax.harmonic import compute_cl_bb, compute_cl_tt, compute_cl_ee, compute_cl_te, compute_cls_all, compute_cls_all_fast  # noqa: F401
 from clax.transfer import compute_pk_from_perturbations, compute_linear_matter_pk_from_perturbations  # noqa: F401
+from clax.lensing import compute_cl_pp, lens_cls  # noqa: F401
 
 import functools
 from dataclasses import dataclass, replace as dataclass_replace