Modified selection function to be updatable

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,64 +29,115 @@ 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
 
     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)

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)

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

setup.py

@@ -1,5 +1,4 @@
-"""Setup script for firecrown.
-"""
+"""Setup script for firecrown."""
 
 from setuptools import setup