cosmicjellyfish.LSSsurvey.jaxcosmo_photo#

JAX-native photometric (GCph + WL) angular power spectra for the jaxcosmo backend.

This module re-implements, with jax.numpy only, the Limber computation done by cosmicjellyfish.LSSsurvey.photo_obs.ComputeCls so that the full chain

(cosmology, bias, IA, photo-z) -> C_ell

is traceable by jax.jit / jax.jacfwd. The existing NumPy/SciPy pipeline is left untouched; this is a parallel observable used by cosmicjellyfish.fishermatrix.jaxcosmo_derivatives.JaxAutodiffProvider.

Scope (explicitly rejected otherwise)#

  • observables: GCph and WL only (no spectroscopic, no CMB);

  • total-matter tracer (GCph_Tracer = "matter"), no modified gravity;

  • flat, massless-neutrino cosmology (see cosmicjellyfish.cosmology.jaxcosmo_theory);

  • intrinsic alignment model eNLA; bias models sqrt, binned, binned_constant, flagship.

Numerical differences from ComputeCls#

  • P(k, z) is evaluated directly (Halofit) at the Limber wavenumbers kappa = (ell + 1/2) / chi(z) on the Limber redshift grid, with kappa capped at kmax (the contribution above kmax is masked to zero anyway). ComputeCls instead reads a bicubic spline built on a coarser z / k table, which clamps at the table edges. Evaluating the model directly keeps the result smooth in the parameters, which both autodiff and finite differences need.

  • The luminosity-ratio table is interpolated with numpy.interp (clamped), whereas the original code uses scipy.interpolate.interp1d (raises outside the table).

Everything else (z grid, ell grid, kernels, trapezoid rules, masks, dictionary keys and their order) follows ComputeCls literally.

Module attributes

COSMO_FREE_NAMES

CosmicJellyfish names of the cosmological parameters that can be differentiated.

Functions

is_bias_key

Whether name looks like a galaxy-bias nuisance parameter (b1, b0, ...).

Classes

JaxPhotoObservable

Differentiable photometric angular power spectra.