Skip to content

refactor: vmap-safe transfer + nonlinear; drop _halofit_modulator inline duplication - #15

Open
MinhMPA wants to merge 6 commits into
smsharma:mainfrom
MinhMPA:refactor/transfer-vmap-safe
Open

refactor: vmap-safe transfer + nonlinear; drop _halofit_modulator inline duplication#15
MinhMPA wants to merge 6 commits into
smsharma:mainfrom
MinhMPA:refactor/transfer-vmap-safe

Conversation

@MinhMPA

@MinhMPA MinhMPA commented May 3, 2026

Copy link
Copy Markdown
Collaborator

Summary

Small follow-up on top of #14 that removes the two vmap obstacles in clax/transfer.py + the bypassable σ check in clax/nonlinear.py, and uses the cleaned-up helpers to delete an inline duplication in _halofit_modulator.

Net diff: 3 files changed, 44 insertions(+), 63 deletions(-).

Stacked on top of #14 (refactor/cl_pp_consolidation_v2). Once #14 merges, the diff against main is just the changes here.

Why

_halofit_modulator (added in #11/#14) couldn't reuse compute_linear_matter_pk_from_perturbations because that function had two vmap-unfriendly choices, so the modulator inlined the τ-axis δ_m interpolation, the 2π²/k³ · P_R · δ_m² formula, and a manual log-log power-law extension. Similarly, compute_pk_nonlinear's σ check used a try/except ConcretizationTypeError that silently bypassed under jit/vmap, so the modulator replicated the σ check inline with jnp.where.

This PR fixes the helpers so the modulator stops needing workarounds.

Changes

clax/transfer.py

  • compute_pk_from_perturbations / compute_linear_matter_pk_from_perturbations gain validate: bool = True. Default keeps the existing solved-support check; internal callers extending the k-grid past pt.k_grid[-1] for σ(R) convergence pass validate=False. The validation is still on for all public table-API callers.
  • The Python-level if z == 0.0: fast path is replaced with always-interpolate along τ, with tau_z clamped to pt.tau_grid[-1]. At z=0 the spline reproduces pt.delta_m[:, -1] exactly (boundary), so accuracy is unchanged. The path is now vmap-safe over z.

clax/nonlinear.py

  • compute_pk_nonlinear's σ(R) ≥ 1 fallback is rewritten as jnp.where(sigma_R(R_min, lnk, pk_lin) >= 1.0, pk_nl, pk_lin). Vmap-safe by construction. The narrow-grid fallback to pk_lin is preserved (test_halofit_sigma_guard).
  • _sigma_convergence_check is kept for direct unit tests, but is no longer in the path of compute_pk_nonlinear.

clax/lensing.py::_halofit_modulator

  • Calls compute_linear_matter_pk_from_perturbations directly (no inline τ-interp + 2π²/k³·P_R·δ_m² duplication).
  • Drops the inline σ check / jnp.where; trusts the helper's now-vmap-safe σ guard.
  • log-log power-law k-extension stays in the modulator (δ_m extrapolation past pt.k_grid[-1] is non-physical, so we extend P_lin in log-log space rather than passing k_ext through the helper).

Test plan

Verified locally on refactor/transfer-vmap-safe:

  • pytest tests/test_cl_pp.py -v → 13 passed in 95s
  • pytest tests/test_lensing.py tests/test_clpp_halofit_ratio.py -v → 15 passed in 204s
  • pytest tests/test_halofit_sigma_guard.py tests/test_halofit_z_cutoff.py tests/test_nonlinear.py -v → 20 passed in 2s
  • pytest tests/test_pk_accuracy.py -v --fast → forthcoming, will paste results here

No new tests added — the change is internal refactoring with the same external behavior covered by the existing suite.

Migration

None — validate is a new keyword-only kwarg with the same default behavior. No public API breakage.

MinhMPA added 6 commits May 3, 2026 12:54
@MinhMPA
refactor: collapse 6 compute_cl_pp_* into single compute_cl_pp(... no…
aa392d2 
…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.
@MinhMPA
docs(README): add C_l^phiphi accuracy table for the consolidated comp…
7863871 
…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.
@MinhMPA
docs(README): tighten linear C_l^phiphi accuracy claim — <1% at all l…
bfbb39a 
…<=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.
@MinhMPA
perf+test(lensing): default k_max_extend=10 (was 0); tighten Halofit …
1e28e12 
…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.
@MinhMPA
perf(lensing): k_max_extend default 10->20 (sub-percent across all l<…
17c7e5a 
…=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.
@MinhMPA
refactor: vmap-safe transfer + nonlinear; drop _halofit_modulator inl…
1d9a70d 
…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).
@MinhMPA

MinhMPA commented May 3, 2026

Copy link
Copy Markdown
Collaborator Author

Depends on #14 — diff shows 12 files / 604+/1055- because it stacks on refactor/cl_pp_consolidation_v2. Once #14 merges, this rebases cleanly to just clax/transfer.py, clax/nonlinear.py, and clax/lensing.py (3 files / +44/-63).

This was referenced May 3, 2026
Draft
Open
Draft
Open
Open
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

No reviews

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

1 participant

@MinhMPA