"""
Models for computing the transfer function.
Note that these are not transfer function "frameworks". The framework is found
in :mod:`~hmf.density_field.transfer`.
"""
import pickle
import warnings
from copy import deepcopy
from typing import Any, ClassVar
import numpy as np
from astropy import cosmology
from scipy.interpolate import InterpolatedUnivariateSpline as Spline
from .._internals._framework import Component, pluggable
try:
import camb
HAVE_CAMB = True
except ImportError: # pragma: no cover
HAVE_CAMB = False
_allfits = ["CAMB", "FromFile", "EH_BAO", "EH_NoBAO", "BBKS", "BondEfs"]
[docs]
@pluggable
class TransferComponent(Component):
r"""
Base class for transfer models.
The only necessary function to specify is ``lnt``, which returns the log
transfer given ``lnk``.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters :
Any model-specific parameters.
"""
_defaults: ClassVar[dict[str, Any]] = {}
[docs]
def __init__(self, cosmo, **model_parameters):
self.cosmo = cosmo
super().__init__(**model_parameters)
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
[docs]
class FromFile(TransferComponent):
r"""
Import a transfer function from file.
.. note:: The file should be in the same format as output from CAMB,
or else in two-column ASCII format (k,T).
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, available
parameters are the following. To see their default values,
check the :attr:`_defaults` class attribute.
:fname: str
Location of the file to import.
"""
_defaults: ClassVar[dict[str, str]] = {"fname": ""}
def _check_low_k(self, lnk, lnT, lnkmin):
"""
Check convergence of transfer function at low k.
Unfortunately, some versions of CAMB produce a transfer which has a
turn-up at low k, which we cut out here.
Parameters
----------
lnk : array_like
Value of log(k)
lnT : array_like
Value of log(transfer)
"""
start = 0
for i in range(len(lnk) - 1):
if abs((lnT[i + 1] - lnT[i]) / (lnk[i + 1] - lnk[i])) < 0.0001:
start = i
break
lnT = lnT[start:-1]
lnk = lnk[start:-1]
lnk[0] = lnkmin
return lnk, lnT
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
try:
T = np.log(np.genfromtxt(self.params["fname"])[:, [0, 6]].T)
except IndexError:
T = np.log(np.genfromtxt(self.params["fname"])[:, [0, 1]].T)
if lnk[0] < T[0, 0]:
lnkout, lnT = self._check_low_k(T[0, :], T[1, :], lnk[0])
else:
lnkout = T[0, :]
lnT = T[1, :]
return Spline(lnkout, lnT, k=1)(lnk)
if HAVE_CAMB:
[docs]
class CAMB(FromFile):
r"""
Transfer function computed by CAMB.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model.
**camb_params:** An instantiated ``CAMBparams`` object, pre-set with desired
accuracy options etc.
**dark_energy_params:** A dictionary of values passed to CAMB's `
`set_dark_energy`` method. Values include
`sound_speed` and `dark_energy_model`.
**extrapolate_with_eh:** Whether to extrapolate past the intrinsic CAMB
kmax by using an EH model. Can cause some problems
if kmax is high, since CAMB diverges from the EH
approximation.
Notes
-----
Neutrino masses are passed to CAMB via the ``mnu`` parameter (sum of neutrino
masses in eV) and CAMB internally computes the neutrino physical density
``omnuh2``. The cold dark matter density ``omch2`` is set from astropy's
``Odm0`` attribute (i.e. ``cosmo.Odm0 * cosmo.h**2``), which represents
CDM-only density (excluding neutrinos and baryons). In astropy, ``Om0``
represents the sum of CDM and baryonic matter only (not massive neutrinos),
so ``Odm0 = Om0 - Ob0``. The neutrino contribution to the matter budget is
accounted for separately by CAMB from the ``mnu`` parameter, ensuring the
total matter density (CDM + baryons + neutrinos) is correctly captured.
"""
_defaults: ClassVar[dict[str, Any]] = {
"camb_params": None,
"dark_energy_params": {},
"extrapolate_with_eh": None,
"kmax": None,
}
[docs]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if not isinstance(self.cosmo, (cosmology.LambdaCDM, cosmology.wCDM, cosmology.w0waCDM)):
raise ValueError("CAMB will only work with LCDM or wCDM cosmologies")
# Save the CAMB object properly for use
# Set the cosmology
if self.params["camb_params"] is None:
self.params["camb_params"] = camb.CAMBparams(
DoLensing=False,
Want_CMB=False,
Want_CMB_lensing=False,
WantCls=False,
WantDerivedParameters=False,
)
self.params["camb_params"].Transfer.high_precision = False
self.params["camb_params"].Transfer.k_per_logint = 0
# If extrapolating with EH, use a lower value of kmax so that the
# calculation is faster.
if self.params["kmax"]:
self.params["camb_params"].Transfer.kmax = self.params["kmax"]
if self.cosmo.Ob0 is None or self.cosmo.Ob0 == 0.0:
raise ValueError(
"To use CAMB, you must set the baryon density in the cosmology explicitly."
)
if self.cosmo.Tcmb0.value == 0:
raise ValueError(
"If using CAMB, the CMB temperature must be set explicitly in the cosmology."
)
self.params["camb_params"].set_cosmology(
H0=self.cosmo.H0.value,
ombh2=self.cosmo.Ob0 * self.cosmo.h**2,
omch2=self.cosmo.Odm0 * self.cosmo.h**2,
mnu=sum(self.cosmo.m_nu.value),
neutrino_hierarchy="degenerate",
omk=self.cosmo.Ok0,
nnu=self.cosmo.Neff,
standard_neutrino_neff=self.cosmo.Neff,
TCMB=self.cosmo.Tcmb0.value,
)
self.params["camb_params"].WantTransfer = True
# Set the DE equation of state. We only support constant w.
if isinstance(self.cosmo, cosmology.wCDM):
self.params["camb_params"].set_dark_energy(w=self.cosmo.w0)
elif isinstance(self.cosmo, cosmology.w0waCDM):
self.params["camb_params"].set_dark_energy(w=self.cosmo.w0, wa=self.cosmo.wa)
if self.params["extrapolate_with_eh"] is None:
warnings.warn(
"'extrapolate_with_eh' was not set. Defaulting to True, which is "
"different behaviour than versions <=3.4.4. This warning may be "
"removed in v4.0. Silence it by setting extrapolate_with_eh explicitly.",
stacklevel=2,
)
self.params["extrapolate_with_eh"] = True
if self.params["extrapolate_with_eh"]:
# Create an EH transfer to extrapolate to at high k.
self._eh = EH(self.cosmo)
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
camb_transfers = camb.get_transfer_functions(self.params["camb_params"])
T = camb_transfers.get_matter_transfer_data().transfer_data
T = np.log(T[[0, 6], :, 0])
if lnk[0] < T[0, 0]:
lnkout, lnT = self._check_low_k(T[0, :], T[1, :], lnk[0])
else:
lnkout = T[0, :]
lnT = T[1, :]
lnT -= lnT[0]
if not self.params["extrapolate_with_eh"]:
return Spline(lnkout, lnT, k=1)(lnk)
# Now add a point one e-fold above the max, with an EH-generated transfer
lnkout = np.concatenate((lnkout, [lnkout[-1] + 1]))
# normalise EH at the final CAMB point.
norm = self._eh.lnt(lnkout[-2]) - lnT[-1]
lnT = np.concatenate((lnT, [self._eh.lnt(lnkout[-1]) - norm]))
lnkmin = lnkout.min()
lnkmax = lnkout.max()
inner_spline = Spline(lnkout, lnT, k=3)
out = np.zeros_like(lnk)
out[lnk < lnkmin] = 0
out[(lnkmin <= lnk) & (lnk <= lnkmax)] = inner_spline(
lnk[(lnkmin <= lnk) & (lnk <= lnkmax)]
)
out[lnk >= lnkmax] = self._eh.lnt(lnk[lnk >= lnkmax]) - norm
return out
def __getstate__(self):
"""Get the state of the object, including converting the CAMBparams object to a dict."""
# We need to get rid of the CAMBparams() object, as it cannot be pickled.
p = self.params["camb_params"]
potential_keys = [ # From https://camb.readthedocs.io/en/latest/model.html
"WantCls",
"WantTransfer",
"WantScalars",
"WantTensors",
"WantVectors",
"WantDerivedParameters",
"Want_cl_2D_array",
"Want_CMB",
"Want_CMB_lensing",
"DoLensing",
"NonLinear",
"Transfer",
"want_zstar",
"want_zdrag",
"min_l",
"max_l",
"max_l_tensor",
"max_eta_k",
"max_eta_k_tensor",
"ombh2",
"omch2",
"omk",
"omnuh2",
"H0",
"TCMB",
"YHe",
"num_nu_massless",
"num_nu_massive",
"nu_mass_eigenstates",
"share_delta_neff",
"InitPower",
"Recomb",
"Reion",
"DarkEnergy",
"NonLinearModel",
"Accuracy",
"SourceTerms",
"z_outputs",
"scalar_initial_condition",
"InitialConditionVector",
"OutputNormalization",
"Alens",
"MassiveNuMethod",
"DoLateRadTruncation",
"Evolve_baryon_cs",
"Evolve_delta_xe",
"Evolve_delta_Ts",
"Do21cm",
"transfer_21cm_cl",
"Log_lvalues",
"use_cl_spline_template",
"SourceWindows",
]
# Unsaveable parameters:
# "nu_mass_degeneracies", "nu_mass_fractions", "nu_mass_numbers", "CustomSources"
dct = {}
for pk in potential_keys:
try:
pickle.dumps(getattr(p, pk))
dct[pk] = getattr(p, pk)
except AttributeError:
warnings.warn(
f"CAMB key '{pk}' is not an attribute. If you provided a "
f"custom CAMBparams, results may be inconsistent. Available: "
f"{dir(p)}",
stacklevel=2,
)
except Exception:
warnings.warn(f"CAMB key {pk} is not pickle-able.", stacklevel=2)
# Deepcopy self
this = {}
for key, val in self.__dict__.items():
if key != "params":
this[key] = deepcopy(val)
this["params"] = {"camb_params": dct}
return this
def __setstate__(self, state):
"""Set the state of the object, including reconstructing the CAMBparams object."""
self.__dict__ = state
self.params["camb_params"] = camb.CAMBparams(**self.params["camb_params"])
[docs]
class FromArray(FromFile):
r"""
Use a spline over a given array to define the transfer function.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, available
parameters are the following. To see their default values,
check the :attr:`_defaults` class attribute.
:k: array
Wavenumbers, in [h/Mpc]
:T: array
Transfer function
"""
_defaults: ClassVar[dict[str, Any]] = {"k": None, "T": None}
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
k = self.params["k"]
T = self.params["T"]
if k is None or T is None:
raise ValueError("You must supply an array for both k and T for this Transfer Model")
if len(k) != len(T):
raise ValueError("k and T must have same length")
if lnk[0] < np.log(k.min()):
lnkout, lnT = self._check_low_k(np.log(k), np.log(T), lnk[0])
else:
lnkout = np.log(k)
lnT = np.log(T)
return Spline(lnkout, lnT, k=1)(lnk)
[docs]
class EH_BAO(TransferComponent):
r"""
Eisenstein & Hu (1998) fitting function with BAO wiggles.
From EH1998, Eqs. 26,28-31. Code adapted from CHOMP.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, there
are no model parameters.
"""
[docs]
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self._set_params()
def _set_params(self):
"""Port of ``TFset_parameters`` from original EH code."""
self.Obh2 = self.cosmo.Ob0 * self.cosmo.h**2
self.Omh2 = self.cosmo.Om0 * self.cosmo.h**2
self.f_baryon = self.cosmo.Ob0 / self.cosmo.Om0
self.theta_cmb = self.cosmo.Tcmb0.value / 2.7
self.z_eq = 2.5e4 * self.Omh2 * self.theta_cmb ** (-4) # really 1+z
self.k_eq = 7.46e-2 * self.Omh2 * self.theta_cmb ** (-2) # units Mpc^-1 (no h!)
self.z_drag_b1 = 0.313 * self.Omh2**-0.419 * (1.0 + 0.607 * self.Omh2**0.674)
self.z_drag_b2 = 0.238 * self.Omh2**0.223
self.z_drag = (
1291.0
* self.Omh2**0.251
/ (1.0 + 0.659 * self.Omh2**0.828)
* (1.0 + self.z_drag_b1 * self.Obh2**self.z_drag_b2)
)
self.r_drag = 31.5 * self.Obh2 * self.theta_cmb**-4 * (1000.0 / (1 + self.z_drag))
self.r_eq = 31.5 * self.Obh2 * self.theta_cmb**-4 * (1000.0 / self.z_eq)
self.sound_horizon = (
(2.0 / (3.0 * self.k_eq))
* np.sqrt(6.0 / self.r_eq)
* np.log(
(np.sqrt(1.0 + self.r_drag) + np.sqrt(self.r_drag + self.r_eq))
/ (1.0 + np.sqrt(self.r_eq))
)
)
self.k_silk = (
1.6 * self.Obh2**0.52 * self.Omh2**0.73 * (1.0 + (10.4 * self.Omh2) ** (-0.95))
)
alpha_c_a1 = (46.9 * self.Omh2) ** 0.670 * (1.0 + (32.1 * self.Omh2) ** (-0.532))
alpha_c_a2 = (12.0 * self.Omh2) ** 0.424 * (1.0 + (45.0 * self.Omh2) ** (-0.582))
self.alpha_c = alpha_c_a1 ** (-self.f_baryon) * alpha_c_a2 ** (-(self.f_baryon**3))
beta_c_b1 = 0.944 / (1.0 + (458.0 * self.Omh2) ** -0.708)
beta_c_b2 = (0.395 * self.Omh2) ** -0.0266
self.beta_c = 1.0 / (1.0 + beta_c_b1 * ((1 - self.f_baryon) ** beta_c_b2 - 1))
y = self.z_eq / (1 + self.z_drag)
alpha_b_G = y * (
-6 * np.sqrt(1 + y) + (2 + 3 * y) * np.log((np.sqrt(1 + y) + 1) / (np.sqrt(1 + y) - 1))
)
self.alpha_b = (
2.07 * self.k_eq * self.sound_horizon * (1.0 + self.r_drag) ** -0.75 * alpha_b_G
)
self.beta_node = 8.41 * self.Omh2**0.435
self.beta_b = (
0.5
+ self.f_baryon
+ (3.0 - 2.0 * self.f_baryon) * np.sqrt((17.2 * self.Omh2) ** 2 + 1.0)
)
@property
def k_peak(self):
"""The k_peak parameter from EH1998, Eq. 29, in units of Mpc^-1 (no h)."""
return 2.5 * np.pi * (1 + 0.217 * self.Omh2) / self.sound_horizon
@property
def sound_horizon_fit(self):
"""Sound horizon in Mpc/h."""
return self.cosmo.h * 44.5 * np.log(9.83 / self.Omh2) / np.sqrt(1 + 10 * (self.Obh2**0.75))
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
# Get k in Mpc^-1
k = np.exp(lnk) * self.cosmo.h
q = k / (13.41 * self.k_eq)
ks = k * self.sound_horizon
T_c_ln_beta = np.log(np.e + 1.8 * self.beta_c * q)
T_c_ln_nobeta = np.log(np.e + 1.8 * q)
T_c_C_alpha = (14.2 / self.alpha_c) + 386.0 / (1.0 + 69.9 * q**1.08)
T_c_C_noalpha = 14.2 + 386.0 / (1.0 + 69.9 * q**1.08)
T_c_f = 1.0 / (1.0 + (ks / 5.4) ** 4)
def term(a, b):
return a / (a + b * q**2)
T_c = T_c_f * term(T_c_ln_beta, T_c_C_noalpha) + (1 - T_c_f) * term(
T_c_ln_beta, T_c_C_alpha
)
s_tilde = self.sound_horizon / (1.0 + (self.beta_node / ks) ** 3) ** (1.0 / 3.0)
ks_tilde = k * s_tilde
T_b_T0 = term(T_c_ln_nobeta, T_c_C_noalpha)
Tb1 = T_b_T0 / (1.0 + (ks / 5.2) ** 2)
Tb2 = (self.alpha_b / (1.0 + (self.beta_b / ks) ** 3)) * np.exp(-((k / self.k_silk) ** 1.4))
T_b = np.sin(ks_tilde) / ks_tilde * (Tb1 + Tb2)
return np.log(self.f_baryon * T_b + (1 - self.f_baryon) * T_c)
[docs]
class EH_NoBAO(EH_BAO):
r"""
Eisenstein & Hu (1998) fitting function without BAO wiggles.
From EH 1998 Eqs. 26,28-31. Code adapted from CHOMP project.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, there are
no model parameters.
"""
@property
def alpha_gamma(self):
"""The alpha_gamma parameter from EH1998, Eq. 31."""
return (
1
- 0.328 * np.log(431 * self.Omh2) * self.f_baryon
+ 0.38 * np.log(22.3 * self.Omh2) * self.f_baryon**2
)
[docs]
def lnt(self, lnk):
r"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
k = np.exp(lnk) * self.cosmo.h
ks = k * self.sound_horizon_fit / self.cosmo.h # need sound horizon in Mpc here
gamma_eff = self.Omh2 * (self.alpha_gamma + (1 - self.alpha_gamma) / (1 + (0.43 * ks) ** 4))
q = k / (13.4 * self.k_eq)
q_eff = q * self.Omh2 / gamma_eff
L0 = np.log(2 * np.e + 1.8 * q_eff)
C0 = 14.2 + 731.0 / (1 + 62.5 * q_eff)
return np.log(L0 / (L0 + C0 * q_eff * q_eff))
[docs]
class BBKS(TransferComponent):
r"""
BBKS (1986) transfer function.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, available
parameters are the following: **a, b, c, d, e**.
To see their default values, check the :attr:`_defaults`
class attribute.
Notes
-----
The fit is given as
.. math:: T(k) = \frac{\ln(1+aq)}{aq}
\left(1 + bq + (cq)^2 + (dq)^3 + (eq)^4\right)^{-1/4},
where
.. math:: q = \frac{k}{\Gamma}
and :math:`\Gamma = \Omega_{m,0} h`. Note that here *k* is in units of h/Mpc, which
accounts for the extra *h* in the equations in BBKS.
These equations are taken from BBKS 1986, Eq. G3.
Further modifications can be made in the presence of baryons. Sugiyama 1995, Eq. 3.9
gives
.. math:: \Gamma \rightarrow \Gamma
\exp\left(-\Omega_{b,0}(1 + 1/\Omega_{m,0})\right)
and Liddle and Lythe (2000) Eq. 5.14 give a slight extra:
.. math:: \Gamma \rightarrow \Gamma
\exp\left(-\Omega_{b,0}(1 + \sqrt{2h}/\Omega_{m,0})\right).
"""
_defaults: ClassVar[dict[str, Any]] = {
"a": 2.34,
"b": 3.89,
"c": 16.1,
"d": 5.46,
"e": 6.71,
"use_sugiyama_baryons": False,
"use_liddle_baryons": True,
}
[docs]
def lnt(self, lnk):
"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [h/Mpc]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
a = self.params["a"]
b = self.params["b"]
c = self.params["c"]
d = self.params["d"]
e = self.params["e"]
Gamma = self.cosmo.Om0 * self.cosmo.h
if self.params["use_sugiyama_baryons"]:
Gamma *= np.exp(-self.cosmo.Ob0 * (1 + 1 / self.cosmo.Om0))
elif self.params["use_liddle_baryons"]:
Gamma *= np.exp(
-self.cosmo.Ob0 * (1 + np.sqrt(self.cosmo.Ob0 * self.cosmo.h) / self.cosmo.Om0)
)
q = np.exp(lnk) / Gamma
return np.log(
np.log(1.0 + a * q)
/ (a * q)
* (1 + b * q + (c * q) ** 2 + (d * q) ** 3 + (e * q) ** 4) ** (-0.25)
)
[docs]
class BondEfs(TransferComponent):
r"""
Transfer function of Bond and Efstathiou.
Parameters
----------
cosmo : :class:`astropy.cosmology.FLRW` instance
The cosmology used in the calculation
\*\*model_parameters : unpack-dict
Parameters specific to this model. In this case, available
parameters are the following: **a, b, c, nu**.
To see their default values, check the :attr:`_defaults`
class attribute.
Notes
-----
The fit is given as
.. math:: T(k) = \left[1 + (\tilde{a}k + (\tilde{b}k)^{3/2} +
(\tilde{c}k)^2)^\nu\right]^{-1/\nu}
where :math:`\tilde{x} = x\alpha` and
.. math:: \alpha = \frac{0.3\times 0.75^2}{\Omega_{m,0} h^2}.
"""
_defaults: ClassVar[dict[str, float]] = {"a": 37.1, "b": 21.1, "c": 10.8, "nu": 1.12}
[docs]
def lnt(self, lnk):
"""
Natural log of the transfer function.
Parameters
----------
lnk : array_like
Wavenumbers [Mpc/h]
Returns
-------
lnt : array_like
The log of the transfer function at lnk.
"""
scale = (0.3 * 0.75**2) / (self.cosmo.Om0 * self.cosmo.h)
a = self.params["a"] * scale
b = self.params["b"] * scale
c = self.params["c"] * scale
nu = self.params["nu"]
k = np.exp(lnk)
return np.log((1 + (a * k + (b * k) ** 1.5 + (c * k) ** 2) ** nu) ** (-1 / nu))
[docs]
class EH(EH_BAO):
"""Alias of :class:`EH_BAO`."""