feat: clax-pt — differentiable one-loop EFTofLSS galaxy power spectra in JAX - #9
Draft
MinhMPA wants to merge 19 commits into
Draft
feat: clax-pt — differentiable one-loop EFTofLSS galaxy power spectra in JAX#9MinhMPA wants to merge 19 commits into
MinhMPA wants to merge 19 commits into
Conversation
Collaborator
Author
|
@smsharma I don't feel so strongly about having this feature as a separate fork, so I'll let you decide whether this feature fits into the current scope of I think the biggest advantage/upside is if we ever decide to do a joint release (and write a code paper) for |
Collaborator
Author
|
To faciliate the review, I have added this notebook in particular: notebooks/clax-pt_full_validation.ipynb. One should compare the notebook outputs directly to their counterparts from this |
…nlinear=...) — redo of smsharma#11 with bug fixes Supersedes PR smsharma#11 (which was reverted via PR smsharma#13 due to post-merge test failures). This PR ships the same API consolidation with the three fixes that smsharma#11 was missing folded in from the start. Public API: compute_cl_pp(pt, params, bg, th, l_max, *, nonlinear="none") # nonlinear in {"none", "halofit"}; ValueError for anything else. # "ept" is reserved for a follow-up clax-pt PR. 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; ~20% over at l=2500) Renamed: - compute_cl_pp_source_limber -> compute_cl_pp (sole public entry) - compute_cl_pp_transfer -> _compute_cl_pp_full_bessel (private oracle) Halofit injection (new private _halofit_modulator): - vmap(compute_pk_nonlinear) over a 100-point z-grid (CLASS-aligned density). - Inline `sigma_R(R_min, ...) >= 1` check + jnp.where to force R=1 where Halofit isn't applicable — replicates CLASS fourier.c:1706-1716 in a vmap-safe way (the Python-level check inside compute_pk_nonlinear gets bypassed under vmap-over-z, which is why smsharma#11 over-corrected ~25% at high l). - k_max_extend = 0 default (no power-law extension; matches CLASS's no-extrapolation behavior). Pass a positive value to enable log-log extension for narrow k-grids. - 2D-interpolates R(k, z) onto (pt.k_grid, pt.tau_grid) via bg.loga_of_tau. - Multiplies sqrt(R) directly into S_transfer (CLASS source-multiplication). 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. - tests/test_clpp_halofit_ratio.py: rewritten for new API; <7% NL/lin ratio agreement at l<=500, <10% at l>=1000 vs CLASS Halofit reference. - tests/test_lensing.py: callers updated; uses local `pipeline` fixture matching upstream-main convention. - Deleted: tests/test_cl_pp_implementations.py, test_cl_pp_limber.py. Verified locally: 32/32 lensing+halofit tests pass, 36/36 EPT tests pass, 104/104 remaining tests pass.
2 tasks
…ute_cl_pp Documents both paths in the new unified API: - nonlinear="none" (linear source-Limber): <1% vs CLASS at all l<=2500 - nonlinear="halofit" (z-aware Halofit injection via source multiplication): <7% at l=500 and <10% at l>=1000 vs CLASS Halofit reference Existing tables (unlensed C_l, lensed C_l with CLASS-generated cl_pp, matter P(k), pipeline, performance, multi-cosmology) are unaffected by PR smsharma#14 — they don't depend on compute_cl_pp. Note on lensed C_l accuracy table: it tests lens_cls in isolation by feeding CLASS-generated C_l^pp; the reported <0.2% is about lens_cls's correlation-function method, not about our compute_cl_pp.
…<=2500
The previous table understated linear accuracy. test_clpp_halofit_ratio.py
::TestClppLinear asserts <1% at l in {100, 500, 1000, 2000, 2500} with
pt_k_max_cl=5.0. Updated the linear column accordingly and added the
prerequisite (pt_k_max_cl >= 5) to the section preamble.
…ratio test thresholds Sweep across modulator settings on default cosmology (l_max=2500): | l | n_z=100, no extension | n_z=100, extend to k=10 | | 100 | 0.01% | 0.01% | | 200 | 0.04% | 0.04% | | 500 | 0.42% | 0.21% | | 1000 | 1.47% | 0.63% | | 1500 | 2.30% | 0.79% | | 2000 | 3.74% | 1.40% | | 2500 | 3.97% | 0.96% | The k-extension restores accuracy at high l. CLASS itself uses a dedicated nonlinear k-grid that extends past the perturbation k-range; our log-log power-law extension provides equivalent coverage with σ(R) bisection well-resolved. Doubling n_z to 200 didn't help (slight degradation at l=1000, l=2500 from interpolation noise), so 100 is the right z-density. Changes: - _halofit_modulator default k_max_extend: 0 → 10 Mpc^-1 - test_ratio_at_low_l threshold: 7% → 1% (measured: 0.21% at l=500) - test_ratio_at_high_l threshold: 10% → 2% (measured: 1.40% worst case) - README accuracy table: replace bounds with measured residuals - Docstrings updated with rationale and measured numbers User asked 'who set the 10% threshold' — it was carried over from the deleted Poisson-Limber test, never measured against the actual source- multiplication recipe. Now it reflects what the code actually achieves.
…=2500); fair absolute-accuracy table k_max_extend sweep on default cosmology, n_z=100 fixed: | l | k=10 | k=20 | k=30 | | 500 | 0.21% | 0.09% | 0.09% | | 1000 | 0.63% | 0.04% | 0.17% | | 1500 | 0.79% | 0.11% | 0.11% | | 2000 | 1.40% | 0.05% | 0.28% | | 2500 | 0.96% | 0.76% | 0.76% | (metric: |R_clax / R_CLASS - 1|, R = NL/lin) k=20 is the sweet spot — far enough for sigma(R) bisection to fully converge, not so far that the local power-law slope drifts from the true asymptote. Going to k=30+ degrades because power-law extrapolation diverges from the true transfer-function shape. Timing benchmark (5-run mean): k=10 110ms, k=20 110ms, k=30 111ms. Accuracy boost is essentially free. Test thresholds: - low_l (l<=500): 1% -> 0.5% - high_l (l>=1000): 2% -> 1% README accuracy table replaced with apples-to-apples absolute residuals for both 'none' and 'halofit' (was mixing absolute-vs-ratio metrics). Both modes share the source-Limber kernel and track each other within ~0.1% on absolute scale; the ratio (third column) isolates the R(k,z) quality from the common linear-kernel systematic.
…ine duplication
clax/transfer.py
- compute_pk_from_perturbations / compute_linear_matter_pk_from_perturbations
gain a `validate=True` kwarg. Default keeps the existing solved-support
check; internal callers that need P(k) past pt.k_grid[-1] (the Halofit
modulator's k-grid extension) pass `validate=False`.
- Replace the Python-level `if z == 0.0:` fast path with always-interpolate
along τ. Clamping τ_z to pt.tau_grid[-1] makes the spline reproduce
pt.delta_m[:, -1] exactly at z=0, so accuracy is unchanged. The path is
now vmap-safe over z.
clax/nonlinear.py
- compute_pk_nonlinear's σ(R)≥1 fallback used a try/except on
ConcretizationTypeError, which silently bypassed the check under jit/vmap
and produced spurious R values when sigma<1. Replace with an inline
`jnp.where(sigma_R(R_min, lnk, pk_lin) >= 1.0, pk_nl, pk_lin)` that is
vmap-safe by construction. The narrow-grid fallback to pk_lin still
holds (test_halofit_sigma_guard).
clax/lensing.py::_halofit_modulator
- Now calls compute_linear_matter_pk_from_perturbations directly (no inline
τ-interp + 2π²/k³·P_R·δ_m² duplication) and trusts compute_pk_nonlinear's
own σ check.
The duplication was a workaround for the two transfer.py obstacles fixed
above. Behavior on the existing tests is unchanged (test_cl_pp,
test_clpp_halofit_ratio, test_lensing, test_halofit_sigma_guard,
test_halofit_z_cutoff, test_nonlinear all pass).
clax/perturbations.py
- _perturbation_solve_setup gains a `tau_max_factor` kwarg (default 0.999
for backward compatibility with the C_l / source-function pipeline).
- The matter-power paths route `tau_max_factor=1.0`:
* `_perturbations_solve_mpk_impl` (used by `compute_pk_table`)
* `_matter_delta_m_single_k_impl` (used by `compute_pk`)
- The 0.999 safety margin removed from these paths missed ~14 Mpc near
today, costing ~0.33% in P(k) via linear growth (D ∝ a, ΔD/D ≈ 0.16%
over 14 Mpc near z=0, P ∝ D² so doubles to ~0.33%).
- The C_l / `perturbations_solve` path is unchanged (still uses 0.999).
Measured P(k) accuracy at planck_fast preset, single-k `compute_pk` vs CLASS:
| k [Mpc^-1] | before | after |
|------------|---------|----------|
| 3e-3 | -0.51% | -0.17% |
| 1e-2 | -0.08% | +0.25% |
| 5e-2 | -0.29% | +0.04% |
| 1e-1 | -0.27% | +0.07% |
| 3e-1 | -0.36% | -0.03% |
| 1.0 | -0.53% | -0.20% |
Worst-case |err| drops from 0.53% to 0.25%; median |err| drops from
~0.30% to ~0.07%. Five of six k values improve.
Background, perturbation evolution, and δ_cdm at the synchronous-gauge
boundary all match CLASS to <0.05% — the residual was almost entirely
the hardcoded 14-Mpc growth deficit.
Tests verified locally:
- tests/test_cl_pp.py + tests/test_lensing.py: 24 passed in 207s
- tests/test_clpp_halofit_ratio.py + test_halofit_sigma_guard.py +
test_nonlinear.py: 20 passed in 89s
- C_l-pipeline-touching tests unchanged because `perturbations_solve`
keeps the default 0.999 factor.
The previous table (1.3–3.5% errors) was measured at an earlier version and has been stale since multiple subsequent improvements to the perturbation pipeline. Replace with planck_fast measurements at default Planck 2018 LCDM, post-`tau_max_factor=1.0` fix, vs CLASS v3.3.4. Median |error| now ~0.07%, worst-case 0.25% (k=0.01).
Cherry-picked from feat/clax-pt onto post-smsharma#11 + smsharma#12 base: clax module: - clax/ept.py — one-loop EFTofLSS power spectra (FFTLog, IR resummation, RSD multipoles) Tests: - tests/test_ept.py, test_ept_accuracy.py, test_ept_assembly.py, test_ept_gradients.py - tests/conftest.py — session-scoped pipeline_fast_cl[_k5] fixtures - tests/test_harmonic.py, test_high_l.py — switched to shared fixtures Docs: - docs/CLASS-PT-summary.md, FFTLog_PT.md, clax-pt.md - docs/accuracy_comparison.png, classpt_reference_table.npz Notebooks: - notebooks/clax-pt_full_validation.ipynb (§11 still uses legacy API; rewrite next commit) - notebooks/clax-pt_pkmm_validation.ipynb - notebooks/figures/fig{1..11}_*.png/.pdf Reference data: - reference_data/classpt_z0.38_fullrange.npz (CLASS-PT z=0.38 spectra) - reference_data/classpt_clpp_1loop.npz (CLASS-PT 1-loop C_l^phiphi) Scripts: - scripts/accuracy_classpt.py, generate_classpt_reference.py - scripts/diag_ploop_components.py, diag_ptree_vs_plin.py, ept_loop_sanity.py - scripts/validate_and_plot.py, generate_validation_figures.py - scripts/create_validation_notebook.py, diagnostic_notes.md Intentionally excluded from feat/clax-pt: - clax/lensing.py legacy plumbing (compute_cl_pp_fast/_vmap/_limber, compute_nl_correction_*) - clax/rosenbrock.py GRKT4 reintroduction (PR smsharma#10 removed it) - tests/test_cl_pp_implementations.py, test_cl_pp_limber.py, test_cl_pp_source_limber.py (deleted by PR smsharma#11) - tests/test_clpp_limber_accuracy.py, test_lensing_nonlinear.py (covered functions deleted) - tests/test_rosenbrock.py GRKT4 tests - paper artifacts (drafts/, docs/clax-pt_human_revision.md) — on paper_drafts branch Next commit will add EPT injection (_ept_modulator + nonlinear="ept") and rewrite §11 of the validation notebook for the new API.
Extends the compute_cl_pp(... nonlinear=...) enum to accept "ept" in
addition to "none" and "halofit". When selected, the lensing source is
multiplied by sqrt(R(k, z(tau))) where R = P_mm,1-loop / P_lin from
clax-pt's compute_ept_from_clax run once at z=0, then growth-rescaled to
other redshifts via the leading-order EFT scaling
R(k, z) - 1 = (R(k, 0) - 1) * [D(z)/D(0)]^2
(loop ~ D^4, ratio minus one ~ D^2). Sufficient at ~1% for k <= 0.3 h/Mpc;
subleading EFT counterterm time dependence neglected — flagged in the
docstring.
- Replace test_ept_not_yet_supported (expected ValueError) with test_ept_accepted (smoke): nonlinear='ept' should not raise. - Add TestEPT class mirroring TestHalofit: positivity, finiteness, NL boost at high l, no boost at low l.
…ed API
Cell 23 (markdown), 24 (perturbation solves + imports), 25 (C_l^phiphi
computation): replace the old triple-path (compute_cl_pp_source_limber +
compute_cl_pp_limber(nonlinear=True) + manual 1-loop ratio applied to
linear) with three calls to the new unified API:
cl_pp_lin = compute_cl_pp(... nonlinear="none")
cl_pp_hf = compute_cl_pp(... nonlinear="halofit")
cl_pp_1loop = compute_cl_pp(... nonlinear="ept")
Plots 8-9 (cell 26+) unchanged — they consume cl_pp_{lin,hf,1loop}.
…ensing path CHANGELOG: prepend two May 3 entries (consolidation + EPT injection) onto feat/clax-pt's clax-pt history block. Final ordering: EPT injection -> consolidation -> clax-pt accuracy log -> pre-clax-pt history. README: add 'One-loop EFTofLSS (clax.ept) accuracy vs CLASS-PT' section with the same accuracy table from feat/clax-pt; update lensing.py table row to mention clax-pt NL injection; add ept.py entry; add CLASS-PT and EFT-with-FFT references.
…lt _halofit_modulator to CLASS-aligned settings
- clax/ept.py compute_ept_from_clax: 'spline(lnk_out)' is not callable on
clax.interpolation.CubicSpline; use spline.evaluate(lnk_out). Also
primordial_scalar_pk takes (k, params), not (params, k) — args were
swapped.
- clax/lensing.py _halofit_modulator defaults updated:
n_z=100 (was 50) — matches CLASS's halofit z-table density.
k_max_extend=0 (was 10) — disables power-law extension by default.
With k_max_extend=0, Halofit operates on pt.k_grid as-is, matching
CLASS's behavior of computing Halofit on its internal P(k,z) table
without extrapolating beyond the perturbation k-range. The inline
sigma_R check (replicating CLASS fourier.c:1706-1716) still forces
R=1 when sigma(R_min)<1, so narrow k-grids gracefully degrade to
'no NL correction' rather than relying on extrapolated P_lin.
Power-law extension remains available by passing k_max_extend>0
explicitly (e.g., for tests with very narrow k-grids).
…of δ_m)
Drops the 'TODO: proper z interpolation' shortcut where compute_ept_from_clax
always used pt.delta_m[:, -1] (= z≈0 snapshot) regardless of the z argument,
and replaces it with vmap-safe τ-interpolation:
tau_z = bg.tau_of_loga.evaluate(log(1/(1+z)))
delta_m_at_z = vmap(lambda dm_k: CubicSpline(pt.tau_grid, dm_k).evaluate(tau_z))(pt.delta_m)
At z=0 this matches the previous behavior up to spline edge effects (last
τ-grid point is τ_0). For z > 0 it gives the correct δ_m(k, z), so loop
integrals reflect the requested redshift.
This unblocks per-z evaluation in clax/lensing.py::_ept_modulator (currently
uses leading-order D² growth scaling around z=0; on-the-fly per-z is now
viable as a future drop-in once the bridge is verified end-to-end).
Standard formula: P_m(k) = 2π² / k³ · P_R(k) · δ_m²(k) The previous compute_ept_from_clax had P_R · δ_m² / k³ — missing the 2π² ≈ 19.74 factor. Result: the input P_lin fed to compute_ept was ~20x too small, contaminating ALL downstream EPT spectra (P_tree, P_loop, multipoles, etc.) that depend on the linear input. This is the root cause of the EPT lensing path failure: at low k, my _ept_modulator was computing R0 = pk_mm_real/pk_lin where pk_mm_real (using EPT internal pipeline) is correctly normalized but pk_lin from compute_ept_from_clax was not, so the ratio was misweighted and the modulator multiplied lensing source by sqrt(R) ≈ 0.7 instead of ≈1. Now matches clax/transfer.py::compute_linear_matter_pk_from_perturbations exactly.
Collaborator
Author
|
Rebased
Once #14/#15/#16 merge upstream, this PR's diff against |
…4 fit
Adds clax.background.sound_horizon_drag(params), implementing the Aubourg
et al. 2015 (arXiv:1411.1074) Eq. (17) Neff-aware fit. Mirrors the
ps_1loop_jax convention and matches it at machine precision.
At Planck 2018 fiducial, agrees with CLASS pth->rs_d to 0.002%.
compute_ept_from_clax now passes sound_horizon_drag(params) * params.h
(the Mpc·h product convention used internally by ept.py) into compute_ept,
so cosmology variations propagate to the IR-resummation Σ²_BAO j₂-filter
instead of using the fiducial-Planck literal.
Direct compute_ept callers keep rs_h=99.0 as a Planck-fiducial fallback;
docstrings updated to clarify the misleading "rs_h" name (the dimensions
are r_s × h in Mpc, NOT r_s / h in Mpc/h).
Validation: scripts/accuracy_classpt.py — all 9 spectra still pass at
fiducial. Cross-validated against ps_1loop_jax across +/-20% omega_b,
+/-20% omega_cdm, mnu in {0, 0.06, 0.15, 0.30} eV, N_ur in {1.5, 2.0328,
3.5} — relative diff = 0.0e+00 in all cases.
Closes the rs_h known caveat in CHANGELOG (Apr-04 PT Known Caveats #2).
This was referenced May 3, 2026
np.trapz was removed in NumPy 2.0 (deprecated since 1.23). The clax env runs Python 3.14 / NumPy 2.x, causing AttributeError in _ir_resummation_numpy and _ir_resummation_gaussian. np.trapezoid is a drop-in replacement with identical semantics. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
float() on JAX-traced values (params.h, sound_horizon_drag, Omega_m_of_z) caused ConcretizationTypeError when computing jax.grad through compute_ept_from_clax. Fix: - Keep params.h as a JAX scalar (h); use stop_gradient for h_conc (numpy-bound ops: unit conversions, rs_h, k_mpc) - Wrap sound_horizon_drag and Omega_m_of_z in stop_gradient before float() - Pre-compute IR resummation arrays with stop_gradient(pk_h) and pass as _ir_precomputed, activating the gradient-safe path in compute_ept where pk_w = pk_lin_h - pk_nw remains JAX-traced, giving non-zero d(pk)/d(params) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Adds clax-pt: a JAX-native, fully differentiable one-loop EFTofLSS pipeline built on top of clax. End-to-end AD from cosmological parameters to P_gg(k, ℓ=0,2,4) via the FFTLog decomposition of CLASS-PT (Chudaykin, Ivanov, Philcox & Simonović 2020). Also extends
compute_cl_pp(... nonlinear=...)to accept"ept"for one-loop nonlinear corrections to CMB lensing C_ℓ^φφ via source multiplication.Dependencies
This PR sits on top of three smaller upstream-bound PRs:
refactor/cl_pp_consolidation_v2compute_cl_pp_*impls into singlecompute_cl_pp(... nonlinear="none"|"halofit"). Redo of the reverted #11 with the three post-merge test failures fixed in a single PR.refactor/transfer-vmap-safevalidate=Falsekwarg tocompute_pk_from_perturbations; replacetry/except ConcretizationTypeErrorσ check with vmap-safejnp.where; drop the inline τ-interp + log-log extension duplication in_halofit_modulator.fix/pk_tau_endtau_max = 0.999 × bg.conformal_agesafety margin in matter-power paths. The 14-Mpc gap near today contributed ~0.33% bias incompute_pk/compute_pk_tablevia linear growth.feat/clax-ptis rebased on top ofupstream/main + #14 + #15 + #16so that once those land,feat/clax-pt's diff against the new upstream/main is just the 9 clax-pt-only commits below.What's new (clax-pt-only commits)
1.
clax/ept.py— one-loop EFTofLSS module (~2,040 LOC)Mirrors
CLASS-PT/source/nonlinear_pt.c::nonlinear_pt_loop(). Layout:_fftlog_decompose,_compute_p22,_compute_p13_load_matrices(M22, M22basic, M13, IFG2; cosmology-independent, loaded once)_ir_resummation_numpy(DST-II nw/w split),_ir_resummation_gaussian(precomputed-Σ² differentiable path)_compute_bias_spectra(52-componentEPTComponents)compute_ept,compute_ept_from_clax(z-aware bridge from claxPerturbationResult)pk_mm_real,pk_gm_real,pk_gg_real,pk_{mm,gg}_l{0,2,4}Public dataclasses:
EPTPrecisionParams,EPTComponents. JAX-differentiable w.r.t. bias/EFT parameters andpk_lin; the DST nw/w split runs in NumPy at setup, and a precomputed-Σ² path keepsjax.gradavailable through the full pipeline (used intests/test_ept_gradients.py).Three EPT bug fixes (caught during the rebase test runs):
compute_ept_from_clax:'CubicSpline' object is not callable→ use.evaluate(). Alsoprimordial_scalar_pkarg order swapped ((params, k)→(k, params)).compute_ept_from_clax: missing2π²factor inP(k) = 2π²/k³ · P_R · δ_m². Was producing P_lin ~20× too small, contaminating downstream EPT spectra.compute_ept_from_clax: was always readingpt.delta_m[:, -1](z≈0) regardless of thezargument. Now properly τ-interpolates δ_m to τ(z) soz > 0queries return the correct redshift.Cosmology-consistent BAO sound horizon for IR resummation (
c09b6aa):rs_h = 99.0(a fiducial-Planck approximation ofr_s_drag × hin Mpc) withclax.background.sound_horizon_drag(params) * params.h, plumbed intocompute_ept_from_clax. The new helper implements Aubourg et al. 2015 (arXiv:1411.1074) Eq. (17), Neff-aware — quoted accuracy 0.119% across0 < Σm_ν < 0.6 eVand3 < N_eff < 5.ps_1loop_jax.background.sound_horizon_drag_aubourg2014_neffat machine precision over a 9-point parameter scan (omega_b ±20%, omega_cdm ±20%, m_ncdm ∈ {0, 0.06, 0.15, 0.30} eV, N_ur ∈ {1.5, 2.0328, 3.5}); matches CLASSpth->rs_dto 0.002% at fiducial Planck.nonlinear_pt.c:5596: rbao = pth->rs_d), not from a hardcoded constant — this fix removes a CLASS-PT mismatch outside fiducial Planck instead of introducing one.compute_ept(... rs_h=...)keeprs_h=99.0as a Planck-fiducial fallback; docstrings updated to clarify the historically misleading variable name (rs_his dimensionallyr_s × hin Mpc, NOTr_s / hin Mpc/h).rs_hknown caveat (PT Known Caveats clax-PT #2 in CHANGELOG).2.
clax/lensing.py— EPT injection for C_ℓ^φφExtends the unified
compute_cl_pp(... nonlinear="none"|"halofit"|"ept")API (introduced by #14) to accept"ept":_ept_modulator(pt, params, bg, th)runscompute_ept_from_clax(z=0.0)once, formsR0(k) = pk_mm_real / pk_lin, and growth-rescales to other redshifts via the leading-order EFT scalingR(k, z) − 1 = (R(k, 0) − 1) · [D(z)/D(0)]²(loop ~D⁴, ratio minus one ~D²). Sufficient at ~1% fork ≤ 0.3 h/Mpc; subleading EFT counterterm time dependence neglected — flagged in the docstring.3.
tests/tests/test_ept.py,test_ept_accuracy.py,test_ept_assembly.py,test_ept_gradients.py— the EPT suite (already validates 9 spectra against CLASS-PT NPZ).tests/test_cl_pp.py::TestEPT—nonlinear="ept"end-to-end tests (positivity, finiteness, NL boost at high ℓ, no boost at low ℓ).4.
notebooks/clax-pt_full_validation.ipynbReproduces the CLASS-PT example notebook plot-for-plot at the same cosmology (Planck 2018, one massive ν), nuisance choices, and z=0.38:
§11 uses the unified
compute_cl_pp(... nonlinear=...)API for all three lensing paths (linear / Halofit / 1-loop PT).tests/test_ept_*.pyand the validation notebook reference data live inreference_data/classpt_z0.38_fullrange.npzandclasspt_clpp_1loop.npz.5.
docs/docs/clax-pt.md— module overview and design.docs/FFTLog_PT.md— derivation: FFTLog decomposition, master integral I(α,β), P13/P22 kernels, RSD μ-channel table, IR resummation formulas, GL projection, units.docs/CLASS-PT-summary.md— 52-componentEPTComponentsindex and all 9 assembly formulas including EFT counterterms.6.
scripts/generate_classpt_reference.py— generatereference_data/classpt_z0.38_fullrange.npzfrom CLASS-PT (classy).accuracy_classpt.py— non-pytest accuracy script (exits 0 if all 9 spectra pass <1% at k<0.3 h/Mpc).validate_and_plot.py,generate_validation_figures.py— produce the figures innotebooks/figures/.diag_ploop_components.py,diag_ptree_vs_plin.py,ept_loop_sanity.py— diagnostics.Accuracy
EPT spectra at Planck 2018 ΛCDM, z=0.38, k<0.30 h/Mpc, vs CLASS-PT
CMB lensing C_ℓ^φφ (set on the new compute_cl_pp API)
compute_cl_pp(nonlinear="halofit")matches CLASS Halofit reference within <1% at all ℓ ≤ 2500 withpt.k_grid[-1] >= 5 Mpc⁻¹(the source-Limber kernel + 100-point z-grid + log-log k-extension to 20 Mpc⁻¹; absolute residual table in README).compute_cl_pp(nonlinear="ept")produces sensible NL boost at high ℓ (>1.001× linear at ℓ=500–2000) and ~1 at low ℓ (tests/test_cl_pp.py::TestEPT).Linear matter P(k)
Single-mode
compute_pkat planck_fast vs CLASS, with PR #16'stau_max_factor=1.0fix: median |err| ~0.07%, worst 0.25% (k=1e-2). See README accuracy table.Test plan
pytest tests/test_ept_*.py -v— 36 passedpytest tests/test_cl_pp.py tests/test_lensing.py tests/test_clpp_halofit_ratio.py tests/test_halofit_sigma_guard.py tests/test_nonlinear.py -v— 48 passed (post-rebase verification)pytest tests/test_ept_assembly.py tests/test_ept_accuracy.py -v— 12 passed (post rs_h plumbing)scripts/accuracy_classpt.py— all 9 spectra pass at k<0.3 h/Mpc (post rs_h plumbing)pytest tests/ -v --fast -q— full quick suite (still need to run end-to-end)notebooks/clax-pt_full_validation.ipynbtop-to-bottom; confirm figures match CLASS-PTnonlinear_pt.ipynbMigration / dependencies
upstream/mainis the cumulative diff (51 files / ~+9.2k / -1.1k); after they merge, the diff cleans down to the 9 clax-pt-only commits listed at the top.clax.background.sound_horizon_drag(params)is a small, additive public helper; technically usable independently of clax-pt, but introduced here to support the ept.py call site (no other current upstream call path needs it).Status
Draft. Will move to ready-for-review after the three dependency PRs merge and the full-suite test plan is checked off.