"""Linear matter power spectrum via CAMB (Lewis, Challinor & Lasenby 2000)."""
import numpy as np
import jax
import jax.numpy as jnp
[docs]
def rho_critical_0() -> float:
"""Critical matter density at z=0 for H₀ = 100 km/s/Mpc, in h-units.
.. math::
\\rho_{\\mathrm{crit},0} = \\frac{3H_{100}^2}{8\\pi G}
\\approx 2.775\\times10^{11}\\;(M_\\odot/h)\\,(\\mathrm{Mpc}/h)^{-3}
In h-unit conventions the h² from :math:`H_0 = 100h` km/s/Mpc cancels the
:math:`h^{-3}` from the comoving volume, so this quantity is independent of h.
The mean matter density follows as
:math:`\\bar{\\rho}_m = \\Omega_m\\,\\rho_{\\mathrm{crit},0}`.
Physical constants used:
:math:`G = 6.67430\\times10^{-11}` m³ kg⁻¹ s⁻²,
1 Mpc = 3.085677581×10²² m,
1 M⊙ = 1.989×10³⁰ kg.
Returns
-------
rho_crit0 : float, (Msun/h) / (Mpc/h)³
"""
G_SI = 6.67430e-11 # m³ kg⁻¹ s⁻²
Mpc_m = 3.085677581e22 # m Mpc⁻¹
Msun_kg = 1.989e30 # kg Msun⁻¹
H100_SI = 1e5 / Mpc_m # 100 km/s/Mpc in s⁻¹
rho_SI = 3.0 * H100_SI**2 / (8.0 * np.pi * G_SI) # kg m⁻³
return float(rho_SI * Mpc_m**3 / Msun_kg)
[docs]
class LinearPowerSpectrum:
"""Linear P(k, z) computed with CAMB.
Parameters
----------
(none — no pre-trained weights required)
"""
def __init__(self):
try:
import camb
except ImportError as e:
raise ImportError("camb not installed — pip install camb") from e
self._camb = camb
def _camb_results(self, z: float, theta: dict):
"""Run CAMB and return results object. Internal helper for pk_linear*."""
h = float(theta["h"])
lnAs = float(theta["ln10^{10}A_s"])
w0 = float(theta.get("w0", -1.0))
wa = float(theta.get("wa", 0.0))
pars = self._camb.CAMBparams()
pars.set_cosmology(
H0=100.0 * h,
ombh2=float(theta["Omega_b"]) * h**2,
omch2=float(theta["Omega_cdm"]) * h**2,
)
pars.InitPower.set_params(
ns=float(theta["n_s"]),
As=np.exp(lnAs) * 1e-10,
)
pars.set_dark_energy(w=w0, wa=wa, dark_energy_model="ppf")
pars.set_matter_power(redshifts=[float(z)], kmax=200.0)
return self._camb.get_results(pars)
def _interp_pk(self, k, kh, pk2d):
"""Log-log interpolate a CAMB P(k) table onto the requested k grid."""
pk_arr = jnp.asarray(pk2d[0])
k_arr = jnp.asarray(kh)
return jnp.power(
10.0,
jnp.interp(jnp.log(k), jnp.log(k_arr), jnp.log10(jnp.maximum(pk_arr, 1e-50))),
)
[docs]
def pk_linear(self, k: jnp.ndarray, z: float, theta: dict) -> jnp.ndarray:
"""Linear P(k) [(Mpc/h)^3] at redshift z (total matter).
Supports CPL dark energy via ``w0`` and ``wa`` keys in ``theta``
(defaults to ΛCDM if absent). CAMB uses the PPF dark energy model
(Hu & Sawicki 2007) which remains accurate for :math:`w < -1`.
Parameters
----------
k : array_like, h/Mpc
z : float
theta : dict — keys: h, Omega_b, Omega_cdm, n_s, ln10^{10}A_s,
w0 (default -1), wa (default 0)
"""
results = self._camb_results(float(z), theta)
kh, _, pk2d = results.get_matter_power_spectrum(
minkh=1e-4, maxkh=200.0, npoints=1024
)
return self._interp_pk(k, kh, pk2d)
[docs]
def pk_linear_cdm(self, k: jnp.ndarray, z: float, theta: dict) -> jnp.ndarray:
"""CDM auto-power spectrum P_CDM(k) [(Mpc/h)^3] at redshift z.
.. math::
P_{\\rm CDM}(k) = T_{\\rm cdm}^2(k)\\,P_{\\rm prim}(k)
where :math:`T_{\\rm cdm}` is the CAMB CDM transfer function.
To recover total matter: :math:`f_b P_b + f_c P_{\\rm CDM} \\approx P_{\\rm tot}`
(exact in linear theory when cross-spectrum terms dominate, valid within ~5%).
Parameters
----------
k : array_like, h/Mpc
z : float
theta : dict — same keys as :meth:`pk_linear`
"""
results = self._camb_results(float(z), theta)
kh, _, pk2d = results.get_matter_power_spectrum(
minkh=1e-4, maxkh=200.0, npoints=1024,
var1="delta_cdm", var2="delta_cdm",
)
return self._interp_pk(k, kh, pk2d)
[docs]
def pk_linear_b(self, k: jnp.ndarray, z: float, theta: dict) -> jnp.ndarray:
"""Baryon auto-power spectrum P_b(k) [(Mpc/h)^3] at redshift z.
.. math::
P_b(k) = T_b^2(k)\\,P_{\\rm prim}(k)
Baryonic BAO wiggles are more pronounced here than in the total spectrum.
Parameters
----------
k : array_like, h/Mpc
z : float
theta : dict — same keys as :meth:`pk_linear`
"""
results = self._camb_results(float(z), theta)
kh, _, pk2d = results.get_matter_power_spectrum(
minkh=1e-4, maxkh=200.0, npoints=1024,
var1="delta_baryon", var2="delta_baryon",
)
return self._interp_pk(k, kh, pk2d)
[docs]
@staticmethod
def default_cosmology() -> dict:
"""Planck 2018 TT,TE,EE+lowE best-fit values (flat ΛCDM)."""
return {
"h": 0.6736,
"Omega_b": 0.0493,
"Omega_cdm": 0.2607,
"Omega_m": 0.3100,
"n_s": 0.9649,
"ln10^{10}A_s": 3.044,
"w0": -1.0,
"wa": 0.0,
}
[docs]
@jax.jit
def eisenstein_hu_pk(k: jnp.ndarray, theta: dict) -> jnp.ndarray:
"""Eisenstein & Hu (1998) transfer function with BAO wiggles.
Implements eqs. (2)–(7), (10)–(24) of EH98. The total transfer function is
:math:`T(k) = f_b T_b(k) + f_c T_c(k)` and the power spectrum is
:math:`P(k) \\propto k^{n_s} T(k)^2`, normalised to unity at
:math:`k = 0.05\\,h\\,\\mathrm{Mpc}^{-1}`.
Parameters
----------
k : array_like, h/Mpc
theta : dict — keys: h, Omega_m, Omega_b, n_s; optional T_cmb (K, default 2.7255)
Accuracy
--------
Normalised to P(k=0.05 h/Mpc) = 1.0 by construction. Shape agrees with
CAMB P(k) to < 10% rms for k ∈ [0.01, 0.3] h/Mpc (Planck 2018, z=0).
Large-scale slope d log P / d log k ≈ n_s to < 0.15 for k < 0.003 h/Mpc
(2026-04-23).
Timing
------
~ 242 µs / call (JIT-compiled, N=200 wavenumbers, CPU x86-64, 2026-04-23).
"""
h = theta["h"]
om = theta["Omega_m"]
ob = theta["Omega_b"]
ns = theta["n_s"]
Tcmb = theta.get("T_cmb", 2.7255)
oc = theta.get("Omega_cdm", om - ob)
fb = ob / om # baryon fraction
fc = oc / om # CDM fraction
omh2 = om * h * h
obh2 = ob * h * h
th27 = Tcmb / 2.7
th272 = th27 * th27
th274 = th272 * th272
kh = k * h # Mpc^-1 (paper uses k in Mpc^-1 internally)
# Eq 2: matter-radiation equality redshift
z_eq = 2.5e4 * omh2 / th274
# Eq 3: equality wavenumber [Mpc^-1]
k_eq = 7.46e-2 * omh2 / th272
# Eq 4: drag epoch redshift
b1d = 0.313 * omh2 ** (-0.419) * (1.0 + 0.607 * omh2 ** 0.674)
b2d = 0.238 * omh2 ** 0.223
z_d = 1291.0 * omh2 ** 0.251 / (1.0 + 0.659 * omh2 ** 0.828) * (1.0 + b1d * obh2 ** b2d)
# Eq 5: baryon-to-photon momentum ratio at drag epoch and equality
R_d = 31.5e3 * obh2 / th274 / z_d
R_eq = 31.5e3 * obh2 / th274 / z_eq
# Eq 6: sound horizon at drag epoch [Mpc]
s = (2.0 / (3.0 * k_eq)) * jnp.sqrt(6.0 / R_eq) * jnp.log(
(jnp.sqrt(1.0 + R_d) + jnp.sqrt(R_d + R_eq)) / (1.0 + jnp.sqrt(R_eq))
)
# Eq 7: Silk damping scale [Mpc^-1]
k_silk = 1.6 * obh2 ** 0.52 * omh2 ** 0.73 * (1.0 + (10.4 * omh2) ** (-0.95))
# Eq 10: scaled wavenumber
q = kh / (13.41 * k_eq)
# Eq 11: CDM suppression alpha_c
a1 = (46.9 * omh2) ** 0.670 * (1.0 + (32.1 * omh2) ** (-0.532))
a2 = (12.0 * omh2) ** 0.424 * (1.0 + (45.0 * omh2) ** (-0.582))
alpha_c = a1 ** (-fb) * a2 ** (-(fb ** 3))
# Eq 12: CDM shift beta_c
b1c = 0.944 / (1.0 + (458.0 * omh2) ** (-0.708))
b2c = (0.395 * omh2) ** (-0.0266)
beta_c = 1.0 / (1.0 + b1c * (fc ** b2c - 1.0))
# Eq 15: G(y) function
y = (1.0 + z_eq) / (1.0 + z_d)
G_y = y * (-6.0 * jnp.sqrt(1.0 + y) + (2.0 + 3.0 * y) * jnp.log(
(jnp.sqrt(1.0 + y) + 1.0) / (jnp.sqrt(1.0 + y) - 1.0)
))
# Eq 14: baryon suppression alpha_b
alpha_b = 2.07 * k_eq * s * (1.0 + R_d) ** (-0.75) * G_y
# Eq 24: beta_b
beta_b = 0.5 + fb + (3.0 - 2.0 * fb) * jnp.sqrt((17.2 * omh2) ** 2 + 1.0)
# Eq 23: node shift parameter
beta_node = 8.41 * omh2 ** 0.435
# Eq 22: effective sound horizon
s_tilde = s / (1.0 + (beta_node / (kh * s)) ** 3) ** (1.0 / 3.0)
def _T0(q_val, ac, bc):
"""Eq 19–20: pressureless T0 with suppression (ac, bc)."""
C = 14.2 / ac + 386.0 / (1.0 + 69.9 * q_val ** 1.08)
L = jnp.log(jnp.e + 1.8 * bc * q_val)
return L / (L + C * q_val ** 2)
# Eq 18: CDM interpolation weight
f = 1.0 / (1.0 + (kh * s / 5.4) ** 4)
# Eq 17: CDM transfer function
T_c = f * _T0(q, 1.0, beta_c) + (1.0 - f) * _T0(q, alpha_c, beta_c)
# Eq 21: baryon transfer function
T0_11 = _T0(q, 1.0, 1.0)
j0_tilde = jnp.sinc(kh * s_tilde / jnp.pi) # sin(x)/x
T_b = (
T0_11 / (1.0 + (kh * s / 5.2) ** 2)
+ alpha_b / (1.0 + (beta_b / (kh * s)) ** 3) * jnp.exp(-(kh / k_silk) ** 1.4)
) * j0_tilde
# Eq 16: density-weighted total transfer function
T = fb * T_b + fc * T_c
pk = k ** ns * T ** 2
pk0 = jnp.interp(jnp.log(jnp.array(0.05)), jnp.log(k), pk)
return pk / pk0
[docs]
@jax.jit
def eisenstein_hu_pk_phys(k: jnp.ndarray, theta: dict) -> jnp.ndarray:
"""Eisenstein & Hu (1998) matter power spectrum in physical (Mpc/h)³ units.
Same transfer function as :func:`eisenstein_hu_pk` but returns
:math:`P(k)` in physical :math:`(\\mathrm{Mpc}/h)^3` units with the
correct amplitude derived from the primordial curvature spectrum via the
Poisson equation in conformal-Newtonian gauge:
.. math::
P(k_h, z=0) = D^2(z=0)\\,\\frac{8\\pi^2}{25}
\\frac{(c/H_{100})^4}{\\Omega_m^2}
A_s \\left(\\frac{h}{k_*}\\right)^{n_s-1}
k_h^{n_s}\\,T^2(k_h)
where :math:`c/H_{100} = 2997.924\\;\\mathrm{Mpc}/h`,
:math:`k_* = 0.05\\;\\mathrm{Mpc}^{-1}` (CAMB pivot, physical),
and :math:`A_s = e^{\\ln10^{10}A_s}\\times10^{-10}`.
:math:`D(z=0)` is the linear growth factor normalised so that
:math:`D \\to a` during matter domination (EdS limit). For
:math:`\\Omega_m < 1` flat :math:`\\Lambda`\\CDM, :math:`D(z=0) < 1`
because dark energy suppresses growth after matter–:math:`\\Lambda` equality.
This factor is **not** encoded in the EH98 transfer function shape and
must be included in the amplitude. It is computed via the exact numerical
integral
.. math::
D(z=0) = \\frac{5\\Omega_m}{2}
\\int_0^1 \\frac{\\mathrm{d}a}{[a\\,H(a)/H_0]^3}
with :math:`H(a) = H_0\\sqrt{\\Omega_m a^{-3} + 1 - \\Omega_m}`.
Parameters
----------
k : array_like, h/Mpc
theta : dict — keys: h, Omega_m, Omega_b, n_s, ln10^{10}A_s;
optional Omega_cdm (defaults to Omega_m − Omega_b),
optional T_cmb [K] (default 2.7255)
Returns
-------
P(k) : jnp.ndarray, (Mpc/h)³
Accuracy
--------
:math:`\\sigma_8` computed from this spectrum agrees with CAMB to
:math:`< 1\\%` for Planck 2018 parameters. Residual discrepancies
reflect the EH98 transfer-function shape error (not the amplitude
formula).
"""
h = theta["h"]
om = theta["Omega_m"]
ob = theta["Omega_b"]
ns = theta["n_s"]
Tcmb = theta.get("T_cmb", 2.7255)
lnAs = theta["ln10^{10}A_s"]
# A_s from ln10^{10}A_s: theta["ln10^{10}A_s"] = ln(10^{10} A_s), so A_s = exp(lnAs)*1e-10
A_s = jnp.exp(lnAs) * 1e-10
# Use Omega_cdm directly when present so jax.grad flows through it
oc = theta.get("Omega_cdm", om - ob)
fb = ob / om
fc = oc / om
omh2 = om * h * h
obh2 = ob * h * h
th27 = Tcmb / 2.7
th272 = th27 * th27
th274 = th272 * th272
kh = k * h # Mpc^-1
z_eq = 2.5e4 * omh2 / th274
k_eq = 7.46e-2 * omh2 / th272
b1d = 0.313 * omh2 ** (-0.419) * (1.0 + 0.607 * omh2 ** 0.674)
b2d = 0.238 * omh2 ** 0.223
z_d = 1291.0 * omh2 ** 0.251 / (1.0 + 0.659 * omh2 ** 0.828) * (1.0 + b1d * obh2 ** b2d)
R_d = 31.5e3 * obh2 / th274 / z_d
R_eq = 31.5e3 * obh2 / th274 / z_eq
s = (2.0 / (3.0 * k_eq)) * jnp.sqrt(6.0 / R_eq) * jnp.log(
(jnp.sqrt(1.0 + R_d) + jnp.sqrt(R_d + R_eq)) / (1.0 + jnp.sqrt(R_eq))
)
k_silk = 1.6 * obh2 ** 0.52 * omh2 ** 0.73 * (1.0 + (10.4 * omh2) ** (-0.95))
q = kh / (13.41 * k_eq)
a1 = (46.9 * omh2) ** 0.670 * (1.0 + (32.1 * omh2) ** (-0.532))
a2 = (12.0 * omh2) ** 0.424 * (1.0 + (45.0 * omh2) ** (-0.582))
alpha_c = a1 ** (-fb) * a2 ** (-(fb ** 3))
b1c = 0.944 / (1.0 + (458.0 * omh2) ** (-0.708))
b2c = (0.395 * omh2) ** (-0.0266)
beta_c = 1.0 / (1.0 + b1c * (fc ** b2c - 1.0))
y = (1.0 + z_eq) / (1.0 + z_d)
G_y = y * (-6.0 * jnp.sqrt(1.0 + y) + (2.0 + 3.0 * y) * jnp.log(
(jnp.sqrt(1.0 + y) + 1.0) / (jnp.sqrt(1.0 + y) - 1.0)
))
alpha_b = 2.07 * k_eq * s * (1.0 + R_d) ** (-0.75) * G_y
beta_b = 0.5 + fb + (3.0 - 2.0 * fb) * jnp.sqrt((17.2 * omh2) ** 2 + 1.0)
beta_node = 8.41 * omh2 ** 0.435
s_tilde = s / (1.0 + (beta_node / (kh * s)) ** 3) ** (1.0 / 3.0)
def _T0(q_val, ac, bc):
C = 14.2 / ac + 386.0 / (1.0 + 69.9 * q_val ** 1.08)
L = jnp.log(jnp.e + 1.8 * bc * q_val)
return L / (L + C * q_val ** 2)
f = 1.0 / (1.0 + (kh * s / 5.4) ** 4)
T_c = f * _T0(q, 1.0, beta_c) + (1.0 - f) * _T0(q, alpha_c, beta_c)
T0_11 = _T0(q, 1.0, 1.0)
j0_tilde = jnp.sinc(kh * s_tilde / jnp.pi)
T_b = (
T0_11 / (1.0 + (kh * s / 5.2) ** 2)
+ alpha_b / (1.0 + (beta_b / (kh * s)) ** 3) * jnp.exp(-(kh / k_silk) ** 1.4)
) * j0_tilde
T = fb * T_b + fc * T_c
# Growth factor D(z=0) normalized so D→a as a→0 (EdS limit).
# For Ω_m<1 flat ΛCDM, D(z=0)<1 — this suppression is absent from the EH98
# transfer function (which is shape-only, not growth-suppressed) and must be
# folded into the amplitude explicitly.
# Integral: D(z=0) = (5Ω_m/2) × ∫₀¹ da / [a H(a)/H₀]³
a_g = jnp.linspace(0.001, 1.0, 500)
ol = 1.0 - om
H_over_H0 = jnp.sqrt(om * a_g ** (-3.0) + ol)
D_z0 = (5.0 * om / 2.0) * jnp.trapezoid(1.0 / (a_g * H_over_H0) ** 3, a_g)
# Physical amplitude: P = D²(z=0) × (8π²/25) × (c/H₁₀₀)⁴/Ω_m² × A_s × (h/k_*)^{n_s−1}
# c/H₁₀₀ = 2997.924 Mpc/h; k_* = 0.05 Mpc^{-1} (physical pivot)
_C_H100 = 2997.924 # Mpc/h
_K_PIVOT = 0.05 # Mpc^{-1}
A_amp = D_z0 ** 2 * (8.0 * jnp.pi ** 2 / 25.0) * (_C_H100 ** 4) / om ** 2 * A_s * (h / _K_PIVOT) ** (ns - 1.0)
return A_amp * k ** ns * T ** 2
[docs]
@jax.jit
def eisenstein_hu_pk_nowiggle(k: jnp.ndarray, theta: dict) -> jnp.ndarray:
"""Eisenstein & Hu (1998) no-wiggle (smooth) power spectrum.
Implements eqs. (26), (28)–(31) of EH98. Captures the baryon-induced shape
suppression through an effective shape parameter :math:`\\Gamma_{\\rm eff}(k)`,
without acoustic oscillations. Useful as a smooth reference spectrum.
:math:`P(k) \\propto k^{n_s} T_0(q_{\\rm eff})^2`, normalised to unity at
:math:`k = 0.05\\,h\\,\\mathrm{Mpc}^{-1}`.
Parameters
----------
k : array_like, h/Mpc
theta : dict — keys: h, Omega_m, Omega_b, n_s; optional T_cmb (K, default 2.7255)
"""
h = theta["h"]
om = theta["Omega_m"]
ob = theta["Omega_b"]
ns = theta["n_s"]
Tcmb = theta.get("T_cmb", 2.7255)
fb = ob / om
omh2 = om * h * h
obh2 = ob * h * h
th27 = Tcmb / 2.7
th272 = th27 * th27
kh = k * h # Mpc^-1
# Eq 26: approximate sound horizon [Mpc]
s = 44.5 * jnp.log(9.83 / omh2) / jnp.sqrt(1.0 + 10.0 * obh2 ** 0.75)
# Eq 31: alpha_Gamma (shape suppression amplitude)
alpha_gamma = (
1.0
- 0.328 * jnp.log(431.0 * omh2) * fb
+ 0.38 * jnp.log(22.3 * omh2) * fb ** 2
)
# Eq 30: effective shape parameter [h/Mpc]
Gamma_eff = om * h * (alpha_gamma + (1.0 - alpha_gamma) / (1.0 + (0.43 * kh * s) ** 4))
# Eq 28: scaled wavenumber (dimensionless)
q = k * th272 / Gamma_eff
# Eq 29: zero-baryon transfer function
L0 = jnp.log(2.0 * jnp.e + 1.8 * q)
C0 = 14.2 + 731.0 / (1.0 + 62.5 * q)
T = L0 / (L0 + C0 * q ** 2)
pk = k ** ns * T ** 2
pk0 = jnp.interp(jnp.log(jnp.array(0.05)), jnp.log(k), pk)
return pk / pk0
[docs]
class CsstLinearPowerSpectrum:
"""Linear P(k, z) via the CSST CEmulator (Chen+2025, v2.0).
The emulator is initialised once on instantiation. Cosmology is set via
``set_cosmos`` before each call to ``pk_linear``. k is in h/Mpc and
output P(k) is in (Mpc/h)^3.
Parameter ranges (will raise ValueError if exceeded):
* Omega_b ∈ [0.04, 0.06]
* Omega_m ∈ [0.24, 0.40]
* H0 ∈ [60, 80] (inferred as h * 100)
* n_s ∈ [0.92, 1.00]
* A_s ∈ [1.7e-9, 2.5e-9]
* w0 ∈ [−1.3, −0.7]
* wa ∈ [−0.5, 0.5]
* m_nu ∈ [0, 0.3] eV (default 0.06 eV when not in theta)
"""
def __init__(self):
try:
from CEmulator.Emulator import CBaseEmulator
except ImportError as e:
raise ImportError("CEmulator not installed — pip install CEmulator") from e
self._emu = CBaseEmulator()
@staticmethod
def _set_cosmos(emu, theta: dict) -> None:
"""Push hod_mod theta dict into the CEmulator cosmology state."""
import numpy as np
emu.set_cosmos(
Omegab=float(theta["Omega_b"]),
Omegac=float(theta["Omega_cdm"]),
H0=float(theta["h"]) * 100.0,
As=np.exp(float(theta["ln10^{10}A_s"])) * 1e-10,
ns=float(theta["n_s"]),
w=float(theta.get("w0", -1.0)),
wa=float(theta.get("wa", 0.0)),
mnu=float(theta.get("mnu", 0.06)),
)
[docs]
def pk_linear(self, k: jnp.ndarray, z: float, theta: dict) -> jnp.ndarray:
"""Linear P(k) [(Mpc/h)^3] at redshift z via the CSST emulator.
Parameters
----------
k : array_like, h/Mpc — interpolated onto emulator k-grid
z : float — must lie in [0, 3]
theta : dict — hod_mod cosmological parameter dict
"""
import numpy as np
self._set_cosmos(self._emu, theta)
k_np = np.asarray(k)
pk2d = self._emu.get_pklin(z=float(z), k=k_np) # shape (1, len(k))
return jnp.asarray(pk2d[0])