From 4265074a581243e271db4cb70e2777a686bb4405 Mon Sep 17 00:00:00 2001 From: sambit-giri Date: Wed, 18 Jun 2025 08:27:40 +0200 Subject: [PATCH 1/3] init custom baryon suppression --- pyccl/baryons/__init__.py | 1 + pyccl/baryons/baryons_custom.py | 89 +++++++++++++++++++++++++++++++++ 2 files changed, 90 insertions(+) create mode 100644 pyccl/baryons/baryons_custom.py 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..f8e76947b --- /dev/null +++ b/pyccl/baryons/baryons_custom.py @@ -0,0 +1,89 @@ +__all__ = ("BaryonsCustom",) + +import numpy as np + +from .. import Pk2D +from . import Baryons + + +class BaryonsCustom(Baryons): + """The "BCM" model boost factor for baryons. BCM stands for the + "baryonic correction model" of `Schneider & Teyssier 2015 + `_. See the + + 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`): Array containing the boost factor 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, filename): + self.filename = filename + + def boost_factor(self, cosmo, k, a): + """The BCM model boost factor for baryons. + + 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_use, k_use = map(np.atleast_1d, [a, k]) + a_use, k_use = a_use[:, None], k_use[None, :] + + z = 1/a_use - 1 + kh = k_use / cosmo['h'] + b0 = 0.105*self.log10Mc - 1.27 + bfunc = b0 / (1. + (z/2.3)**2.5) + kg = 0.7 * (1-bfunc)**4 * self.eta_b**(-1.6) + gf = bfunc / (1 + (kh/kg)**3) + 1. - bfunc + scomp = 1 + (kh / self.k_s)**2 + fka = gf * scomp + + if np.ndim(k) == 0: + fka = np.squeeze(fka, axis=-1) + if np.ndim(a) == 0: + fka = np.squeeze(fka, axis=0) + return fka + + def update_parameters(self, log10Mc=None, eta_b=None, k_s=None): + """Update BCM parameters. All parameters set to ``None`` will + be left untouched. + + Args: + log10Mc (:obj:`float`): logarithmic mass scale of hot + gas suppression. + eta_b (:obj:`float`): ratio of escape to ejection radii. + k_s (:obj:`float`): Characteristic scale (wavenumber) of + the stellar component. + """ + if log10Mc is not None: + self.log10Mc = log10Mc + if eta_b is not None: + self.eta_b = eta_b + if k_s is not None: + self.k_s = k_s + + 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) From a42c96e42a840ab7e1a7057a4609bd896ffa71c1 Mon Sep 17 00:00:00 2001 From: sambit-giri Date: Thu, 19 Jun 2025 10:47:28 +0200 Subject: [PATCH 2/3] working for flamingo data --- pyccl/baryons/baryons_custom.py | 79 ++++++++++++++++++++------------- 1 file changed, 47 insertions(+), 32 deletions(-) diff --git a/pyccl/baryons/baryons_custom.py b/pyccl/baryons/baryons_custom.py index f8e76947b..334e83898 100644 --- a/pyccl/baryons/baryons_custom.py +++ b/pyccl/baryons/baryons_custom.py @@ -4,12 +4,10 @@ from .. import Pk2D from . import Baryons - +from scipy.interpolate import RegularGridInterpolator class BaryonsCustom(Baryons): - """The "BCM" model boost factor for baryons. BCM stands for the - "baryonic correction model" of `Schneider & Teyssier 2015 - `_. See the + """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)`. @@ -22,11 +20,21 @@ class BaryonsCustom(Baryons): name = 'BaryonsCustom' __repr_attrs__ = __eq_attrs__ = ("boost_data", "k_data", "a_data") - def __init__(self, filename): - self.filename = filename + 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): - """The BCM model boost factor for baryons. + """Interpolated baryonic boost factor. Args: cosmo (:class:`~pyccl.cosmology.Cosmology`): Cosmological parameters. @@ -37,25 +45,23 @@ def boost_factor(self, cosmo, k, a): :obj:`float` or `array`: Correction factor to apply to \ the power spectrum. """ # noqa - a_use, k_use = map(np.atleast_1d, [a, k]) - a_use, k_use = a_use[:, None], k_use[None, :] - - z = 1/a_use - 1 - kh = k_use / cosmo['h'] - b0 = 0.105*self.log10Mc - 1.27 - bfunc = b0 / (1. + (z/2.3)**2.5) - kg = 0.7 * (1-bfunc)**4 * self.eta_b**(-1.6) - gf = bfunc / (1 + (kh/kg)**3) + 1. - bfunc - scomp = 1 + (kh / self.k_s)**2 - fka = gf * scomp - - if np.ndim(k) == 0: - fka = np.squeeze(fka, axis=-1) - if np.ndim(a) == 0: - fka = np.squeeze(fka, axis=0) - return fka - - def update_parameters(self, log10Mc=None, eta_b=None, k_s=None): + 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) + + # Match output shape to inputs + if np.ndim(a) == 0 and np.ndim(k) == 0: + return boost_vals[0, 0] + elif np.ndim(a) == 0: + return boost_vals[0] + elif np.ndim(k) == 0: + return boost_vals[:, 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. @@ -66,12 +72,21 @@ def update_parameters(self, log10Mc=None, eta_b=None, k_s=None): k_s (:obj:`float`): Characteristic scale (wavenumber) of the stellar component. """ - if log10Mc is not None: - self.log10Mc = log10Mc - if eta_b is not None: - self.eta_b = eta_b - if k_s is not None: - self.k_s = k_s + 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=(a_data, k_data), + values=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 From 34dea1351f56a9905c10f397b5aeff3fda782534 Mon Sep 17 00:00:00 2001 From: sambit-giri Date: Wed, 18 Mar 2026 15:48:34 +0100 Subject: [PATCH 3/3] Address review comments: fix docstrings, simplify boost_factor, fix update_parameters bug --- pyccl/baryons/baryons_custom.py | 41 +++++++++++++++++---------------- 1 file changed, 21 insertions(+), 20 deletions(-) diff --git a/pyccl/baryons/baryons_custom.py b/pyccl/baryons/baryons_custom.py index 334e83898..7d97bebfa 100644 --- a/pyccl/baryons/baryons_custom.py +++ b/pyccl/baryons/baryons_custom.py @@ -6,14 +6,17 @@ from . import Baryons from scipy.interpolate import RegularGridInterpolator + class BaryonsCustom(Baryons): - """The custom baryonic model boost factor for 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`): Array containing the boost factor data. + 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. """ @@ -37,28 +40,27 @@ 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}`). + 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 + """ # 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) - - # Match output shape to inputs - if np.ndim(a) == 0 and np.ndim(k) == 0: - return boost_vals[0, 0] - elif np.ndim(a) == 0: - return boost_vals[0] - elif np.ndim(k) == 0: - return boost_vals[:, 0] + + 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): @@ -66,11 +68,10 @@ def update_parameters(self, boost_data=None, k_data=None, a_data=None): be left untouched. Args: - log10Mc (:obj:`float`): logarithmic mass scale of hot - gas suppression. - eta_b (:obj:`float`): ratio of escape to ejection radii. - k_s (:obj:`float`): Characteristic scale (wavenumber) of - the stellar component. + 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 @@ -82,8 +83,8 @@ def update_parameters(self, boost_data=None, k_data=None, a_data=None): # 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=(a_data, k_data), - values=boost_data, + 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 )