smsharma · MinhMPA · May 3, 2026 · May 3, 2026 · May 3, 2026 · May 3, 2026
diff --git a/CHANGELOG.md b/CHANGELOG.md
diff --git a/README.md b/README.md
@@ -14,7 +14,27 @@ The goal is a differentiable alternative to [CLASS](https://github.com/lesgourg/
 
 ## Status
 
-Sub-0.2% unlensed C_l^TT/EE at l=20-1200. Full lensed C_l^TT/EE/TE/BB (sub-0.2% at l=10-2000). Multi-cosmology validated (10 LCDM parameter points). Full ncdm Boltzmann hierarchy. `fit_cl` preset for HMC / fitting; `planck_cl` for science-grade accuracy (487s on H100). See [CHANGELOG.md](CHANGELOG.md) for details.
+Sub-0.2% unlensed C_l^TT/EE at l=20-1200. Full lensed C_l^TT/EE/TE/BB (sub-0.2% at l=10-2000). Multi-cosmology validated (10 LCDM parameter points). Full ncdm Boltzmann hierarchy. `fit_cl` preset for HMC / fitting; `planck_cl` for science-grade accuracy (487s on H100). One-loop EFTofLSS module (`clax.ept`) ports the CLASS-PT FFTLog pipeline to JAX with sub-percent accuracy on monopoles/quadrupoles. See [CHANGELOG.md](CHANGELOG.md) for details.
+
+## One-loop EFTofLSS (`clax.ept`) accuracy vs CLASS-PT
+
+All comparisons at Planck 2018 best-fit LCDM, z=0.38, k < 0.30 h/Mpc:
+
+| Spectrum            | Max error | Mean error | Notes                              |
+|---------------------|-----------|------------|------------------------------------|
+| P_mm real-space     | **0.33%** | 0.10%      | Matter-matter, IR resummed         |
+| P_gg real-space     | **0.35%** | 0.12%      | Galaxy-galaxy (b1=2)               |
+| P_gm real-space     | **0.28%** | 0.09%      | Galaxy-matter cross                |
+| P_mm ell=0          | **0.21%** | 0.08%      | Matter monopole                    |
+| P_mm ell=2          | **0.34%** | 0.11%      | Matter quadrupole                  |
+| P_mm ell=4          | **1.2%**  | 0.45%      | Matter hexadecapole                |
+| P_gg ell=0          | **0.42%** | 0.15%      | Galaxy monopole                    |
+| P_gg ell=2          | **0.89%** | 0.50%      | Galaxy quadrupole                  |
+| P_gg ell=4          | **1.8%**  | 0.72%      | Galaxy hexadecapole                |
+
+Sub-percent on all monopoles and quadrupoles. Hexadecapole within 2% (limited by the small signal amplitude and zero-crossings — see project memory `project_ept_hexadecapole_floor`).
+
+The `clax.ept` module additionally feeds CMB lensing C_l^phiphi corrections via `compute_cl_pp(... nonlinear="ept")`.
 
 ## Accuracy comparison against CLASS v3.3.4
 
@@ -53,6 +73,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 +115,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 +194,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,8 +314,9 @@ 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              |
-| `nonlinear.py`      | HaloFit (Takahashi 2012)                         |
+| `lensing.py`        | C_l^phiphi (source-Limber, optional Halofit / clax-pt NL) and lensed C_l (correlation-function method) |
+| `nonlinear.py`      | HaloFit (Takahashi 2012 + Bird 2012 ν correction)|
+| `ept.py`            | One-loop EFTofLSS via FFTLog (CLASS-PT algorithm in JAX); real-space and RSD multipoles, IR resummation, bias |
 | `shooting.py`       | theta_s -> H0 via Newton + `custom_vjp`          |
 
 ## Precision presets
@@ -287,11 +361,15 @@ Default parameters correspond to Planck 2018 best-fit LCDM:
 ## References
 
 - **CLASS v3.3.4**: Blas, Lesgourgues & Tram (2011). [arXiv:1104.2933](https://arxiv.org/abs/1104.2933)
+- **CLASS-PT**: Chudaykin, Ivanov, Philcox & Simonović (2020). Non-linear perturbation theory extension of CLASS. [arXiv:2004.10607](https://arxiv.org/abs/2004.10607)
+- **EFT-with-FFT**: Nguyen (2026). Pedagogical notes on FFTLog computation of one-loop PT integrals. [GitHub](https://github.com/MinhMPA/EFT-with-FFT)
 - **DISCO-DJ**: Hahn, List & Porqueres (2023). Differentiable Einstein-Boltzmann solver in JAX. [arXiv:2311.03291](https://arxiv.org/abs/2311.03291)
 - **Bolt.jl**: Li, Sullivan & Millea (2023). Differentiable Boltzmann solver in Julia. [Zenodo: 10.5281/zenodo.10065126](https://zenodo.org/records/10065126)
 - Seljak & Zaldarriaga (1996). Line-of-sight integration approach. [arXiv:astro-ph/9603033](https://arxiv.org/abs/astro-ph/9603033)
 - Ma & Bertschinger (1995). Cosmological perturbation theory. [arXiv:astro-ph/9506072](https://arxiv.org/abs/astro-ph/9506072)
 
+For more on the `clax.ept` module, see `docs/clax-pt.md`, `docs/CLASS-PT-summary.md`, and `docs/FFTLog_PT.md`.
+
 ## License
 
 MIT
diff --git a/clax/__init__.py b/clax/__init__.py
@@ -19,12 +19,13 @@
 
 from clax.constants import *  # noqa: F401,F403
 from clax.params import CosmoParams, PrecisionParams  # noqa: F401
-from clax.background import background_solve, BackgroundResult, H_of_z, angular_diameter_distance  # noqa: F401
+from clax.background import background_solve, BackgroundResult, H_of_z, angular_diameter_distance, sound_horizon_drag  # noqa: F401
 from clax.thermodynamics import thermodynamics_solve, ThermoResult  # noqa: F401
 from clax.perturbations import MatterPerturbationResult, PerturbationResult, TensorPerturbationResult, perturbations_solve, perturbations_solve_mpk, tensor_perturbations_solve  # noqa: F401
 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

diff --git a/clax/background.py b/clax/background.py
@@ -785,3 +785,33 @@ def luminosity_distance(bg: BackgroundResult, z: float) -> float:
     D_L(z) = chi(z) * (1 + z)  (for flat universe)
     """
     return comoving_distance(bg, z) * (1.0 + z)
+
+
+def sound_horizon_drag(params: CosmoParams) -> float:
+    """Sound horizon at baryon drag epoch r_s(z_drag) in Mpc.
+
+    Aubourg et al. 2015 (arXiv:1411.1074) Eq. (17), the Neff-aware variant.
+    Quoted accuracy 0.119% across 0 < sum m_nu < 0.6 eV, 3 < N_eff < 5.
+    Cross-checked against CLASS pth->rs_d at Planck 2018 fiducial: 0.002%.
+    Mirrors ps_1loop_jax sound_horizon_drag_aubourg2014_neff.
+
+    The hardcoded literal `99.0` in clax/ept.py corresponded to r_s_drag * h
+    at fiducial Planck 2018 (147.05 Mpc * 0.6736 = 99.05). This function
+    computes the cosmology-varying value used everywhere CLASS-PT uses
+    `pth->rs_d` (nonlinear_pt.c:5596).
+
+    Returns: r_s(z_drag) in Mpc (multiply by h for the Mpc*h convention
+    used by clax/ept.py IR resummation).
+    """
+    omega_cb = params.omega_b + params.omega_cdm
+    omega_nu = params.m_ncdm / 93.14
+    neff = params.N_ur + float(params.N_ncdm)
+    return (
+        56.067
+        * jnp.exp(-49.7 * (omega_nu + 0.002) ** 2)
+        / (
+            omega_cb ** 0.2436
+            * params.omega_b ** 0.128876
+            * (1.0 + (neff - 3.046) / 30.60)
+        )
+    )