cosmicjellyfish.cosmology.jax_theory#

Backend-neutral, differentiable theory interface for JAX observables.

CosmicJellyfish separates the theory backend (which produces distances, growth and matter power spectra) from the observables (which turn those quantities into angular power spectra or cluster counts). The same separation is useful for the JAX/differentiable path: a JAX observable should not know which theory backend produced its inputs.

This module defines that small contract:

  • TheoryTables is the bundle of quantities an observable needs, always in CosmicJellyfish units (chi in Mpc, h_over_c in 1/Mpc, P in Mpc^3);

  • JaxTheory is the protocol a differentiable backend implements;

  • get_jax_theory() resolves the backend selected in options["code"].

Only backends whose computations are written with jax.numpy can provide exact automatic derivatives. A NumPy/SciPy backend can still drive the plain CosmicJellyfish observables through finite differences, but it cannot be used here.

Nothing in this module imports JAX at module level, so it stays importable without the optional jaxcosmo extra.

Functions

get_jax_theory

Return the differentiable theory implementation selected by options["code"].

Classes

JaxTheory

A theory backend that can be traced and differentiated by JAX.

TheoryTables

Differentiable theory quantities on a redshift/wavenumber grid.