### Install Jax-Cosmo using pip

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/index.rst

This command installs the jax-cosmo library using pip. It's a straightforward installation process as the library is written in plain Python.

```bash
$ pip install jax-cosmo
```

--------------------------------

### Build Documentation with Sphinx

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Navigates to the 'docs' directory and builds the HTML documentation using Sphinx. Requires Sphinx and its dependencies to be installed.

```bash
cd docs && make html
```

--------------------------------

### Install with Test Dependencies

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Installs the jax-cosmo package in editable mode along with extra dependencies required for testing, such as pyccl.

```bash
pip install -e .[test]
```

--------------------------------

### Import Libraries and Setup for jax_cosmo and CCL Comparison

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Imports necessary libraries from pyccl and jax_cosmo, and sets up the environment for 64-bit precision in JAX. This is a prerequisite for running the comparison.

```python
%pylab inline
import os
os.environ['JAX_ENABLE_X64']='True'

import pyccl as ccl
import jax
from jax_cosmo import Cosmology, background
```

--------------------------------

### Install Black and Pre-commit for Python Formatting

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CONTRIBUTING.md

Installs the Black code formatter and pre-commit hooks for automatic Python code formatting and import reordering. Pre-commit ensures that code is formatted correctly before each commit.

```bash
pip install --user black pre-commit reorder_python_imports
pre-commit install
```

--------------------------------

### Install in Development Mode

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Installs the jax-cosmo package in editable mode, allowing for direct changes to the source code without reinstallation. Useful during development.

```bash
pip install -e .
```

--------------------------------

### Jax-cosmo Container Class Example

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Demonstrates the definition of a custom tracer redshift distribution ('gaussian_nz') inheriting from the jax-cosmo 'container' class. This pattern is used to store traceable parameters and static configurations for JAX compatibility.

```python
from jax_cosmo.core import container
from jax import numpy as jnp

class gaussian_nz(container):
    def __init__(self, z0, sigma, zmax=10, **kwargs):
        super().__init__(z0, sigma, zmax=zmax, **kwargs)
    
    def __call__(self, z):
        z0, sigma = self.params
        return jnp.exp(-0.5 * (z - z0)**2 / sigma**2)
```

--------------------------------

### Prepare for Galaxy Clustering Comparison

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Imports the necessary module for galaxy clustering bias from jax_cosmo. This sets up the next step for comparing galaxy clustering calculations.

```python
# Let's try galaxy clustering now
from jax_cosmo.bias import constant_linear_bias
```

--------------------------------

### Access Cosmology Parameters

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Demonstrates how to access individual cosmological parameters, such as the Hubble parameter 'h', directly from the cosmology object.

```python
# Parameters can be easily accessed from the cosmology object
cosmo.h
```

--------------------------------

### Import JAX and jax-cosmo

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Imports essential JAX libraries, including jax, jax_cosmo, and jax.numpy. It also prints the versions of JAX and jax-cosmo, which is useful for debugging and ensuring compatibility.

```python
import jax
import jax_cosmo as jc
import jax.numpy as np

print("JAX version:", jax.__version__)
print("jax-cosmo version:", jc.__version__)
```

--------------------------------

### Vectorized Likelihood Evaluation over Parameter Grid

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Demonstrates how to use `jax.vmap` to efficiently evaluate a likelihood function over a grid of cosmological parameters. This is useful for tasks like parameter space exploration or generating data for analysis.

```python
import jax.numpy as np
import jax

# Assume 'likelihood' is a defined JAX-compatible function
# def likelihood(p):
#     # ... implementation ...
#     return log_likelihood

@jax.vmap
def likelihood_grid(p):
    return likelihood(p)

param_grid = np.stack([
    np.linspace(0.2, 0.3, 10),
    np.linspace(0.75, 0.85, 10)
], axis=1)
logL_grid = likelihood_grid(param_grid)
```

--------------------------------

### Define Equivalent Cosmologies in CCL and jax_cosmo

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Defines two cosmology objects with identical parameters, one using pyccl and the other using jax_cosmo. This ensures a fair comparison of subsequent calculations.

```python
# We first define equivalent CCL and jax_cosmo cosmologies
cosmo_ccl = ccl.Cosmology(
    Omega_c=0.3, Omega_b=0.05, h=0.7, sigma8 = 0.8, n_s=0.96, Neff=0,
    transfer_function='eisenstein_hu', matter_power_spectrum='halofit')

cosmo_jax = Cosmology(Omega_c=0.3, Omega_b=0.05, h=0.7, sigma8 = 0.8, n_s=0.96,
                      Omega_k=0., w0=-1., wa=0.)
```

--------------------------------

### Run All Tests with pytest

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Executes all tests within the project using the pytest framework. This is a standard command for verifying code integrity.

```bash
pytest
```

--------------------------------

### Cosmological Probes (Weak Lensing, Galaxy Clustering) in Python

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Illustrates the definition of cosmological probes, including Weak Lensing and Number Counts (Galaxy Clustering), using redshift distributions and bias models. It shows how to configure probes with intrinsic alignments, shear calibration biases, and per-bin properties.

```python
import jax_cosmo as jc

# Define redshift bins
nz_bins = [
    jc.redshift.smail_nz(1., 2., 0.4, gals_per_arcmin2=10.),
    jc.redshift.smail_nz(1., 2., 0.8, gals_per_arcmin2=10.),
]

# Weak Lensing probe
# sigma_e: intrinsic ellipticity dispersion (shape noise)
probe_wl = jc.probes.WeakLensing(
    nz_bins,
    sigma_e=0.26  # Typical value for ground-based surveys
)

# Weak Lensing with intrinsic alignments (NLA model)
ia_bias = jc.bias.des_y1_ia_bias(0.5, -1.0, 0.62)
probe_wl_ia = jc.probes.WeakLensing(
    nz_bins,
    ia_bias=ia_bias,         # Enable NLA intrinsic alignments
    multiplicative_bias=0.0,  # Shear calibration bias (1+m)
    sigma_e=0.26
)

# Per-bin multiplicative bias
probe_wl_mbias = jc.probes.WeakLensing(
    nz_bins,
    multiplicative_bias=[0.01, -0.02],  # Per-bin m values
    sigma_e=[0.26, 0.28]                # Per-bin sigma_e
)

# Galaxy Clustering (Number Counts) probe
bias = jc.bias.constant_linear_bias(1.5)
probe_nc = jc.probes.NumberCounts(
    nz_bins,
    bias,
    has_rsd=False  # Redshift space distortions (not yet implemented)
)

# Per-bin galaxy bias
bin_biases = [
    jc.bias.constant_linear_bias(1.2),
    jc.bias.constant_linear_bias(1.6),
]
probe_nc_binned = jc.probes.NumberCounts(nz_bins, bin_biases)

# Access probe properties
print(f"Number of tracers: {probe_wl.n_tracers}")  # 2
print(f"Maximum redshift: {probe_wl.zmax}")
```

--------------------------------

### Sparse Matrix Operations for Covariance Matrices in jax_cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Showcases the use of `jax_cosmo.sparse` module for efficient handling of sparse covariance matrices, common in Gaussian approximations. It covers checking sparsity, conversion to dense format, matrix inversion, matrix-vector products, matrix-matrix products, and log-determinant calculation.

```python
import jax_cosmo as jc
import jax.numpy as np

cosmo = jc.Planck15()
nzs = [jc.redshift.smail_nz(1., 2., 0.5, gals_per_arcmin2=10.)]
probes = [jc.probes.WeakLensing(nzs, sigma_e=0.26)]
ell = np.logspace(1, 3, 30)

# Get sparse covariance
mu, cov_sparse = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo, ell, probes, sparse=True
)

# Check if matrix is sparse
is_sparse = jc.sparse.is_sparse(cov_sparse)  # True

# Convert sparse to dense
cov_dense = jc.sparse.to_dense(cov_sparse)

# Sparse matrix inverse
cov_inv = jc.sparse.inv(cov_sparse)

# Sparse matrix-vector product
x = np.ones(len(mu))
y = jc.sparse.sparse_dot_vec(cov_inv, x)

# Sparse matrix-matrix product: A @ B @ C
result = jc.sparse.dot(cov_sparse, cov_inv, cov_sparse)

# Log determinant
sign, logdet = jc.sparse.slogdet(cov_sparse)
```

--------------------------------

### Format Code with Black

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Applies the Black code formatter to all Python files in the current directory and its subdirectories. Ensures consistent code style.

```bash
black .
```

--------------------------------

### Galaxy Bias Models in Python

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Demonstrates the implementation of various galaxy bias models within the jax_cosmo library. It covers constant linear bias, inverse growth bias, and the DES Y1 intrinsic alignment bias model, showing how to apply them as a function of redshift.

```python
import jax_cosmo as jc
import jax.numpy as np

cosmo = jc.Planck15()
z = np.linspace(0, 2, 50)

# Constant linear bias b(z) = b
bias_const = jc.bias.constant_linear_bias(1.5)
b = bias_const(cosmo, z)  # Returns array of 1.5

# Inverse growth bias: b(z) = b0 / D(z)
# Useful for conserved tracers
bias_inv = jc.bias.inverse_growth_linear_bias(1.0)
b = bias_inv(cosmo, z)  # Increases with redshift

# DES Y1 intrinsic alignment bias model
# A_IA * ((1+z)/(1+z0))^eta
ia_bias = jc.bias.des_y1_ia_bias(
    0.5,   # A: amplitude
    -1.0,  # eta: redshift slope
    0.62   # z0: pivot redshift
)
A_IA = ia_bias(cosmo, z)

# Per-bin bias (list of bias objects)
bin_biases = [
    jc.bias.constant_linear_bias(1.2),
    jc.bias.constant_linear_bias(1.4),
    jc.bias.constant_linear_bias(1.6),
]
```

--------------------------------

### Defining Gaussian Likelihood with JAX JIT

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Defines a Gaussian likelihood function that depends on cosmological parameters (Omega_c, sigma8). The function is decorated with `@jax.jit` for just-in-time compilation, enabling fast execution on GPUs. It computes the mean and covariance of angular Cls and returns the negative log-likelihood.

```python
data = mu # We create some fake data from the fiducial cosmology
# Let's define a parameter vector for Omega_cdm, sigma8, which we initialize 
# at the fiducial cosmology used to produce the data vector.
params = np.array([cosmo.Omega_c, cosmo.sigma8]) 

# Note the `jit` decorator for just in time compilation, this makes your code
# run fast on GPU :-)
@jax.jit
def likelihood(p):
  # Create a new cosmology at these parameters
  cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])

  # Compute mean and covariance of angular Cls
  m, C = jc.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, probes, sparse=True)

  # Return likelihood value assuming constant covariance, so we stop the gradient
  # at the level of the precision matrix, and we will not include the logdet term
  # in the likelihood
  P = jc.sparse.inv(jax.lax.stop_gradient(C))
  r = data - m
  return -0.5 * r.T @ jc.sparse.sparse_dot_vec(P, r)
```

--------------------------------

### Evaluate Jacobian and Plot Contours with Jax_cosmo

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This snippet shows how to evaluate the Jacobian of the mean power spectrum with respect to cosmological parameters using `jac_mean` and then compose the Fisher matrix. It also includes plotting contours derived from the Fisher matrix for visualization and comparison.

```python
dmu = jac_mean(params)

# For fun, we can also time it
%timeit jac_mean(params).block_until_ready()

# Now we can compose the Fisher matrix:
F_2 = jc.sparse.dot(dmu.T, jc.sparse.inv(cov), dmu)

# We can now plot contours obtained with this 
plot_contours(F, params, fill=False,color='black',lw=4);
plot_contours(F_2, params, fill=False, color='red', lw=4, linestyle='dashed');
xlabel('Omega_m')
ylabel('sigma8');
```

--------------------------------

### Create and Access Cosmology Objects in jax-cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Demonstrates how to instantiate the Cosmology object using predefined cosmological models (e.g., Planck15) or by providing custom parameters. It also shows how to access derived cosmological quantities.

```python
import jax_cosmo as jc
import jax.numpy as np

# Create cosmology with Planck 2015 parameters
cosmo = jc.Planck15()

# Override specific parameters
cosmo_modified = jc.Planck15(Omega_c=0.27, sigma8=0.8, h=0.7)

# Access cosmological parameters
print(f"Hubble parameter h: {cosmo.h}")           # 0.6774
print(f"Matter density Omega_m: {cosmo.Omega_m}") # Omega_b + Omega_c
print(f"Dark energy density: {cosmo.Omega_de}")   # 1 - Omega_m - Omega_k
print(f"Spectral index n_s: {cosmo.n_s}")         # 0.9667
print(f"sigma8: {cosmo.sigma8}")                  # 0.8159

# Create custom cosmology from scratch
from jax_cosmo.core import Cosmology
custom_cosmo = Cosmology(
    Omega_c=0.25,    # Cold dark matter density
    Omega_b=0.05,    # Baryon density
    h=0.67,          # Hubble constant / 100 km/s/Mpc
    n_s=0.96,        # Scalar spectral index
    sigma8=0.81,     # Matter fluctuation amplitude
    Omega_k=0.0,     # Curvature
    w0=-1.0,         # Dark energy equation of state
    wa=0.0           # Dark energy evolution parameter
)
```

--------------------------------

### Accessing Documentation for Radial Comoving Distance

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This code snippet shows how to access the documentation for the `radial_comoving_distance` function within the `jax_cosmo.background` module. This is useful for understanding the units and parameters of the function.

```python
# Not sure what are the units of the comoving distance? just ask:
jc.background.radial_comoving_distance?
```

--------------------------------

### Compare Growth Factor Calculation

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Compares the calculation of the growth factor as a function of scale factor between CCL and jax_cosmo. The growth factor describes the evolution of matter density perturbations in the universe.

```python
# Comparing the growth factor
plot(a, ccl.growth_factor(cosmo_ccl,a), label='CCL')
plot(a, background.growth_factor(cosmo_jax, a), '--', label='jax_cosmo')
legend()
xlabel('a')
ylabel('Growth factor')
```

--------------------------------

### Define Default Cosmology

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Creates a cosmology object using the default Planck 2015 parameters. This is the simplest way to initialize a cosmology for calculations.

```python
# Create a cosmology with default parameters
cosmo = jc.Planck15()
```

--------------------------------

### Evaluate n(z) and Redshift Bins in Python

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

This snippet demonstrates how to evaluate a normalized redshift distribution n(z) using the smail_nz function and how to create multiple redshift bins for tomographic analysis. It also shows how to check the normalization of the n(z) distribution.

```python
import jax.numpy as np
import jax_cosmo as jc

# Evaluate n(z) - returns normalized distribution
z = np.linspace(0, 3, 100)
pz = nz_smail(z)  # Shape: (100,)

# Check normalization (should be ~1.0)
norm = jc.scipy.integrate.romb(nz_smail, 0., 3.)

# Multiple redshift bins for tomography
nz_bins = [
    jc.redshift.smail_nz(1., 2., 0.3, gals_per_arcmin2=8.0),
    jc.redshift.smail_nz(1., 2., 0.6, gals_per_arcmin2=8.0),
    jc.redshift.smail_nz(1., 2., 0.9, gals_per_arcmin2=8.0),
]
```

--------------------------------

### Computing Jacobian of Mean Function with JAX

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Computes the Jacobian of the `mean_fn` with respect to its input parameters using JAX's forward-mode differentiation (`jax.jacfwd`). The resulting Jacobian function is JIT-compiled for performance.

```python
# We compute it's jacobian with JAX, and we JIT it for efficiency
jac_mean = jax.jit(jax.jacfwd(mean_fn))
```

--------------------------------

### Compare Linear Matter Power Spectrum

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Compares the linear matter power spectrum P(k) at a=1.0 for both CCL and jax_cosmo. The matter power spectrum describes the density fluctuations in the universe as a function of wavenumber k.

```python
from jax_cosmo.power import linear_matter_power

k = np.logspace(-3,-0.5)

#Let's have a look at the linear power
pk_ccl = ccl.linear_matter_power(cosmo_ccl, k, 1.0)
pk_jax = linear_matter_power(cosmo_jax, k/cosmo_jax.h, a=1.0)

loglog(k,pk_ccl,label='CCL')
loglog(k,pk_jax/cosmo_jax.h**3, '--', label='jax_cosmo')
legend()
xlabel('k [Mpc]')
ylabel('pk')
```

--------------------------------

### Define and Plot Redshift Distribution using Smail Profile

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Defines a redshift distribution n(z) using the Smail profile from jax_cosmo and plots it. This is commonly used for galaxy surveys.

```python
from jax_cosmo.redshift import smail_nz

# Let's define a redshift distribution
# with a Smail distribution with a=1, b=2, z0=1
nz = smail_nz(1.,2., 1.)

z = linspace(0,4,1024)
plot(z, nz(z))
xlabel(r'Redshift $z$');
title('Normalized n(z)')
```

--------------------------------

### Define Galaxy Tracers and Calculate Auto-Spectrum (CCL vs Jax Cosmo)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

This snippet defines number count tracers for both CCL and Jax Cosmo, incorporating a linear bias model. It then calculates the angular power spectrum for the auto-correlation of these tracers using both libraries for comparison.

```python
import jax_cosmo.constants as cst
from jax_cosmo import probes
from jax_cosmo.angular_cl import angular_cl
import pyccl as ccl

# Assuming cosmo_ccl, cosmo_jax, z, nz, ell, bias are defined elsewhere
# bias = constant_linear_bias(1.)

tracer_ccl_n = ccl.NumberCountsTracer(cosmo_ccl,
                                      has_rsd=False,
                                      dndz=(z, nz(z)),
                                      bias=(z, bias(cosmo_jax, z)))
tracer_jax_n = probes.NumberCounts([nz], bias)

cl_ccl = ccl.angular_cl(cosmo_ccl, tracer_ccl_n, tracer_ccl_n, ell)
cl_jax = angular_cl(cosmo_jax, ell, [tracer_jax_n])
```

--------------------------------

### Compute Fisher Matrix using Jax Auto-Differentiation

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Demonstrates how to compute the Fisher matrix for cosmological parameters using JAX's automatic differentiation capabilities within jax-cosmo. It involves defining a likelihood function, computing its Hessian, and optionally calculating it from the Jacobian of the mean power spectrum. This allows for analytical gradient and Hessian computation, enabling efficient parameter constraint estimation.

```python
import jax
import jax_cosmo as jc
import jax.numpy as np

# Setup
nzs = [
    jc.redshift.smail_nz(1., 2., 0.5, gals_per_arcmin2=10.),
    jc.redshift.smail_nz(1., 2., 1.0, gals_per_arcmin2=10.),
]
probes = [jc.probes.WeakLensing(nzs, sigma_e=0.26)]
ell = np.logspace(1, 3, 30)

# Fiducial cosmology
cosmo_fid = jc.Planck15()
params_fid = np.array([cosmo_fid.Omega_c, cosmo_fid.sigma8])

# Create mock data at fiducial cosmology
data, cov = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo_fid, ell, probes, sparse=True
)

# Define likelihood as function of parameters
@jax.jit  # JIT compile for GPU acceleration
def likelihood(p):
    cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])
    mu, C = jc.angular_cl.gaussian_cl_covariance_and_mean(
        cosmo, ell, probes, sparse=True
    )
    # Stop gradient on covariance for Fisher matrix
    P = jc.sparse.inv(jax.lax.stop_gradient(C))
    r = data - mu
    return -0.5 * r.T @ jc.sparse.sparse_dot_vec(P, r)

# Evaluate likelihood
logL = likelihood(params_fid)

# Compute gradient of likelihood
grad_logL = jax.grad(likelihood)(params_fid)

# Compute Fisher matrix via Hessian
# F = -d^2 logL / d theta^2
fisher_fn = jax.jit(jax.hessian(likelihood))
F = -fisher_fn(params_fid)

# Parameter constraints from Fisher matrix
cov_params = np.linalg.inv(F)
sigma_Omega_c = np.sqrt(cov_params[0, 0])
sigma_sigma8 = np.sqrt(cov_params[1, 1])

# Alternative: Compute Fisher from Jacobian of mean
# F_ab = sum_ij (dmu_i/dtheta_a) C^-1_ij (dmu_j/dtheta_b)
def mean_fn(p):
    cosmo = jc.Planck15(Omega_c=p[0], sigma8=p[1])
    return jc.angular_cl.angular_cl(cosmo, ell, probes).flatten()

# Jacobian of mean w.r.t. parameters
jac_mean = jax.jit(jax.jacfwd(mean_fn))
dmu = jac_mean(params_fid)  # Shape: (n_data, n_params)

# Fisher matrix from Jacobian
F_alt = jc.sparse.dot(dmu.T, jc.sparse.inv(cov), dmu)
```

--------------------------------

### Compute Background Cosmological Quantities in jax-cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Illustrates the use of functions within the `jax_cosmo.background` module to calculate key cosmological quantities such as comoving distance, Hubble parameter, and growth factors as a function of scale factor.

```python
import jax_cosmo as jc
import jax.numpy as np

cosmo = jc.Planck15()

# Define scale factors (a=1 is today, a=0.5 is z=1)
a = np.linspace(0.1, 1.0, 100)

# Radial comoving distance in Mpc/h
chi = jc.background.radial_comoving_distance(cosmo, a)
# chi at a=0.5 (z=1) is approximately 2300 Mpc/h

# Hubble parameter H(a) in km/s/(Mpc/h)
H = jc.background.H(cosmo, a)
# H(a=1) = H0 = 100 km/s/(Mpc/h)

# Angular diameter distance in Mpc/h
d_A = jc.background.angular_diameter_distance(cosmo, a)

# Transverse comoving distance (accounts for curvature)
f_k = jc.background.transverse_comoving_distance(cosmo, a)

# Linear growth factor D(a), normalized to D(a=1)=1
D = jc.background.growth_factor(cosmo, a)

# Growth rate f = dlnD/dlna
f = jc.background.growth_rate(cosmo, a)

# Matter density evolution Omega_m(a)
Omega_m_a = jc.background.Omega_m_a(cosmo, a)

# Dark energy equation of state w(a) = w0 + wa*(1-a)
w = jc.background.w(cosmo, a)

# E^2(a) where H(a) = H0 * E(a)
Esqr = jc.background.Esqr(cosmo, a)
```

--------------------------------

### Differentiating a Function with a JAX Container

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/design.md

Demonstrates how to compute the gradient of a function with respect to a JAX container object. This showcases the power of JAX's automatic differentiation on custom data structures.

```python
# Define a likelihood, function of the redshift distribution
def likelihood(nz):
  ... # some computation that uses this nz
  return likelihood_value

nz = gaussian_nz(1., 0.1)
print(jax.grad(likelihood)(nz))
```

--------------------------------

### Compare Non-linear Matter Power Spectrum

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Compares the non-linear matter power spectrum P(k) at a=1.0 for both CCL and jax_cosmo. This comparison includes non-linear effects which are important at smaller scales.

```python
from jax_cosmo.power import nonlinear_matter_power

k = np.logspace(-3,1)
#Let's have a look at the non linear power
pk_ccl = ccl.nonlin_matter_power(cosmo_ccl, k, 1.0)
pk_jax = nonlinear_matter_power(cosmo_jax, k/cosmo_jax.h, a=1.0)

loglog(k,pk_ccl,label='CCL')
loglog(k,pk_jax/cosmo_jax.h**3, '--', label='jax_cosmo')
legend()
xlabel('k [Mpc]')
ylabel('pk')
```

--------------------------------

### Define Redshift Distributions in jax-cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Shows how to define and use various parameterized redshift distribution functions (n(z)) for galaxy samples, including Smail, Fu et al., KDE-based, and delta function distributions. These are fully differentiable with respect to their parameters.

```python
import jax_cosmo as jc
import jax.numpy as np

# Smail distribution: n(z) = z^a * exp(-(z/z0)^b)
# Parameters: (a, b, z0)
nz_smail = jc.redshift.smail_nz(
    1.0,   # a: power law index
    2.0,   # b: exponential index
    0.5,   # z0: characteristic redshift
    gals_per_arcmin2=10.0,  # Galaxy number density
    zmax=3.0                # Maximum redshift
)

# Fu et al. distribution: n(z) = (z^a + z^(ab)) / (z^b + c)
nz_fu = jc.redshift.fu_nz(
    0.4710,  # a parameter
    5.1843,  # b parameter
    0.7259,  # c parameter
    gals_per_arcmin2=30.0
)

# KDE-based distribution from catalog
redshift_catalog = np.array([0.1, 0.3, 0.5, 0.7, 0.9, 1.1])
weights = np.ones(len(redshift_catalog))
nz_kde = jc.redshift.kde_nz(
    redshift_catalog,
    weights,
    bw=0.1,  # Bandwidth for KDE
    gals_per_arcmin2=5.0
)

# Delta function distribution (single redshift plane)
nz_delta = jc.redshift.delta_nz(1.0)  # Sources at z=1
```

--------------------------------

### Compute Likelihood, Gradients, and Hessians with jax-cosmo

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/README.md

This snippet demonstrates how to compute the mean and covariance of angular Cls using jax-cosmo, then calculate the Gaussian log-likelihood. It further shows how to obtain the gradients and Hessians of the likelihood function with respect to cosmological parameters using JAX's autodiff capabilities.

```python
import jax
import jax_cosmo

def likelihood(cosmo):
  # Compute mean and covariance of angular Cls, for specific probes
  mu, cov = jax_cosmo.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, probes)
  # Return likelihood value
  return jax_cosmo.likelihood.gaussian_log_likelihood(data, mu, cov)

# Compute derivatives of the likelihood with respect to cosmological parameters
g = jax.grad(likelihood)(cosmo)

# Compute Fisher matrix of cosmological parameters
F = - jax.hessian(likelihood)(cosmo)
```

--------------------------------

### Compare Linear Growth Rate Calculation

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Compares the calculation of the linear growth rate (f) as a function of scale factor between CCL and jax_cosmo. The growth rate is the logarithmic derivative of the growth factor with respect to log(a).

```python
# Comparing linear growth rate
plot(a, ccl.growth_rate(cosmo_ccl,a), label='CCL')
plot(a, background.growth_rate(cosmo_jax, a), '--', label='jax_cosmo')
legend()
xlabel('a')
ylabel('growth rate')
```

--------------------------------

### Define Galaxy Tracers and Calculate Cross-Spectrum (CCL vs Jax Cosmo)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

This snippet calculates the angular cross-power spectrum between galaxy density and shape-position observables. It defines the necessary tracers for both CCL and Jax Cosmo and then computes the spectra.

```python
# Assuming cosmo_ccl, cosmo_jax, tracer_ccl, tracer_ccl_n, ell, tracer_jax, tracer_jax_n are defined elsewhere

cl_ccl = ccl.angular_cl(cosmo_ccl, tracer_ccl, tracer_ccl_n, ell)
cl_jax = angular_cl(cosmo_jax, ell, [tracer_jax, tracer_jax_n])
```

--------------------------------

### Compare Weak Lensing Angular Power Spectra (Cls)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Computes and compares the angular power spectra (Cls) for weak lensing between CCL and jax_cosmo. This involves defining lensing tracers and calculating Cls over a range of multipoles (ell).

```python
from jax_cosmo.angular_cl import angular_cl
from jax_cosmo import probes

# Let's first compute some Weak Lensing cls
tracer_ccl = ccl.WeakLensingTracer(cosmo_ccl, (z, nz(z)), use_A_ia=False)
tracer_jax = probes.WeakLensing([nz])

ell = np.logspace(0.1,3)

cl_ccl = ccl.angular_cl(cosmo_ccl, tracer_ccl, tracer_ccl, ell)
cl_jax = angular_cl(cosmo_jax, ell, [tracer_jax])

loglog(ell, cl_ccl,label='CCL')
loglog(ell, cl_jax[0], '--', label='jax_cosmo')

legend()
xlabel(r'$ell$')
ylabel(r'Lensing angular $C_ell$')
```

--------------------------------

### Run Specific Test File with pytest

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/CLAUDE.md

Executes tests from a specific file within the 'tests' directory. Useful for targeted debugging or development.

```bash
pytest tests/test_*.py
```

--------------------------------

### Defining Lensing and Clustering Probes

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This code defines a list of cosmological probes, including WeakLensing and NumberCounts, using previously defined redshift distributions (`nzs`). It specifies parameters like intrinsic shear dispersion (`sigma_e`) for weak lensing and galaxy bias for number counts.

```python
nzs = [nz1, nz2]

probes = [ jc.probes.WeakLensing(nzs, sigma_e=0.26), 
           jc.probes.NumberCounts(nzs, jc.bias.constant_linear_bias(1.)) ]
```

--------------------------------

### Converting Sparse Covariance to Dense Representation

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This snippet demonstrates the initial step of converting a sparse covariance matrix to a dense representation for visualization or further computation. It sets up a matplotlib figure for plotting the dense covariance matrix.

```python
figure(figsize=(10,10))
# Here we convert the covariance matrix from sparse to dense reprensentation
```

--------------------------------

### Compute Angular Power Spectra (Cls) in Python

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

This section details how to compute tomographic angular power spectra (C_l) for defined cosmological probes using the Limber approximation in jax_cosmo. It covers calculating signal, noise, and Gaussian covariances, including sparse and dense representations.

```python
import jax_cosmo as jc
import jax.numpy as np

cosmo = jc.Planck15()

# Define probes
nzs = [
    jc.redshift.smail_nz(1., 2., 0.5, gals_per_arcmin2=10.),
    jc.redshift.smail_nz(1., 2., 1.0, gals_per_arcmin2=10.),
]
probes = [
    jc.probes.WeakLensing(nzs, sigma_e=0.26),
    jc.probes.NumberCounts(nzs, jc.bias.constant_linear_bias(1.))
]

# Define multipole range
ell = np.logspace(1, 3, 50)  # l = 10 to 1000

# Compute angular Cls
# Returns shape (n_cls, n_ell) where n_cls = n_tracers*(n_tracers+1)/2
cls = jc.angular_cl.angular_cl(cosmo, ell, probes)
# For 2 probes with 2 bins each: 4 tracers -> 10 Cls

# Cls ordering: first by probe, then by bin
# [WL1-WL1, WL1-WL2, WL1-NC1, WL1-NC2, WL2-WL2, WL2-NC1, WL2-NC2, NC1-NC1, NC1-NC2, NC2-NC2]

# Compute noise power spectrum (shape noise for WL, shot noise for NC)
noise_cls = jc.angular_cl.noise_cl(ell, probes)

# Compute signal, noise, and Gaussian covariance together
mu, cov = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo, ell, probes,
    f_sky=0.25,   # Sky fraction
    sparse=True   # Use sparse representation for memory efficiency
)
# mu: flattened signal vector, shape (n_cls * n_ell,)
# cov: sparse covariance matrix

# Non-sparse covariance (dense matrix)
mu_dense, cov_dense = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo, ell, probes, f_sky=0.25, sparse=False
)
# cov_dense shape: (n_cls * n_ell, n_cls * n_ell)
```

--------------------------------

### Computing Fisher Matrix using JAX Hessian

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Computes the Fisher matrix by taking the negative Hessian of the likelihood function with respect to the cosmological parameters. The `jax.hessian` function is used, and the entire operation is JIT-compiled for efficiency.

```python
# Compile a function that computes the Hessian of the likelihood
hessian_loglik = jax.jit(jax.hessian(likelihood))

# Evalauate the Hessian at fiductial cosmology to retrieve Fisher matrix
# This is a bit slow at first....
F = - hessian_loglik(params)
```

--------------------------------

### Plotting Redshift Distributions

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This code plots two Smail redshift distributions, `nz1` and `nz2`, against redshift `z`. It uses `numpy.linspace` to create the redshift array and matplotlib for plotting. The distributions are normalized and callable.

```python
# And let's plot it
z = np.linspace(0,5,256)

# Redshift distributions are callable, and they return the normalized distribution
plot(z, nz1(z), label='z0=1.')
plot(z, nz2(z), label='z0=0.5')
legend();
xlabel('Redshift $z$');
```

--------------------------------

### Plotting Angular Power Spectra Comparison (CCL vs Jax Cosmo)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

This code snippet visualizes the comparison between angular power spectra calculated by CCL and Jax Cosmo. It uses log-log plots to display the clustering power spectra and includes labels for clarity.

```python
from matplotlib.pyplot import loglog, legend, xlabel, ylabel

# Assuming ell, cl_ccl, cl_jax are defined from previous snippet
loglog(ell, cl_ccl,label='CCL')
loglog(ell, cl_jax[0], '--', label='jax_cosmo')

legend()
xlabel(r'$\ell$')
ylabel(r'Clustering angular $C_\ell$')
```

--------------------------------

### Computing Angular Power Spectra (Cls)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This snippet computes the tomographic angular power spectra (`Cls`) for the defined cosmological probes and a given cosmology (`cosmo`). It uses the `jax_cosmo.angular_cl.angular_cl` function with a specified range of multipole moments (`ell`). Computations are performed under the Limber approximation.

```python
ell = np.logspace(1,3) # Defines a range of \ell 

# And compute the data vector
cls = jc.angular_cl.angular_cl(cosmo, ell, probes)
```

--------------------------------

### Compute Gaussian Log-Likelihood with Jax_Cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Calculates the Gaussian log-likelihood for cosmological parameter inference using `jax_cosmo.likelihood`. This function requires observed data, theoretical predictions (mean), and a covariance matrix. It can optionally include the log-determinant term of the covariance matrix and supports different methods for inverting dense covariance matrices.

```python
import jax_cosmo as jc
import jax.numpy as np

cosmo = jc.Planck15()

# Setup probes
nzs = [jc.redshift.smail_nz(1., 2., 0.5, gals_per_arcmin2=10.)]
probes = [jc.probes.WeakLensing(nzs, sigma_e=0.26)]
ell = np.logspace(1, 3, 30)

# Compute theory prediction and covariance
mu, cov = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo, ell, probes, sparse=True
)

# Simulate data (here using theory as mock data)
data = mu

# Gaussian log-likelihood with sparse covariance
logL = jc.likelihood.gaussian_log_likelihood(
    data, mu, cov,
    include_logdet=True  # Include -0.5*log|C| term
)

# Without log determinant (faster, for fixed covariance)
logL_no_det = jc.likelihood.gaussian_log_likelihood(
    data, mu, cov,
    include_logdet=False
)

# Dense covariance matrix version
_, cov_dense = jc.angular_cl.gaussian_cl_covariance_and_mean(
    cosmo, ell, probes, sparse=False
)
logL_dense = jc.likelihood.gaussian_log_likelihood(
    data, mu, cov_dense,
    inverse_method='cholesky'  # or 'inverse'
)
```

--------------------------------

### Computing Gaussian Angular Cl Covariance and Mean

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This snippet computes the Gaussian approximation for the angular power spectrum covariance matrix and the mean power spectra. It uses `jax_cosmo.angular_cl.gaussian_cl_covariance_and_mean` with the `sparse=True` option for efficiency. The output `mu` represents the mean and `cov` represents the covariance.

```python
mu, cov = jc.angular_cl.gaussian_cl_covariance_and_mean(cosmo, ell, probes, sparse=True);
```

--------------------------------

### Calculate Matter Power Spectra with Jax_Cosmo

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Computes linear and non-linear matter power spectra using the `jax_cosmo.power` module. It supports various options for transfer functions (e.g., Eisenstein_Hu) and non-linear prescriptions (e.g., halofit). Inputs include a cosmology object, wave numbers (k), and optionally scale factors (a).

```python
import jax_cosmo as jc
import jax.numpy as np
import jax_cosmo.power as power
import jax_cosmo.transfer as transfer

cosmo = jc.Planck15()

# Wave numbers in h/Mpc
k = np.logspace(-3, 1, 100)  # k = 0.001 to 10 h/Mpc

# Scale factors
a = np.array([1.0, 0.5, 0.33])  # z = 0, 1, 2

# Linear matter power spectrum P(k) in (Mpc/h)^3
pk_linear = power.linear_matter_power(cosmo, k, a=1.0)

# With custom transfer function
pk_linear_eh = power.linear_matter_power(
    cosmo, k, a=1.0,
    transfer_fn=transfer.Eisenstein_Hu
)

# Nonlinear power spectrum using halofit
pk_nonlinear = power.nonlinear_matter_power(
    cosmo, k, a=1.0,
    transfer_fn=transfer.Eisenstein_Hu,
    nonlinear_fn=power.halofit
)

# Direct halofit call with prescription choice
pk_halofit = power.halofit(
    cosmo, k, a=1.0,
    transfer_fn=transfer.Eisenstein_Hu,
    prescription='takahashi2012'  # or 'smith2003'
)

# Primordial power spectrum P(k) = k^n_s
pk_primordial = power.primordial_matter_power(cosmo, k)

# Transfer function T(k)
Tk = transfer.Eisenstein_Hu(cosmo, k, type='eisenhu_osc')
# type options: 'eisenhu' (no BAO), 'eisenhu_osc' (with BAO wiggles)
```

--------------------------------

### JAX Container Class for Differentiable Parameters

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/design.md

Defines a base 'container' class to hold parameters for JAX tracing. It separates traceable parameters into `self.params` and static configuration into `self.config`, enabling automatic differentiation of functions using these containers.

```python
class container(object):
  def __init__(self, *args, **kwargs):
    self.params = list(args)
    self.config = kwargs

class gaussian_nz(container):

  def __init__(self, z0, sigma, zmax=10, **kwargs):
    super(gaussian_nz, self).__init__(z0, sigma, # Traceable parameters
                                      zmax=zmax, **kwargs) # non-traceable configuration

  def __call__(self, z):
    z0, sigma = self.params
    return np.clip(exp(-0.5*( z - z0)**2/sigma**2),
                   0., self.config['zmax'])
```

--------------------------------

### Plotting Shape-Position Angular Power Spectra Comparison (CCL vs Jax Cosmo)

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

This code snippet visualizes the comparison of angular power spectra for shape-position observables between CCL and Jax Cosmo. It uses log-log plots and appropriate labels to represent the calculated spectra.

```python
import numpy as np
from matplotlib.pyplot import loglog, legend, xlabel, ylabel

ell = np.logspace(1,3)

# Assuming cl_ccl, cl_jax are defined from previous snippet
loglog(ell, cl_ccl,label='CCL')
loglog(ell, cl_jax[1], '--', label='jax_cosmo')

legend()
xlabel(r'$\ell$')
ylabel(r'shape-position angular $C_\ell$')
```

--------------------------------

### Plotting Logarithm of Covariance Matrix with Matplotlib

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

Visualizes the logarithm of a sparse covariance matrix, converted to a dense format, using Matplotlib's imshow function. It adds a small constant to avoid issues with log(0).

```python
imshow(np.log10(jc.sparse.to_dense(cov)+1e-11),cmap='gist_stern');
```

--------------------------------

### Test Radial Comoving Distance Calculation

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/CCL_comparison.ipynb

Calculates and plots the radial comoving distance as a function of scale factor using both CCL and jax_cosmo. This verifies the consistency of distance calculations between the two libraries.

```python
# Test array of scale factors
a = np.linspace(0.01, 1.)

# Testing the radial comoving distance
chi_ccl = ccl.comoving_radial_distance(cosmo_ccl, a)
chi_jax = background.radial_comoving_distance(cosmo_jax, a)/cosmo_jax.h

plot(a, chi_ccl, label='CCL')
plot(a, chi_jax, '--', label='jax_cosmo')
legend()
xlabel('a')
ylabel('radial comoving distance [Mpc]')
```

--------------------------------

### Compute Angular Power Spectra with Custom Functions

Source: https://context7.com/differentiableuniverseinitiative/jax_cosmo/llms.txt

Calculates angular power spectra (cls) using custom transfer and nonlinear functions. It requires a cosmology object, angular power spectrum bins (ell), and probe definitions. The function `angular_cl` from `jax_cosmo.angular_cl` is used, allowing specification of `transfer_fn` and `nonlinear_fn`.

```python
import jax_cosmo as jc
import jax_cosmo.transfer as tklib
import jax_cosmo.power as power

# Assuming cosmo, ell, and probes are defined elsewhere
cls_custom = jc.angular_cl.angular_cl(
    cosmo, ell, probes,
    transfer_fn=tklib.Eisenstein_Hu,
    nonlinear_fn=power.halofit
)
```

--------------------------------

### Plotting Angular Power Spectrum

Source: https://github.com/differentiableuniverseinitiative/jax_cosmo/blob/master/docs/notebooks/jax-cosmo-intro.ipynb

This code plots the first angular power spectrum (`cls[0]`) against the multipole moment `ell` on a log-log scale. It labels the axes and sets a title for clarity, visualizing the lensing auto-spectrum of the first redshift bin.

```python
# This is for instance the first bin auto-spectrum 
loglog(ell, cls[0])
ylabel(r'$C_\ell$')
xlabel(r'$\ell$');
title(r'Angular $C_\ell$');
```