diff --git a/pyccl/baryons/__init__.py b/pyccl/baryons/__init__.py index 344b2bf61..8a35d889c 100644 --- a/pyccl/baryons/__init__.py +++ b/pyccl/baryons/__init__.py @@ -2,3 +2,4 @@ from .schneider15 import * from .baccoemu_baryons import * from .vandaalen19 import * +from .baryons_custom import * diff --git a/pyccl/baryons/baryons_custom.py b/pyccl/baryons/baryons_custom.py new file mode 100644 index 000000000..7d97bebfa --- /dev/null +++ b/pyccl/baryons/baryons_custom.py @@ -0,0 +1,105 @@ +__all__ = ("BaryonsCustom",) + +import numpy as np + +from .. import Pk2D +from . import Baryons +from scipy.interpolate import RegularGridInterpolator + + +class BaryonsCustom(Baryons): + """The custom baryonic model boost factor for baryons. + + The boost factor is applied multiplicatively so that + :math:`P_{\\rm bar.}(k, a) = P_{\\rm DMO}(k, a)\\, f_{\\rm BCM}(k, a)`. + + Args: + boost_data (:obj:`array`): 2D array of shape ``(n_a, n_k)`` + containing the boost factor, with rows corresponding to + entries in ``a_data`` and columns to entries in ``k_data``. + k_data (:obj:`array`): Wavenumber (in :math:`{\\rm Mpc}^{-1}`). + a_data (:obj:`array`): Scale factor. + """ + name = 'BaryonsCustom' + __repr_attrs__ = __eq_attrs__ = ("boost_data", "k_data", "a_data") + + def __init__(self, boost_data, k_data, a_data): + self.boost_data = boost_data + self.k_data = k_data + self.a_data = a_data + + # Create interpolator in (a, k) space + self._interpolator = RegularGridInterpolator( + points=(a_data, k_data), + values=boost_data, + bounds_error=False, + fill_value=1.0 # default to 1 (no boost) outside bounds + ) + + def boost_factor(self, cosmo, k, a): + """Interpolated baryonic boost factor. + + Args: + cosmo (:class:`~pyccl.cosmology.Cosmology`): Cosmological + parameters. + k (:obj:`float` or `array`): Wavenumber + (in :math:`{\\rm Mpc}^{-1}`). + a (:obj:`float` or `array`): Scale factor. + + Returns: + :obj:`float` or `array`: Correction factor to apply to \ + the power spectrum. + """ # noqa + a_arr = np.atleast_1d(a) + k_arr = np.atleast_1d(k) + a_mesh, k_mesh = np.meshgrid(a_arr, k_arr, indexing='ij') + query_points = np.column_stack([a_mesh.ravel(), k_mesh.ravel()]) + + boost_vals = self._interpolator(query_points).reshape(a_mesh.shape) + + if np.ndim(k) == 0: + boost_vals = np.squeeze(boost_vals, axis=-1) + if np.ndim(a) == 0: + boost_vals = np.squeeze(boost_vals, axis=0) + return boost_vals + + def update_parameters(self, boost_data=None, k_data=None, a_data=None): + """Update BCM parameters. All parameters set to ``None`` will + be left untouched. + + Args: + boost_data (:obj:`array`): 2D array of shape ``(n_a, n_k)`` + containing the boost factor. + k_data (:obj:`array`): Wavenumber (in :math:`{\\rm Mpc}^{-1}`). + a_data (:obj:`array`): Scale factor. + """ + if boost_data is not None: + self.boost_data = boost_data + if k_data is not None: + self.k_data = k_data + if a_data is not None: + self.a_data = a_data + + # Update interpolator in (a, k) space + if boost_data is not None or k_data is not None or a_data is not None: + self._interpolator = RegularGridInterpolator( + points=(self.a_data, self.k_data), + values=self.boost_data, + bounds_error=False, + fill_value=1.0 # default to 1 (no boost) outside bounds + ) + + def _include_baryonic_effects(self, cosmo, pk): + # Applies boost factor + a_arr, lk_arr, pk_arr = pk.get_spline_arrays() + k_arr = np.exp(lk_arr) + fka = self.boost_factor(cosmo, k_arr, a_arr) + pk_arr *= fka + + if pk.psp.is_log: + np.log(pk_arr, out=pk_arr) # in-place log + + return Pk2D(a_arr=a_arr, lk_arr=lk_arr, pk_arr=pk_arr, + is_logp=pk.psp.is_log, + extrap_order_lok=pk.extrap_order_lok, + extrap_order_hik=pk.extrap_order_hik)