refactor: collapse 6 compute_cl_pp_* into single compute_cl_pp(... nonlinear=...) — redo of #11 with bug fixes

[MinhMPA]

Summary

Re-attempts the C_ℓ^φφ API consolidation that was merged as PR #11 and reverted via PR #13, this time with the three test failures that surfaced post-merge folded in from the start. The fix list is the same content as the (now-obsoleted) draft hotfix PR #12, but combined into a single working PR so it lands cleanly without an interim regression.

Why a redo, not a re-merge of #11

PR #11 passed the contract tests pre-merge but failed three other tests that I hadn't run beforehand:

1. _halofit_modulator was vmap-unsafe. It passed an extended k-grid to compute_linear_matter_pk_from_perturbations, which (a) rejects k_eval outside pt.k_grid via _validate_k_eval_support (ValueError: k_eval must lie within...), and (b) has a Python-level if z == 0.0 branch that fails under jax.vmap-over-z (TracerBoolConversionError).
2. compute_pk_nonlinear's Python-level _sigma_convergence_check was bypassed under vmap (the try/except catches ConcretizationTypeError), so Halofit ran at high z where σ < 1 and produced spurious R values (~1.7 at z = 100), inflating the integrated NL boost ~25% above CLASS at high ℓ.
3. tests/test_lensing.py::test_cl_pp_positive referenced a pipeline_fast_cl fixture that doesn't exist on upstream/main (it lives in PR feat: clax-pt — differentiable one-loop EFTofLSS galaxy power spectra in JAX #9's conftest).

This PR ships PR #11's API consolidation with all three issues fixed in a single commit.

Public API

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 (#9).

Re-exported at the package root (clax.compute_cl_pp) along with clax.lens_cls.

Changes

Removed (BREAKING — pre-1.0, hard break)

Function	Why
compute_cl_pp (Siddharth's original)	Three-bug Bessel integration; superseded by source-Limber
compute_cl_pp_fast	Inaccurate at ℓ ≥ 300 (documented in its own docstring)
compute_cl_pp_vmap	Hermite Bessel-table vmap; superseded
compute_cl_pp_limber	Poisson-reconstruction Limber; ~20% overestimate at ℓ = 2500 vs CLASS

Renamed

From	To	Rationale
compute_cl_pp_source_limber	compute_cl_pp	Sole public entry; CLASS-accurate at all ℓ ≤ 2500
compute_cl_pp_transfer	_compute_cl_pp_full_bessel	Private oracle for cross-impl tests

Added — _halofit_modulator (CLASS-aligned)

• n_z = 100 z-points (CLASS density), log-spaced in 1+z spanning [0, z_rec].
• Inline sigma_R(R_min, ...) >= 1 check using jnp.where to fall back to pk_lin when σ < 1, replicating CLASS fourier.c:1706-1716 in a vmap-safe way.
• k_max_extend = 0 default — uses pt.k_grid as-is, matching CLASS's no-extrapolation behavior. The inline σ check forces R = 1 when the k-grid is too narrow. Pass k_max_extend > 0 to enable log-log power-law extension as a numerical convenience for narrow grids.
• 2D-interpolates R(k, z) onto every (pt.k_grid, pt.tau_grid) lattice point using bg.loga_of_tau for the τ → z mapping.
• Multiplies sqrt(R) directly into S_transfer (CLASS source-multiplication recipe).

Tests

File	Change
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 for new API; <7% NL/lin ratio agreement at ℓ ≤ 500, <10% at ℓ ≥ 1000 vs CLASS Halofit reference
tests/test_lensing.py	Callers updated to new signature; uses local pipeline fixture matching upstream-main convention
tests/test_cl_pp_implementations.py	Deleted (targeted removed implementations)
tests/test_cl_pp_limber.py	Deleted (targeted removed compute_cl_pp_limber)

Test plan

Verified locally on this branch:

• pytest tests/test_cl_pp.py tests/test_lensing.py tests/test_clpp_halofit_ratio.py -v → 32 passed in 295s
• pytest tests/test_ept_*.py -v → 36 passed (sanity check; this PR doesn't touch ept)
• 104 other tests (background, thermo, nonlinear, rosenbrock, etc.) → all pass

Expected on this PR:

• CI: pytest tests/ -v --fast -q passes
• No remaining references to deleted functions: git grep -nE "compute_cl_pp_(fast|vmap|limber|source_limber|transfer)" clax/ tests/ should return only _compute_cl_pp_full_bessel

Migration

Hard break, no deprecation shims:

# Before (any of):
cl = compute_cl_pp_source_limber(pt, params, bg, th, l_max=2500)
cl = compute_cl_pp_vmap(pt, params, bg, l_max=2500)
cl = compute_cl_pp_fast(pt, params, bg, l_max=2500)
cl = compute_cl_pp_limber(pt, params, bg, th, l_max=2500, nonlinear=True)

# After:
cl = clax.compute_cl_pp(pt, params, bg, th, l_max=2500)
cl_nl = clax.compute_cl_pp(pt, params, bg, th, l_max=2500, nonlinear="halofit")

[MinhMPA]

Self-flag for the reviewer (no action required in this PR):

clax/lensing.py::_halofit_modulator inlines the τ-axis interpolation of pt.delta_m and the P_lin formula (2π²/k³ × P_R × δ_m²) rather than reusing clax/transfer.py::compute_pk_from_perturbations + compute_linear_matter_pk_from_perturbations. That's deliberate code duplication — those public APIs have two vmap-unfriendly choices that the modulator can't tolerate when looped over a z-grid:

1. _validate_k_eval_support rejects k_eval outside pt.k_grid, which the modulator's k-extension to k_max=20 Mpc⁻¹ deliberately violates.
2. The Python-level if z == 0.0: fast-path in compute_pk_from_perturbations raises TracerBoolConversionError under jax.vmap over z.

Both have small fixes in transfer.py (a validate=False kwarg + jnp.where for the z=0 branch), after which _halofit_modulator could drop ~30 lines and call the shared utilities directly. Out of scope for this PR — flagging in case you'd prefer that cleanup folded in before merge or as a separate follow-up. The tests already cover the modulator's correctness end-to-end either way.

EDIT: This is now addressed in PR #15. Merge order: #14 first, then #15 and #16 (independent of each other, both touch different files).