JaxTheory#

class cosmicjellyfish.cosmology.jax_theory.JaxTheory(*args, **kwargs)[source]#

Bases: Protocol

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

Implementations must be pure functions of the parameters: no NumPy conversion and no SciPy splines on the path between parameters and returned arrays.

cosmo_names: tuple[str, ...]#: Canonical cosmological parameter names expected in the params mapping.

free_name_map: Mapping[str, str]#: Mapping from CosmicJellyfish parameter names to this backend’s canonical names.

resolve_cosmology(params)[source]#

Validate a CosmicJellyfish cosmology mapping and return canonical float values.

Return type:: dict[str, float]

tables(params, z, k_hmpc, nonlinear=True)[source]#

Evaluate TheoryTables at z and k_hmpc for params.

Return type:: TheoryTables

power(params, z, k_hmpc, nonlinear=True)[source]#

Matter power in Mpc^3 at per-redshift wavenumbers k_hmpc (h/Mpc).

z has shape (nz,). k_hmpc may have shape (nk,) (the same wavenumbers for every redshift) or (nz, nk) to request a different wavenumber per redshift, which is what a Limber projection needs. The return value has shape (nz, nk).

Return type:: Any