fix(compute_pk): integrate to full bg.conformal_age + recover dtau_end gradient via Taylor correction

[MinhMPA]

Summary

Two coupled fixes to compute_pk / compute_pk_table, both rooted in the integration endpoint tau_end = bg.conformal_age:

1. Primal accuracy. Integrate the matter-power paths to the full bg.conformal_age (1.0×) instead of the historical 0.999× safety margin. Closes a ~0.33% bias in P(k) from the missing 14 Mpc of linear growth near today.
2. Reverse-mode AD correctness. Recover the chain-rule term sacrificed by the necessary stop_gradient(bg.conformal_age) via a first-order Taylor expansion around the frozen endpoint. Closes a 65–84% AD-vs-FD gradient error on h, omega_b, omega_cdm (parameters that flow through conformal_age via H(z)).

The two changes are coupled: routing the matter-power paths to tau_max_factor=1.0 makes tau_end parameter-dependent (it was a fixed 0.999 × age before), which exposes the latent gradient bug. Both must land together.

Independent of #15 — they touch different files. No longer stacked on #14; rebased onto main after #14 was redone.

Root cause — primal (1.0× vs 0.999×)

Bottom-up diagnostic at planck_fast precision (60 k/decade, l_max=50, q=5, rtol=1e-6):

• Background scalars (H₀, Ω_m, conformal_age) match CLASS to <0.0001%.
• δ_cdm and δ_b match CLASS synchronous-gauge perturbations to <0.05% when compared at the same τ. (The +1.7% I first measured was an np.interp clamping artifact past CLASS's saved tau_max=13998 vs clax's tau_end=14139.)
• P(k) showed -0.5% bias because tau_end = 0.999 × conformal_age = 14139 Mpc while CLASS evaluates pk(z=0) at the full conformal_age (14153 Mpc).

Linear growth from τ=14139 to τ=14153: ΔD/D ≈ ½ × Δa/a ≈ ½ × aH × Δτ ≈ 0.16%, doubling in P ∝ D² to 0.33%.

Root cause — gradient (Taylor correction)

After the primal fix, the single-mode path needs tau_end to be parameter-dependent (any θ that changes H(z) changes conformal_age). But tau_end = bg.conformal_age cannot flow as a traced t1 into diffrax's ODE solve: RecursiveCheckpointAdjoint + PIDController raise RuntimeError: Unexpected tangent (the controller's accepted-step factor is marked eqxi.nondifferentiable). The fix uses stop_gradient(bg.conformal_age) as t1 and reintroduces the chain-rule term analytically.

Chain rule for δ_m(τ_end(θ), y(τ_end(θ); θ)) decomposes into:

Term	Captured by frozen-t1 AD?	Magnitude vs total
∂δ_m/∂y · dy(τ_end; θ)/dθ	yes (RecursiveCheckpointAdjoint)	~30%
dδ_m/dτ · dτ_end/dθ	no (stop_gradient drops it)	~70%

This perfectly explains the diagnostic pattern from scripts/benchmark_gradients.py:

Param	AD (before)	FD	|AD/FD-1|
h	-4894	-13912	64.8% ❌
omega_b	-17064	-105290	83.8% ❌
omega_cdm	+213990	+122410	74.8% ❌
ln10A_s	+10526	+10526	0.0% ✅
n_s	+7296	+7296	0.0% ✅

ln10A_s and n_s enter P(k) as pure scalar prefactors of the primordial spectrum — they don't change τ_end — so the missing chain-rule term is zero for them.

Fix. First-order Taylor expansion around the frozen endpoint:

tau_traced = bg.conformal_age
tau_end    = jax.lax.stop_gradient(tau_traced)
sol = diffrax.diffeqsolve(... t1=tau_end ...)

dy_dtau = _perturbation_rhs(tau_end, sol.ys[-1], ode_args)
delta_m_frozen, ddelta_m_dtau = jax.jvp(
    _extract_delta_m_partial,
    (sol.ys[-1], tau_end),
    (dy_dtau, jnp.ones_like(tau_end)),
)
return delta_m_frozen + ddelta_m_dtau * (tau_traced - tau_end)

Two properties make this exact at fiducial:

• (tau_traced − tau_end) = 0 identically → primal value unchanged.
• jax.jvp with tangent (dy/dτ, 1) captures both the implicit (state) and explicit (background-density-at-τ) contributions to dδ_m/dτ, so AD picks up the correct linear coefficient.

Changes

clax/perturbations.py

Primal fix (existing in PR):

• _perturbation_solve_setup gains a tau_max_factor kwarg, default 0.999 (preserves the C_l / source-function pipeline).
• _perturbations_solve_mpk_impl (used by compute_pk_table) and _matter_delta_m_single_k_impl (used by compute_pk) both now pass tau_max_factor=1.0 to integrate to the full bg.conformal_age.
• The C_l-driving path (perturbations_solve) is unchanged.

Gradient fix (new):

• _matter_delta_m_single_k_impl now applies the Taylor correction to recover dτ_end/dθ. Inline comment documents the diffrax structural constraint forcing the stop_gradient.

tests/test_pk_gradients.py

• New TestPkScalarDensityGradients regression test: asserts dP/d{h, omega_b, omega_cdm} matches centered FD to <5% at PK_CONTRACT_PREC (rtol=1e-6). Fails loudly on the original stop_gradient-only code (where the same params were 65-84% off); the 5% threshold accommodates the AD/FD precision floor at rtol=1e-6 (empirically ~3% on omega_b at k=3e-4).

Measured improvement

Primal (P(k) value)

k [Mpc⁻¹]	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 drops from ~0.30% to ~0.07%.

Gradient (dP/dθ vs centered FD), trajectory across pt_ode_rtol

Param	rtol=1e-3 (was)	rtol=1e-5	rtol=1e-7
h	6.41%	1.08%	0.32%
omega_b	8.91%	3.18%	1.23%
omega_cdm	4.83%	0.85%	0.61%
ln10A_s, n_s	0%	0%	0%

Clean precision-floor convergence (~3-4× per rtol decade), consistent with the FD oracle's own step-noise floor at each rtol (FD step-spread on h is 2.2% at rtol=1e-7). Tightening rtol further would be exercise without information gain — both AD and FD converge to the same value.

Test plan

• pytest tests/test_pk_gradients.py::TestPkScalarDensityGradients -v → 3/3 pass (h: 0.6%, omega_b: 3.1%, omega_cdm: 0.7%; all under 5% threshold).
• pytest tests/test_thermodynamics.py -v --fast → 9/9 pass.
• pytest 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 → unchanged from prior PR fix(compute_pk): integrate to full bg.conformal_age + recover dtau_end gradient via Taylor correction #16 state (44 passed in 296s).
• Existing TestPkScalarGradients (primordial-only direct-path coverage) and TestPkPublicTableGradients (table-API density coverage) unchanged.

Notes

• Forward-mode AD (jax.jvp) is not yet enabled end-to-end. Three structural blockers remain: RecursiveCheckpointAdjoint's internal custom_vjp (the checkpointed_while_loop), clax/shooting.py's shoot_fn custom_vjp, and clax/thermodynamics.py's _find_z_reio custom_vjp. None affect reverse-mode (HMC, optimization). The Taylor correction in this PR is mode-agnostic — once those blockers are converted, jax.jvp(compute_pk) will pick up the chain-rule term automatically. This is a follow-up workstream; tracked separately.
• Residual sub-percent biases at moderate rtol come from the AD/FD precision floor (both adjoint and FD oracle inherit ODE solver tolerance), not from a missing term. At pt_ode_rtol=1e-7 all five params pass <1.3%.