Source code for wale.CosmologyModel

import numpy as np
from scipy.integrate import quad
from scipy.optimize import brentq
import pyccl as ccl
from scipy.interpolate import interp1d
from scipy.integrate import trapezoid

from pyccl.halos.pk_4pt import halomod_Tk3D_4h


[docs] class Cosmology_function: """ A class for initializing and handling cosmological computations using PyCCL. This class encapsulates key cosmological parameters and provides methods to compute quantities such as the Hubble parameter, comoving distances, redshift inversions, non-linear matter power spectrum, and lensing weights. Internally, it constructs a `pyccl.Cosmology` object for use with the Core Cosmology Library (CCL). """ def __init__(self, h, Oc, Ob, w=-1.0, wa=0.0, **kwargs): """ Initialize a cosmological model with the given parameters. Parameters ---------- h : float Dimensionless Hubble parameter (i.e., H0 = 100 * h km/s/Mpc). Oc : float Cold dark matter density fraction, Ω_c. Ob : float Baryon density fraction, Ω_b. w : float, optional Dark energy equation-of-state parameter at present (default is -1.0). wa : float, optional Evolution parameter of dark energy equation-of-state (default is 0.0). **kwargs : dict, optional Additional optional parameters: - 'sigma8' : float, RMS density fluctuations at 8 Mpc/h. - 'As' : float, amplitude of primordial scalar fluctuations. - 'ns' : float, scalar spectral index (default is 0.973). - 'kmin' : float, minimum wavenumber for calculations (default is 1e-4). - 'kmax' : float, maximum wavenumber (default is 10.0). - 'dk' : float, step size for k-array (default is 0.1). Raises ------ ValueError If neither 'sigma8' nor 'As' is specified. """ self.h = h self.Oc = Oc self.Ob = Ob self.w = w self.wa = wa self.H0 = 100.0 * h # Hubble constant in km/s/Mpc self.Om = Oc + Ob self.ns = kwargs.get("ns", 0.973) # scalar spectral index # power normalization (either sigma8 or As must be set) self.sig8 = kwargs.get("sigma8", None) self.As = kwargs.get("As", None) if self.sig8 is None and self.As is None: raise ValueError("Please specify either sigma8 or As.") self.speed_light = 299792.458 self.kmin = kwargs.get("kmin", 1e-4) self.kmax = kwargs.get("kmax", 10.0) self.dk = kwargs.get("dk", 0.1) self.k = np.arange(self.kmin, self.kmax, self.dk) self.nk = len(self.k) def _set_params(self): if self.sig8 is not None: return ccl.Cosmology( h=self.h, Omega_c=self.Oc, Omega_b=self.Ob, sigma8=self.sig8, n_s=self.ns, w0=self.w, wa=self.wa, transfer_function="boltzmann_camb", ) elif self.As is not None: return ccl.Cosmology( h=self.h, Omega_c=self.Oc, Omega_b=self.Ob, A_s=self.As, n_s=self.ns, w0=self.w, wa=self.wa, transfer_function="boltzmann_camb", ) self.cosmoccl = _set_params(self) def _E(self, z): """ Compute the dimensionless Hubble expansion rate E(z) = H(z)/H0. This method uses the CPL (Chevallier-Polarski-Linder) parameterization for the dark energy equation of state: w(a) = w0 + wa * (1 - a), where a = 1 / (1 + z) is the scale factor. The expansion rate is given by: E(z)^2 = Ω_m (1 + z)^3 + Ω_k (1 + z)^2 + Ω_de (1 + z)^{3(1 + w0 + wa)} exp[-3wa z / (1 + z)] Parameters ---------- z : float or array_like Redshift value(s) at which to compute the dimensionless expansion rate. Returns ------- E : float or ndarray The dimensionless Hubble parameter E(z) = H(z)/H0. Notes ----- - Assumes a flat or curved universe with Ω_k = 0 by default. - This function does not use the pyccl Cosmology object and computes E(z) analytically. """ self.Ok = 0.0 self.Ode = 1.0 - self.Om - self.Ok # CPL parameterization for w(a)=w0 + wa(1−a) return np.sqrt( self.Om * (1 + z) ** 3 + self.Ok * (1 + z) ** 2 + self.Ode * (1 + z) ** (3 * (1 + self.w + self.wa)) * np.exp(-3 * self.wa * z / (1 + z)) )
[docs] def get_H(self, z): """ Compute the Hubble parameter H(z) in km/s/Mpc. Parameters ---------- z : float or array_like Redshift(s) at which to evaluate H(z). Returns ------- H : float or ndarray Hubble parameter at redshift z, in units of km/s/Mpc. """ return self.H0 * self._E(z)
[docs] def get_chi(self, z): """ Compute the comoving radial distance χ(z) in Mpc. This function numerically integrates the inverse of the dimensionless Hubble parameter E(z) = H(z)/H0 to compute: χ(z) = ∫₀ᶻ (c / H(z')) dz' Parameters ---------- z : float or array_like Redshift(s) at which to evaluate the comoving distance. Returns ------- chi : float or ndarray Comoving radial distance(s) to redshift z, in units of Mpc. Notes ----- - Uses `scipy.integrate.quad` for numerical integration. - Assumes a flat universe with constant c = 299792.458 km/s. - The integration is done individually for each redshift if an array is given. """ H0 = self.H0 # km/s/Mpc c = self.speed_light # speed of light in km/s def integrand(zp): return 1.0 / self._E(zp) # scalar case if np.isscalar(z): integral, _ = quad(integrand, 0.0, z) return (c / H0) * integral # array case z = np.asanyarray(z) chi = np.empty_like(z, dtype=float) for i, zi in enumerate(z): integral, _ = quad(integrand, 0.0, zi) chi[i] = (c / H0) * integral return chi
[docs] def get_z_from_chi(self, chi_target, z_max=10.0, tol=1e-8): """ Invert χ(z) to find z such that χ(z)=chi_target. Parameters ---------- chi_target : float or array_like Comoving distance(s) in Mpc. z_max : float Maximum redshift to search within (must be large enough that χ(z_max) > chi_target). tol : float Desired precision in z. Returns ------- z : float or ndarray Redshift(s) corresponding to chi_target. """ # root function for a given target χ def f(z, chi_t): return self.get_chi(z) - chi_t def find_root(chi_t): # ensure bracket: χ(0)=0, χ(z_max)>chi_t chi_max = self.get_chi(z_max) if chi_t < 0 or chi_t > chi_max: raise ValueError( f"chi_target={chi_t:.3f} outside [0, {chi_max:.3f}] Mpc (z_max={z_max})" ) return brentq(f, 0.0, z_max, args=(chi_t,), xtol=tol) if np.isscalar(chi_target): return find_root(chi_target) chi_arr = np.asanyarray(chi_target) z_arr = np.empty_like(chi_arr, dtype=float) for i, chi_t in enumerate(chi_arr): z_arr[i] = find_root(chi_t) return z_arr
[docs] def get_nonlinear_pk(self, z, ks=None): """ Compute the non-linear matter power spectrum P(k, z) using HALOFIT. Parameters ---------- z : float Redshift at which to evaluate the power spectrum. ks : array_like, optional Wavenumber values in h/Mpc. If None, uses the default self.k array. Returns ------- Pnl : ndarray Non-linear matter power spectrum at redshift z for the given ks, in (Mpc/h)^3. """ # default k-grid if ks is None: ks = self.k # np.logspace(np.log10(self.kmin), # np.log10(self.kmax), # self.nk) else: ks = np.atleast_1d(ks) a = 1.0 / (1.0 + z) # compute HALOFIT non-linear power Pnl = ccl.nonlin_matter_power(self.cosmoccl, self.k, a) return Pnl
[docs] def get_lensing_weight_array(self, chis, chi_source): """ Compute lensing weight array W(chi) for a single source plane at chi_source. Parameters ---------- chis : array_like Comoving radial distances (chi) at which to compute lensing weights. chi_source : float Comoving distance to the source plane (single redshift). Returns ------- z_values : ndarray Redshift values corresponding to chis, computed via inversion. lensing_weight : ndarray Lensing weight W(chi) evaluated at each input chi. """ z_values = self.get_z_from_chi(chis) lensing_weight = np.zeros_like(chis) for i in range(len(chis)): z = z_values[i] lensing_weight[i] = ( 1.5 * self.Om * (self.speed_light**-2.0) * ((self.H0) ** 2.0) * chis[i] * (1 - (chis[i] / chi_source)) * (1 + z) ) return z_values, lensing_weight
[docs] def get_lensing_weight_array_nz(self, chis, z_nz, n_z): """ Compute lensing weight W(chi) using a redshift distribution n(z), without relying on pyccl. Parameters ---------- chis : array_like Comoving radial distances (chi) at which to compute lensing weights. z_nz : array_like Redshift values at which the distribution n(z) is defined. n_z : array_like Source redshift distribution n(z), not necessarily normalized. Returns ------- z_values : ndarray Redshift values corresponding to the input chis, estimated via interpolation. lensing_weight : ndarray Lensing weight values W(chi) computed using the n(z) distribution. """ n_norm = n_z / trapezoid(n_z, z_nz) # normalized n(z) # Get chi(z) using self.get_chi chi_nz = self.get_chi(z_nz) # in Mpc z_of_chi_interp = interp1d(chi_nz, z_nz, bounds_error=False, fill_value=0.0) # Convert n(z) → n(chi) dz_dchi = np.gradient(z_nz, chi_nz) n_chi = n_norm * dz_dchi n_chi_interp = interp1d(chi_nz, n_chi, bounds_error=False, fill_value=0.0) # Prepare lensing weights lensing_weight = np.zeros_like(chis) z_values = z_of_chi_interp(chis) for i, chi in enumerate(chis): # Integrate over chi' > chi chi_prime = chi_nz[chi_nz > chi] if chi_prime.size == 0: lensing_weight[i] = 0.0 continue integrand = n_chi_interp(chi_prime) * (chi_prime - chi) / chi_prime integral = trapezoid(integrand, chi_prime) lensing_weight[i] = ( (1.5 * self.Om * (self.H0) ** 2 / self.speed_light**2) * chi * (1 + z_values[i]) * integral ) return z_values, lensing_weight
[docs] def get_lensing_weight(self, chis, chisource, **kwargs): """ Dispatch method to compute lensing weights depending on source distribution. Parameters ---------- chis : array_like Comoving radial distances where lensing weights are evaluated. chisource : float Source-plane comoving distance (used only if no n(z) is provided). **kwargs : Optional arguments: - nz_file : str, optional If provided, uses redshift distribution from file (not yet implemented). Returns ------- z_values : ndarray Redshifts corresponding to chis. lensing_weight : ndarray Lensing weights W(chi) using either a single plane or n(z). """ if kwargs.get("nz_file") is not None: return self.get_lensing_weight_array_nz(chis) else: return self.get_lensing_weight_array(chis, chisource)