JAXcosmo and automatic differentiation#

The optional jaxcosmo backend supplies differentiable cosmological theory from jax-cosmo. Use options["code"] = "jaxcosmo" to select it, and pass a derivative provider to FisherMatrix to choose automatic differentiation (AD) or finite differences (FD) of the JAX-native observables.

Installation#

From a source checkout containing the JAXcosmo implementation:

uv sync --extra jaxcosmo
# Alternatively:
pip install -e '.[jaxcosmo]'

The extra installs JAX, jaxlib, jax-cosmo and its compatibility dependencies. The shipped backend profile enables JAX float64 (NUMERICS.x64: True), which is important for small finite-difference steps and Fisher inversions. Initial JIT compilation contributes to the first evaluation’s runtime; subsequent evaluations with the same shapes reuse compiled functions.

A photometric forecast#

This example keeps both the neutrino mass and the spectral index fixed. Parameters omitted from freepars retain their fiducial values; the usual survey nuisance parameters are added by CosmicJellyfish.

from cosmicjellyfish.fishermatrix.cosmicfish import FisherMatrix
from cosmicjellyfish.fishermatrix.jaxcosmo_derivatives import JaxAutodiffProvider

fiducial = {
    "10^9As": 2.1,
    "Omegam": 0.32,
    "Omegab": 0.05,
    "h": 0.67,
    "ns": 0.96,
    "mnu": 0.0,
}
free = {name: 0.01 for name in ("10^9As", "Omegam", "Omegab", "h")}
options = {
    "code": "jaxcosmo",
    "accuracy": 1,
    "nonlinear": True,
    "derivatives": "3PT",
    "survey_name_photo": "Euclid-Photometric-ISTF-Pessimistic",
    "results_dir": "results/jaxcosmo/",
    "outroot": "photo_ad",
}

forecast = FisherMatrix(
    fiducialpars=fiducial,
    freepars=free,
    options=options,
    observables=["GCph", "WL"],
    cosmoModel="LCDM",
    surveyName="Euclid",
    derivative_provider=JaxAutodiffProvider(),
)
fisher = forecast.compute()

JaxAutodiffProvider uses jax.jacfwd to differentiate the implemented observable; the values in freepars select the parameters but do not set an AD step size. For finite differences they are relative steps (absolute steps when the fiducial value is zero), as described in Derivative methods.

Three routes with the same theory#

Route	`derivative_provider`	Observable used for derivatives
JAX AD	`JaxAutodiffProvider()`	JAX-native CosmicJellyfish observable
JAX FD	`JaxFiniteDifferenceProvider()`	The same JAX-native observable
Standard FD	Omitted, or `None`	Existing NumPy/SciPy observable

Import both providers from cosmicjellyfish.fishermatrix.jaxcosmo_derivatives. All three routes can use options["code"] = "jaxcosmo".

AD versus JAX FD isolates the derivative method. Decrease the FD step to check convergence. JAX FD versus standard FD checks the observable port, including its interpolation, integration and boundary choices. The JAX observables reproduce CosmicJellyfish’s photometric Limber and spectroscopic galaxy-power calculations; they do not call jax_cosmo.angular_cl. The existing covariance and Fisher assembly are retained; selecting AD does not add covariance-derivative terms.

The normal backend adapter returns NumPy/SciPy tables for the existing pipeline. Those tables do not preserve JAX gradients, which is why choosing the backend alone does not enable autodiff. Use a fresh options dictionary for each forecast: configuration initialisation fills defaults into it.

Comparison notebook#

Download 07-JAXcosmo-autodiff-vs-finite-differences.ipynb and run it from the repository checkout (or open it directly in notebooks/). Select that checkout’s Python environment, for example .venv/bin/python, as the Jupyter kernel. Install jupyterlab and ipykernel into that environment if needed.

The notebook compares the three JAXcosmo routes above with the standard CosmicJellyfish SYREN-NEW forecast. It includes:

linear and nonlinear matter power, and the response to the spectral index;
angular spectra from the JAX and standard observable implementations;
FD step-size convergence, marginalized errors and a triangle overlay;
an optional CLASS reference, disabled by default.

The fiducial model is flat LCDM with massless neutrinos. The Fisher runs keep ns = 0.96 fixed, varying 10^9As, Omegam, Omegab and h, along with the default photometric nuisance parameters. This avoids using the known SYREN-NEW nonlinear spectral-index response discrepancy in the forecast comparison; see SYREN-NEW spectral-tilt response and photometric Fisher discrepancies. The diagnostic response plot still varies ns to show the theory response, independently of the Fisher parameter list. Fixing ns does not make the two backends’ predictions identical.

The notebook writes figures and its summary into results/jaxcosmo_ad_vs_fd/. These are local outputs. Re-run all cells after changing the free parameters so the displayed tables and figures share the same configuration.

Why SYREN is a separate theory comparison#

JAXcosmo uses an Eisenstein–Hu transfer function (with baryon wiggles by default) and Halofit nonlinear power. SYREN-NEW uses an Eisenstein–Hu-based linear fit with its own corrections and growth model, and a separate nonlinear emulator. A shared Eisenstein–Hu ingredient does not make their power spectra or derivatives equal. Compare AD against FD with the same JAX observable to assess differentiation; compare SYREN against the standard JAXcosmo route to assess theory dependence.

Where the SYREN amplitude difference comes from#

With ns and mnu fixed, the marginalized 10^9As error differs by about 10% between SYREN-NEW and JAXcosmo. The autodiff and finite-difference routes through the same JAX observable agree to 1.8e-5 in max|dF|/max|F| (the largest Fisher change divided by the largest Fisher entry), so the difference is not a differentiation problem. The following decomposition locates it.

`10^9As` error, fiducial 1	SYREN-NEW	JAXcosmo AD	AD/SYREN
unmarginalized (`1/sqrt(F_AsAs)`)	0.00191324	0.00186387	-2.6%
cosmology only	0.0223482	0.0237162	+6.1%
cosmology + 10 galaxy biases	0.0294102	0.0263472	-10.4%
cosmology + biases + IA (all)	0.0294396	0.0264081	-10.3%

The two models agree closely on the unmarginalized amplitude error and on dlnP_nl/dlnAs (ratio 0.98-1.00 over k = 10^-3 to 1 h/Mpc). The gap appears only once the galaxy-bias parameters are marginalized: SYREN’s amplitude error grows by a factor 1.32, JAXcosmo’s by 1.11. The amplitude-bias correlations differ in shape, for example r(10^9As, b8) = -0.37 for SYREN and -0.28 for JAXcosmo, while the amplitude response of the spectra themselves is nearly identical.

The decisive test replaces the matter power spectrum inside the JAX observable with SYREN-NEW’s, keeping the JAX kernels, integration, covariance and Fisher assembly. That hybrid reproduces the SYREN result:

route	marginalized `10^9As` error	ratio to SYREN
SYREN-NEW, standard pipeline	0.0294396	1.00000
JAXcosmo, autodiff	0.0264081	0.89703
SYREN-NEW power inside the JAX observable	0.0294160	0.99920

The hybrid’s fiducial angular spectra match the standard SYREN ones to max|dC|/max|C| = 8.3e-4. The same comparison at a second fiducial gives -10.3% and -11.3%, so the effect is stable rather than a single-point accident. The difference is therefore a genuine matter-power-model difference, not the observable implementation, the covariance, the Fisher assembly or the autodiff. It cannot be removed without changing the theory.

Both models are “Eisenstein-Hu”, so why do they differ?#

A common objection is that both backends are Eisenstein-Hu based, so the spectra should agree. They are not the same formula:

jax_cosmo evaluates the Eisenstein-Hu 1998 fit including baryon wiggles (eisenhu_osc) and normalises it to sigma8.
SYREN-NEW starts from the Eisenstein-Hu no-wiggle fit (get_eisensteinhu_nw) and multiplies it by exp(logF), an emulated correction of the true linear spectrum relative to the Eisenstein-Hu fit (Bartlett et al. 2023), a growth correction R and a further S(k) correction with its own coefficients. In other words SYREN-NEW’s “linear” model is calibrated to reproduce a full Boltzmann spectrum with Eisenstein-Hu only as a backbone.

At equal As the two linear spectra differ by a few percent with a k-dependent sign, while the growth factors agree to 4e-4:

band [`h/Mpc`]	`P_lin`(SYREN)/`P_lin`(EH), z=0
0.001-0.005	1.0032
0.005-0.02	1.0164
0.02-0.05	1.0093
0.05-0.1	0.9872
0.1-0.2	0.9914
0.2-0.5	1.0093
0.5-1	1.0161

Those percent-level shape changes sit exactly where the galaxy-bias marginalization operates. Feeding Halofit with SYREN-NEW’s linear shape, while keeping the JAX observable and the Halofit nonlinear prescription, gives a marginalized 10^9As error of 0.0301323, i.e. +2.4% with respect to SYREN (0.0294396) and +14% with respect to plain Eisenstein-Hu plus Halofit (0.0264081). The linear input therefore carries essentially all of the difference: the two nonlinear prescriptions (Halofit emulator versus the SYREN-NEW emulator) contribute at the percent level in the opposite direction. This is why the gap cannot be removed by changing the derivative method, the grids or the marginalization; only a common power-spectrum model can.

Reproduce the decomposition with:

uv run --extra jaxcosmo python scripts/validation/compare_theory_amplitude_response.py \
    --outdir results/theory_amplitude_response

The script prints the theory response table, the nuisance-marginalization decomposition, the hybrid test, the linear-shape comparison, the Halofit-on-SYREN-linear bridge, both fiducials and the AD-versus-FD control, and writes summary.json. The response-optimised SYREN coefficient sets from the ns-response study are not a fix here: they increase this amplitude error.

Theory and supported configurations#

Component	Current implementation
Cosmology	Flat `LCDM` and `w0waCDM`, massless neutrinos (`mnu = 0`)
Matter tracer	Total matter; clustering/CB tracers are unsupported
Background and growth	`jax-cosmo` distances, expansion and scale-independent growth
Linear power	Eisenstein–Hu, with or without baryon wiggles
Nonlinear power	Halofit, `takahashi2012` or `smith2003`
Amplitude	JAX port of the `symbolic_pofk` maximum-precision `As`–`sigma8` fit
JAX observables	Photometric `GCph`/`WL`; spectroscopic `GCsp` default modelling path
Outside this implementation	CMB, intensity-mapping autodiff, massive neutrinos and curvature

The amplitude fit is re-evaluated inside the traced function, retaining its cosmology dependence. Supply one amplitude coordinate, for example 10^9As or sigma8, rather than both. The AD parameter mapping currently exposes 10^9As, Omegam, Omegab, h, ns, w0 and wa; an accepted backend input alias is not necessarily an available AD coordinate. Use the coordinates in the example for a controlled comparison.

The GCsp JAX path requires fix_cosmo_nl_terms=True, bfs8terms=False and a supported bias model (linear, linear_log or linear_Qbias). Photometry supports the implemented photo-z distributions, galaxy-bias models and eNLA intrinsic alignments. Unsupported modelling choices raise explicit errors.

For the two-stage GCsp forecast in shape, background, growth and bias parameters, see GCsp forecasts through intermediate parameters. It reuses the GCsp JAX observable and adds paper-style finite differences, autodiff comparisons and a Fisher-space Jacobian projection.

Backend configuration lives in cosmicjellyfish/configs/default_boltzmann_yaml_files/jaxcosmo/default.yaml. To customise it, copy the profile and set options["jaxcosmo_config_yaml"] to the copy’s path. Key controls are ACCURACY.transfer, ACCURACY.halofit, ACCURACY.distance_steps, and the NUMERICS power/background grids and x64 flag. Public CosmicJellyfish power uses k in inverse Mpc and P(k) in cubic Mpc; the theory adapter converts the internal jax-cosmo h-scaled units.

Smooth Halofit and table boundaries#

The theory layer includes a locally maintained Halofit nonlinear-scale solver. It uses a converged Newton solve with an implicit derivative to avoid the fixed-grid root-interpolation derivative artefacts in jax-cosmo 0.1.0. It is used by both the standard backend adapter and the JAX-native theory path. The upstream proposal and numerical comparisons are recorded in jaxcosmo backend: review log.

The GCsp JAX observable evaluates power directly by default. Set NUMERICS.legacy_pk_clamp: True in a custom profile to reproduce the standard spline’s clamp at its tabulated k boundaries for parity studies. This introduces a derivative kink at the boundary. Spectroscopic Savitzky–Golay window changes can also affect finite differences at large steps; inspect step convergence rather than relying on a single step size.

Reproducible parity checks and development#

For a straight-line terminal example timing SYREN finite differences against JAXcosmo autodiff with ns fixed:

uv run --extra jaxcosmo python scripts/validation/time_syren_vs_jaxcosmo.py

It prints setup, fiducial spectra, covariance, derivatives and Fisher assembly times, followed by marginalized errors. The first AD derivative call includes compilation and execution; a second call measures the compiled Jacobian’s warm execution. The two backends use different theory, so this timing comparison is distinct from the same-theory parity checks below.

From the checkout, with the extra installed:

uv run --extra jaxcosmo python scripts/validation/compare_jax_photo_fishers.py \
    --outdir results/jax_photo_parity
uv run --extra jaxcosmo python scripts/validation/compare_jax_spectro_fishers.py \
    --outdir results/jax_spectro_parity --fd-step 0.003

These scripts compare the three same-theory routes and save reports and numerical summaries. They are separate from the SYREN comparison notebook. Existing parity measurements cover a limited Euclid, massless-neutrino LCDM setup; the jaxcosmo backend: review log records metrics, normalisations and remaining review items.

The backend-neutral contract in cosmology/jax_theory.py separates the traceable theory from observables. A future differentiable backend can implement that contract; a backend returning NumPy arrays alone cannot supply cosmological AD. The standard likelihood/context interface remains available to sampling code. Derivative providers cache compiled observables locally and discard that cache when pickled; create them within workers. Nautilus serial-versus-spawn parity with JAX remains a separate validation item, rather than a demonstrated result of these Fisher comparisons.