classy_szlite

Fast, differentiable cosmology in pure JAX.

A standalone Python package providing JIT-compiled, jax.grad-friendly access to:

• CMB angular power spectra (TT, TE, EE)

• Linear and non-linear matter power spectrum

• Distances (H(z), comoving, angular-diameter)

• Derived parameters (σ8, Ω_m, S8)

• Halo-model tSZ Cl^yy (Arnaud 2010 GNFW pressure profile)

Backed by the high-accuracy v2 CosmoPower emulators — the same emulators used in the ACT DR6 extended-cosmology analysis (2025) and the ACT DR6 + DESI DR2 analysis by Poulin et al. (2025). They match the CAMB-based Jense et al. (2024) emulators to well under 0.1 \(\sigma\) in \(\Lambda\)CDM (where \(\sigma\) is the typical 68% CL on parameter constraints from ACT DR6).

Runtime dependencies: just jax, numpy, and mcfit. See Installation for the emulator-coverage details (LCDM, \(m_\nu\)-LCDM, \(w\)CDM, \(N_{\rm eff}\)-LCDM, EDE).

Contents

• Installation
  • From PyPI (recommended)
  • From source
  • Emulator data
  • Verification
• Quick start
  • The cosmology container
  • The five calculations
  • The fast path for MCMC
  • JAX gradients
• API reference
  • Parameter containers
  • Derived parameters
  • CMB angular power spectra
  • Matter power spectrum
  • Distances
  • Halo-model tSZ Cl^yy
  • Utility
• Throughput
  • Reproduce
• JAX gradients
  • Why this matters
  • Verification: autodiff vs finite difference
  • Recommended path: gradient via the factory closure
  • Gradient through the full pipeline (cosmology + profile)
  • Gradient w.r.t. CosmoParams as a pytree
  • Jacobians (Fisher matrix building block)
  • Caveats
• Cl^yy convergence
  • What was measured
  • Numerical results
  • Recommended settings
  • Caveats and notes
  • How to reproduce
• Tutorials & examples
  • CMB angular power spectra
  • Matter Pk (linear + nonlinear)
  • Cosmological distances
  • Linear growth σ₈(z)
  • Halo-model tSZ Cl^yy (1h + 2h decomposition)
  • Battaglia-2012 profile: validation against classy_sz
  • Bandpower covariance: Gaussian + 1-halo trispectrum
  • Bestfit + NUTS + RW-MH on synthetic Cl^yy bandpowers (baseline cosmology)
  • Posterior bands on the GNFW pressure profile
  • Fisher matrix in one autodiff sweep
  • End-to-end MCMC pattern (cobaya Theory)
  • Cosmology scan
  • Exploring EDE space
  • Pre-compiling a forward + gradient function
  • EDE NUTS demo — recovering an MCMC posterior in 3 hours

Note

Inference recipes shipped as runnable scripts: examples/nuts_clyy_profile.py (NumPyro NUTS + cobaya RW-MH overlay), examples/profile_bands.py (GNFW profile posterior bands), examples/fisher_clyy_profile.py (Fisher matrix via jax.jacfwd). See Tutorials & examples for full walkthroughs.

Quick example

import jax.numpy as jnp
import classy_szlite as csl

cosmo = csl.CosmoParams()

# Derived
print(csl.derived(cosmo))            # → {'sigma_8': 0.812, 'Omega_m': 0.311, 'S8': 0.827, ...}

# CMB
cls = csl.cl_TTTEEE(cosmo)           # → {'ell','tt','te','ee'}, dimensionless D_ℓ

# Matter Pk
k, pk = csl.Pk(cosmo, [0., 0.5, 1., 2.])

# Distances
Hz, chi, Da = csl.distances(cosmo, [0.1, 0.5, 1.0])

# tSZ Cl^yy
profile = csl.ProfileParamsA10(P0=8.13, beta=5.48, B=1.25)
ell = jnp.geomspace(2, 9000, 80)
cl_1h, cl_2h = csl.cl_yy(cosmo, profile, ell)

# MCMC fast path: precompute cosmology + halo grids once, then ~5 ms/call
ev = csl.cl_yy_factory(cosmo, ell)
cl_1h, cl_2h = ev(profile)

Repository

github.com/CLASS-SZ/classy_szlite