feat(thermo/shooting): forward-mode AD via custom_jvp for _find_z_reio and shoot_fn - #21
Open
MinhMPA wants to merge 15 commits into
Open
feat(thermo/shooting): forward-mode AD via custom_jvp for _find_z_reio and shoot_fn#21MinhMPA wants to merge 15 commits into
MinhMPA wants to merge 15 commits into
Conversation
…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.
…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.
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).
The single-mode `_matter_delta_m_single_k_impl` (used by `compute_pk`)
wraps `bg.conformal_age` in `stop_gradient` before passing it as `t1` to
the diffrax ODE solve. That stop_gradient is structurally required:
RecursiveCheckpointAdjoint + PIDController cannot reverse-mode AD through
a traced `t1` (the controller's accepted-step factor is marked
`eqxi.nondifferentiable`, so any tangent flowing through `t1` raises
`RuntimeError: Unexpected tangent`). Removing it crashes the adjoint.
But stop_gradient sacrifices a chain-rule term. The full derivative of
`delta_m(tau_end(theta), y(tau_end(theta); theta))` has two parts:
(1) d_y delta_m * dy(tau_end; theta)/dtheta (frozen-t1 AD captures)
(2) ddelta_m/dtau * dtau_end/dtheta (stop_gradient drops)
For h/omega_b/omega_cdm — all of which change `conformal_age` through H(z)
— part (2) is ~70% of the gradient, making `jax.grad(compute_pk)` 65-84%
off vs centered FD on this branch. ln10A_s/n_s pass at 0% off because
they don't enter the ODE and don't change tau_end.
Recovery: first-order Taylor expansion around the frozen endpoint,
delta_m(tau_traced) ~ delta_m(tau_frozen)
+ (ddelta_m/dtau)|_{tau_end} * (tau_traced - tau_frozen)
Two properties make this exact at fiducial:
- `(tau_traced - tau_frozen) = 0` identically -> primal value unchanged.
- `jax.jvp` on `_extract_delta_m` with tangent `(dy/dtau, 1)` captures
both the implicit (state) and explicit (background-density-at-tau)
contributions to `ddelta_m/dtau`, so AD picks up the correct linear
coefficient.
Verified at PK_CONTRACT_PREC (rtol=1e-6) and at rtol=1e-7:
rtol=1e-3: h 6.41% / omega_b 8.91% / omega_cdm 4.83%
rtol=1e-5: h 1.08% / omega_b 3.18% / omega_cdm 0.85%
rtol=1e-7: h 0.32% / omega_b 1.23% / omega_cdm 0.61%
Clean precision-floor convergence (~3-4x per rtol decade), consistent with
the FD oracle's own step-noise floor at each rtol. ln10A_s/n_s remain
exact at 0% off across all rtol values.
Adds `TestPkScalarDensityGradients` regression test on the contract path.
The 5% threshold fails loudly on the original stop_gradient-only code
(>50% off) while accommodating the FD precision floor at rtol=1e-6
(empirically ~3% on omega_b at k=3e-4).
Fails until _find_z_reio is converted from custom_vjp to custom_jvp. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
…ga gradient
Forward-mode AD:
- Reformulate C_He as (inv_A + Lambda)/(inv_A + Lambda + Rup) to avoid
inf-inf NaN tangents in JVP when He_Boltz is huge at low T.
- Replace jnp.maximum with jnp.where in _tau_reio_for_zreio so the
boundary tangent is zero (not NaN) during jax.jvp.
- Convert _find_z_reio from custom_vjp to custom_jvp using the implicit
function theorem: z_reio_dot = -F_dot / dF_dz. stop_gradient on
z_reio_sg and dF_dz prevents JAX from adding spurious second-order terms
when it derives VJP by transposition.
GREEN: test_find_z_reio_forward_mode_matches_fd
Reverse-mode AD:
- Fix dkappa_dot_dloga_of_loga gradient: replace _first_derivative_table
with sg(kappa_dot_grid)*(n_H_0/sg(n_H_0)) before spline-differentiating.
kappa_dot ∝ n_H_0 exactly, so this preserves values while routing the
gradient only through n_H_0 ∝ omega_b. The accumulated Friedmann-scan
gradient d(a^{-2})/d(params) (which blows up to ~10^17 in 1/h) is stopped.
At loga=-8 (Saha regime): d(dkd/dloga)/d(omega_b) ≈ -2*kd/omega_b ≈ -334,
matching FD to <1%.
GREEN: test_opacity_logderivative_gradient_matches_fd_for_omega_b
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
jax.jvp through a custom_vjp function raises an error; this test documents the requirement that shoot_fn must define a custom_jvp rule instead. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Enables forward-mode AD (jax.jvp) through the theta_s shooting method. JVP rule applies the implicit function theorem: dh/d(theta_s) = 1/dtheta_dh. custom_jvp also provides VJP via JAX transposition, preserving all existing reverse-mode gradient tests. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
diag_jvp_nan_source.py: pinpoints which ThermoResult tangent leaf carries NaN when jax.jvp is called through thermodynamics_solve. Used to isolate the He C_He inf-inf cancellation at z~15 (Task B3). STATUS.md: task-claiming log per plan etiquette. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
This was referenced May 12, 2026
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
Enables `jax.jvp` / `jax.jacfwd` through the thermodynamics shooting loop by converting `_find_z_reio` and `shoot_fn` from `custom_vjp` to `custom_jvp`.
Why: The implicit-function-theorem wrappers were implemented with `custom_vjp`, which is only differentiable in reverse mode. Forward-mode passes (`jax.jvp`, `jax.jacfwd`, FWDPASS inside `jax.hessian`) raised `NotImplementedError: Forward differentiation rule not defined`.
Fix: Replace `@custom_vjp` + `fwd/bwd` with `@custom_jvp` + `jvp` implementing the implicit-function-theorem tangent:
```
∂output/∂θ · v = −(∂F/∂output)⁻¹ · (∂F/∂θ · v)
```
where `F(output; θ) = 0` is the shooting residual. Analytically equivalent to the old VJP, works for both forward and reverse mode.
Changes
Merge order (4-PR AD-correctness series)
Test plan