Refactoring the integration / integrand workflow in the cluster count calculation

examples/cluster_number_counts/cluster_redshift_richness.py

@@ -1,24 +1,22 @@
 """Likelihood factory function for cluster number counts."""
 
 import os
+from typing import Tuple
 
 import pyccl as ccl
 import sacc
 
-from firecrown.models.cluster.integrator.numcosmo_integrator import NumCosmoIntegrator
 from firecrown.likelihood.gauss_family.gaussian import ConstGaussian
 from firecrown.likelihood.gauss_family.statistic.binned_cluster_number_counts import (
     BinnedClusterNumberCounts,
 )
-from firecrown.models.cluster.properties import ClusterProperty
+from firecrown.likelihood.likelihood import Likelihood, NamedParameters
 from firecrown.modeling_tools import ModelingTools
 from firecrown.models.cluster.abundance import ClusterAbundance
-from firecrown.models.cluster.kernel import (
-    SpectroscopicRedshift,
+from firecrown.models.cluster.properties import ClusterProperty
+from firecrown.models.cluster.recipes.murata_binned_spec_z import (
+    MurataBinnedSpecZRecipe,
 )
-from firecrown.models.cluster.mass_proxy import MurataBinned
-from firecrown.likelihood.likelihood import NamedParameters, Likelihood
-from typing import Tuple
 
 
 def get_cluster_abundance() -> ClusterAbundance:
@@ -27,14 +25,6 @@ def get_cluster_abundance() -> ClusterAbundance:
     min_z, max_z = 0.2, 0.8
     cluster_abundance = ClusterAbundance(min_mass, max_mass, min_z, max_z, hmf)
 
-    # Create and add the kernels you want in your cluster abundance
-    pivot_mass, pivot_redshift = 14.625862906, 0.6
-    mass_observable_kernel = MurataBinned(pivot_mass, pivot_redshift)
-    cluster_abundance.add_kernel(mass_observable_kernel)
-
-    redshift_proxy_kernel = SpectroscopicRedshift()
-    cluster_abundance.add_kernel(redshift_proxy_kernel)
-
     return cluster_abundance
 
 
@@ -44,19 +34,17 @@ def build_likelihood(
     """
     Here we instantiate the number density (or mass function) object.
     """
-    integrator = NumCosmoIntegrator()
-    # integrator = ScipyIntegrator()
 
     # Pull params for the likelihood from build_parameters
-    average_properties = ClusterProperty.NONE
+    average_on = ClusterProperty.NONE
     if build_parameters.get_bool("use_cluster_counts", True):
-        average_properties |= ClusterProperty.COUNTS
+        average_on |= ClusterProperty.COUNTS
     if build_parameters.get_bool("use_mean_log_mass", False):
-        average_properties |= ClusterProperty.MASS
+        average_on |= ClusterProperty.MASS
 
     survey_name = "numcosmo_simulated_redshift_richness"
     likelihood = ConstGaussian(
-        [BinnedClusterNumberCounts(average_properties, survey_name, integrator)]
+        [BinnedClusterNumberCounts(average_on, survey_name, MurataBinnedSpecZRecipe())]
     )
 
     # Read in sacc data

examples/cluster_number_counts/compare_integration_methods.py

@@ -4,14 +4,14 @@
 
 import pyccl
 import numpy as np
-from firecrown.models.cluster.mass_proxy import MurataBinned
-from firecrown.models.cluster.kernel import Kernel
 from firecrown.models.cluster.integrator.numcosmo_integrator import (
-    NumCosmoIntegrator,
     NumCosmoIntegralMethod,
 )
 from firecrown.models.cluster.abundance import ClusterAbundance
-from firecrown.models.cluster.kernel import SpectroscopicRedshift
+from firecrown.models.cluster.recipes.murata_binned_spec_z import (
+    MurataBinnedSpecZRecipe,
+)
+from firecrown.models.cluster.binning import TupleBin
 
 
 def get_cosmology() -> pyccl.Cosmology:
@@ -44,64 +44,45 @@ def get_cosmology() -> pyccl.Cosmology:
     return cosmo_ccl
 
 
-def get_mass_richness() -> Kernel:
-    pivot_mass = 14.625862906
-    pivot_redshift = 0.6
-
-    mass_richness = MurataBinned(pivot_mass, pivot_redshift)
-
-    mass_richness.mu_p0 = 3.0
-    mass_richness.mu_p1 = 0.86
-    mass_richness.mu_p2 = 0.0
-    mass_richness.sigma_p0 = 3.0
-    mass_richness.sigma_p1 = 0.7
-    mass_richness.sigma_p2 = 0.0
-
-    return mass_richness
-
-
 def compare_integration() -> None:
     """Compare integration methods."""
     hmf = pyccl.halos.MassFuncTinker08()
     abundance = ClusterAbundance(13, 15, 0, 4, hmf)
+    cluster_recipe = MurataBinnedSpecZRecipe()
 
-    mass_richness = get_mass_richness()
-    abundance.add_kernel(mass_richness)
-
-    redshift_proxy_kernel = SpectroscopicRedshift()
-    abundance.add_kernel(redshift_proxy_kernel)
+    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
 
     cosmo = get_cosmology()
     abundance.update_ingredients(cosmo)
 
     sky_area = 360**2
-    integrand = abundance.get_integrand()
 
     for method in NumCosmoIntegralMethod:
         counts_list = []
         t_start = time.time()
 
-        nc_integrator = NumCosmoIntegrator(method=method)
-        nc_integrator.set_integration_bounds(abundance, 496, (0, 4), (13, 15))
+        cluster_recipe.integrator.method = method
+
+        # nc_integrator.set_integration_bounds(abundance, 496, (0, 4), (13, 15))
 
         z_bins = np.linspace(0.0, 1.0, 4)
         mass_proxy_bins = np.linspace(1.0, 2.5, 5)
 
         for z_idx, mass_proxy_idx in itertools.product(range(3), range(4)):
-            z_proxy_limits = (z_bins[z_idx], z_bins[z_idx + 1])
-            mass_proxy_limits = (
+            mass_limits = (
                 mass_proxy_bins[mass_proxy_idx],
                 mass_proxy_bins[mass_proxy_idx + 1],
             )
+            z_limits = (z_bins[z_idx], z_bins[z_idx + 1])
 
-            nc_integrator.set_integration_bounds(
-                abundance,
-                sky_area,
-                z_proxy_limits,
-                mass_proxy_limits,
-            )
+            bin = TupleBin([mass_limits, z_limits])
 
-            counts = nc_integrator.integrate(integrand)
+            counts = cluster_recipe.evaluate_theory_prediction(abundance, bin, sky_area)
             counts_list.append(counts)
 
         t_stop = time.time()

firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py

@@ -4,20 +4,23 @@
 clusters within a single redshift and mass bin.
 """
 from __future__ import annotations
+
 from typing import List, Optional
-import sacc
+
 import numpy as np
-from firecrown.models.cluster.integrator.integrator import Integrator
-from firecrown.models.cluster.abundance_data import AbundanceData
-from firecrown.models.cluster.binning import SaccBin
-from firecrown.models.cluster.properties import ClusterProperty
+import sacc
+
+from firecrown.likelihood.gauss_family.statistic.source.source import SourceSystematic
 from firecrown.likelihood.gauss_family.statistic.statistic import (
-    Statistic,
     DataVector,
+    Statistic,
     TheoryVector,
 )
-from firecrown.likelihood.gauss_family.statistic.source.source import SourceSystematic
 from firecrown.modeling_tools import ModelingTools
+from firecrown.models.cluster.abundance_data import AbundanceData
+from firecrown.models.cluster.binning import SaccBin
+from firecrown.models.cluster.properties import ClusterProperty
+from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe
 
 
 class BinnedClusterNumberCounts(Statistic):
@@ -31,15 +34,15 @@ def __init__(
         self,
         cluster_properties: ClusterProperty,
         survey_name: str,
-        integrator: Integrator,
+        cluster_recipe: ClusterRecipe,
         systematics: Optional[List[SourceSystematic]] = None,
     ):
         super().__init__()
         self.systematics = systematics or []
         self.theory_vector: Optional[TheoryVector] = None
         self.cluster_properties = cluster_properties
         self.survey_name = survey_name
-        self.integrator = integrator
+        self.cluster_recipe = cluster_recipe
         self.data_vector = DataVector.from_list([])
         self.sky_area = 0.0
         self.bins: List[SaccBin] = []
@@ -110,23 +113,17 @@ def get_binned_cluster_property(
         the clusters in each bin."""
         assert tools.cluster_abundance is not None
 
-        cluster_masses = []
-        for bin_edge, counts in zip(self.bins, cluster_counts):
-            integrand = tools.cluster_abundance.get_integrand(
-                average_properties=cluster_properties
-            )
-            self.integrator.set_integration_bounds(
-                tools.cluster_abundance,
-                self.sky_area,
-                bin_edge.z_edges,
-                bin_edge.mass_proxy_edges,
+        mean_values = []
+        for this_bin, counts in zip(self.bins, cluster_counts):
+            total_observable = self.cluster_recipe.evaluate_theory_prediction(
+                tools.cluster_abundance, this_bin, self.sky_area, cluster_properties
             )
+            cluster_counts.append(counts)
 
-            total_mass = self.integrator.integrate(integrand)
-            mean_mass = total_mass / counts
-            cluster_masses.append(mean_mass)
+            mean_observable = total_observable / counts
+            mean_values.append(mean_observable)
 
-        return cluster_masses
+        return mean_values
 
     def get_binned_cluster_counts(self, tools: ModelingTools) -> List[float]:
         """Computes the number of clusters in each bin
@@ -137,16 +134,10 @@ def get_binned_cluster_counts(self, tools: ModelingTools) -> List[float]:
         assert tools.cluster_abundance is not None
 
         cluster_counts = []
-        for bin_edge in self.bins:
-            self.integrator.set_integration_bounds(
-                tools.cluster_abundance,
-                self.sky_area,
-                bin_edge.z_edges,
-                bin_edge.mass_proxy_edges,
+        for this_bin in self.bins:
+            counts = self.cluster_recipe.evaluate_theory_prediction(
+                tools.cluster_abundance, this_bin, self.sky_area
             )
-
-            integrand = tools.cluster_abundance.get_integrand()
-            counts = self.integrator.integrate(integrand)
             cluster_counts.append(counts)
 
         return cluster_counts

firecrown/models/cluster/abundance.py

@@ -4,29 +4,13 @@
 and phenomenological predictions.  This module contains the classes and
 functions that produce those predictions.
 """
-from typing import List, Callable, Optional, Dict, Tuple
+from typing import Dict, Tuple
 import numpy as np
 import numpy.typing as npt
 import pyccl
 import pyccl.background as bkg
 from pyccl.cosmology import Cosmology
 from firecrown.updatable import Updatable, UpdatableCollection
-from firecrown.models.cluster.kernel import Kernel
-from firecrown.models.cluster.properties import ClusterProperty
-
-
-AbundanceIntegrand = Callable[
-    [
-        npt.NDArray[np.float64],
-        npt.NDArray[np.float64],
-        float,
-        npt.NDArray[np.float64],
-        npt.NDArray[np.float64],
-        Tuple[float, float],
-        Tuple[float, float],
-    ],
-    npt.NDArray[np.float64],
-]
 
 
 class ClusterAbundance(Updatable):
@@ -43,23 +27,6 @@ def cosmo(self) -> Cosmology:
         """The cosmology used to predict the cluster number count."""
         return self._cosmo
 
-    @property
-    def analytic_kernels(self) -> List[Kernel]:
-        """The kernels in in the integrand which have an analytic solution."""
-        return [x for x in self.kernels if x.has_analytic_sln]
-
-    @property
-    def dirac_delta_kernels(self) -> List[Kernel]:
-        """The kernels in in the integrand which are dirac delta functions."""
-        return [x for x in self.kernels if x.is_dirac_delta]
-
-    @property
-    def integrable_kernels(self) -> List[Kernel]:
-        """The kernels in in the integrand which must be integrated."""
-        return [
-            x for x in self.kernels if not x.is_dirac_delta and not x.has_analytic_sln
-        ]
-
     def __init__(
         self,
         min_mass: float,
@@ -78,10 +45,6 @@ def __init__(
         self._hmf_cache: Dict[Tuple[float, float], float] = {}
         self._cosmo: Cosmology = None
 
-    def add_kernel(self, kernel: Kernel) -> None:
-        """Add a kernel to the cluster abundance integrand"""
-        self.kernels.append(kernel)
-
     def update_ingredients(self, cosmo: Cosmology) -> None:
         """Update the cluster abundance calculation with a new cosmology."""
         self._cosmo = cosmo
@@ -126,43 +89,3 @@ def mass_function(
             return_vals.append(val)
 
         return np.asarray(return_vals, dtype=np.float64)
-
-    def get_integrand(
-        self, *, average_properties: Optional[ClusterProperty] = None
-    ) -> AbundanceIntegrand:
-        """Returns a callable that evaluates the complete integrand."""
-
-        def integrand(
-            mass: npt.NDArray[np.float64],
-            z: npt.NDArray[np.float64],
-            sky_area: float,
-            mass_proxy: npt.NDArray[np.float64],
-            z_proxy: npt.NDArray[np.float64],
-            mass_proxy_limits: Tuple[float, float],
-            z_proxy_limits: Tuple[float, float],
-        ) -> npt.NDArray[np.float64]:
-            integrand = self.comoving_volume(z, sky_area) * self.mass_function(mass, z)
-
-            for kernel in self.kernels:
-                assert isinstance(kernel, Kernel)
-                integrand *= kernel.distribution(
-                    mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits
-                )
-
-            if average_properties is None:
-                return integrand
-
-            for cluster_prop in ClusterProperty:
-                include_prop = cluster_prop & average_properties
-                if not include_prop:
-                    continue
-                if cluster_prop == ClusterProperty.MASS:
-                    integrand *= mass
-                elif cluster_prop == ClusterProperty.REDSHIFT:
-                    integrand *= z
-                else:
-                    raise NotImplementedError(f"Average for {cluster_prop}.")
-
-            return integrand
-
-        return integrand