Source code for cosmofit.cosmology.numerics.integrals
"""
Numerical integration utilities.
This module computes the dimensionless comoving distance
χ(z) = ∫ dz / E(z)
with a corrected (Euler-Maclaurin) trapezoid rule on a uniform
redshift grid -- fourth order, using the model's own dE/dz -- and
interpolates it with a cubic Hermite spline whose slopes are the
exact 1/E(z).
The integral is evaluated only once for each cosmology,
making MCMC analyses orders of magnitude faster than
repeated scipy.integrate.quad evaluations.
"""
from __future__ import annotations
import numpy as np
from cosmofit.cosmology.numerics.hermite import hermite_spline
# ============================================================
# Adaptive integration grid
# ============================================================
def adaptive_grid(zmax):
"""
Grid size for the chi(z) = int dz/E(z) interpolation table.
chi(z) is an extremely smooth function of z (E(z) itself is
smooth for every model in this library), so PCHIP interpolation
converges fast: tested against a 20000-point reference across
randomized LCDM/CPL parameters, even 300 points reach ~5e-5
relative accuracy in chi(z) at zmax=5 -- a fractional distance
error around 100x smaller than the ~0.1-1% measurement
precision of any dataset this grid is used for (CC, BAO, SN).
The 400-1300 range below keeps a comfortable safety margin
over that while still shrinking the grid (and the per-MCMC-step
cost of rebuilding it -- see ``DistanceIntegrator.rebuild``)
several-fold versus a much denser grid.
"""
return int(
np.clip(
350 + 200 * np.log10(zmax + 1),
400,
1300,
)
)
# ============================================================
# Distance Integrator
# ============================================================
[docs]
class DistanceIntegrator:
"""
Fast evaluator for the dimensionless comoving distance.
The integral
χ(z) = ∫ dz / E(z)
is computed only once on a dense redshift grid and
subsequently evaluated through a cubic Hermite interpolant.
Asked for a redshift beyond the grid, it extends the grid
rather than returning NaN.
Parameters
----------
cosmology : Cosmology
Cosmological model.
zmax : float, optional
Maximum redshift of the interpolation grid.
ngrid : int, optional
Number of grid points. If None, an adaptive
grid size is automatically selected.
interpolator : str, optional
Interpolation method.
Currently only "hermite" is implemented.
"""
[docs]
def __init__(
self,
cosmology,
zmax: float = 5.0,
ngrid: int | None = None,
interpolator: str = "hermite",
):
self.cosmo = cosmology
self.zmax = float(zmax)
self.interpolator = interpolator.lower()
if ngrid is None:
ngrid = adaptive_grid(self.zmax)
self.ngrid = int(ngrid)
self._build()
# --------------------------------------------------------
def _build(self) -> None:
"""
Build the interpolation table.
"""
self.z_grid = np.linspace(
0.0,
self.zmax,
self.ngrid,
)
if self.interpolator != "hermite":
raise ValueError(
f"Unknown interpolator '{self.interpolator}'."
)
Ez = self.cosmo.E(self.z_grid)
invE = 1.0 / Ez
# d(1/E)/dz, analytically. Having the integrand's own
# derivative is what makes both steps below exact to
# fourth order instead of second.
dinvE = -self.cosmo.dEdz(self.z_grid) / Ez ** 2
h = np.diff(self.z_grid)
# Corrected (Euler-Maclaurin) trapezoid: the plain rule
# plus the end-derivative term, which is the integral of
# the cubic Hermite interpolant of 1/E over each interval.
# Fourth order, for one extra array evaluation.
segments = (
0.5 * h * (invE[:-1] + invE[1:])
+ h ** 2 / 12.0 * (dinvE[:-1] - dinvE[1:])
)
chi = np.concatenate([[0.0], np.cumsum(segments)])
# Hermite rather than a shape-preserving cubic: chi'(z) is
# exactly 1/E(z), so there is nothing to estimate. Pchip
# spent most of the build time deriving slopes that were
# already known, and got them slightly wrong -- against a
# quadrature reference on the default grid, the maximum
# relative error in chi(z) falls from 4.5e-06 to 1.3e-10.
#
# Monotonicity, which is why a shape-preserving spline was
# chosen originally, is not at risk: chi is the integral of
# a strictly positive function, and the interpolant is
# built from that function's own positive values.
self._chi = hermite_spline(
self.z_grid,
chi,
invE,
extrapolate=False,
)
# --------------------------------------------------------
def rebuild(self) -> None:
"""
Recompute the chi(z) interpolation table from the
current state of the associated cosmology.
This MUST be called whenever the cosmology's parameters
change (e.g. at every MCMC step) -- the table is built
once from ``cosmo.E(z)`` and is NOT automatically kept
in sync with parameter updates, since ``Cosmology``
objects are designed to be mutated in place (via
``params.update(...)``) for performance rather than
reconstructed every step.
See also :meth:`cosmology.core.base.Cosmology.refresh`,
which is the method user code should normally call.
"""
self._build()
# --------------------------------------------------------
def prepare(
self,
zmax: float,
ngrid: int | None = None,
) -> None:
"""
Prepare the integration grid.
The grid is rebuilt only if a larger redshift
range is required.
"""
if zmax <= self.zmax:
return
self.zmax = float(zmax)
if ngrid is None:
self.ngrid = adaptive_grid(self.zmax)
else:
self.ngrid = int(ngrid)
self._build()
# --------------------------------------------------------
def chi(
self,
z,
):
"""
Evaluate the dimensionless comoving distance.
Parameters
----------
z : float or ndarray
Returns
-------
float or ndarray
χ(z)
Notes
-----
A redshift beyond the table's ``zmax`` (5 by default) used to
come back NaN -- out of the spline's range -- which a
likelihood turned into chi2 = NaN and a sampler into a silent
rejection of every point. No bundled dataset reaches past
z = 2.33, so nothing showed it; a quasar or GRB Hubble
diagram would. The table now grows to cover the request, and
keeps that range on every later rebuild.
"""
top = float(np.max(z)) if np.size(z) else 0.0
if top > self.zmax:
self.prepare(1.1 * top)
return self._chi(z)