fix(lensing): post-#11 hotfix — vmap-safe Halofit modulator + test corrections - #12
Closed
MinhMPA wants to merge 2 commits into
Closed
fix(lensing): post-#11 hotfix — vmap-safe Halofit modulator + test corrections#12MinhMPA wants to merge 2 commits into
MinhMPA wants to merge 2 commits into
Conversation
…to low l Three fixes to issues in the C_l^phiphi consolidation merged in smsharma#11: 1. _halofit_modulator was passing 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 and (b) has a Python-level `if z == 0.0` branch that fails under vmap-over-z. Both surface as test failures. Fix: inline the τ-axis interpolation of pt.delta_m, build P_lin on pt.k_grid only, then power-law-extrapolate in log-log to k_max_extend. Pure JAX, vmap-safe, no validation conflict. 2. tests/test_lensing.py::test_cl_pp_positive used the `pipeline_fast_cl` fixture, which is defined in PR smsharma#9's conftest but not on upstream/main. Replaced with a local `pipeline` fixture matching upstream-main convention. 3. tests/test_cl_pp.py::TestCrossImplAgreement compared source-Limber against _compute_cl_pp_full_bessel at l up to 500. The oracle uses clax.bessel.spherical_jl (upward recurrence) which is unstable in the classically-forbidden region at high l; ratios at l=200 hit ~0.5x. Restricted the cross-impl test to l <= 100 where the oracle is reliable, and documented why. Higher-l accuracy is verified against an external CLASS reference in TestLinearAccuracy. These three failures emerged from a post-merge sanity check and were not caught before smsharma#11 because only TestContract had been run.
compute_pk_nonlinear has a Python-level sigma_convergence_check that bypasses Halofit when sigma(R_min)<1 (matching CLASS fourier.c:1706-1716, where nl_corr_density=1.0). Under vmap-over-z, that check is bypassed because the try/except catches ConcretizationTypeError, and Halofit then runs on linearly-extrapolated pk_lin at high z where it produces spurious values (e.g., R~1.7 at z=100 in the test fixture). These spurious R values contaminate the lensing source-multiplication and make the integrated NL boost ~25% larger than CLASS. Diagnostic showed R(k,z=0.5..1) matches CLASS to <0.01%, but R(k, z>=2) diverges where Halofit shouldn't apply. Fix: call sigma_R inline (vmap-safe pure JAX) and use jnp.where to fall back to pk_lin when sigma<1, replicating CLASS's behavior. Verified to fix the over-correction surfaced by tests/test_clpp_halofit_ratio.py.
Collaborator
Author
Owner
|
@MinhMPA no worries, reverted for now. |
6 tasks
Collaborator
Author
|
Thanks! That works. Let me also close this and consolidate PR #11 with this into a new PR then. |
2 tasks
MinhMPA
added a commit
to MinhMPA/clax-pt
that referenced
this pull request
May 3, 2026
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.
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
2 participants
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
Four fixes to issues introduced by the C_l^phiphi consolidation merged in #11. The contract tests passed pre-merge, but a fuller post-merge sanity check exposed failures that this PR addresses.
Changes
1.
_halofit_modulatoris now vmap-safe (P_lin path) —clax/lensing.py_halofit_modulatororiginally passed an extended k-grid tocompute_linear_matter_pk_from_perturbations, which has two issues for our use case:k_evaloutsidept.k_gridvia_validate_k_eval_support(so the extended grid triggersValueError: k_eval must lie within...).if z == 0.0branch that fails underjax.vmap-over-z withTracerBoolConversionError.Fix: inline the τ-axis interpolation of
pt.delta_m, build P_lin onpt.k_gridonly, then power-law-extrapolate in log-log space tok_max_extend. Pure JAX, vmap-safe, no validation conflict.2. Inline
sigma(R) >= 1check in_halofit_modulator—clax/lensing.pycompute_pk_nonlinearhas a Python-level_sigma_convergence_checkthat bypasses Halofit whensigma(R_min) < 1(matching CLASSfourier.c:1706-1716, wherenl_corr_density = 1.0). Under vmap-over-z that check is bypassed (the try/except catchesConcretizationTypeError), so Halofit then runs at high z wheresigma < 1and produces spurious R values (e.g., R ≈ 1.7 at z=100 in the test fixture). These contaminate the lensing source-multiplication and inflate the integrated NL boost ~25% above CLASS at high l.Fix: call
sigma_Rinline (vmap-safe pure JAX) andjnp.whereto fall back topk_linwhensigma < 1, replicating CLASS's behavior. Diagnostic confirmedR(k, z=0..1)now matches CLASS to <0.01% and high-z R is forced to 1 as it should be.3.
tests/test_lensing.py::test_cl_pp_positive— local fixtureThe post-merge file referenced a
pipeline_fast_clfixture that doesn't exist onupstream/main(it lives in PR #9's conftest). Replaced with a localpipelinefixture matching the upstream-main convention.4.
tests/test_cl_pp.py::TestCrossImplAgreement— restrict to low lThe cross-impl test compared source-Limber against
_compute_cl_pp_full_bessel(the slow Bessel oracle) at l up to 500. The oracle usesclax.bessel.spherical_jl(upward recurrence), which is unstable in the classically-forbidden region (x < 0.7l) at high l; the ratio at l=200 was ~0.5x.The Limber approximation, in contrast, becomes more accurate as l increases — CLASS uses
l_switch_limberbetween 9 and 40 for the lensing path because Limber is the correct answer there. Restricted the cross-impl test to l ≤ 100 where the oracle is reliable, and documented why. Higher-l accuracy is already verified against an external CLASS reference inTestLinearAccuracy.Why this wasn't caught before #11
Only the cheap
TestContract(4 shape/signature tests) was run pre-merge — the other tests need the full perturbation solve (~6 min/fixture) and I cut myself off after the contract tests passed. The post-merge sanity check ran the full lensing test suite and surfaced these issues.Test plan
Expected on hotfix branch: all passing.