diff --git a/firecrown/models/cluster/kernel.py b/firecrown/models/cluster/kernel.py index 3b7b6d78..3e0763b1 100644 --- a/firecrown/models/cluster/kernel.py +++ b/firecrown/models/cluster/kernel.py @@ -8,6 +8,15 @@ import numpy.typing as npt import numpy as np +from firecrown import parameters +from firecrown.updatable import Updatable + + +REDMAPPER_DEFAULT_AC_NC = 0.38 +REDMAPPER_DEFAULT_BC_NC = 1.2634 +REDMAPPER_DEFAULT_AC_MC = 13.31 +REDMAPPER_DEFAULT_BC_MC = 0.2025 + class KernelType(Enum): """The kernels that can be included in the cluster abundance integrand.""" @@ -20,47 +29,98 @@ class KernelType(Enum): PURITY = 6 -class Completeness: +class Completeness(Updatable): """The completeness kernel for the numcosmo simulated survey. This kernel will affect the integrand by accounting for the incompleteness of a cluster selection. """ + def __init__( + self, + ): + super().__init__() + # Updatable parameters + self.ac_nc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_AC_NC + ) + self.bc_nc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_BC_NC + ) + self.ac_mc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_AC_MC + ) + self.bc_mc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_BC_MC + ) + + def _mc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + ac_mc = self.ac_mc + bc_mc = self.bc_mc + log_mc = ac_mc + bc_mc * (1.0 + z) + mc = 10.0**log_mc + return mc.astype(np.float64) + + def _nc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + ac_nc = self.ac_nc + bc_nc = self.bc_nc + nc = ac_nc + bc_nc * (1.0 + z) + assert isinstance(nc, np.ndarray) + return nc + def distribution( self, - mass: npt.NDArray[np.float64], + log_mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], ) -> npt.NDArray[np.float64]: """Evaluates and returns the completeness contribution to the integrand.""" - a_nc = 1.1321 - b_nc = 0.7751 - a_mc = 13.31 - b_mc = 0.2025 - log_mc = a_mc + b_mc * (1.0 + z) - nc = a_nc + b_nc * (1.0 + z) - completeness = (mass / log_mc) ** nc / ((mass / log_mc) ** nc + 1.0) + mc = self._mc(z) + mass = 10.0**log_mass + nc = self._nc(z) + completeness = (mass / mc) ** nc / ((mass / mc) ** nc + 1.0) assert isinstance(completeness, np.ndarray) return completeness -class Purity: +REDMAPPER_DEFAULT_AP_NC = 3.9193 +REDMAPPER_DEFAULT_BP_NC = -0.3323 +REDMAPPER_DEFAULT_AP_RC = 1.1839 +REDMAPPER_DEFAULT_BP_RC = -0.4077 + + +class Purity(Updatable): """The purity kernel for the numcosmo simulated survey. This kernel will affect the integrand by accounting for the purity of a cluster selection. """ - def _ln_rc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - a_rc = 2.2183 - b_rc = -0.6592 - ln_rc = a_rc + b_rc * (1.0 + z) - return ln_rc.astype(np.float64) + def __init__(self): + super().__init__() + self.ap_nc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_AP_NC + ) + self.bp_nc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_BP_NC + ) + self.ap_rc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_AP_RC + ) + self.bp_rc = parameters.register_new_updatable_parameter( + default_value=REDMAPPER_DEFAULT_BP_RC + ) + + def _rc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + ap_rc = self.ap_rc + bp_rc = self.bp_rc + log_rc = ap_rc + bp_rc * (1.0 + z) + rc = 10**log_rc + return rc.astype(np.float64) def _nc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: - b_nc = np.log(10) * 0.3527 - a_nc = np.log(10) * 0.8612 - nc = a_nc + b_nc * (1.0 + z) + bp_nc = self.bp_nc + ap_nc = self.ap_nc + nc = ap_nc + bp_nc * (1.0 + z) assert isinstance(nc, np.ndarray) return nc @@ -68,16 +128,16 @@ def distribution( self, z: npt.NDArray[np.float64], mass_proxy: npt.NDArray[np.float64], - mass_proxy_limits: tuple[float, float], + mass_proxy_limits: tuple[float, float] = None, ) -> npt.NDArray[np.float64]: """Evaluates and returns the purity contribution to the integrand.""" if all(mass_proxy == -1.0): mean_mass = (mass_proxy_limits[0] + mass_proxy_limits[1]) / 2 - ln_r = np.log(10**mean_mass) + r = 10**mean_mass else: - ln_r = np.log(10**mass_proxy) + r = np.power(10.0, mass_proxy) - r_over_rc = ln_r / self._ln_rc(z) + r_over_rc = r / self._rc(z) purity = (r_over_rc) ** self._nc(z) / (r_over_rc ** self._nc(z) + 1.0) assert isinstance(purity, np.ndarray) diff --git a/firecrown/models/cluster/mass_proxy.py b/firecrown/models/cluster/mass_proxy.py index 3fc105cf..4a0a1e1b 100644 --- a/firecrown/models/cluster/mass_proxy.py +++ b/firecrown/models/cluster/mass_proxy.py @@ -93,7 +93,7 @@ def _distribution_unbinned( normalization = 1 / np.sqrt(2 * np.pi * proxy_sigma**2) result = normalization * np.exp( - -0.5 * ((mass_proxy - proxy_mean) / proxy_sigma) ** 2 + -0.5 * ((mass_proxy * np.log(10) - proxy_mean) / proxy_sigma) ** 2 ) assert isinstance(result, np.ndarray) diff --git a/firecrown/models/cluster/recipes/murata_binned_spec_z_selection.py b/firecrown/models/cluster/recipes/murata_binned_spec_z_selection.py new file mode 100644 index 00000000..4fff475b --- /dev/null +++ b/firecrown/models/cluster/recipes/murata_binned_spec_z_selection.py @@ -0,0 +1,147 @@ +"""Module for defining the classes used in the MurataBinnedSpecZ cluster recipe.""" + +from typing import Callable + +import numpy as np +import numpy.typing as npt + +from firecrown.models.cluster.abundance import ClusterAbundance +from firecrown.models.cluster.binning import NDimensionalBin +from firecrown.models.cluster.integrator.numcosmo_integrator import NumCosmoIntegrator +from firecrown.models.cluster.kernel import SpectroscopicRedshift +from firecrown.models.cluster.mass_proxy import MurataUnbinned +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe +from firecrown.models.cluster.kernel import Completeness, Purity + + +class MurataBinnedSpecZSelectionRecipe(ClusterRecipe): + """Cluster recipe with Murata19 mass-richness and spec-zs. + + This recipe uses the Murata 2019 binned mass-richness relation and assumes + perfectly measured spec-zs. + """ + + def __init__(self) -> None: + super().__init__() + + self.integrator = NumCosmoIntegrator() + self.redshift_distribution = SpectroscopicRedshift() + pivot_mass, pivot_redshift = 14.3, 0.5 + self.mass_distribution = MurataUnbinned(pivot_mass, pivot_redshift) + self.completeness_distribution = Completeness() + self.purity_distribution = Purity() + self.my_updatables.append(self.mass_distribution) + self.my_updatables.append(self.completeness_distribution) + self.my_updatables.append(self.purity_distribution) + + def get_theory_prediction( + self, + cluster_theory: ClusterAbundance, + average_on: None | ClusterProperty = None, + ) -> Callable[ + [ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + float, + ], + npt.NDArray[np.float64], + ]: + """Get a callable that evaluates a cluster theory prediction. + + Returns a callable function that accepts mass, redshift, mass proxy limits, + and the sky area of your survey and returns the theoretical prediction for the + expected number of clusters. + """ + + def theory_prediction( + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + mass_proxy: npt.NDArray[np.float64], + sky_area: float, + ): + prediction = ( + cluster_theory.comoving_volume(z, sky_area) + * cluster_theory.mass_function(mass, z) + * self.redshift_distribution.distribution() + * self.mass_distribution.distribution(mass, z, mass_proxy / np.log(10)) + * self.completeness_distribution.distribution(mass, z) + / self.purity_distribution.distribution(z, mass_proxy=mass_proxy) + ) + + if average_on is None: + return prediction + + for cluster_prop in ClusterProperty: + include_prop = cluster_prop & average_on + if not include_prop: + continue + if cluster_prop == ClusterProperty.MASS: + prediction *= mass + elif cluster_prop == ClusterProperty.REDSHIFT: + prediction *= z + else: + raise NotImplementedError(f"Average for {cluster_prop}.") + + return prediction + + return theory_prediction + + def get_function_to_integrate( + self, + prediction: Callable[ + [ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + float, + ], + npt.NDArray[np.float64], + ], + ) -> Callable[[npt.NDArray, npt.NDArray], npt.NDArray]: + """Returns a callable function that can be evaluated by an integrator. + + This function is responsible for mapping arguments from the numerical integrator + to the arguments of the theoretical prediction function. + """ + + def function_mapper( + int_args: npt.NDArray, extra_args: npt.NDArray + ) -> npt.NDArray[np.float64]: + mass = int_args[:, 0] + z = int_args[:, 1] + mass_proxy = int_args[:, 2] + + sky_area = extra_args[0] + + return prediction(mass, z, mass_proxy, sky_area) + + return function_mapper + + def evaluate_theory_prediction( + self, + cluster_theory: ClusterAbundance, + this_bin: NDimensionalBin, + sky_area: float, + average_on: None | ClusterProperty = None, + ) -> float: + """Evaluate the theory prediction for this cluster recipe. + + Evaluate the theoretical prediction for the observable in the provided bin + using the Murata 2019 binned mass-richness relation and assuming perfectly + measured redshifts. + """ + self.integrator.integral_bounds = [ + (cluster_theory.min_mass, cluster_theory.max_mass), + this_bin.z_edges, + np.log(10) * np.array(this_bin.mass_proxy_edges), + ] + self.integrator.extra_args = np.array([sky_area]) + + theory_prediction = self.get_theory_prediction(cluster_theory, average_on) + prediction_wrapper = self.get_function_to_integrate(theory_prediction) + + counts = self.integrator.integrate(prediction_wrapper) + + return counts diff --git a/setup.py b/setup.py index 177abd0a..970808c8 100644 --- a/setup.py +++ b/setup.py @@ -1,5 +1,4 @@ -"""Setup script for firecrown. -""" +"""Setup script for firecrown.""" from setuptools import setup diff --git a/tests/cluster_recipes/test_murata_binned_spec_z_selection.py b/tests/cluster_recipes/test_murata_binned_spec_z_selection.py new file mode 100644 index 00000000..07a1d804 --- /dev/null +++ b/tests/cluster_recipes/test_murata_binned_spec_z_selection.py @@ -0,0 +1,199 @@ +"""Tests for the cluster abundance module.""" + +from unittest.mock import Mock + +import numpy as np +import pyccl +import pytest + +from firecrown.models.cluster.abundance import ClusterAbundance +from firecrown.models.cluster.binning import NDimensionalBin +from firecrown.models.cluster.integrator.numcosmo_integrator import NumCosmoIntegrator +from firecrown.models.cluster.kernel import SpectroscopicRedshift +from firecrown.models.cluster.mass_proxy import MurataBinned +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe +from firecrown.models.cluster.recipes.murata_binned_spec_z_selection import ( + MurataBinnedSpecZSelectionRecipe, +) + + +@pytest.fixture(name="cluster_abundance") +def fixture_cluster_abundance() -> ClusterAbundance: + hmf = pyccl.halos.MassFuncBocquet16() + cl_abundance = ClusterAbundance( + min_z=0, + max_z=2, + min_mass=13, + max_mass=17, + halo_mass_function=hmf, + ) + cl_abundance.update_ingredients(pyccl.CosmologyVanillaLCDM()) + return cl_abundance + + +@pytest.fixture(name="murata_binned_spec_z_selection") +def fixture_murata_binned_spec_z_selection() -> MurataBinnedSpecZSelectionRecipe: + cluster_recipe = MurataBinnedSpecZSelectionRecipe() + cluster_recipe.mass_distribution.mu_p0 = 3.0 + cluster_recipe.mass_distribution.mu_p1 = 0.86 + cluster_recipe.mass_distribution.mu_p2 = 0.0 + cluster_recipe.mass_distribution.sigma_p0 = 3.0 + cluster_recipe.mass_distribution.sigma_p1 = 0.7 + cluster_recipe.mass_distribution.sigma_p2 = 0.0 + cluster_recipe.purity_distribution.ap_rc = 1.1839 + cluster_recipe.purity_distribution.bp_rc = -0.4077 + cluster_recipe.purity_distribution.ap_nc = 3.9193 + cluster_recipe.purity_distribution.bp_nc = -0.3323 + cluster_recipe.completeness_distribution.ac_mc = 13.31 + cluster_recipe.completeness_distribution.bc_mc = 0.2025 + cluster_recipe.completeness_distribution.ac_nc = 0.38 + cluster_recipe.completeness_distribution.bc_nc = 1.2634 + return cluster_recipe + + +def test_murata_binned_spec_z_selection_init(): + recipe = MurataBinnedSpecZSelectionRecipe() + + assert recipe is not None + assert isinstance(recipe, ClusterRecipe) + assert recipe.integrator is not None + assert isinstance(recipe.integrator, NumCosmoIntegrator) + assert recipe.redshift_distribution is not None + assert isinstance(recipe.redshift_distribution, SpectroscopicRedshift) + assert recipe.mass_distribution is not None + assert isinstance(recipe.mass_distribution, MurataBinned) + assert recipe.completeness_distribution is not None + assert recipe.purity_distribution is not None + assert recipe.my_updatables is not None + assert len(recipe.my_updatables) == 3 + assert recipe.my_updatables[0] is recipe.mass_distribution + assert recipe.my_updatables[1] is recipe.completeness_distribution + assert recipe.my_updatables[2] is recipe.purity_distribution + + +def test_get_theory_prediction_returns_value( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z_selection: MurataBinnedSpecZSelectionRecipe, +): + prediction = murata_binned_spec_z_selection.get_theory_prediction(cluster_abundance) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2, dtype=np.float64) + z = np.linspace(0.1, 1, 2, dtype=np.float64) + mass_proxy_limits = (0, 5) + sky_area = 360**2 + + result = prediction(mass, z, mass_proxy_limits, sky_area) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 2 + assert np.all(result > 0) + + +def test_get_theory_prediction_with_average_returns_value( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z_selection: MurataBinnedSpecZSelectionRecipe, +): + mass = np.linspace(13, 17, 2, dtype=np.float64) + z = np.linspace(0.1, 1, 2, dtype=np.float64) + mass_proxy_limits = (0, 5) + sky_area = 360**2 + + prediction = murata_binned_spec_z_selection.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.MASS + ) + + assert prediction is not None + assert callable(prediction) + + result = prediction(mass, z, mass_proxy_limits, sky_area) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 2 + assert np.all(result > 0) + + prediction = murata_binned_spec_z_selection.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.REDSHIFT + ) + + assert prediction is not None + assert callable(prediction) + + result = prediction(mass, z, mass_proxy_limits, sky_area) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 2 + assert np.all(result > 0) + + prediction = murata_binned_spec_z_selection.get_theory_prediction( + cluster_abundance, average_on=(ClusterProperty.REDSHIFT | ClusterProperty.MASS) + ) + + assert prediction is not None + assert callable(prediction) + + result = prediction(mass, z, mass_proxy_limits, sky_area) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 2 + assert np.all(result > 0) + + +def test_get_theory_prediction_throws_with_nonimpl_average( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z_selection: MurataBinnedSpecZSelectionRecipe, +): + prediction = murata_binned_spec_z_selection.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.SHEAR + ) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2, dtype=np.float64) + z = np.linspace(0.1, 1, 2, dtype=np.float64) + mass_proxy_limits = (0, 5) + sky_area = 360**2 + + with pytest.raises(NotImplementedError): + _ = prediction(mass, z, mass_proxy_limits, sky_area) + + +def test_get_function_to_integrate_returns_value( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z_selection: MurataBinnedSpecZSelectionRecipe, +): + prediction = murata_binned_spec_z_selection.get_theory_prediction(cluster_abundance) + function_to_integrate = murata_binned_spec_z_selection.get_function_to_integrate( + prediction + ) + + assert function_to_integrate is not None + assert callable(function_to_integrate) + + int_args = np.array([[13.0, 0.1], [17.0, 1.0]]) + extra_args = np.array([0, 5, 360**2]) + + result = function_to_integrate(int_args, extra_args) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 2 + assert np.all(result > 0) + + +def test_evaluates_theory_prediction_returns_value( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z_selection: MurataBinnedSpecZSelectionRecipe, +): + mock_bin = Mock(spec=NDimensionalBin) + mock_bin.mass_proxy_edges = (0, 5) + mock_bin.z_edges = (0, 1) + + prediction = murata_binned_spec_z_selection.evaluate_theory_prediction( + cluster_abundance, mock_bin, 360**2 + ) + + assert prediction > 0 diff --git a/tests/test_cluster_kernels.py b/tests/test_cluster_kernels.py index 36e83d4e..5740d88c 100644 --- a/tests/test_cluster_kernels.py +++ b/tests/test_cluster_kernels.py @@ -23,12 +23,28 @@ def test_create_mass_kernel(): def test_create_completeness_kernel(): ck = Completeness() + ck.ac_mc = 13.31 + ck.bc_mc = 0.2025 + ck.ac_nc = 0.38 + ck.bc_nc = 1.2634 assert ck is not None + assert ck.ac_mc == 13.31 + assert ck.bc_mc == 0.2025 + assert ck.ac_nc == 0.38 + assert ck.bc_nc == 1.2634 def test_create_purity_kernel(): pk = Purity() + pk.ap_nc = 3.9193 + pk.bp_nc = -0.3323 + pk.ap_rc = 1.1839 + pk.bp_rc = -0.4077 assert pk is not None + assert pk.ap_nc == 3.9193 + assert pk.bp_nc == -0.3323 + assert pk.ap_rc == 1.1839 + assert pk.bp_rc == -0.4077 def test_spec_z_distribution(): @@ -44,56 +60,64 @@ def test_true_mass_distribution(): @pytest.mark.precision_sensitive def test_purity_distribution(): pk = Purity() - mass_proxy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + pk.ap_nc = 3.9193 + pk.bp_nc = -0.3323 + pk.ap_rc = 1.1839 + pk.bp_rc = -0.4077 + mass_proxy = np.linspace(0.0, 2.5, 10) z = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) mass_proxy_limits = (1.0, 10.0) truth = np.array( [ - 0.77657274, - 0.96966127, - 0.99286409, - 0.99780586, - 0.999224, - 0.99970302, - 0.99988111, - 0.99995125, - 0.99997982, - 0.99999166, - ] + 0.00242882, + 0.03294582, + 0.3122527, + 0.85213252, + 0.98584893, + 0.99875485, + 0.99988632, + 0.99998911, + 0.99999891, + 0.99999988, + ], + dtype=np.float64, ) purity = pk.distribution(z, mass_proxy, mass_proxy_limits) assert isinstance(purity, np.ndarray) for ref, true in zip(purity, truth): - assert ref == pytest.approx(true, rel=1e-7, abs=0.0) + assert ref == pytest.approx(true, rel=1e-5, abs=0.0) @pytest.mark.precision_sensitive def test_purity_distribution_uses_mean(): pk = Purity() + pk.ap_nc = 3.9193 + pk.bp_nc = -0.3323 + pk.ap_rc = 1.1839 + pk.bp_rc = -0.4077 z = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) - mass_proxy = np.ones_like(z) * -1.0 - mass_proxy_limits = (1.0, 10.0) + mass_proxy = np.array([-1.0]) + mass_proxy_limits = (0.0, 2.0) truth = np.array( [ - 0.9978693724040568, - 0.9984319673134954, - 0.9988620014089232, - 0.9991864843696077, - 0.9994279315032029, - 0.999604893383804, - 0.9997324678841709, - 0.9998227843987537, - 0.9998854531462606, - 0.9999279749997235, + 0.89705651, + 0.92238419, + 0.94154163, + 0.95593305, + 0.96670586, + 0.97476117, + 0.98078884, + 0.98530847, + 0.98870753, + 0.99127329, ], dtype=np.float64, ) - - purity = pk.distribution(z, mass_proxy.astype(np.float64), mass_proxy_limits) + purity = pk.distribution(z, mass_proxy, mass_proxy_limits) assert isinstance(purity, np.ndarray) for ref, true in zip(purity, truth): assert ref == pytest.approx(true, rel=1e-7, abs=0.0) @@ -102,21 +126,25 @@ def test_purity_distribution_uses_mean(): @pytest.mark.precision_sensitive def test_completeness_distribution(): ck = Completeness() - mass = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) + ck.ac_mc = 13.31 + ck.bc_mc = 0.2025 + ck.ac_nc = 0.38 + ck.bc_nc = 1.2634 + mass = np.linspace(13.0, 15.0, 10) z = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) truth = np.array( [ - 0.0056502277493542, - 0.01896566878380423, - 0.03805597500308377, - 0.06224888967250564, - 0.09124569979282898, - 0.12486247682690908, - 0.16290218589569144, - 0.20507815091349266, - 0.2509673905442634, - 0.2999886170051561, + 0.10239024, + 0.19090539, + 0.35438466, + 0.58952617, + 0.80866296, + 0.93327968, + 0.98115635, + 0.99543348, + 0.99902667, + 0.99981606, ] )