DOCSEARCHbyAlgolia
Submit Docs

jax-cosmo

repository·master·Indexed 42 days ago

https://github.com/differentiableuniverseinitiative/jax_cosmo

A differentiable cosmology library built on JAX for computing cosmological observables, likelihoods, and their derivatives. It enables efficient parameter estimation and Fisher matrix analysis using automatic differentiation (gradients and Hessians) without finite differences. Key features include the computation of angular Cls via Limber approximation, background cosmology functions, and support for probes such as Weak Lensing and Number Counts.

Tokens
4.2K
Snippets
16
Records
18
Agent score
53%

What's inside jax-cosmo

  1. Understand the jax-cosmo module structure

    master

    The library is organized into several core modules that handle different aspects of cosmological computation:

    • Cosmology: A class that stores cosmological parameters; it is an instance of the container class.
    • background: Functions for computing various background-related quantities.
    • transfer: A library of transfer functions (e.g., EisensteinHu).
    • probes: Definitions of various probes (e.g., lensing, number counts). Probes are container objects that can be differentiated and typically gather redshift distributions and other parameters.
    • angular_cl: Contains Limber integration code and covariance tools.
    • non_linear: (Planned) For non-linear corrections like halofit.

  2. Use the container class for differentiable parameters

    master

    The container class is a generic structure used to hold parameters that need to be traced by JAX (differentiable parameters) alongside static configuration arguments.

    • Traceable parameters: These are passed as positional arguments to __init__ and are stored in self.params. Anything that needs to be differentiable (e.g., mean redshift, cosmological parameters) should be placed here.
    • Configuration arguments: These are passed as keyword arguments to __init__ and are stored in self.config. These are non-traceable/static (e.g., zmax, numerical precision flags).

    By inheriting from container, your objects become compatible with JAX pytrees, allowing you to compute gradients with respect to the object's parameters using jax.grad.

    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'])

# Usage for differentiation
def likelihood(nz):
  # ... computation using nz ...
  return likelihood_value

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

  3. How 2pt functions are computed

    master

    In jax-cosmo, all two-point (2pt) functions are currently computed using the Limber approximation.

    Following the structure used by CCL, each tracer (such as lensing or number counts) is characterized by three components:

    1. A radial kernel function.
    2. An $\ell$-dependent prefactor.
    3. A transfer function.

    Probes are implemented as container objects, meaning they can be differentiated and can aggregate multiple redshift distributions and other necessary parameters.

  4. Compute likelihoods and derivatives with jax-cosmo

    master

    jax-cosmo is designed to work with JAX to provide differentiable cosmology. You can use it to compute the mean and covariance of angular $C_ ext{l}$s, calculate Gaussian log-likelihoods, and then use JAX's automatic differentiation to compute gradients and Hessians (Fisher matrices) without finite differences.

    Typical workflow:

    1. Use jax_cosmo.angular_cl.gaussian_cl_covariance_and_mean to get the mean ($\mu$) and covariance ($\text{cov}$) of angular $C_ ext{l}$s.
    2. Use jax_cosmo.likelihood.gaussian_log_likelihood to compute the likelihood value.
    3. Use jax.grad to compute derivatives of the likelihood with respect to cosmological parameters.
    4. Use jax.hessian to compute the Fisher matrix.
    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)

  7. Define a Cosmology and compute background quantities

    master

    Cosmologies can be instantiated using predefined models (like jc.Planck15()) or by overriding default parameters. Background quantities are computed using the jax_cosmo.background module, which typically takes a cosmology object and a scale factor a as arguments.

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

# Alternatively, override defaults
cosmo_modified = jc.Planck15(h=0.7)

# Access parameters directly
print(cosmo.h)

# Compute background quantities (e.g., radial comoving distance)
import jax.numpy as np
a = np.linspace(0.01, 1.)
chi = jc.background.radial_comoving_distance(cosmo, a)

  8. Define probes and compute angular $C_\ell$

    master

    Probes (e.g., lensing or galaxy clustering) are defined in jax_cosmo.probes. They contain parameters like redshift distributions and galaxy bias. Angular power spectra ($C_\ell$) are computed using jax_cosmo.angular_cl.angular_cl, which takes a cosmology, an $\ell$ vector, and a list of probes. It computes all auto- and cross-correlations between all redshift bins of all probes.

    # 1. Define redshift bins
nzs = [nz1, nz2]

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

# 3. Compute angular Cls
ell = np.logspace(1, 3)
cls = jc.angular_cl.angular_cl(cosmo, ell, probes)

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

  11. Compute Fisher matrix via Jacobian of the mean

    master

    An alternative way to compute the Fisher matrix is using the formula $F_{\alpha, \beta} = \sum_{i,j} \frac{d \mu_i}{d \theta_\alpha} C^{-1}{i,j} \frac{d \mu_j}{d \theta\beta}$. You can use jax.jacfwd to compute the Jacobian of the mean signal with respect to the parameters.

    # 1. Define a function that returns the mean signal vector
def mean_fn(p):
    cosmo_p = jc.Planck15(Omega_c=p[0], sigma8=p[1])
    m = jc.angular_cl.angular_cl(cosmo_p, ell, probes)
    return m.flatten()

# 2. Compute the Jacobian using JAX
jac_mean = jax.jit(jax.jacfwd(mean_fn))
dmu = jac_mean(params)

# 3. Compose the Fisher matrix using sparse operations
# F = dmu.T @ inv(cov) @ dmu
F_2 = jc.sparse.dot(dmu.T, jc.sparse.inv(cov), dmu)

  12. Compute Fisher matrices using Automatic Differentiation

    master

    Because jax-cosmo is built on JAX, you can analytically compute Fisher matrices using the Hessian of a likelihood function without resorting to finite differences. This is done by applying jax.hessian to a likelihood function that depends on cosmological parameters.

    # Define a likelihood function that depends on parameters p
@jax.jit
def likelihood(p):
    # Create a new cosmology at these parameters
    cosmo_p = jc.Planck15(Omega_c=p[0], sigma8=p[1])
    
    # Compute mean and covariance
    m, C = jc.angular_cl.gaussian_cl_covariance_and_mean(cosmo_p, ell, probes, sparse=True)
    
    # Use stop_gradient on the covariance to avoid differentiating through it
    P = jc.sparse.inv(jax.lax.stop_gradient(C))
    r = data - m
    return -0.5 * r.T @ jc.sparse.sparse_dot_vec(P, r)

# Compute the Hessian (Fisher Matrix) analytically
hessian_loglik = jax.jit(jax.hessian(likelihood))
params = np.array([cosmo.Omega_c, cosmo.sigma8])
F = - hessian_loglik(params)