Skip to content
BranchesTags
 
 

Repository files navigation

clax-pt

clax-pt

A differentiable cosmological Boltzmann solver and one-loop EFTofLSS engine in JAX.

Development

This codebase was written entirely by Claude Code (Opus and Sonnet models). The development process -- including architecture decisions, bug hunting through the CLASS and CLASS-PT C source code, and numerical validation -- is documented in CHANGELOG.md and CLAUDE.md.

Overview

clax-pt extends clax with a JAX-native one-loop EFTofLSS engine, exposing real-space and redshift-space galaxy power-spectrum multipoles, IR resummation, bias operators, and EFT counterterms — all end-to-end differentiable via JAX automatic differentiation.

While clax mirrors CLASS, clax-pt mirrors CLASS-PT. Both clax and clax-pt were written by Claude Code (Opus and Sonnet models) under human supervision; see CHANGELOG.md and CLAUDE.md for the development log.

clax solves the coupled Einstein-Boltzmann equations for cosmological perturbations from first principles: background cosmology, hydrogen recombination, the full photon-baryon-neutrino Boltzmann hierarchy in synchronous gauge, line-of-sight integration for CMB angular power spectra, HaloFit for nonlinear matter power, gravitational lensing, and a shooting method for theta_s parametrization. The entire pipeline -- from cosmological parameters to P(k), C_l^TT/EE/TE/BB, and lensed C_l -- is end-to-end differentiable via JAX automatic differentiation.

The goal is a differentiable alternative to CLASS and CLASS-PT that enables gradient-based cosmological inference (HMC, variational methods) on CMB and large-scale structure data.

Status

Sub-0.2% unlensed C_l^TT/EE at l=20-1200. Full lensed C_l^TT/EE/TE/BB (sub-0.2% at l=10-2000). Multi-cosmology validated (10 LCDM parameter points). Full ncdm Boltzmann hierarchy. fit_cl preset for HMC / fitting; planck_cl for science-grade accuracy (487s on H100). See CHANGELOG.md for details.

Accuracy comparison against CLASS v3.3.4

All comparisons at Planck 2018 best-fit LCDM. GPU: H100-80GB.

Unlensed C_l angular power spectra

planck_cl preset: k_max=1.0, 60 k/decade (300 modes), l_max=50, full ncdm Psi_l(q) hierarchy (5 q-bins x 18 multipoles), source-interpolated to 10000 fine k-points.

Multipole l C_l^TT error C_l^EE error C_l^TE error
20 -0.08% -0.21% -0.3%
30 -0.05% -0.11% -0.5%
50 -0.05% -0.05% +0.8%
100 -0.02% +0.02% -0.03%
200 -0.05% -0.04% -0.05%
300 -0.06% -0.02% -0.04%
500 -0.15% -0.15% -0.01%
700 -0.23% -0.11% +0.08%
1000 -0.57% -0.26% +1.7%
1200 -0.07% +0.03% —

Bold = sub-percent. TT sub-0.1% at l=20-300 and l=1200. EE sub-0.3% from l=20 to l=1200. TE zero crossings near l=52 and l=400 cause large relative errors.

Lensed C_l angular power spectra

Full spin-2 correlation function lensing with Cgl2 corrections, 12 Wigner d-functions:

Multipole l Lensed TT err Lensed EE err BB ratio
50 -0.000% +0.000% 1.000
200 +0.000% -0.001% 1.000
500 +0.002% -0.003% 0.999
1000 +0.006% +0.005% 0.996
2000 -0.199% -0.166% 0.937

Lensing algorithm sub-0.2% at all l=10-2000 for TT and EE (tested with CLASS unlensed+pp as input).

Lensing potential C_l^phiphi

clax.compute_cl_pp(pt, params, bg, th, l_max, *, nonlinear="none") uses the source-based Limber kernel (CLASS transfer.c + harmonic.c):

With pt_k_max_cl >= 5.0 Mpc^-1 (required for Halofit's σ(R) bisection), measured residuals at the default cosmology:

Absolute |C_l^phiphi_clax / C_l^phiphi_CLASS - 1| at the default cosmology:

Multipole l nonlinear="none" nonlinear="halofit" NL/lin ratio agreement
100 0.18% 0.16% 0.02%
200 0.14% 0.10% 0.05%
500 0.23% 0.15% 0.09%
1000 0.12% 0.08% 0.04%
1500 0.49% 0.61% 0.11%
2000 0.71% 0.76% 0.05%
2500 0.39% 1.16% 0.76%

Both nonlinear modes share the same source-Limber kernel, so they track each other closely — the absolute Halofit residual is dominated by the common linear-kernel systematic, with a small extra contribution from the R(k, z) accuracy that the third column isolates. Halofit injection multiplies the lensing source by sqrt(R(k, z(τ))) where R = P_NL/P_lin is computed on a 100-point z-grid via vmap(compute_pk_nonlinear) with log-log k-extension to k_max=20 Mpc^-1 (CLASS itself uses a dedicated nonlinear k-grid that extends past the perturbation k-range; the extension provides the equivalent coverage). Sub-percent absolute residual across all ℓ ≤ 2000; ~1% at ℓ=2500 where linear-kernel and R-residual systematics add constructively. Narrower k-grids (pt_k_max_cl < 5) gracefully degrade to no NL correction (R = 1, per CLASS fourier.c:1706-1716).

Multi-cosmology validation

Validated at 10 LCDM parameter variations (omega_b, omega_cdm, h, n_s, tau_reio at +/-20%):

  • TT: sub-0.5% at ALL l for ALL 10 cosmologies
  • EE: sub-0.3% at l>=100 for all cosmologies

Matter power spectrum P(k)

Single-mode compute_pk(params, prec, k) at PrecisionParams.planck_fast() (60 k/decade, l_max=50, ncdm_q_size=5, rtol=1e-6) vs CLASS v3.3.4 linear pk_lin(k, z=0) at default Planck 2018 LCDM:

k [Mpc^-1] clax / CLASS Error
0.003 0.9983 -0.17%
0.010 1.0025 +0.25%
0.050 1.0004 +0.04%
0.100 1.0007 +0.07%
0.300 0.9997 -0.03%
1.000 0.9980 -0.20%

Median |error| ≈ 0.07%, worst-case 0.25%. Underlying perturbation evolution (δ_cdm and δ_b at synchronous gauge) matches CLASS to <0.05% — most of any remaining residual is from photon-hierarchy and IC corrections at next order.

Pipeline accuracy

Module Accuracy Notes
Background (H, D_A, r_s) < 0.01% 6+ significant digits
Thermodynamics (x_e) 0.25% at z_star RECFAST with Heun stepping
Visibility function g(tau) 0.04% Bisection reionization, corrected kappa
Perturbation ODE (Phi, Psi) 0.01-0.25% Gauge-corrected at recombination
AD gradients (dP(k)/d(params)) 0.03% vs finite differences

Performance

Preset Background Thermo Perturbations Harmonic Total GPU
planck_cl 1s 53s 401s 33s 487s H100

fit_cl uses 20 k/decade, l_max=17, table-based Bessel (all l at once), rtol=1e-3, max_steps=1024, and defaults to the Rodas5 Rosenbrock solver. Suitable for HMC/fitting at <1.5% TT/EE accuracy (l<=500). planck_cl gives sub-0.2% science-grade accuracy with Kvaerno5. All times are JIT-cached (second call onwards).

Quick start

import jax
import clax
from clax import CosmoParams, PrecisionParams

# Background + thermodynamics
result = clax.compute(CosmoParams(h=0.6736, omega_b=0.02237))
print(result.bg.H0)            # Hubble constant [Mpc^-1] (CLASS internal units)
print(result.bg.conformal_age) # conformal age in Mpc

# Practical dense-spectrum / reusable-table workflow: solve once on an
# internal k-grid, then interpolate P(k) to the requested points.
pk_table = clax.compute_pk_table(CosmoParams(), k_eval=[1e-3, 5e-2, 1.0])
print(pk_table.pk_grid)
print(pk_table.pk(0.05))

# Exact single-k diagnostic / spot-check path
pk = clax.compute_pk(CosmoParams(), k=0.05)

# Local gradient at one k-mode
grad = jax.grad(lambda p: clax.compute_pk(p, k=0.05))(CosmoParams())
print(grad.omega_cdm)  # dP(k)/d(omega_cdm)

# Practical multi-k gradient workflow: differentiate one reusable table solve
# instead of many repeated exact single-k solves.
grad_table = jax.grad(
    lambda p: jax.numpy.sum(
        clax.compute_pk_table(p, k_eval=[1.5e-2, 3.0e-2, 9.0e-2, 2.5e-1]).pk_grid
    )
)(CosmoParams())
print(grad_table.h)  # d(sum_i P(k_i))/dh

# Angular power spectra (C_l) -- science quality
from clax.perturbations import perturbations_solve
from clax.harmonic import compute_cl_tt_interp, compute_cl_ee_interp
prec = PrecisionParams.planck_cl()
params = CosmoParams()
bg = clax.background_solve(params, prec)
th = clax.thermodynamics_solve(params, prec, bg)
pt = perturbations_solve(params, prec, bg, th)
cl_tt = compute_cl_tt_interp(pt, params, bg, [30, 100, 200])
cl_ee = compute_cl_ee_interp(pt, params, bg, [30, 100, 200])

# CMB lensing potential C_l^phiphi (linear or Halofit-corrected)
cl_pp_lin = clax.compute_cl_pp(pt, params, bg, th, l_max=2500)
cl_pp_nl  = clax.compute_cl_pp(pt, params, bg, th, l_max=2500,
                                nonlinear="halofit")

# Lens the unlensed CMB spectra
cl_tt_lensed, cl_ee_lensed, cl_te_lensed, cl_bb_lensed = clax.lens_cls(
    cl_tt_unlensed, cl_ee_unlensed, cl_te_unlensed, cl_bb_unlensed,
    cl_pp_nl, l_max=2500)

Installation

Requires Python >= 3.10.

pip install jax jaxlib diffrax equinox jaxtyping

# Clone and install
git clone https://github.com/MinhMPA/clax-pt.git
cd clax-pt
pip install -e .

# Run tests
pip install pytest
pytest tests/ -v
pytest tests/ -v -m "not slow"  # skip slow integration tests

Architecture

A sequential pipeline of pure functions, each returning a frozen PyTree:

CosmoParams --> background --> thermodynamics --> perturbations --> primordial
                                                                       |
                                             lensing <-- harmonic <-- transfer <-- nonlinear

Key design choices:

  • CosmoParams fields are JAX-traced for automatic differentiation. PrecisionParams fields are static (control array shapes, not traced).
  • Full Boltzmann hierarchy with smooth RSA damping post-recombination. Full ncdm Psi_l(q) phase-space hierarchy (5 q-bins x 18 multipoles). TCA (tight-coupling approximation) with CLASS-matching dual criteria for numerical stability.
  • Perturbation ODE solved via Diffrax with either Kvaerno5 (ESDIRK, default) or Rodas5 (8-stage Rosenbrock, default in fit_cl); Rodas5Batched shares time-stepping across k-modes for batched compute_pk_table solves.
  • compute_pk_table() / compute_pk_interpolator() are the practical dense-spectrum and reusable-table P(k) APIs, especially on GPU; exact compute_pk() remains the single-mode reference path.
  • Practical reverse-mode P(k) objectives should usually differentiate one table-backed multi-k solve, not many repeated exact compute_pk() calls.
  • Backend-aware batching over k-modes for GPU/CPU perturbation solves.
  • RecursiveCheckpointAdjoint for memory-efficient reverse-mode AD through the ODE solve.
  • Synchronous gauge throughout (matching CLASS default).
  • RECFAST recombination matching CLASS wrap_recfast.c (Peebles C with fudge, Gaussian K correction, Heun stepping).
  • Source-interpolated C_l integration: source functions S(k,tau) interpolated to fine k-grid (10000 points) before line-of-sight integration, resolving the oscillatory T_l(k) transfer function robustly.
  • JIT-compiled: all solve functions + per-l harmonic functions cached for 2x speedup on repeated calls.

Choosing an adjoint mode

PrecisionParams.ode_adjoint controls Diffrax reverse-mode AD through the ODE solves. clax-pt currently supports:

  • "recursive_checkpoint": production default and recommended choice
  • "direct": optional alternate reverse pass for diagnostics or environment-specific benchmarking

Use "recursive_checkpoint" by default on both CPU and GPU. That includes:

  • laptop / local CPU development
  • gating gradient tests
  • production P(k) / C_l gradient workflows on GPU

Use "direct" only after validating it on the exact environment and problem size you care about. It is most appropriate for:

  • reduced or short ODE solves where memory is easy to bound
  • adjoint debugging, where you want to compare reverse-pass implementations
  • target-specific benchmarking after AD-vs-FD checks pass

Do not assume a CPU validation transfers to GPU, or vice versa. In particular, the single-mode perturbation solve's DirectAdjoint path has shown unstable density-parameter gradients on a CPU/macOS checkout even when the forward solve and thermodynamics AD were healthy.

Before switching from the default, validate on the target setup:

  • compare AD vs central finite differences on representative P(k) probes
  • include density parameters (h, omega_b, omega_cdm) plus amplitude/tilt probes
  • check runtime and memory, not just gradient agreement
  • keep recursive_checkpoint as the reference for gating tests unless the alternate mode is revalidated

Example:

from dataclasses import replace
from clax import CosmoParams, PrecisionParams, compute_pk

params = CosmoParams()
prec = PrecisionParams()  # defaults to ode_adjoint="recursive_checkpoint"
pk = compute_pk(params, prec, k=0.05)

# Optional diagnostic / benchmark mode; validate before trusting it.
prec_direct = replace(prec, ode_adjoint="direct")
pk_direct = compute_pk(params, prec_direct, k=0.05)

Source modules

Module Description
constants.py Physical constants matching CLASS
params.py CosmoParams (traced) and PrecisionParams (static)
interpolation.py Pytree-registered CubicSpline
ode.py Diffrax ODE solver wrappers
rosenbrock.py Rodas5 / Rodas5Batched Rosenbrock ODE solver
background.py Friedmann equation, distances, growth factor
thermodynamics.py RECFAST recombination, visibility function
perturbations.py Full scalar + tensor Boltzmann hierarchy
primordial.py Power-law scalar and tensor spectra
bessel.py Spherical Bessel functions j_l(x)
transfer.py Linear matter P(k) from perturbation solve
harmonic.py C_l^TT/EE/TE/BB from line-of-sight integration
lensing.py C_l^phiphi (source-Limber, optional Halofit NL) and lensed C_l (correlation-function method)
nonlinear.py HaloFit (Takahashi 2012)
shooting.py theta_s -> H0 via Newton + custom_vjp

Precision presets

PrecisionParams provides six presets controlling the accuracy/speed tradeoff:

Preset k/decade l_max k_max Use case
fit_cl() 20 17 1.0 HMC / fitting (<1.5%)
planck_fast() 60 50 1.0 Science-grade, optimized
fast_cl() 15 25 0.15 Quick iteration, testing
medium_cl() 20 50 0.3 Moderate accuracy
planck_cl() 60 50 1.0 Science-grade, original
science_cl() 200 50 0.35 Sub-percent C_l

For science-grade results, use compute_cl_tt_interp / compute_cl_ee_interp which interpolate source functions to a fine k-grid (10000 points) before computing the transfer integral. This is robust regardless of the perturbation k-density.

Cosmological parameters

Default parameters correspond to Planck 2018 best-fit LCDM:

Parameter Default Description
h 0.6736 H0 / (100 km/s/Mpc)
omega_b 0.02237 Physical baryon density
omega_cdm 0.1200 Physical CDM density
ln10A_s 3.0445 Log primordial amplitude (A_s = 2.1e-9)
n_s 0.9649 Scalar spectral index
tau_reio 0.0544 Reionization optical depth
m_ncdm 0.06 eV Neutrino mass (single species)
w0, wa -1.0, 0 CPL dark energy equation of state

Known limitations

  • Speed: fit_cl preset targets HMC / fitting; planck_cl runs in 487s on H100 for science-grade work. Further speedup requires fewer k-modes or float32 (currently infeasible with jax_enable_x64=True).
  • TT l=400-800: +0.10-0.18% residual from SW+Doppler source amplitude (0.06% excess at k0.03). Comparable to CAMB-CLASS inter-code variation (~0.07%).
  • TT l>1200: Degrades due to k-integration under-resolution (Bessel oscillation period constant in k, but log-uniform grid spacing grows). Hybrid linear/log k-grid would fix this.
  • EE l=20-30: ~0.2% from RECFAST visibility function bias. HyRec recombination would improve to sub-0.1%.
  • BB tensor modes: Lensing BB is accurate (<0.5% at l<=1000), but primordial BB still ~2x off CLASS.
  • TE zero crossings: Large relative errors near l=52 and l=400 where C_l^TE crosses zero.

References

License

MIT

About

Differentiable Boltzmann solver + one-loop EFTofLSS engine in JAX (mirrors CLASS and CLASS-PT).

Resources

Readme
License
Activity

Stars

18 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages