import numpy as np # type: ignore
from scipy.optimize import root_scalar # type: ignore
from scipy.integrate import simpson # type: ignore
from scipy.interpolate import splev # type: ignore
import camb # type: ignore
from camb import model # type: ignore
import logging
from importlib import resources
from pathlib import Path
from urllib import request
from appdirs import user_cache_dir
import skops.io as skio
# Use relative imports to find the other modules in this package
from .cosmology import Cosmology
from .growth import GrowthCalculator
# Get a logger for this module
log = logging.getLogger(__name__)
[docs]
class AletheiaEmu:
"""Emulator for the non-linear matter power spectrum and the halo mass function.
This class provides predictions for the non-linear matter power spectrum,
P_NL(k) and the halo multiplicity function f(nu) for a given cosmology and
redshift. It is based on a set of Gaussian Process (GP) models trained
on high-fidelity N-body simulations.
The P_NL(k) emulation method :meth:`get_pnl` combines a de-wiggled linear
power spectrum with a GP prediction for the non-linear boost factor, B(k),
and a response function, dR/dxtide, that captures the effects of different
growth of structure histories.
The emulation of f(nu) is handled by :meth:`get_fnu`, using the same emulation
strategy and the same simulation suite, and reuse the cosmology validation, the
reference-redshift solver and the growth-history machinery of P_NL; only the
trained models, the memory-kernel width and the region of validity differ.
The method :meth:`get_fnu` returns the differential mass function dn/dlnM
based on the f(nu) emulation.
Attributes
----------
gp_B : sklearn.gaussian_process.GaussianProcessRegressor
The trained GP model for the non-linear boost factor.
gp_dRdx : sklearn.gaussian_process.GaussianProcessRegressor
The trained GP model for the response to xtilde.
correction_function : scipy.interpolate.RectBivariateSpline
A 2D spline object for correcting resolution effects in the prediction.
planck_means : np.ndarray
Mean values of [omega_b, omega_c, n_s] from Planck 2018.
eigenvecs : np.ndarray
Eigenvectors of the Planck 2018 covariance matrix.
planck_sigmas : np.ndarray
Standard deviations along each eigenvector direction.
gp_fnu : sklearn.gaussian_process.GaussianProcessRegressor
The trained GP model for ln[f(nu) / f_ref(nu)].
gp_dfdx : sklearn.gaussian_process.GaussianProcessRegressor
The trained GP model for the response of f(nu) to xtilde.
reference_spline_hmf : tuple
A `splrep` representation of ln f_ref as a function of ln nu.
nu_range : tuple of float
Bounding box of the region of validity in nu: the loosest (min, max)
peak height reached anywhere in :attr:`SIGMA12_RANGE`. It is a summary
only -- the check applied to a prediction is the sigma12-dependent
window returned by :meth:`nu_range_at`.
SIGMA12_RANGE : tuple of float
The (min, max) clustering amplitude the emulator is trained for.
REFERENCE_COSPAR : dict
Parameters of the reference cosmology that defines the evolution-mapping
reference growth history.
MEMORY_SCALE_HMF : float
Width of the mass-function memory kernel, in units of ln D.
"""
# --- Define model URLs pointing to .skops files ---
MODEL_URLS = {
"gp_B": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/Aletheia_GP_B_skl1.7.skops?ref_type=heads&inline=false",
"gp_dRdx": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/Aletheia_GP_dRdxt_skl1.7.skops?ref_type=heads&inline=false",
"correction": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/resolution_correction_skl1.7.skops?ref_type=heads&inline=false",
"gp_fnu": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/Aletheia_GP_fnu_skl1.7.skops?ref_type=heads&inline=false",
"gp_dfdx": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/Aletheia_GP_dfdxt_skl1.7.skops?ref_type=heads&inline=false",
"reference": "https://gitlab.mpcdf.mpg.de/arielsan/aletheia/-/raw/main/src/aletheiacosmo/data/fnu_reference_skl1.7.skops?ref_type=heads&inline=false",
}
# --- Define cache directory for downloaded models ---
CACHE_DIR = Path(user_cache_dir("AletheiaCosmo", "AletheiaTeam"))
# --- Define trusted types for skops loading ---
TRUSTED_TYPES = ['numpy.ndarray',
'builtins.dict',
'builtins.tuple',
'builtins.list',
'builtins.str',
'numpy.float64',
'numpy.random.mtrand.RandomState',
'sklearn.gaussian_process.kernels.Matern',
'sklearn.gaussian_process._gpr.GaussianProcessRegressor',
'scipy.interpolate._fitpack2.RectBivariateSpline',
'builtins.int' # needed by the f(nu) reference spline
]
# --- Planck covariance matrix ---
PLANCK_FILE = "planck_2018_lcdm_parcov.dat"
def __init__(self):
log.info("AletheiaEmu instance created. Checking for models...")
# --- Validity range of the emulator ---
self.SIGMA12_RANGE = (0.2, 1.0)
# --- Parameters defining the reference growth history ---
# These are the parameters of the "i0" nodes of the AletheiaEmu suite, i.e. the
# nodes on which the f(nu) emulator itself was trained. get_pnl applies the
# same values inline.
self.REFERENCE_COSPAR = {'w_0': -1.0,
'w_a': 0.0, 'A_s': 2.101e-9, 'h': 0.673}
# --- Definition of xtilde for the mass function ---
# Width of the Gaussian memory kernel in units of ln D, fitted for M200b.
self.MEMORY_SCALE_HMF = 0.2884
# Ensure the cache directory exists
# Use self.CACHE_DIR to access the class attribute
self.CACHE_DIR.mkdir(parents=True, exist_ok=True)
# --- Define local file paths for LARGE models ---
gp_b_path = self.CACHE_DIR / "Aletheia_GP_B_skl1.7.skops"
gp_drdx_path = self.CACHE_DIR / "Aletheia_GP_dRdxt_skl1.7.skops"
correction_path = self.CACHE_DIR / "resolution_correction_skl1.7.skops"
gp_fnu_path = self.CACHE_DIR / "Aletheia_GP_fnu_skl1.7.skops"
gp_dfdx_path = self.CACHE_DIR / "Aletheia_GP_dfdxt_skl1.7.skops"
reference_path = self.CACHE_DIR / "fnu_reference_skl1.7.skops"
# --- Download models if they are missing ---
self._download_if_missing(self.MODEL_URLS["gp_B"], gp_b_path)
self._download_if_missing(self.MODEL_URLS["gp_dRdx"], gp_drdx_path)
self._download_if_missing(
self.MODEL_URLS["correction"], correction_path)
self._download_if_missing(self.MODEL_URLS["gp_fnu"], gp_fnu_path)
self._download_if_missing(self.MODEL_URLS["gp_dfdx"], gp_dfdx_path)
self._download_if_missing(self.MODEL_URLS["reference"], reference_path)
# --- Use skio.load to load emulator files and resolution correction ---
self.gp_B = skio.load(gp_b_path, trusted=self.TRUSTED_TYPES)
self.gp_dRdx = skio.load(gp_drdx_path, trusted=self.TRUSTED_TYPES)
self.correction_function = skio.load(
correction_path, trusted=self.TRUSTED_TYPES)
self.gp_fnu = skio.load(gp_fnu_path, trusted=self.TRUSTED_TYPES)
self.gp_dfdx = skio.load(gp_dfdx_path, trusted=self.TRUSTED_TYPES)
self.reference_spline_hmf = skio.load(
reference_path, trusted=self.TRUSTED_TYPES)
# --- Load the bundled Planck data file ---
try:
# This uses importlib.resources to find the file *inside* the package
with resources.files('aletheiacosmo.data').joinpath(self.PLANCK_FILE).open('rb') as f:
covmat = np.loadtxt(f)
except FileNotFoundError as e:
raise FileNotFoundError(
"Planck covariance matrix file (planck_2018_lcdm_parcov.dat) was not found. "
"This file should be bundled with the package. Please reinstall AletheiaCosmo."
)
# Construct Planck 2018 parameter covariance
# eigenvectors and eigenvalues
# required data for parameter validation
self.planck_means = np.array([0.02236164, 0.12071002, 0.96479956])
eigenvals, eigenvecs = np.linalg.eigh(covmat)
self.eigenvecs = eigenvecs
self.planck_sigmas = np.sqrt(eigenvals)
log.info("Emulator models and data loaded successfully.")
# --- Derive the mass-function region of validity from the training data ---
# The boundary is that of the C(nu) GP, the leading term of the
# prediction; the response GP is allowed to extrapolate beyond its own coverage.
# The fit is rebuilt from the training inputs carried
# in the model file, so it stays in step with it if the GP is retrained.
self.nu_lo_coeffs = self._fit_nu_boundary(self.gp_fnu, 'lower')
self.nu_hi_coeffs = self._fit_nu_boundary(self.gp_fnu, 'upper')
# Bounding box of that region, for reference only.
s12_grid = np.linspace(*self.SIGMA12_RANGE, 801)
self.nu_range = (float(np.polyval(self.nu_lo_coeffs, s12_grid).min()),
float(np.polyval(self.nu_hi_coeffs, s12_grid).max()))
log.info(f"Emulator covers nu in [{self.nu_range[0]:.3f}, "
f"{self.nu_range[1]:.3f}] across sigma12 in "
f"[{self.SIGMA12_RANGE[0]}, {self.SIGMA12_RANGE[1]}], with the "
f"nu window narrowing as sigma12 falls.")
# --- One-slot cache for the CAMB-backed cosmology engine ---
# Unlike P_nl, mass-function calls are usually repeated over many masses
# and redshifts for the same cosmology, and each engine costs a CAMB run.
self._cached_cospar = None
self._cached_engine = None
def _download_if_missing(self, url, filepath):
"""Helper function to download an emulator skop file if it doesn't exist locally."""
if not filepath.exists():
log.warning(
f"Data file not found. Downloading {filepath.name} to cache: {filepath}")
try:
opener = request.build_opener()
opener.addheaders = [
('User-agent', 'AletheiaCosmo-Downloader')]
request.install_opener(opener)
request.urlretrieve(url, filepath)
log.info("Download complete.")
except Exception as e:
log.error(f"Failed to download model from {url}. Error: {e}")
raise RuntimeError(
f"Could not download model data. Please check your internet connection or the model URL.")
[docs]
def get_pnl(self, kvec, cospar, z):
"""Calculates the non-linear matter power spectrum.
This is the main method of the emulator. It takes a set of wavenumbers,
a cosmology, and a redshift, and returns the emulated P_NL(k).
Parameters
----------
kvec : array_like
Array of wavenumbers, k, in units of **1/Mpc**.
Must be within the emulator's valid range [0.006, 2.0] 1/Mpc.
cospar : dict
A dictionary of cosmological parameters such as the one created
by `create_cosmo_dict`.
z : float
The redshift at which to calculate the power spectrum.
Returns
-------
np.ndarray
The non-linear matter power spectrum, P_NL(k), in units of **Mpc^3**.
Raises
------
ValueError
If any k-values in `kvec` are outside the emulator's valid
training range [0.006, 2.0] 1/Mpc.
ValueError
If the input cosmology fails the validation checks (e.g., sigma12
is out of range [0.2, 1.0] or shape parameters are out of
the 5-sigma Planck box).
Notes
-----
The calculation involves several steps:
1. Input parameters are validated (k-range, sigma12, shape).
2. A `Cosmology` object is created to compute linear spectra.
3. The non-linear boost `B(k)` and response `dR/dxi` are predicted by GPs.
4. The parameter `xtilde` is computed for the target and a reference cosmology.
5. All components are combined and a final resolution correction is applied.
6. A warning is logged if the resolution correction is > 1% at any scale.
"""
log.info(f"Received request for P_nl at z={z:.2f}")
kvec = np.atleast_1d(kvec)
z = float(z)
# --- k-vector validation (units are 1/Mpc) ---
if np.any(kvec < 0.006) or np.any(kvec > 2.0):
msg = (f"k-values are outside the valid emulator range [0.006, 2.0] 1/Mpc. "
f"Found min={np.min(kvec):.4f}, max={np.max(kvec):.4f}")
log.error(msg)
raise ValueError(msg)
log.info("Initializing cosmology engine for target cosmology.")
cosmology_engine = Cosmology(cospar)
sigma12 = cosmology_engine.get_sigma12(z)
self._validate_params(cospar, sigma12)
# --- Correction factor for resolution effects ---
correction = self.correction_function(kvec, sigma12).flatten()
correction_factors = 1.0 / correction
max_correction = np.max(correction_factors)
if max_correction > 1.03:
k_at_max_corr = kvec[np.argmax(correction_factors)]
log.warning(f"Resolution correction is > 3% (max: {max_correction:.3f}x "
f"at k={k_at_max_corr:.2f} 1/Mpc). ")
cosmology_engine.compute_all_spectra(sigma12)
pdw_use = cosmology_engine.get_dewiggled_pk(kvec)
log.info("Computing emulator predictions from GPs.")
x_combined = self._build_gp_input(kvec, cospar, sigma12)
B = np.exp(self._gp_predict(self.gp_B, x_combined))
dRdx = self._gp_predict(self.gp_dRdx, x_combined)
log.info("Computing xtilde for target and reference cosmologies.")
xtilde = self._get_xtilde(z, cosmology_engine.growth)
cospar_ref = cospar.copy()
# cospar_ref.update({'w_0':-1.0, 'w_a':0.0, 'A_s':2.101e-9, 'h':0.673})
cospar_ref.update(self.REFERENCE_COSPAR)
growth_ref = GrowthCalculator(cospar_ref)
z0 = self._get_redshift(
sigma12, # The value to match
growth_ref, # The reference growth engine
cosmology_engine, # The target cosmology engine
cospar_ref # The reference parameter dict
)
xtilde_0 = self._get_xtilde(z0, growth_ref)
log.debug(
f"Target xtilde={xtilde:.4f}. Reference z0={z0:.4f}, xtilde_0={xtilde_0:.4f}")
log.info("Combining components for final prediction.")
dxtilde = xtilde - xtilde_0
Pnl_uncorrected = pdw_use * B * (1. + dRdx * dxtilde)
# Apply final resolution correction
Pnl_final = Pnl_uncorrected * correction_factors
log.info("P_nl calculation complete.")
return Pnl_final
def _build_gp_input(self, xvec, cospar, sigma12):
"""Constructs the 2D input array for the Gaussian Process models."""
# xvec is the array of k-values (1/Mpc) or nu-values, depending on the GP being called
lxvec = np.log(xvec)
# Reshape all inputs to be (N, 1) column vectors
x1 = lxvec.reshape(-1, 1)
x2 = np.full_like(x1, cospar['omega_b'])
x3 = np.full_like(x1, cospar['omega_c'])
x4 = np.full_like(x1, cospar['n_s'])
x5 = np.full_like(x1, sigma12)
return np.concatenate((x1, x2, x3, x4, x5), axis=1)
def _gp_predict(self, gp_model, x_input):
"""Predicts outputs from a Gaussian Process model, bypassing Scikit-Learn's predict().
This method computes the GP prediction using the kernel and training data
directly, which can be more efficient for large datasets.
Parameters
----------
gp_model : sklearn.gaussian_process.GaussianProcessRegressor
The trained GP model to use for prediction.
x_input : np.ndarray
The input array for which to predict outputs.
Returns
-------
np.ndarray
The predicted outputs from the GP model.
"""
K_trans = gp_model.kernel_(x_input, gp_model.X_train_)
y_pred = K_trans.dot(gp_model.alpha_)
if gp_model.normalize_y:
y_pred = y_pred * gp_model._y_train_std + gp_model._y_train_mean
return y_pred
def _validate_params(self, cospar, sigma12):
"""
Checks if the input cosmology is within the valid range of the emulator.
This function performs two checks:
1. It verifies if sigma12 falls within the emulator's trained range
of [0.2, 1.0].
2. It transforms the shape parameters (omega_b, omega_c, n_s) into the
eigenvector basis of the Planck 2018 covariance matrix and checks
that each projected component is within a +/- 5-sigma box.
Args:
cospar (dict): The input cosmology dictionary.
sigma12 (array_like): The input value of sigma12 (as a 1-element array).
Raises:
ValueError: If any parameter falls outside the valid range.
"""
log.info("Validating input parameters...")
# Extract the scalar value for the check and logging
sigma12_val = sigma12[0]
# --- Validate sigma12 ---
log.debug(f"Checking sigma12 = {sigma12_val:.4f}")
if not (0.2 <= sigma12_val <= 1.0):
raise ValueError(
f"Validation failed: Calculated sigma12 ({sigma12_val:.4f}) is outside "
f"the valid emulator range of [0.2, 1.0]."
)
# --- Use the scalar value for logging ---
log.info(f" sigma12 = {sigma12_val:.4f} (OK)")
# --- Validate shape parameters in Planck eigenbasis ---
log.debug("Checking shape parameters against Planck prior box...")
input_params = np.array(
[cospar['omega_b'], cospar['omega_c'], cospar['n_s']])
centered_params = input_params - self.planck_means
projected_params = self.eigenvecs.T @ centered_params
deviations = np.abs(projected_params) / self.planck_sigmas
log.debug(f"Parameter deviations (in sigmas): {deviations}")
if np.any(deviations > 5.0):
failed_axis = np.argmax(deviations)
msg = (
f"Validation failed: Shape parameters are outside the 5-sigma Planck prior box.\n"
f" - Problem is in eigenvector direction {failed_axis}.\n"
f" - Deviation is {deviations[failed_axis]:.2f} sigma (limit is 5.0 sigma)."
)
log.error(msg)
raise ValueError(msg)
log.info("Shape parameters are within 5-sigma box (OK)")
def _get_xtilde(self, z, growth_obj, memory_scale=0.12, n_efolds=2.0, n_eta=1000):
"""Calculates the smoothed growth-dependent parameter xtilde.
The growth history X(tau) is convolved with a Gaussian memory kernel
looking back over a stretch of the history ending at the requested
redshift. The defaults are the P_nl values; the mass function uses a
wider kernel and correspondingly a longer, better-sampled window, which
:meth:`get_fnu` passes explicitly.
Parameters
----------
z : float
Redshift at which xtilde is evaluated.
growth_obj : GrowthCalculator
Growth history to integrate over.
memory_scale : float, optional
Width of the Gaussian memory kernel, in units of ln D. Default 0.12,
the value fitted for P_nl; :attr:`MEMORY_SCALE_HMF` is the mass-function
counterpart.
n_efolds : float, optional
Length of the integration window, in e-folds of the growth factor.
Default 2.
n_eta : int, optional
Number of samples across that window. Default 1000.
Returns
-------
float
The smoothed parameter xtilde.
"""
eta = np.log(growth_obj.Dgrowth(z))
eta_vec = np.linspace(eta - n_efolds, eta, n_eta)
x_vec = growth_obj.X_tau(eta_vec)
xtilde = simpson(
self.gaussian_kernel(eta, eta_vec, memory_scale) * x_vec, eta_vec)
return xtilde
def _get_redshift(self, target_sigma12, growth_ref, cosmo_engine_target, cospar_ref):
"""Finds the redshift z0 in a reference cosmology with the same sigma12.
This is an internal root-finding method. It finds the redshift `z0`
at which a standard reference LCDM cosmology has a sigma_12 value
equal to the `target_sigma12` of the user's input cosmology.
Parameters
----------
target_sigma12 : float
The sigma_12 value of the target cosmology that we want to match.
growth_ref : GrowthCalculator
An initialized GrowthCalculator instance for the reference LCDM cosmology.
cosmo_engine_target : Cosmology
The initialized Cosmology instance for the target cosmology.
cospar_ref : dict
The cosmology dictionary for the reference LCDM cosmology.
Returns
-------
float
The redshift, z0, in the reference cosmology.
"""
# Get D(z=0) for the reference cosmology
D_ref_0 = growth_ref.Dgrowth(0.)
# Calculate the expected sigma12 at z=0 for the reference cosmology.
# This is done by taking the sigma12(z=0) from the target cosmology's CAMB run
# and rescaling it by the ratio of sqrt(A_s) values.
# The ratio of D(0) values is a small correction for different normalizations.
sigma12_ref_0 = (cosmo_engine_target.sigma12_0 *
np.sqrt(cospar_ref['A_s'] / cosmo_engine_target.parameters.InitPower.As) *
(D_ref_0 / cosmo_engine_target.D0))
# This nested function describes the evolution of sigma12 in the reference model
def sig12_in_ref_cosmology(z):
return sigma12_ref_0 * growth_ref.Dgrowth(z) / D_ref_0
# The function whose root we want to find: f(z) = sigma12_ref(z) - target = 0
def delta_sigma12(z):
return float(np.ravel(sig12_in_ref_cosmology(z) - target_sigma12)[0])
# Use a robust root-finding algorithm to find z0
try:
result = root_scalar(
delta_sigma12, bracket=[-0.9, 4.4], method='brentq')
if not result.converged:
raise RuntimeError("Root-finding for z0 did not converge.")
return result.root
except ValueError as e:
# This can happen if delta_sigma12 has the same sign at both ends of the bracket
raise ValueError(f"Could not find a valid z0 for target sigma12={target_sigma12}. "
f"The value may be outside the solvable range. Original error: {e}")
# ==================================================================
# Halo mass function
# ==================================================================
[docs]
def get_fnu(self, nu, cospar, z):
"""Calculates the halo multiplicity function f(nu).
This is the main method of the HMF emulator. It takes a set of peak heights,
a cosmology, and a redshift, and returns the emulated f(nu), defined such
that dn/dlnM = f(nu) * rho_b / M * dln(nu)/dlnM.
Parameters
----------
nu : array_like
Array of peak heights, nu = delta_c / sigma(M, z).
Must lie within the window the emulator is trained for at the
sigma12 of this cosmology and redshift, :meth:`nu_range_at`.
cospar : dict
A dictionary of cosmological parameters such as the one created
by `create_cosmo_dict`.
z : float
The redshift at which to calculate the multiplicity function.
Returns
-------
np.ndarray
The halo multiplicity function, f(nu), dimensionless.
Raises
------
ValueError
If the input cosmology fails the validation checks (sigma12 outside
[0.2, 1.0] or shape parameters outside the 5-sigma Planck box), or if
any nu-value falls outside the sigma12-dependent window the emulator
is trained for.
Notes
-----
The calculation involves several steps:
1. A `Cosmology` object is created to compute sigma12(z).
2. Inputs are validated: shape parameters and sigma12 first, then nu
against the window that sigma12 implies.
3. The ratio f/f_ref and the response dR/dxtilde are predicted by GPs.
4. `xtilde` is computed for the target cosmology and for the reference
cosmology at the redshift z0 where it has the same sigma12.
5. All components are combined into the final predicted f(nu).
"""
log.info(f"Received request for f(nu) at z={z:.2f}")
nu = np.atleast_1d(nu).astype(float)
z = float(z)
log.info("Initializing cosmology engine for target cosmology.")
cosmology_engine = self._get_cosmology_engine(cospar)
sigma12 = cosmology_engine.get_sigma12(z)
# Input parameter validation: shape params
# and sigma12 range first
self._validate_params(cospar, sigma12)
# then check the (nu, sigma12) pair is within the boundary of the training data
self._validate_nu(nu, sigma12)
log.info("Computing emulator predictions from GPs.")
x_combined = self._build_gp_input(nu, cospar, sigma12)
# --- Bypass Scikit-Learn predict() ---
ratio = np.exp(self._gp_predict(self.gp_fnu, x_combined))
dfdx = self._gp_predict(self.gp_dfdx, x_combined)
log.info("Computing xtilde for target and reference cosmologies.")
# Memory kernel of the mass function: wider than the P_nl one, so the
# window is longer and more finely sampled. Both evaluations below must
# use the same recipe as the training set.
kernel = dict(memory_scale=self.MEMORY_SCALE_HMF,
n_efolds=2.0, n_eta=3000)
xtilde = self._get_xtilde(z, cosmology_engine.growth, **kernel)
cospar_ref = cospar.copy()
cospar_ref.update(self.REFERENCE_COSPAR)
growth_ref = GrowthCalculator(cospar_ref)
z0 = self._get_redshift(
sigma12, # The value to match
growth_ref, # The reference growth engine
cosmology_engine, # The target cosmology engine
cospar_ref # The reference parameter dict
)
xtilde_0 = self._get_xtilde(z0, growth_ref, **kernel)
log.debug(
f"Target xtilde={xtilde:.4f}. Reference z0={z0:.4f}, xtilde_0={xtilde_0:.4f}")
log.info("Combining components for final prediction.")
dxtilde = xtilde - xtilde_0
fnu_reference = self._reference_fnu(nu)
fnu = fnu_reference * ratio * (1. + dfdx * dxtilde)
log.info("f(nu) calculation complete.")
return fnu
[docs]
def nM_dlogM(self, M, cospar, z):
"""Computes the differential mass function dn/dlnM at redshift z.
Parameters
----------
M : array_like
Halo masses at which the mass function is evaluated, in Msun.
cospar : dict
A dictionary of cosmological parameters such as the one created
by `create_cosmo_dict`.
z : float
The redshift at which to calculate the mass function.
Returns
-------
np.ndarray
dn/dlnM, in Mpc^-3.
Raises
------
ValueError
If the masses map onto peak heights outside the window the emulator
is trained for at this sigma12, or if the cosmology fails validation.
The accessible mass range shrinks towards high redshift.
Notes
-----
The mass-to-peak-height conversion is :meth:`Cosmology.mass_to_nu`:
sigma(R) from CAMB at z=0 for the input cosmology, scaled to the
requested redshift with the linear growth factor, exactly as in
:class:`HMF_Fiorilli25`. Call it on a `Cosmology` built from `cospar` to
pick masses that land inside :meth:`nu_range_at`, or to label a mass axis
in peak height. Masses are always M200b, the single mass definition the
emulator is trained for.
"""
M = np.atleast_1d(M).astype(float)
z = float(z)
cosmology_engine = self._get_cosmology_engine(cospar)
# --- Convert M to nu ---
nu = np.atleast_1d(cosmology_engine.mass_to_nu(M, z))
# dln sigma / dln R at each Lagrangian radius, needed to turn f(nu) into
# dn/dlnM. Redshift independent, and it reuses the engine's cached splines.
_, spline_log_sigma0, _ = cosmology_engine.sigma0_splines()
lagrangian_radius = (
3 * M / (4 * np.pi * cosmology_engine.rho_bar))**(1. / 3.) # Mpc
dlogsigma_dlogR = spline_log_sigma0.derivative()(np.log(lagrangian_radius))
# Shape parameters and sigma12 first, then nu against the window that
# sigma12 implies. get_fnu repeats both checks.
sigma12 = cosmology_engine.get_sigma12(z)
self._validate_params(cospar, sigma12)
self._validate_nu(nu, sigma12)
fnu = self.get_fnu(nu, cospar, z)
# --- Convert f(nu) to dn/dlnM ---
nM_dlogM = (fnu * cosmology_engine.rho_bar * -1.0 / 3.0
* dlogsigma_dlogR / M)
if nM_dlogM.size == 1:
return nM_dlogM[0]
return nM_dlogM
def _fit_nu_boundary(self, gp, side):
"""Fits one nu boundary of a GP's training coverage as a function of sigma12.
Each training node occupies a contiguous stretch of nu bins whose extent
is set by the resolution and volume of the box it came from, and both
limits move with sigma12. Grouping a GP's stored inputs by node recovers
the (nu_min, nu_max) pair at each node's sigma12; the raw limits snap to
discrete nu bin edges, so a low-order polynomial through them is a better
estimate of the underlying boundary than the steps.
Parameters
----------
gp : sklearn.gaussian_process.GaussianProcessRegressor
A trained GP whose `X_train_` holds columns
(ln nu, omega_b, omega_c, n_s, sigma12).
side : {'lower', 'upper'}
Which boundary of the nu coverage to fit.
Returns
-------
np.ndarray
Polynomial coefficients in `np.polyfit` order, for nu as a function
of sigma12.
"""
X = gp.X_train_
# Columns 1..4 identify the training node; all its rows share a sigma12.
nodes, index = np.unique(X[:, 1:5], axis=0, return_inverse=True)
sigma12 = nodes[:, 3]
reduce = np.min if side == 'lower' else np.max
nu_edge = np.array([np.exp(reduce(X[index == j, 0]))
for j in range(len(nodes))])
# Degree 3, as in plot_parameter_range.ipynb: the raw per-node limits snap
# to discrete nu bin edges, so a low-order fit is a better estimate of the
# underlying boundary than the steps themselves.
return np.polyfit(sigma12, nu_edge, 3)
[docs]
def nu_range_at(self, sigma12):
"""Peak heights the emulator is trained for at a given sigma12.
The window is the nu coverage of the f(nu) GP, which carries the leading
term of the prediction. It narrows as sigma12 falls, from a width of
about 2.2 at sigma12 = 1 down to less than 0.1 at sigma12 = 0.2.
Parameters
----------
sigma12 : float
The clustering amplitude of the target cosmology at the redshift of
interest. Assumed to lie in :attr:`SIGMA12_RANGE`, which is checked
separately; the polynomials are not meaningful outside it.
Returns
-------
tuple of float
The (min, max) valid peak height at this sigma12.
"""
sigma12 = float(np.atleast_1d(sigma12)[0])
return (float(np.polyval(self.nu_lo_coeffs, sigma12)),
float(np.polyval(self.nu_hi_coeffs, sigma12)))
def _validate_nu(self, nu, sigma12):
"""Checks the requested peak heights against the region of validity.
Unlike the sigma12 and shape checks, this one cannot be made from the
inputs alone: the accessible nu window depends on sigma12, so the
cosmology has to be solved first.
Parameters
----------
nu : np.ndarray
The requested peak heights.
sigma12 : array_like
The clustering amplitude of the target cosmology at the requested
redshift.
Raises
------
ValueError
If any peak height falls outside the window.
"""
sigma12_val = float(np.atleast_1d(sigma12)[0])
nu_lo, nu_hi = self.nu_range_at(sigma12_val)
outside = (nu < nu_lo) | (nu > nu_hi)
if np.any(outside):
msg = (f"Validation failed: {int(outside.sum())} of {nu.size} peak "
f"height(s) fall outside the emulator's region of validity. "
f"At sigma12 = {sigma12_val:.4f} the training set covers "
f"nu in [{nu_lo:.3f}, {nu_hi:.3f}]; the requested values span "
f"[{np.min(nu):.3f}, {np.max(nu):.3f}].")
msg += (" Note that this window narrows as sigma12 falls; see "
"nu_range_at().")
log.error(msg)
raise ValueError(msg)
log.info(f" nu within [{nu_lo:.3f}, {nu_hi:.3f}] at "
f"sigma12 = {sigma12_val:.4f} (OK)")
def _reference_fnu(self, nu):
"""Evaluates the reference multiplicity function at the requested nu."""
return np.exp(splev(np.log(nu), self.reference_spline_hmf))
def _get_cosmology_engine(self, cospar):
"""Returns a `Cosmology` for `cospar`, reusing the last one if unchanged."""
key = tuple(sorted(cospar.items()))
if key != self._cached_cospar:
log.debug("Cosmology cache miss; running CAMB.")
self._cached_engine = Cosmology(cospar)
self._cached_cospar = key
return self._cached_engine
[docs]
@staticmethod
def gaussian_kernel(tau, tau_prime, tau_s):
"""Computes a normalized Gaussian kernel."""
return np.exp(-((tau - tau_prime) ** 2) / (2 * tau_s ** 2)) * 2. / (np.sqrt(2 * np.pi) * tau_s)
[docs]
@staticmethod
def create_cosmo_dict(h, omega_b=None, omega_c=None, omega_k=0., Omega_b=None, Omega_c=None,
Omega_k=0., A_s=2.1e-9, n_s=0.96, w_0=-1.0, w_a=0.0,
model='LCDM', density_type='physical'):
"""
Creates a standardized cosmology dictionary for the Aletheia Emulator.
Args:
h (float): The Hubble parameter. Required for all conversions.
omega_b, omega_c, omega_k (float, optional): Physical baryon/CDM densities.
Omega_b, Omega_c, Omega_nu (float, optional): Fractional baryon/CDM densities.
... (other parameters with defaults)
model (str, optional): Cosmological model, e.g., 'LCDM' or 'W0WACDM'.
density_type (str, optional): 'physical' (little omega) or 'fractional' (big Omega).
Returns:
dict: A validated dictionary ready for the emulator.
"""
cospar = {}
# --- Handle density parameter conversion ---
if density_type == 'physical':
if omega_b is None or omega_c is None:
raise ValueError(
"For 'physical' density_type, 'omega_b', 'omega_c', and 'omega_k' must be provided.")
cospar['omega_b'] = omega_b
cospar['omega_c'] = omega_c
cospar['omega_k'] = omega_k
cospar['omega_nu'] = 0. # The current version assumes massless neutrinos
elif density_type == 'fractional':
if Omega_b is None or Omega_c is None:
raise ValueError(
"For 'fractional' density_type, 'Omega_b', 'Omega_c', and 'Omega_k' must be provided.")
cospar['omega_b'] = Omega_b * h**2
cospar['omega_c'] = Omega_c * h**2
cospar['omega_k'] = Omega_k * h**2
cospar['omega_nu'] = 0. # The current version assumes massless neutrinos
else:
raise ValueError(f"Unknown density_type: '{density_type}'")
# --- Compute derived parameters ---
cospar['omega_de'] = h**2 - cospar['omega_c'] - \
cospar['omega_b'] - cospar['omega_k'] - cospar['omega_nu']
# --- Set remaining parameters ---
cospar['h'] = h
cospar['A_s'] = A_s
cospar['n_s'] = n_s
# --- Handle model-specific defaults ---
if model.upper() == 'LCDM':
cospar['w_0'] = -1.0
cospar['w_a'] = 0.0
elif model.upper() == 'W0WACDM':
cospar['w_0'] = w_0
cospar['w_a'] = w_a
else:
raise ValueError(f"Unknown model: '{model}'")
return cospar