Impure mass distribution exact [version:1.0.1]

crow/__init__.py

@@ -8,4 +8,4 @@
 from .recipes.binned_grid import GridBinnedClusterRecipe
 from .recipes.binned_parent import BinnedClusterRecipe
 
-__version__ = "1.0.0"
+__version__ = "1.0.1"

crow/recipes/binned_exact.py

@@ -60,74 +60,62 @@ def _setup_with_completeness(self):
             self._completeness_distribution = self.completeness.distribution
 
     def _setup_with_purity(self):
-        """Makes mass distribution use additional integral with completeness"""
+        """Makes mass distribution use additional integral with purity"""
         if self.purity is None:
             self._mass_distribution_distribution = self.mass_distribution.distribution
         else:
-            self._mass_distribution_distribution = self._impure_mass_distribution
+            self._mass_distribution_distribution = self._mass_distribution_purity
 
-    def _impure_mass_distribution(self, log_mass, z, log_mass_proxy_limits):
+    def _mass_distribution_purity(self, log_mass, z, log_mass_proxy):
 
-        ##############################
-        # Fix this function, Henrique
-        # Good luck!
-        ##############################
-
-        integrator = NumCosmoIntegrator(
-            relative_tolerance=1e-6,
-            absolute_tolerance=1e-12,
+        return (
+            self.mass_distribution.gaussian_kernel(log_mass, z, log_mass_proxy)
+            / self.purity.distribution(log_mass_proxy, z)
+            * np.log(10)
         )
 
-        def integration_func(int_args, extra_args):
-            ln_mass_proxy = int_args[:, 0]
-            log_mass_proxy = ln_mass_proxy / np.log(10.0)
-            return np.array(
-                [
-                    self.mass_distribution.gaussian_kernel(
-                        log_mass, z, np.array([_log_mass_proxy])
-                    )
-                    / self.purity.distribution(np.array([_log_mass_proxy]), z)
-                    for _log_mass_proxy in log_mass_proxy
-                ]
-            )
-
-        integrator.integral_bounds = [
-            (
-                np.log(10.0) * log_mass_proxy_limits[0],
-                np.log(10.0) * log_mass_proxy_limits[1],
-            )
-        ]
-
-        return integrator.integrate(integration_func)
-
     def _get_theory_prediction_counts(
         self,
         average_on: None | ClusterProperty = None,
     ) -> Callable[
-        [npt.NDArray[np.float64], npt.NDArray[np.float64], tuple[float, float], float],
+        [
+            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,
+        Returns a callable function that accepts mass, redshift, mass_proxy,
         and the sky area of your survey and returns the theoretical prediction for the
-        expected number of clusters.
+        expected number of clusters including false detections.
         """
 
         def theory_prediction(
             mass: npt.NDArray[np.float64],
             z: npt.NDArray[np.float64],
-            mass_proxy_limits: tuple[float, float],
+            mass_proxy: npt.NDArray[np.float64],
             sky_area: float,
         ):
             prediction = (
                 self.cluster_theory.comoving_volume(z, sky_area)
                 * self.cluster_theory.mass_function(mass, z)
                 * self._completeness_distribution(mass, z)
                 * self.redshift_distribution.distribution()
-                * self._mass_distribution_distribution(mass, z, mass_proxy_limits)
             )
 
+            if self.purity == None:
+                assert (
+                    len(mass_proxy) == 2
+                ), "mass_proxy with no purity should be size 2"
+                prediction *= self._mass_distribution_distribution(
+                    mass, z, (mass_proxy[0], mass_proxy[1])
+                )
+            else:
+                prediction *= self._mass_distribution_distribution(mass, z, mass_proxy)
+
             if average_on is None:
                 return prediction
 
@@ -149,7 +137,7 @@ def _get_function_to_integrate_counts(
             [
                 npt.NDArray[np.float64],
                 npt.NDArray[np.float64],
-                tuple[float, float],
+                npt.NDArray[np.float64],
                 float,
             ],
             npt.NDArray[np.float64],
@@ -167,11 +155,14 @@ def function_mapper(
             mass = int_args[:, 0]
             z = int_args[:, 1]
 
-            mass_proxy_low = extra_args[0]
-            mass_proxy_high = extra_args[1]
-            sky_area = extra_args[2]
+            if self.purity == None:
+                mass_proxy = np.array([extra_args[0], extra_args[1]])
+                sky_area = extra_args[2]
+            else:
+                mass_proxy = int_args[:, 2]
+                sky_area = extra_args[0]
 
-            return prediction(mass, z, (mass_proxy_low, mass_proxy_high), sky_area)
+            return prediction(mass, z, mass_proxy, sky_area)
 
         return function_mapper
 
@@ -188,11 +179,22 @@ def evaluate_theory_prediction_counts(
         using the Murata 2019 binned mass-richness relation and assuming perfectly
         measured redshifts.
         """
-        self.integrator.integral_bounds = [
-            self.mass_interval,
-            z_edges,
-        ]
-        self.integrator.extra_args = np.array([*log_proxy_edges, sky_area])
+        assert len(log_proxy_edges) == 2, "log_proxy_edges should be size 2"
+        assert len(z_edges) == 2, "z_edges should be size 2"
+
+        if self.purity == None:
+            self.integrator.integral_bounds = [
+                self.mass_interval,
+                z_edges,
+            ]
+            self.integrator.extra_args = np.array([*log_proxy_edges, sky_area])
+        else:
+            self.integrator.integral_bounds = [
+                self.mass_interval,
+                z_edges,
+                log_proxy_edges,
+            ]
+            self.integrator.extra_args = np.array([sky_area])
 
         theory_prediction = self._get_theory_prediction_counts(average_on)
         prediction_wrapper = self._get_function_to_integrate_counts(theory_prediction)
@@ -208,7 +210,7 @@ def _get_theory_prediction_shear_profile(
         [
             npt.NDArray[np.float64],
             npt.NDArray[np.float64],
-            tuple[float, float],
+            npt.NDArray[np.float64],
             float,
             float,
         ],
@@ -224,7 +226,7 @@ def _get_theory_prediction_shear_profile(
         def theory_prediction(
             mass: npt.NDArray[np.float64],
             z: npt.NDArray[np.float64],
-            mass_proxy_limits: tuple[float, float],
+            mass_proxy: npt.NDArray[np.float64],
             sky_area: float,
             radius_center: float,
         ):
@@ -233,7 +235,6 @@ def theory_prediction(
                 * self.cluster_theory.mass_function(mass, z)
                 * self.redshift_distribution.distribution()
                 * self._completeness_distribution(mass, z)
-                * self.mass_distribution.distribution(mass, z, mass_proxy_limits)
             )
             if average_on is None:
                 # pylint: disable=no-member
@@ -248,6 +249,18 @@ def theory_prediction(
                     z=z,
                     radius_center=radius_center,
                 )
+                if self.purity == None:
+                    assert (
+                        len(mass_proxy) == 2
+                    ), "mass_proxy with no purity should be size 2"
+                    prediction *= self._mass_distribution_distribution(
+                        mass, z, (mass_proxy[0], mass_proxy[1])
+                    )
+                else:
+                    prediction *= self._mass_distribution_distribution(
+                        mass, z, mass_proxy
+                    )
+
             return prediction
 
         return theory_prediction
@@ -258,7 +271,7 @@ def _get_function_to_integrate_shear_profile(
             [
                 npt.NDArray[np.float64],
                 npt.NDArray[np.float64],
-                tuple[float, float],
+                npt.NDArray[np.float64],
                 float,
                 float,
             ],
@@ -274,16 +287,20 @@ def _get_function_to_integrate_shear_profile(
         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_low = extra_args[0]
-            mass_proxy_high = extra_args[1]
-            sky_area = extra_args[2]
-            radius_center = extra_args[3]
-            return prediction(
-                mass, z, (mass_proxy_low, mass_proxy_high), sky_area, radius_center
-            )
+            if self.purity == None:
+                mass_proxy = np.array([extra_args[0], extra_args[1]])
+                sky_area = extra_args[2]
+                radius_center = extra_args[3]
+            else:
+                mass_proxy = int_args[:, 2]
+                sky_area = extra_args[0]
+                radius_center = extra_args[1]
+
+            return prediction(mass, z, mass_proxy, sky_area, radius_center)
 
         return function_mapper
 
@@ -301,15 +318,28 @@ def evaluate_theory_prediction_lensing_profile(
         using the Murata 2019 binned mass-richness relation and assuming perfectly
         measured redshifts.
         """
-        self.integrator.integral_bounds = [
-            self.mass_interval,
-            z_edges,
-        ]
+        assert len(log_proxy_edges) == 2, "log_proxy_edges should be size 2"
+        assert len(z_edges) == 2, "z_edges should be size 2"
+
+        if self.purity == None:
+            self.integrator.integral_bounds = [
+                self.mass_interval,
+                z_edges,
+            ]
+            extra_args = [*log_proxy_edges]
+        else:
+            self.integrator.integral_bounds = [
+                self.mass_interval,
+                z_edges,
+                log_proxy_edges,
+            ]
+            extra_args = []
+
         deltasigma_list = []
+
         for radius_center in radius_centers:
-            self.integrator.extra_args = np.array(
-                [*log_proxy_edges, sky_area, radius_center]
-            )
+            extra_args.extend([sky_area, radius_center])
+            self.integrator.extra_args = np.array(extra_args)
             if self.cluster_theory._beta_parameters is not None:
                 self.cluster_theory.set_beta_s_interp(*z_edges)
             theory_prediction = self._get_theory_prediction_shear_profile(average_on)

crow/recipes/binned_grid.py

@@ -111,7 +111,6 @@ def _get_hmf_grid(
             vol = self.cluster_theory.comoving_volume(z, sky_area)
             # assign
             self._hmf_grid[key] = vol[:, np.newaxis] * mass_function_2d
-
         return self._hmf_grid[key]
 
     def _get_mass_richness_grid(
@@ -147,7 +146,6 @@ def _get_completeness_grid(self, z: npt.NDArray[np.float64], key):
             self._completeness_grid[key] = self._completeness_distribution(
                 self.log_mass_grid[np.newaxis, :], z[:, np.newaxis]
             )
-
         return self._completeness_grid[key]
 
     def _get_purity_grid(