diff --git a/examples/cluster_number_counts/cluster_redshift_richness.py b/examples/cluster_number_counts/cluster_redshift_richness.py index a7acf0b3..2b2c5e02 100644 --- a/examples/cluster_number_counts/cluster_redshift_richness.py +++ b/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 diff --git a/examples/cluster_number_counts/compare_integration_methods.py b/examples/cluster_number_counts/compare_integration_methods.py index fb556340..b8d216bc 100644 --- a/examples/cluster_number_counts/compare_integration_methods.py +++ b/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() diff --git a/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py b/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py index 8d5e9e01..305f4b5d 100644 --- a/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py +++ b/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,7 +34,7 @@ def __init__( self, cluster_properties: ClusterProperty, survey_name: str, - integrator: Integrator, + cluster_recipe: ClusterRecipe, systematics: Optional[List[SourceSystematic]] = None, ): super().__init__() @@ -39,7 +42,7 @@ def __init__( 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 diff --git a/firecrown/models/cluster/abundance.py b/firecrown/models/cluster/abundance.py index 6d836ff1..8c9b9a51 100644 --- a/firecrown/models/cluster/abundance.py +++ b/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 diff --git a/firecrown/models/cluster/binning.py b/firecrown/models/cluster/binning.py index f4183b71..7220fa1c 100644 --- a/firecrown/models/cluster/binning.py +++ b/firecrown/models/cluster/binning.py @@ -35,9 +35,6 @@ def mass_proxy_edges(self) -> Tuple[float, float]: def __str__(self) -> str: return f"[{self.z_edges}, {self.mass_proxy_edges}]\n" - def __repr__(self) -> str: - return f"[{self.z_edges}, {self.mass_proxy_edges}]\n" - class SaccBin(NDimensionalBin[sacc.BaseTracer]): """An implementation of the N dimensional bin using sacc tracers.""" @@ -71,8 +68,7 @@ def __eq__(self, other: object) -> bool: other_bin = [ x for x in other.coordinate_bins if x.tracer_type == my_bin.tracer_type ] - if len(other_bin) != 1: - return False + if my_bin.lower != other_bin[0].lower: return False if my_bin.upper != other_bin[0].upper: @@ -84,3 +80,44 @@ def __hash__(self) -> int: """One bin's hash is determined by the dimension and lower/upper bound.""" bin_bounds = [(bin.lower, bin.upper) for bin in self.coordinate_bins] return hash((self.dimension, tuple(bin_bounds))) + + +class TupleBin(NDimensionalBin[Tuple]): + """An implementation of the N dimensional bin using sacc tracers.""" + + def __init__(self, bins: List[Tuple]): + super().__init__(bins) + + @property + def mass_proxy_edges(self) -> Tuple[float, float]: + mass_bin = self.coordinate_bins[0] + return mass_bin[0], mass_bin[1] + + @property + def z_edges(self) -> Tuple[float, float]: + z_bin = self.coordinate_bins[1] + return z_bin[0], z_bin[1] + + def __eq__(self, other: object) -> bool: + """Two bins are equal if they have the same lower/upper bound.""" + if not isinstance(other, TupleBin): + return False + + if self.dimension != other.dimension: + return False + + for i, my_bin in enumerate(self.coordinate_bins): + other_bin = other.coordinate_bins[i] + if len(my_bin) != len(other_bin): + return False + if my_bin[0] != other_bin[0]: + return False + if my_bin[1] != other_bin[1]: + return False + + return True + + def __hash__(self) -> int: + """One bin's hash is determined by the dimension and lower/upper bound.""" + bin_bounds = [(bin[0], bin[1]) for bin in self.coordinate_bins] + return hash((self.dimension, tuple(bin_bounds))) diff --git a/firecrown/models/cluster/integrator/integrator.py b/firecrown/models/cluster/integrator/integrator.py index 4503753b..7620623f 100644 --- a/firecrown/models/cluster/integrator/integrator.py +++ b/firecrown/models/cluster/integrator/integrator.py @@ -1,64 +1,29 @@ """The integrator module This module holds the classes that define the interface required to -integrate an assembled cluster abundance. +integrate a function. """ -import inspect from abc import ABC, abstractmethod -from typing import Tuple, Dict, get_args, List -from firecrown.models.cluster.abundance import ClusterAbundance, AbundanceIntegrand -from firecrown.models.cluster.kernel import KernelType +from typing import Callable, List, Tuple + +import numpy as np +import numpy.typing as npt class Integrator(ABC): """The integrator base class - This class acts as an adapter around an integration library, and must provides - a specific set of methods to be used to integrate a cluster abundance integral.""" + This class acts as an adapter around an integration library.""" def __init__(self) -> None: - self.z_proxy_limits: Tuple[float, float] = (-1.0, -1.0) - self.mass_proxy_limits: Tuple[float, float] = (-1.0, -1.0) - self.sky_area: float = 360**2 self.integral_bounds: List[Tuple[float, float]] = [] - self.integral_args_lkp: Dict[KernelType, int] = self._default_integral_args() + self.extra_args: npt.NDArray[np.float64] = np.array([], dtype=np.float64) @abstractmethod def integrate( self, - integrand: AbundanceIntegrand, + func_to_integrate: Callable[ + [npt.NDArray[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64] + ], ) -> float: """Call this method to integrate the provided integrand argument.""" - - @abstractmethod - def set_integration_bounds( - self, - cl_abundance: ClusterAbundance, - sky_area: float, - z_proxy_limits: Tuple[float, float], - mass_proxy_limits: Tuple[float, float], - ) -> None: - """Sets the limits of integration used by the integration library. - - This method uses the mass and redshift proxy arguments, along with - the cluster abundance argument to construct the limits of integration - to be used in evaluating the cluster abundance.""" - - def _default_integral_args(self) -> Dict[KernelType, int]: - lkp: Dict[KernelType, int] = {} - lkp[KernelType.MASS] = 0 - lkp[KernelType.Z] = 1 - return lkp - - def _validate_integrand(self, integrand: AbundanceIntegrand) -> None: - expected_args, expected_return = get_args(AbundanceIntegrand) - - signature = inspect.signature(integrand) - params = signature.parameters.values() - param_types = [param.annotation for param in params] - - assert len(params) == len(expected_args) - - assert param_types == list(expected_args) - - assert signature.return_annotation == expected_return diff --git a/firecrown/models/cluster/integrator/numcosmo_integrator.py b/firecrown/models/cluster/integrator/numcosmo_integrator.py index d3105307..e457620b 100644 --- a/firecrown/models/cluster/integrator/numcosmo_integrator.py +++ b/firecrown/models/cluster/integrator/numcosmo_integrator.py @@ -2,13 +2,13 @@ This module holds the NumCosmo implementation of the integrator classes """ -from typing import Tuple, Callable, Sequence, Optional from enum import Enum +from typing import Callable, Optional, Tuple + import numpy as np import numpy.typing as npt from numcosmo_py import Ncm -from firecrown.models.cluster.kernel import KernelType -from firecrown.models.cluster.abundance import ClusterAbundance, AbundanceIntegrand + from firecrown.models.cluster.integrator.integrator import Integrator @@ -31,83 +31,22 @@ def __init__( absolute_tolerance: float = 1e-12, ) -> None: super().__init__() + self.method = method or NumCosmoIntegralMethod.P_V self._relative_tolerance = relative_tolerance self._absolute_tolerance = absolute_tolerance - self._method = method or NumCosmoIntegralMethod.P_V - - def _integral_wrapper( - self, - integrand: AbundanceIntegrand, - ) -> Callable[[npt.NDArray], Sequence[float]]: - self._validate_integrand(integrand) - - # mypy strict issue: npt.NDArray[npt.NDArray[np.float64]] not supported - def ncm_integrand(int_args: npt.NDArray) -> Sequence[float]: - default = np.ones_like(int_args[0]) * -1.0 - # pylint: disable=R0801 - mass = self._get_or_default(int_args, KernelType.MASS, default) - z = self._get_or_default(int_args, KernelType.Z, default) - mass_proxy = self._get_or_default(int_args, KernelType.MASS_PROXY, default) - z_proxy = self._get_or_default(int_args, KernelType.Z_PROXY, default) - - return_val = integrand( - mass, - z, - self.sky_area, - mass_proxy, - z_proxy, - self.mass_proxy_limits, - self.z_proxy_limits, - ).tolist() - assert isinstance(return_val, list) - return return_val - - return ncm_integrand - - def set_integration_bounds( - self, - cl_abundance: ClusterAbundance, - sky_area: float, - z_proxy_limits: Tuple[float, float], - mass_proxy_limits: Tuple[float, float], - ) -> None: - # pylint: disable=R0801 - self.integral_args_lkp = self._default_integral_args() - self.integral_bounds = [ - (cl_abundance.min_mass, cl_abundance.max_mass), - (cl_abundance.min_z, cl_abundance.max_z), - ] - - self.mass_proxy_limits = mass_proxy_limits - self.z_proxy_limits = z_proxy_limits - self.sky_area = sky_area - - for kernel in cl_abundance.dirac_delta_kernels: - if kernel.kernel_type == KernelType.Z_PROXY: - self.integral_bounds[1] = z_proxy_limits - - elif kernel.kernel_type == KernelType.MASS_PROXY: - self.integral_bounds[0] = mass_proxy_limits - - for kernel in cl_abundance.integrable_kernels: - idx = len(self.integral_bounds) - - if kernel.kernel_type == KernelType.Z_PROXY: - self.integral_bounds.append(z_proxy_limits) - self.integral_args_lkp[KernelType.Z_PROXY] = idx - - elif kernel.kernel_type == KernelType.MASS_PROXY: - self.integral_bounds.append(mass_proxy_limits) - self.integral_args_lkp[KernelType.MASS_PROXY] = idx def integrate( self, - integrand: AbundanceIntegrand, + func_to_integrate: Callable[ + [npt.NDArray[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64] + ], ) -> float: Ncm.cfg_init() - ncm_integrand = self._integral_wrapper(integrand) - int_nd = CountsIntegralND(len(self.integral_bounds), ncm_integrand) - int_nd.set_method(self._method.value) + + int_nd = CountsIntegralND( + len(self.integral_bounds), func_to_integrate, self.extra_args + ) + int_nd.set_method(self.method.value) int_nd.set_reltol(self._relative_tolerance) int_nd.set_abstol(self._absolute_tolerance) res = Ncm.Vector.new(1) @@ -117,18 +56,6 @@ def integrate( int_nd.eval(Ncm.Vector.new_array(bl), Ncm.Vector.new_array(bu), res, err) return res.get(0) - def _get_or_default( - self, - # mypy strict issue: npt.NDArray[npt.NDArray[np.float64]] not supported - int_args: npt.NDArray, - kernel_type: KernelType, - default: npt.NDArray[np.float64], - ) -> npt.NDArray[np.float64]: - try: - return int_args[:, self.integral_args_lkp[kernel_type]] - except KeyError: - return default - class CountsIntegralND(Ncm.IntegralND): """Integral subclass used to compute the integrals using NumCosmo.""" @@ -136,11 +63,15 @@ class CountsIntegralND(Ncm.IntegralND): def __init__( self, dim: int, - fun: Callable[[npt.NDArray], Sequence[float]], + fun: Callable[ + [npt.NDArray[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64] + ], + args: npt.NDArray[np.float64], ) -> None: super().__init__() self.dim = dim self.fun = fun + self.extra_args = args # pylint: disable-next=arguments-differ def do_get_dimensions(self) -> Tuple[int, int]: @@ -158,4 +89,4 @@ def do_integrand( ) -> None: """Called by NumCosmo to evaluate the integrand.""" x = np.array(x_vec.dup_array()).reshape(npoints, dim) - fval_vec.set_array(self.fun(x)) + fval_vec.set_array(list(self.fun(x, self.extra_args))) diff --git a/firecrown/models/cluster/integrator/scipy_integrator.py b/firecrown/models/cluster/integrator/scipy_integrator.py index 87bc8420..4869f60c 100644 --- a/firecrown/models/cluster/integrator/scipy_integrator.py +++ b/firecrown/models/cluster/integrator/scipy_integrator.py @@ -2,13 +2,13 @@ This module holds the scipy implementation of the integrator classes """ -from typing import Callable, Dict, Tuple +from typing import Callable + import numpy as np import numpy.typing as npt from scipy.integrate import nquad + from firecrown.models.cluster.integrator.integrator import Integrator -from firecrown.models.cluster.kernel import KernelType -from firecrown.models.cluster.abundance import ClusterAbundance, AbundanceIntegrand class ScipyIntegrator(Integrator): @@ -21,80 +21,16 @@ def __init__( self._relative_tolerance = relative_tolerance self._absolute_tolerance = absolute_tolerance - self.integral_args_lkp: Dict[KernelType, int] = self._default_integral_args() - - def _integral_wrapper( - self, - integrand: AbundanceIntegrand, - ) -> Callable[..., float]: - self._validate_integrand(integrand) - - def scipy_integrand(*int_args: float) -> float: - default = -1.0 - # pylint: disable=R0801 - mass = self._get_or_default(int_args, KernelType.MASS, default) - z = self._get_or_default(int_args, KernelType.Z, default) - mass_proxy = self._get_or_default(int_args, KernelType.MASS_PROXY, default) - z_proxy = self._get_or_default(int_args, KernelType.Z_PROXY, default) - - return_val = integrand( - mass, - z, - self.sky_area, - mass_proxy, - z_proxy, - self.mass_proxy_limits, - self.z_proxy_limits, - )[0] - assert isinstance(return_val, float) - return return_val - - return scipy_integrand - - def set_integration_bounds( - self, - cl_abundance: ClusterAbundance, - sky_area: float, - z_proxy_limits: Tuple[float, float], - mass_proxy_limits: Tuple[float, float], - ) -> None: - # pylint: disable=R0801 - self.integral_args_lkp = self._default_integral_args() - self.integral_bounds = [ - (cl_abundance.min_mass, cl_abundance.max_mass), - (cl_abundance.min_z, cl_abundance.max_z), - ] - - self.mass_proxy_limits = mass_proxy_limits - self.z_proxy_limits = z_proxy_limits - self.sky_area = sky_area - - for kernel in cl_abundance.dirac_delta_kernels: - if kernel.kernel_type == KernelType.Z_PROXY: - self.integral_bounds[1] = z_proxy_limits - - elif kernel.kernel_type == KernelType.MASS_PROXY: - self.integral_bounds[0] = mass_proxy_limits - - for kernel in cl_abundance.integrable_kernels: - idx = len(self.integral_bounds) - - if kernel.kernel_type == KernelType.Z_PROXY: - self.integral_bounds.append(z_proxy_limits) - self.integral_args_lkp[KernelType.Z_PROXY] = idx - - elif kernel.kernel_type == KernelType.MASS_PROXY: - self.integral_bounds.append(mass_proxy_limits) - self.integral_args_lkp[KernelType.MASS_PROXY] = idx - def integrate( self, - integrand: AbundanceIntegrand, + func_to_integrate: Callable[ + [npt.NDArray[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64] + ], ) -> float: - scipy_integrand = self._integral_wrapper(integrand) val = nquad( - scipy_integrand, + func_to_integrate, ranges=self.integral_bounds, + args=self.extra_args, opts={ "epsabs": self._absolute_tolerance, "epsrel": self._relative_tolerance, @@ -102,14 +38,3 @@ def integrate( )[0] assert isinstance(val, float) return val - - def _get_or_default( - self, - int_args: Tuple[float, ...], - kernel_type: KernelType, - default: float, - ) -> npt.NDArray[np.float64]: - try: - return np.atleast_1d(int_args[self.integral_args_lkp[kernel_type]]) - except KeyError: - return np.atleast_1d(default) diff --git a/firecrown/models/cluster/kernel.py b/firecrown/models/cluster/kernel.py index 3e094269..df3fcd08 100644 --- a/firecrown/models/cluster/kernel.py +++ b/firecrown/models/cluster/kernel.py @@ -2,12 +2,10 @@ This module holds the classes that define the kernels that can be included in the cluster abundance integrand.""" -from abc import ABC, abstractmethod from enum import Enum -from typing import List, Tuple, Optional +from typing import Tuple import numpy.typing as npt import numpy as np -from firecrown.updatable import Updatable class KernelType(Enum): @@ -21,56 +19,19 @@ class KernelType(Enum): PURITY = 6 -class Kernel(Updatable, ABC): - """The Kernel base class - - This class defines the common interface that any kernel must implement in order - to be included in the cluster abundance integrand.""" - - def __init__( - self, - kernel_type: KernelType, - is_dirac_delta: bool = False, - has_analytic_sln: bool = False, - integral_bounds: Optional[List[Tuple[float, float]]] = None, - ): - super().__init__() - self.integral_bounds = integral_bounds - self.is_dirac_delta = is_dirac_delta - self.kernel_type = kernel_type - self.has_analytic_sln = has_analytic_sln - - @abstractmethod - def distribution( - self, - mass: npt.NDArray[np.float64], - z: npt.NDArray[np.float64], - 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]: - """The functional form of the distribution or spread of this kernel""" - - -class Completeness(Kernel): - """The completeness kernel +# pylint: disable=too-few-public-methods +class Completeness: + """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) -> None: - super().__init__(KernelType.COMPLETENESS) - def distribution( self, mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], - _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]: + """Evaluates and returns the completeness contribution to the integrand.""" a_nc = 1.1321 b_nc = 0.7751 a_mc = 13.31 @@ -82,15 +43,13 @@ def distribution( return completeness -class Purity(Kernel): - """The purity kernel +# pylint: disable=too-few-public-methods +class Purity: + """The purity kernel for the numcosmo simulated survey This kernel will affect the integrand by accounting for the purity of a cluster selection.""" - def __init__(self) -> None: - super().__init__(KernelType.PURITY) - def _ln_rc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: a_rc = 2.2183 b_rc = -0.6592 @@ -106,13 +65,11 @@ def _nc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: def distribution( self, - _mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], 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]: + """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) @@ -126,42 +83,26 @@ def distribution( return purity -class TrueMass(Kernel): +# pylint: disable=too-few-public-methods +class TrueMass: """The true mass kernel. Assuming we measure the true mass, this will always be 1.""" - def __init__(self) -> None: - super().__init__(KernelType.MASS_PROXY, True) - - def distribution( - self, - _mass: npt.NDArray[np.float64], - _z: npt.NDArray[np.float64], - _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]: + def distribution(self) -> npt.NDArray[np.float64]: + """Evaluates and returns the mass distribution contribution to the integrand. + We have set this to 1.0 (i.e. it does not affect the mass distribution)""" return np.atleast_1d(1.0) -class SpectroscopicRedshift(Kernel): +# pylint: disable=too-few-public-methods +class SpectroscopicRedshift: """The spec-z kernel. Assuming the spectroscopic redshift has no uncertainties, this is akin to multiplying by 1.""" - def __init__(self) -> None: - super().__init__(KernelType.Z_PROXY, True) - - def distribution( - self, - _mass: npt.NDArray[np.float64], - _z: npt.NDArray[np.float64], - _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]: + def distribution(self) -> npt.NDArray[np.float64]: + """Evaluates and returns the z distribution contribution to the integrand. + We have set this to 1.0 (i.e. it does not affect the redshift distribution)""" return np.atleast_1d(1.0) diff --git a/firecrown/models/cluster/mass_proxy.py b/firecrown/models/cluster/mass_proxy.py index f64a4f6a..54170f7d 100644 --- a/firecrown/models/cluster/mass_proxy.py +++ b/firecrown/models/cluster/mass_proxy.py @@ -3,16 +3,18 @@ This module holds the classes that define the mass richness relations that can be included in the cluster abundance integrand. These are implementations of Kernels.""" -from typing import List, Tuple, Optional from abc import abstractmethod +from typing import Tuple + import numpy as np import numpy.typing as npt from scipy import special + from firecrown import parameters -from firecrown.models.cluster.kernel import Kernel, KernelType +from firecrown.updatable import Updatable -class MassRichnessGaussian(Kernel): +class MassRichnessGaussian(Updatable): """The representation of mass richness relations that are of a gaussian form.""" @staticmethod @@ -53,10 +55,7 @@ def _distribution_binned( self, mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], - _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]: proxy_mean = self.get_proxy_mean(mass, z) proxy_sigma = self.get_proxy_sigma(mass, z) @@ -88,9 +87,6 @@ def _distribution_unbinned( mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], 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]: proxy_mean = self.get_proxy_mean(mass, z) proxy_sigma = self.get_proxy_sigma(mass, z) @@ -111,10 +107,8 @@ def __init__( self, pivot_mass: float, pivot_redshift: float, - integral_bounds: Optional[List[Tuple[float, float]]] = None, ): - super().__init__(KernelType.MASS_PROXY, False, True, integral_bounds) - + super().__init__() self.pivot_redshift = pivot_redshift self.pivot_mass = pivot_mass * np.log(10.0) # ln(M) self.log1p_pivot_redshift = np.log1p(self.pivot_redshift) @@ -159,14 +153,10 @@ def distribution( self, mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], - 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]: - return self._distribution_binned( - mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits - ) + """Evaluates and returns the mass-richness contribution to the integrand.""" + return self._distribution_binned(mass, z, mass_proxy_limits) class MurataUnbinned(MassRichnessGaussian): @@ -176,10 +166,8 @@ def __init__( self, pivot_mass: float, pivot_redshift: float, - integral_bounds: Optional[List[Tuple[float, float]]] = None, ): - super().__init__(KernelType.MASS_PROXY, False, True, integral_bounds) - + super().__init__() self.pivot_redshift = pivot_redshift self.pivot_mass = pivot_mass * np.log(10.0) # ln(M) self.log1p_pivot_redshift = np.log1p(self.pivot_redshift) @@ -192,8 +180,6 @@ def __init__( self.sigma_p1 = parameters.register_new_updatable_parameter() self.sigma_p2 = parameters.register_new_updatable_parameter() - # Verify this gets called last or first - def get_proxy_mean( self, mass: npt.NDArray[np.float64], @@ -225,10 +211,6 @@ def distribution( mass: npt.NDArray[np.float64], z: npt.NDArray[np.float64], 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]: - return self._distribution_unbinned( - mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits - ) + """Evaluates and returns the mass-richness contribution to the integrand.""" + return self._distribution_unbinned(mass, z, mass_proxy) diff --git a/firecrown/models/cluster/recipes/__init__.py b/firecrown/models/cluster/recipes/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/firecrown/models/cluster/recipes/cluster_recipe.py b/firecrown/models/cluster/recipes/cluster_recipe.py new file mode 100644 index 00000000..4de2bf49 --- /dev/null +++ b/firecrown/models/cluster/recipes/cluster_recipe.py @@ -0,0 +1,29 @@ +"""Module for defining the ClusterRecipe class""" +from abc import ABC, abstractmethod +from typing import Optional + +from firecrown.models.cluster.abundance import ClusterAbundance +from firecrown.models.cluster.binning import SaccBin +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.updatable import Updatable, UpdatableCollection + + +class ClusterRecipe(Updatable, ABC): + """Abstract class defining a cluster recipe. + + A cluster recipe is a combination of different cluster theoretrical predictions + and models that produces a single prediction for an observable.""" + + def __init__(self, parameter_prefix: Optional[str] = None) -> None: + super().__init__(parameter_prefix) + self.my_updatables: UpdatableCollection = UpdatableCollection() + + @abstractmethod + def evaluate_theory_prediction( + self, + cluster_theory: ClusterAbundance, + this_bin: SaccBin, + sky_area: float, + average_on: Optional[ClusterProperty] = None, + ) -> float: + """Evaluate the theory prediction for this cluster recipe""" diff --git a/firecrown/models/cluster/recipes/murata_binned_spec_z.py b/firecrown/models/cluster/recipes/murata_binned_spec_z.py new file mode 100644 index 00000000..bdfbf9ff --- /dev/null +++ b/firecrown/models/cluster/recipes/murata_binned_spec_z.py @@ -0,0 +1,124 @@ +"""Module for defining the classes used in the MurataBinnedSpecZ cluster recipe.""" +from typing import Callable, Optional, Tuple + +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 MurataBinned +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe + + +class MurataBinnedSpecZRecipe(ClusterRecipe): + """Cluster recipe using the Murata 2019 binned mass-richness relation and assuming + perfectly measured spec-zs.""" + + def __init__(self) -> None: + super().__init__() + + self.integrator = NumCosmoIntegrator() + self.redshift_distribution = SpectroscopicRedshift() + pivot_mass, pivot_redshift = 14.625862906, 0.6 + self.mass_distribution = MurataBinned(pivot_mass, pivot_redshift) + self.my_updatables.append(self.mass_distribution) + + def get_theory_prediction( + self, + cluster_theory: ClusterAbundance, + average_on: Optional[ClusterProperty] = None, + ) -> Callable[ + [npt.NDArray[np.float64], npt.NDArray[np.float64], Tuple[float, float], float], + npt.NDArray[np.float64], + ]: + """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_limits: Tuple[float, float], + 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_limits) + ) + + 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], + Tuple[float, float], + 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_low = extra_args[0] + mass_proxy_high = extra_args[1] + sky_area = extra_args[2] + + return prediction(mass, z, (mass_proxy_low, mass_proxy_high), sky_area) + + return function_mapper + + def evaluate_theory_prediction( + self, + cluster_theory: ClusterAbundance, + this_bin: NDimensionalBin, + sky_area: float, + average_on: Optional[ClusterProperty] = None, + ) -> float: + """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, + ] + self.integrator.extra_args = np.array([*this_bin.mass_proxy_edges, 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/firecrown/models/cluster/recipes/murata_unbinned_spec_z.py b/firecrown/models/cluster/recipes/murata_unbinned_spec_z.py new file mode 100644 index 00000000..472a9b34 --- /dev/null +++ b/firecrown/models/cluster/recipes/murata_unbinned_spec_z.py @@ -0,0 +1,130 @@ +"""Module for defining the classes used in the MurataUnbinnedSpecZ cluster recipe.""" +from typing import Callable, Optional + +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 + + +class MurataUnbinnedSpecZRecipe(ClusterRecipe): + """Cluster recipe using the Murata 2019 unbinned mass-richness relation and assuming + perfectly measured spec-zs.""" + + def __init__(self) -> None: + super().__init__() + + self.integrator = NumCosmoIntegrator() + self.redshift_distribution = SpectroscopicRedshift() + pivot_mass, pivot_redshift = 14.625862906, 0.6 + self.mass_distribution = MurataUnbinned(pivot_mass, pivot_redshift) + self.my_updatables.append(self.mass_distribution) + + def get_theory_prediction( + self, + cluster_theory: ClusterAbundance, + average_on: Optional[ClusterProperty] = None, + ) -> Callable[ + [ + npt.NDArray[np.float64], + npt.NDArray[np.float64], + npt.NDArray[np.float64], + float, + ], + npt.NDArray[np.float64], + ]: + """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.""" + + 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) + ) + + 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[np.float64], npt.NDArray[np.float64]], npt.NDArray[np.float64] + ]: + """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[np.float64], extra_args: npt.NDArray[np.float64] + ) -> 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: Optional[ClusterProperty] = None, + ) -> float: + """Evaluate the theoretical prediction for the observable in the provided bin + using the Murata 2019 unbinned mass-richness relation and assuming perfectly + measured redshifts.""" + self.integrator.integral_bounds = [ + (cluster_theory.min_mass, cluster_theory.max_mass), + this_bin.z_edges, + 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/firecrown/models/pylintrc b/firecrown/models/pylintrc index 83f475be..abd8b0a6 100644 --- a/firecrown/models/pylintrc +++ b/firecrown/models/pylintrc @@ -12,6 +12,9 @@ py-version=3.9 # Discover python modules and packages in the file system subtree. recursive=yes +# Add custom pylint plugins +load-plugins=pylint_plugins.duplicate_code + [MESSAGES CONTROL] # Enable the message, report, category or checker with the given id(s). You can @@ -19,6 +22,7 @@ recursive=yes # multiple time (only on the command line, not in the configuration file where # it should appear only once). See also the "--disable" option for examples. enable=c-extension-no-member +disable=too-few-public-methods [MISCELLANEOUS] diff --git a/pylint_plugins/README.md b/pylint_plugins/README.md new file mode 100644 index 00000000..fe9bb619 --- /dev/null +++ b/pylint_plugins/README.md @@ -0,0 +1,10 @@ +### Custom Pylint Plugins +Pylint does not have any built in functionality to selectively apply rules to subdirectories. + +As an example, if we want to ignore a pylint warning in only one class/directory/module, we would need to create a custom `pylintrc` file and manually exclude this file from the parent pylint runs. + +Custom plugins allow us to selectively ignore warnings. This directory holds the custom plugins we load into pylint through the `pylintrc` files. + +### Table of Contents + +* `duplicate_code`: A custom plugin that will ignore the `duplicate-code` warning. Currently the only classes we allow to ignore this message is `firecrown.models.cluster.recipes`. \ No newline at end of file diff --git a/pylint_plugins/duplicate_code.py b/pylint_plugins/duplicate_code.py new file mode 100644 index 00000000..3e95d734 --- /dev/null +++ b/pylint_plugins/duplicate_code.py @@ -0,0 +1,21 @@ +from astroid import MANAGER +from astroid import scoped_nodes +from pylint.lint import PyLinter + + +def register(linter: PyLinter): + pass + + +def transform(mod): + if "firecrown.models.cluster.recipes." not in mod.name: + return + + c = mod.stream().read() + c = b"# pylint: disable=duplicate-code\n" + c + + # pylint will read from `.file_bytes` attribute later when tokenization + mod.file_bytes = c + + +MANAGER.register_transform(scoped_nodes.Module, transform) diff --git a/pylintrc b/pylintrc index c7fefa57..59b3b0b3 100644 --- a/pylintrc +++ b/pylintrc @@ -12,6 +12,9 @@ py-version=3.9 # Discover python modules and packages in the file system subtree. recursive=yes +# Add custom pylint plugins +load-plugins=pylint_plugins.duplicate_code + [MESSAGES CONTROL] # Enable the message, report, category or checker with the given id(s). You can @@ -19,6 +22,7 @@ recursive=yes # multiple time (only on the command line, not in the configuration file where # it should appear only once). See also the "--disable" option for examples. enable=c-extension-no-member +disable=too-few-public-methods [MISCELLANEOUS] diff --git a/tests/cluster_recipes/test_murata_binned_spec_z.py b/tests/cluster_recipes/test_murata_binned_spec_z.py new file mode 100644 index 00000000..22ebb6d5 --- /dev/null +++ b/tests/cluster_recipes/test_murata_binned_spec_z.py @@ -0,0 +1,184 @@ +"""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 import ( + MurataBinnedSpecZRecipe, +) + + +@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") +def fixture_murata_binned_spec_z() -> MurataBinnedSpecZRecipe: + cluster_recipe = MurataBinnedSpecZRecipe() + 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 + return cluster_recipe + + +def test_murata_binned_spec_z_init(): + recipe = MurataBinnedSpecZRecipe() + + 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.my_updatables is not None + assert len(recipe.my_updatables) == 1 + assert recipe.my_updatables[0] is recipe.mass_distribution + + +def test_get_theory_prediction_returns_value( + cluster_abundance: ClusterAbundance, + murata_binned_spec_z: MurataBinnedSpecZRecipe, +): + prediction = murata_binned_spec_z.get_theory_prediction(cluster_abundance) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + 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: MurataBinnedSpecZRecipe, +): + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + mass_proxy_limits = (0, 5) + sky_area = 360**2 + + prediction = murata_binned_spec_z.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.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.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: MurataBinnedSpecZRecipe, +): + prediction = murata_binned_spec_z.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.SHEAR + ) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + 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: MurataBinnedSpecZRecipe, +): + prediction = murata_binned_spec_z.get_theory_prediction(cluster_abundance) + function_to_integrate = murata_binned_spec_z.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: MurataBinnedSpecZRecipe, +): + mock_bin = Mock(spec=NDimensionalBin) + mock_bin.mass_proxy_edges = (0, 5) + mock_bin.z_edges = (0, 1) + + prediction = murata_binned_spec_z.evaluate_theory_prediction( + cluster_abundance, mock_bin, 360**2 + ) + + assert prediction > 0 diff --git a/tests/cluster_recipes/test_murata_unbinned_spec_z.py b/tests/cluster_recipes/test_murata_unbinned_spec_z.py new file mode 100644 index 00000000..76207c14 --- /dev/null +++ b/tests/cluster_recipes/test_murata_unbinned_spec_z.py @@ -0,0 +1,184 @@ +"""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 MurataUnbinned +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe +from firecrown.models.cluster.recipes.murata_unbinned_spec_z import ( + MurataUnbinnedSpecZRecipe, +) + + +@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_unbinned_spec_z") +def fixture_murata_unbinned_spec_z() -> MurataUnbinnedSpecZRecipe: + cluster_recipe = MurataUnbinnedSpecZRecipe() + 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 + return cluster_recipe + + +def test_murata_unbinned_spec_z_init(): + recipe = MurataUnbinnedSpecZRecipe() + + 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, MurataUnbinned) + assert recipe.my_updatables is not None + assert len(recipe.my_updatables) == 1 + assert recipe.my_updatables[0] is recipe.mass_distribution + + +def test_get_theory_prediction_returns_value( + cluster_abundance: ClusterAbundance, + murata_unbinned_spec_z: MurataUnbinnedSpecZRecipe, +): + prediction = murata_unbinned_spec_z.get_theory_prediction(cluster_abundance) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + mass_proxy = np.linspace(0, 5, 2) + sky_area = 360**2 + + result = prediction(mass, z, mass_proxy, 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_unbinned_spec_z: MurataUnbinnedSpecZRecipe, +): + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + mass_proxy = np.linspace(0, 5, 2) + sky_area = 360**2 + + prediction = murata_unbinned_spec_z.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.MASS + ) + + assert prediction is not None + assert callable(prediction) + + result = prediction(mass, z, mass_proxy, 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_unbinned_spec_z.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.REDSHIFT + ) + + assert prediction is not None + assert callable(prediction) + + result = prediction(mass, z, mass_proxy, 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_unbinned_spec_z.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, 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_unbinned_spec_z: MurataUnbinnedSpecZRecipe, +): + prediction = murata_unbinned_spec_z.get_theory_prediction( + cluster_abundance, average_on=ClusterProperty.SHEAR + ) + + assert prediction is not None + assert callable(prediction) + + mass = np.linspace(13, 17, 2) + z = np.linspace(0.1, 1, 2) + mass_proxy = np.linspace(0, 5, 2) + sky_area = 360**2 + + with pytest.raises(NotImplementedError): + _ = prediction(mass, z, mass_proxy, sky_area) + + +def test_get_function_to_integrate_returns_value( + cluster_abundance: ClusterAbundance, + murata_unbinned_spec_z: MurataUnbinnedSpecZRecipe, +): + prediction = murata_unbinned_spec_z.get_theory_prediction(cluster_abundance) + function_to_integrate = murata_unbinned_spec_z.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, 0.1], [17.0, 1.0, 4.9]]) + extra_args = np.array([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_unbinned_spec_z: MurataUnbinnedSpecZRecipe, +): + mock_bin = Mock(spec=NDimensionalBin) + mock_bin.mass_proxy_edges = (0, 5) + mock_bin.z_edges = (0, 1) + + prediction = murata_unbinned_spec_z.evaluate_theory_prediction( + cluster_abundance, mock_bin, 360**2 + ) + + assert prediction > 0 diff --git a/tests/conftest.py b/tests/conftest.py index ea5b5c33..6944986e 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,7 +15,6 @@ from firecrown.parameters import ParamsMap from firecrown.connector.mapping import MappingCosmoSIS, mapping_builder from firecrown.modeling_tools import ModelingTools -from firecrown.models.cluster.abundance import ClusterAbundance def pytest_addoption(parser): @@ -157,15 +156,3 @@ def fixture_cluster_sacc_data() -> sacc.Sacc: s.add_data_point(mlm, ("not_my_survey", "z_bin_tracer_2", "mass_bin_tracer_2"), 1) return s - - -@pytest.fixture(name="empty_cluster_abundance") -def fixture_empty_cluster_abundance() -> ClusterAbundance: - cl_abundance = ClusterAbundance( - min_z=0, - max_z=2, - min_mass=13, - max_mass=17, - halo_mass_function=None, - ) - return cl_abundance diff --git a/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py b/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py index ff34e4d8..cf542d14 100644 --- a/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py +++ b/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py @@ -3,7 +3,7 @@ import sacc import pytest import pyccl -from firecrown.models.cluster.integrator.integrator import Integrator +from firecrown.models.cluster.recipes.cluster_recipe import ClusterRecipe from firecrown.likelihood.gauss_family.statistic.source.source import SourceSystematic from firecrown.modeling_tools import ModelingTools from firecrown.models.cluster.properties import ClusterProperty @@ -15,9 +15,8 @@ def test_create_binned_number_counts(): - integrator = Mock(spec=Integrator) - integrator.integrate.return_value = 1.0 - bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", integrator) + recipe = Mock(spec=ClusterRecipe) + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", recipe) assert bnc is not None assert bnc.cluster_properties == ClusterProperty.NONE assert bnc.survey_name == "Test" @@ -26,31 +25,27 @@ def test_create_binned_number_counts(): assert len(bnc.data_vector) == 0 bnc = BinnedClusterNumberCounts( - (ClusterProperty.COUNTS | ClusterProperty.MASS), "Test", integrator + (ClusterProperty.COUNTS | ClusterProperty.MASS), "Test", recipe ) assert ClusterProperty.COUNTS in bnc.cluster_properties assert ClusterProperty.MASS in bnc.cluster_properties systematics = [SourceSystematic("mock_systematic")] - bnc = BinnedClusterNumberCounts( - ClusterProperty.NONE, "Test", integrator, systematics - ) + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", recipe, systematics) assert bnc.systematics == systematics def test_get_data_vector(): - integrator = Mock(spec=Integrator) - integrator.integrate.return_value = 1.0 - bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", integrator) + recipe = Mock(spec=ClusterRecipe) + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", recipe) dv = bnc.get_data_vector() assert dv is not None assert len(dv) == 0 -def test_read(cluster_sacc_data: sacc.Sacc): - integrator = Mock(spec=Integrator) - integrator.integrate.return_value = 1.0 - bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "my_survey", integrator) +def test_read_throws_if_no_property(cluster_sacc_data: sacc.Sacc): + recipe = Mock(spec=ClusterRecipe) + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "my_survey", recipe) with pytest.raises( ValueError, @@ -58,7 +53,11 @@ def test_read(cluster_sacc_data: sacc.Sacc): ): bnc.read(cluster_sacc_data) - bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", integrator) + +def test_read_single_property(cluster_sacc_data: sacc.Sacc): + recipe = Mock(spec=ClusterRecipe) + + bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", recipe) bnc.read(cluster_sacc_data) assert bnc.sky_area == 4000 assert len(bnc.bins) == 2 @@ -66,7 +65,7 @@ def test_read(cluster_sacc_data: sacc.Sacc): assert bnc.sacc_indices is not None assert len(bnc.sacc_indices) == 2 - bnc = BinnedClusterNumberCounts(ClusterProperty.MASS, "my_survey", integrator) + bnc = BinnedClusterNumberCounts(ClusterProperty.MASS, "my_survey", recipe) bnc.read(cluster_sacc_data) assert bnc.sky_area == 4000 assert len(bnc.bins) == 2 @@ -74,8 +73,11 @@ def test_read(cluster_sacc_data: sacc.Sacc): assert bnc.sacc_indices is not None assert len(bnc.sacc_indices) == 2 + +def test_read_multiple_properties(cluster_sacc_data: sacc.Sacc): + recipe = Mock(spec=ClusterRecipe) bnc = BinnedClusterNumberCounts( - (ClusterProperty.COUNTS | ClusterProperty.MASS), "my_survey", integrator + (ClusterProperty.COUNTS | ClusterProperty.MASS), "my_survey", recipe ) bnc.read(cluster_sacc_data) assert bnc.sky_area == 4000 @@ -86,8 +88,8 @@ def test_read(cluster_sacc_data: sacc.Sacc): def test_compute_theory_vector(cluster_sacc_data: sacc.Sacc): - integrator = Mock(spec=Integrator) - integrator.integrate.return_value = 1.0 + recipe = Mock(spec=ClusterRecipe) + recipe.evaluate_theory_prediction.return_value = 1.0 tools = ModelingTools() hmf = pyccl.halos.MassFuncBocquet16() @@ -98,20 +100,20 @@ def test_compute_theory_vector(cluster_sacc_data: sacc.Sacc): tools.update(params) tools.prepare(cosmo) - bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", integrator) + bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", recipe) bnc.read(cluster_sacc_data) tv = bnc.compute_theory_vector(tools) assert tv is not None assert len(tv) == 2 - bnc = BinnedClusterNumberCounts(ClusterProperty.MASS, "my_survey", integrator) + bnc = BinnedClusterNumberCounts(ClusterProperty.MASS, "my_survey", recipe) bnc.read(cluster_sacc_data) tv = bnc.compute_theory_vector(tools) assert tv is not None assert len(tv) == 2 bnc = BinnedClusterNumberCounts( - (ClusterProperty.COUNTS | ClusterProperty.MASS), "my_survey", integrator + (ClusterProperty.COUNTS | ClusterProperty.MASS), "my_survey", recipe ) bnc.read(cluster_sacc_data) tv = bnc.compute_theory_vector(tools) diff --git a/tests/pylintrc b/tests/pylintrc index 5c53d424..977e8043 100644 --- a/tests/pylintrc +++ b/tests/pylintrc @@ -19,6 +19,7 @@ recursive=yes # multiple time (only on the command line, not in the configuration file where # it should appear only once). See also the "--disable" option for examples. enable=c-extension-no-member +disable=duplicate-code [MISCELLANEOUS] diff --git a/tests/test_cluster_abundance.py b/tests/test_cluster_abundance.py index f811d873..7dd17178 100644 --- a/tests/test_cluster_abundance.py +++ b/tests/test_cluster_abundance.py @@ -1,11 +1,9 @@ """Tests for the cluster abundance module.""" -from unittest.mock import Mock -import pytest -import pyccl import numpy as np +import pyccl +import pytest + from firecrown.models.cluster.abundance import ClusterAbundance -from firecrown.models.cluster.kernel import Kernel, KernelType -from firecrown.models.cluster.properties import ClusterProperty @pytest.fixture(name="cluster_abundance") @@ -16,54 +14,6 @@ def fixture_cluster_abundance(): return ca -@pytest.fixture(name="dirac_delta_kernel") -def fixture_dirac_delta_kernel(): - mk = Mock( - spec=Kernel, - kernel_type=KernelType.MASS, - is_dirac_delta=True, - has_analytic_sln=False, - ) - mk.distribution.return_value = np.atleast_1d(1.0) - return mk - - -@pytest.fixture(name="integrable_kernel") -def fixture_integrable_kernel(): - mk = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - is_dirac_delta=False, - has_analytic_sln=False, - ) - mk.distribution.return_value = np.atleast_1d(1.0) - return mk - - -@pytest.fixture(name="analytic_sln_kernel") -def fixture_analytic_sln_solved_kernel(): - mk = Mock( - spec=Kernel, - kernel_type=KernelType.MASS, - is_dirac_delta=False, - has_analytic_sln=True, - ) - mk.distribution.return_value = np.atleast_1d(1.0) - return mk - - -@pytest.fixture(name="integrand_args") -def fixture_integrand_args(): - sky_area = 439.78986 - mass = np.linspace(13, 17, 5) - z = np.linspace(0, 1, 5) - mass_proxy = np.linspace(0, 5, 5) - z_proxy = np.linspace(0, 1, 5) - mass_proxy_limits = (0, 5) - z_proxy_limits = (0, 1) - return (mass, z, sky_area, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits) - - def test_cluster_abundance_init(cluster_abundance: ClusterAbundance): assert cluster_abundance is not None assert cluster_abundance.cosmo is None @@ -88,42 +38,6 @@ def test_cluster_update_ingredients(cluster_abundance: ClusterAbundance): assert cluster_abundance._hmf_cache == {} -def test_cluster_add_integrable_kernel( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel -): - cluster_abundance.add_kernel(integrable_kernel) - assert len(cluster_abundance.kernels) == 1 - assert isinstance(cluster_abundance.kernels[0], Kernel) - - assert len(cluster_abundance.analytic_kernels) == 0 - assert len(cluster_abundance.dirac_delta_kernels) == 0 - assert len(cluster_abundance.integrable_kernels) == 1 - - -def test_cluster_add_dirac_delta_kernel( - cluster_abundance: ClusterAbundance, dirac_delta_kernel: Kernel -): - cluster_abundance.add_kernel(dirac_delta_kernel) - assert len(cluster_abundance.kernels) == 1 - assert isinstance(cluster_abundance.kernels[0], Kernel) - - assert len(cluster_abundance.analytic_kernels) == 0 - assert len(cluster_abundance.dirac_delta_kernels) == 1 - assert len(cluster_abundance.integrable_kernels) == 0 - - -def test_cluster_add_analytic_sln_kernel( - cluster_abundance: ClusterAbundance, analytic_sln_kernel: Kernel -): - cluster_abundance.add_kernel(analytic_sln_kernel) - assert len(cluster_abundance.kernels) == 1 - assert isinstance(cluster_abundance.kernels[0], Kernel) - - assert len(cluster_abundance.analytic_kernels) == 1 - assert len(cluster_abundance.dirac_delta_kernels) == 0 - assert len(cluster_abundance.integrable_kernels) == 0 - - def test_abundance_comoving_returns_value(cluster_abundance: ClusterAbundance): cosmo = pyccl.CosmologyVanillaLCDM() cluster_abundance.update_ingredients(cosmo) @@ -147,81 +61,3 @@ def test_abundance_massfunc_returns_value(cluster_abundance: ClusterAbundance): assert np.issubdtype(result.dtype, np.float64) assert len(result) == 5 assert np.all(result > 0) - - -@pytest.mark.slow -def test_abundance_get_integrand( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel, integrand_args -): - cosmo = pyccl.CosmologyVanillaLCDM() - cluster_abundance.update_ingredients(cosmo) - cluster_abundance.add_kernel(integrable_kernel) - - integrand = cluster_abundance.get_integrand() - assert callable(integrand) - result = integrand(*integrand_args) - assert isinstance(result, np.ndarray) - assert np.issubdtype(result.dtype, np.float64) - - -@pytest.mark.slow -def test_abundance_get_integrand_avg_mass( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel, integrand_args -): - cosmo = pyccl.CosmologyVanillaLCDM() - cluster_abundance.update_ingredients(cosmo) - cluster_abundance.add_kernel(integrable_kernel) - - integrand = cluster_abundance.get_integrand(average_properties=ClusterProperty.MASS) - assert callable(integrand) - result = integrand(*integrand_args) - assert isinstance(result, np.ndarray) - assert np.issubdtype(result.dtype, np.float64) - - -@pytest.mark.slow -def test_abundance_get_integrand_avg_redshift( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel, integrand_args -): - cosmo = pyccl.CosmologyVanillaLCDM() - cluster_abundance.update_ingredients(cosmo) - cluster_abundance.add_kernel(integrable_kernel) - - integrand = cluster_abundance.get_integrand( - average_properties=ClusterProperty.REDSHIFT - ) - assert callable(integrand) - result = integrand(*integrand_args) - assert isinstance(result, np.ndarray) - assert np.issubdtype(result.dtype, np.float64) - - -@pytest.mark.slow -def test_abundance_get_integrand_avg_mass_and_redshift( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel, integrand_args -): - cosmo = pyccl.CosmologyVanillaLCDM() - cluster_abundance.update_ingredients(cosmo) - cluster_abundance.add_kernel(integrable_kernel) - - average_properties = ClusterProperty.MASS | ClusterProperty.REDSHIFT - integrand = cluster_abundance.get_integrand(average_properties=average_properties) - assert callable(integrand) - result = integrand(*integrand_args) - assert isinstance(result, np.ndarray) - assert np.issubdtype(result.dtype, np.float64) - - -@pytest.mark.slow -def test_abundance_get_integrand_avg_not_implemented_throws( - cluster_abundance: ClusterAbundance, integrable_kernel: Kernel, integrand_args -): - cosmo = pyccl.CosmologyVanillaLCDM() - cluster_abundance.update_ingredients(cosmo) - cluster_abundance.add_kernel(integrable_kernel) - - average_properties = ClusterProperty.SHEAR - integrand = cluster_abundance.get_integrand(average_properties=average_properties) - assert callable(integrand) - with pytest.raises(NotImplementedError): - _ = integrand(*integrand_args) diff --git a/tests/test_cluster_binning.py b/tests/test_cluster_binning.py index 2304e1a6..ad665603 100644 --- a/tests/test_cluster_binning.py +++ b/tests/test_cluster_binning.py @@ -1,7 +1,18 @@ """Tests for the cluster binning module""" from unittest.mock import Mock + +import pytest import sacc -from firecrown.models.cluster.binning import SaccBin, NDimensionalBin + +from firecrown.models.cluster.binning import NDimensionalBin, SaccBin, TupleBin + + +def test_bin_str(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + + sb = SaccBin([tracer_z, tracer_lambda]) + assert str(sb) == "[(0, 1), (4, 5)]\n" def test_create_sacc_bin_with_correct_dimension(): @@ -25,6 +36,22 @@ def test_sacc_bin_z_edges(): assert sb.z_edges == (0, 1) +def test_sacc_bin_z_edges_throws_when_multiple_z_bins(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_z, tracer_z, tracer_lambda]) + with pytest.raises(ValueError, match="SaccBin must have exactly one z bin"): + print(sb.z_edges) + + +def test_sacc_bin_mass_proxy_edges_throws_when_multiple_mass_proxy_bins(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_z, tracer_lambda, tracer_lambda]) + with pytest.raises(ValueError, match="SaccBin must have exactly one richness bin"): + print(sb.mass_proxy_edges) + + def test_sacc_bin_richness_edges(): tracer_z = sacc.tracers.BinZTracer("", 0, 1) tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) @@ -95,3 +122,55 @@ def test_sacc_bin_must_be_equal_type(): sb = SaccBin([tracer_z, tracer_lambda]) assert sb != other_bin + + +def test_create_tuple_bin(): + tb = TupleBin([(1, 2), (3, 4)]) + assert tb is not None + assert tb.dimension == 2 + assert tb.mass_proxy_edges == (1, 2) + assert tb.z_edges == (3, 4) + + +def test_tuple_bins_are_equal(): + tb1 = TupleBin([(1, 2), (3, 4)]) + tb2 = TupleBin([(1, 2), (3, 4)]) + assert tb1 == tb2 + assert hash(tb1) == hash(tb2) + assert tb1 is not tb2 + + +def test_tuple_bin_neq_sacc_bin(): + tb = TupleBin([(1, 2), (3, 4)]) + tracer_z = sacc.tracers.BinZTracer("", 3, 4) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 1, 2) + sb = SaccBin([tracer_z, tracer_lambda]) + assert tb != sb + + +def test_tuple_bin_different_dimensions_not_equal(): + tb1 = TupleBin([(1, 2), (3, 4), (5, 6)]) + tb2 = TupleBin([(1, 2), (3, 4)]) + assert tb1 != tb2 + + tb2 = TupleBin([(1, 2, 3), (3, 4)]) + tb1 = TupleBin([(1, 2), (3, 4)]) + assert tb1 != tb2 + + +def test_tuple_bin_different_bins_not_equal(): + tb1 = TupleBin([(1, 2), (3, 4)]) + tb2 = TupleBin([(1, 2), (3, 5)]) + assert tb1 != tb2 + + tb1 = TupleBin([(0, 2), (3, 4)]) + tb2 = TupleBin([(1, 2), (3, 4)]) + assert tb1 != tb2 + + tb1 = TupleBin([(0, 2), (0, 4)]) + tb2 = TupleBin([(1, 2), (3, 4)]) + assert tb1 != tb2 + + tb1 = TupleBin([(3, 4), (1, 2)]) + tb2 = TupleBin([(1, 2), (3, 4)]) + assert tb1 != tb2 diff --git a/tests/test_cluster_integrators.py b/tests/test_cluster_integrators.py index acaf31a3..fc1d6b93 100644 --- a/tests/test_cluster_integrators.py +++ b/tests/test_cluster_integrators.py @@ -1,291 +1,49 @@ """Tests for the integrator module.""" -from typing import Tuple -from unittest.mock import Mock + import numpy as np import numpy.typing as npt import pytest -from firecrown.models.cluster.integrator.integrator import Integrator from firecrown.models.cluster.integrator.scipy_integrator import ScipyIntegrator from firecrown.models.cluster.integrator.numcosmo_integrator import NumCosmoIntegrator -from firecrown.models.cluster.kernel import KernelType, Kernel -from firecrown.models.cluster.abundance import ClusterAbundance - - -@pytest.fixture(name="integrator", params=[ScipyIntegrator, NumCosmoIntegrator]) -def fixture_integrator(request) -> Integrator: - return request.param() - - -def test_set_integration_bounds_no_kernels( - empty_cluster_abundance: ClusterAbundance, integrator: Integrator -): - z_array = np.linspace(0, 2, 10) - m_array = np.linspace(13, 17, 10) - z_bins = list(zip(z_array[:-1], z_array[1:])) - m_bins = list(zip(m_array[:-1], m_array[1:])) - sky_area = 100**2 - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [(13, 17), (0, 2)] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } - - -def test_set_integration_bounds_dirac_delta( - empty_cluster_abundance: ClusterAbundance, integrator: Integrator -): - z_array = np.linspace(0, 2, 10) - m_array = np.linspace(13, 17, 10) - z_bins = list(zip(z_array[:-1], z_array[1:])) - m_bins = list(zip(m_array[:-1], m_array[1:])) - sky_area = 100**2 - - dd_kernel = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - is_dirac_delta=True, - has_analytic_sln=False, - ) - dd_kernel.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(dd_kernel) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [mass_limits, (0, 2)] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } - - empty_cluster_abundance.kernels.clear() - dd_kernel = Mock( - spec=Kernel, - kernel_type=KernelType.Z_PROXY, - has_analytic_sln=False, - is_dirac_delta=True, - ) - dd_kernel.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(dd_kernel) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [(13, 17), z_limits] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } - dd_kernel2 = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - has_analytic_sln=False, - is_dirac_delta=True, - ) - dd_kernel2.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(dd_kernel2) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [mass_limits, z_limits] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } - - -def test_set_integration_bounds_integrable_kernels( - empty_cluster_abundance: ClusterAbundance, integrator: Integrator -): - z_array = np.linspace(0, 2, 10) - m_array = np.linspace(13, 17, 10) - z_bins = list(zip(z_array[:-1], z_array[1:])) - m_bins = list(zip(m_array[:-1], m_array[1:])) - sky_area = 100**2 - ig_kernel = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - has_analytic_sln=False, - is_dirac_delta=False, - ) - ig_kernel.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(ig_kernel) - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) +# @pytest.fixture(name="integrator", params=[ScipyIntegrator, NumCosmoIntegrator]) +# def fixture_integrator(request) -> Integrator: +# return request.param() - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 3 - assert integrator.integral_bounds == [(13, 17), (0, 2), mass_limits] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - KernelType.MASS_PROXY: 2, - } - empty_cluster_abundance.kernels.clear() - ig_kernel = Mock( - spec=Kernel, - kernel_type=KernelType.Z_PROXY, - has_analytic_sln=False, - is_dirac_delta=False, - ) - ig_kernel.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(ig_kernel) +def test_numcosmo_integrator_integrate(): + integrator = NumCosmoIntegrator() - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 3 - assert integrator.integral_bounds == [(13, 17), (0, 2), z_limits] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - KernelType.Z_PROXY: 2, - } - - ig_kernel2 = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - has_analytic_sln=False, - is_dirac_delta=False, - ) - ig_kernel2.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(ig_kernel2) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 4 - assert integrator.integral_bounds == [(13, 17), (0, 2), z_limits, mass_limits] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - KernelType.Z_PROXY: 2, - KernelType.MASS_PROXY: 3, - } - - -def test_set_integration_bounds_analytic_slns( - empty_cluster_abundance: ClusterAbundance, integrator: Integrator -): - z_array = np.linspace(0, 2, 10) - m_array = np.linspace(13, 17, 10) - z_bins = list(zip(z_array[:-1], z_array[1:])) - m_bins = list(zip(m_array[:-1], m_array[1:])) - sky_area = 100**2 - - a_kernel = Mock( - spec=Kernel, - kernel_type=KernelType.MASS_PROXY, - has_analytic_sln=True, - is_dirac_delta=False, - ) - a_kernel.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(a_kernel) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) - - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [(13, 17), (0, 2)] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } - a_kernel2 = Mock( - spec=Kernel, - kernel_type=KernelType.Z_PROXY, - has_analytic_sln=True, - is_dirac_delta=False, - ) - a_kernel2.distribution.return_value = np.atleast_1d(1.0) - empty_cluster_abundance.add_kernel(a_kernel2) - - for z_limits, mass_limits in zip(z_bins, m_bins): - integrator.set_integration_bounds( - empty_cluster_abundance, sky_area, z_limits, mass_limits - ) + def integrand( + int_args: npt.NDArray[np.float64], _extra_args: npt.NDArray[np.float64] + ) -> npt.NDArray[np.float64]: + # xy + a = int_args[:, 0] + b = int_args[:, 1] + result = a * b + return result - assert integrator.mass_proxy_limits == mass_limits - assert integrator.z_proxy_limits == z_limits - assert len(integrator.integral_bounds) == 2 - assert integrator.integral_bounds == [(13, 17), (0, 2)] - assert integrator.integral_args_lkp == { - KernelType.MASS: 0, - KernelType.Z: 1, - } + integrator.integral_bounds = [(0, 1), (0, 1)] + integrator.extra_args = np.array([], dtype=np.float64) + result = integrator.integrate(integrand) + # \int_0^1 \int_0^1 xy dx dy = 1/4 + assert result == pytest.approx(0.25, rel=1e-15, abs=0) -def test_integrator_integrate( - empty_cluster_abundance: ClusterAbundance, integrator: Integrator -): - empty_cluster_abundance.min_mass = 0 - empty_cluster_abundance.max_mass = 1 - empty_cluster_abundance.min_z = 0 - empty_cluster_abundance.max_z = 1 - sky_area = 100**2 +def test_scipy_integrator_integrate(): + integrator = ScipyIntegrator() + # TODO: should we just remove this? def integrand( - a: npt.NDArray[np.float64], - b: npt.NDArray[np.float64], - _c: float, - _d: npt.NDArray[np.float64], - _e: npt.NDArray[np.float64], - _f: Tuple[float, float], - _g: Tuple[float, float], + a: npt.NDArray[np.float64], b: npt.NDArray[np.float64] ) -> npt.NDArray[np.float64]: # xy result = a * b return result - integrator.set_integration_bounds( - empty_cluster_abundance, - sky_area, - (0, 1), - (0, 1), - ) + integrator.integral_bounds = [(0, 1), (0, 1)] + integrator.extra_args = np.array([], dtype=np.float64) result = integrator.integrate(integrand) # \int_0^1 \int_0^1 xy dx dy = 1/4 assert result == pytest.approx(0.25, rel=1e-15, abs=0) diff --git a/tests/test_cluster_kernels.py b/tests/test_cluster_kernels.py index 1719bdbf..749426b4 100644 --- a/tests/test_cluster_kernels.py +++ b/tests/test_cluster_kernels.py @@ -1,11 +1,10 @@ """Tests for the cluster kernel module.""" -import pytest import numpy as np +import pytest + from firecrown.models.cluster.kernel import ( Completeness, Purity, - KernelType, - Kernel, SpectroscopicRedshift, TrueMass, ) @@ -13,84 +12,41 @@ def test_create_spectroscopic_redshift_kernel(): srk = SpectroscopicRedshift() - assert isinstance(srk, Kernel) - assert srk.kernel_type == KernelType.Z_PROXY - assert srk.is_dirac_delta is True - assert srk.integral_bounds is None - assert srk.has_analytic_sln is False + assert srk is not None def test_create_mass_kernel(): mk = TrueMass() - assert isinstance(mk, Kernel) - assert mk.kernel_type == KernelType.MASS_PROXY - assert mk.is_dirac_delta is True - assert mk.integral_bounds is None - assert mk.has_analytic_sln is False + assert mk is not None def test_create_completeness_kernel(): ck = Completeness() - assert isinstance(ck, Kernel) - assert ck.kernel_type == KernelType.COMPLETENESS - assert ck.is_dirac_delta is False - assert ck.integral_bounds is None - assert ck.has_analytic_sln is False + assert ck is not None def test_create_purity_kernel(): pk = Purity() - assert isinstance(pk, Kernel) - assert pk.kernel_type == KernelType.PURITY - assert pk.is_dirac_delta is False - assert pk.integral_bounds is None - assert pk.has_analytic_sln is False + assert pk is not None def test_spec_z_distribution(): srk = SpectroscopicRedshift() - - assert ( - srk.distribution( - _mass=np.linspace(13, 17, 5), - _z=np.linspace(0, 1, 5), - _mass_proxy=np.linspace(0, 5, 5), - _z_proxy=np.linspace(0, 1, 5), - _mass_proxy_limits=(0, 5), - _z_proxy_limits=(0, 1), - ) - == 1.0 - ) + assert srk.distribution() == 1.0 def test_true_mass_distribution(): tmk = TrueMass() - - assert ( - tmk.distribution( - _mass=np.linspace(13, 17, 5), - _z=np.linspace(0, 1, 5), - _mass_proxy=np.linspace(0, 5, 5), - _z_proxy=np.linspace(0, 1, 5), - _mass_proxy_limits=(0, 5), - _z_proxy_limits=(0, 1), - ) - == 1.0 - ) + assert tmk.distribution() == 1.0 @pytest.mark.precision_sensitive def test_purity_distribution(): pk = Purity() - - mass = np.linspace(13, 17, 10) mass_proxy = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) z = np.array([0.1, 0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]) - z_proxy = np.linspace(0, 1, 10) - mass_proxy_limits = (1.0, 10.0) - z_proxy_limits = (0.1, 1.0) truth = np.array( [ @@ -107,9 +63,7 @@ def test_purity_distribution(): ] ) - purity = pk.distribution( - mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_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) @@ -118,15 +72,9 @@ def test_purity_distribution(): @pytest.mark.precision_sensitive def test_purity_distribution_uses_mean(): pk = Purity() - - mass = np.linspace(13, 17, 10) - z_proxy = np.linspace(0, 1, 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 = np.ones_like(z) * -1.0 - mass_proxy_limits = (1.0, 10.0) - z_proxy_limits = (0.1, 1.0) truth = np.array( [ @@ -143,9 +91,7 @@ def test_purity_distribution_uses_mean(): ] ) - purity = pk.distribution( - mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_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) @@ -156,11 +102,6 @@ def test_completeness_distribution(): ck = Completeness() mass = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 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 = np.array([1, 2, 3, 4, 5, 6, 7, 8, 9, 10]) - z_proxy = np.linspace(0, 1, 10) - - mass_proxy_limits = (1.0, 10.0) - z_proxy_limits = (0.1, 1.0) truth = np.array( [ @@ -177,9 +118,7 @@ def test_completeness_distribution(): ] ) - comp = ck.distribution( - mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits - ) + comp = ck.distribution(mass, z) assert isinstance(comp, np.ndarray) for ref, true in zip(comp, truth): assert ref == pytest.approx(true, rel=1e-7, abs=0.0) diff --git a/tests/test_cluster_mass_richness.py b/tests/test_cluster_mass_richness.py index 7914cdcc..840c99ca 100644 --- a/tests/test_cluster_mass_richness.py +++ b/tests/test_cluster_mass_richness.py @@ -7,10 +7,7 @@ MurataUnbinned, MassRichnessGaussian, ) -from firecrown.models.cluster.kernel import ( - KernelType, - Kernel, -) + PIVOT_Z = 0.6 PIVOT_MASS = 14.625862906 @@ -54,11 +51,6 @@ def fixture_murata_unbinned() -> MurataUnbinned: def test_create_musigma_kernel(): mb = MurataBinned(1, 1) - assert isinstance(mb, Kernel) - assert mb.kernel_type == KernelType.MASS_PROXY - assert mb.is_dirac_delta is False - assert mb.integral_bounds is None - assert mb.has_analytic_sln is True assert mb.pivot_mass == 1 * np.log(10) assert mb.pivot_redshift == 1 assert mb.log1p_pivot_redshift == np.log1p(1) @@ -91,7 +83,6 @@ def test_cluster_observed_mass(): def test_cluster_murata_binned_distribution(murata_binned_relation: MurataBinned): mass_array = np.linspace(7.0, 26.0, 20) mass_proxy_limits = (1.0, 5.0) - z_proxy_limits = (0.0, 1.0) for z in np.geomspace(1.0e-18, 2.0, 20): flip = False @@ -99,14 +90,12 @@ def test_cluster_murata_binned_distribution(murata_binned_relation: MurataBinned mass1 = np.atleast_1d(mass1) mass2 = np.atleast_1d(mass2) z = np.atleast_1d(z) - z_proxy = np.atleast_1d(0) - mass_proxy = np.atleast_1d(1) probability_0 = murata_binned_relation.distribution( - mass1, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits + mass1, z, mass_proxy_limits ) probability_1 = murata_binned_relation.distribution( - mass2, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits + mass2, z, mass_proxy_limits ) assert probability_0 >= 0 @@ -167,8 +156,6 @@ def test_cluster_murata_binned_variance(murata_binned_relation: MurataBinned): def test_cluster_murata_unbinned_distribution(murata_unbinned_relation: MurataUnbinned): mass_array = np.linspace(7.0, 26.0, 20) - mass_proxy_limits = (1.0, 5.0) - z_proxy_limits = (0.0, 1.0) for z in np.geomspace(1.0e-18, 2.0, 20): flip = False @@ -176,15 +163,10 @@ def test_cluster_murata_unbinned_distribution(murata_unbinned_relation: MurataUn mass1 = np.atleast_1d(mass1) mass2 = np.atleast_1d(mass2) z = np.atleast_1d(z) - z_proxy = np.atleast_1d(0) mass_proxy = np.atleast_1d(1) - probability_0 = murata_unbinned_relation.distribution( - mass1, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits - ) - probability_1 = murata_unbinned_relation.distribution( - mass2, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits - ) + probability_0 = murata_unbinned_relation.distribution(mass1, z, mass_proxy) + probability_1 = murata_unbinned_relation.distribution(mass2, z, mass_proxy) # Probability density should be initially monotonically increasing # and then monotonically decreasing. It should flip only once. @@ -215,12 +197,10 @@ def test_cluster_murata_unbinned_distribution_is_normalized( sigma = murata_unbinned_relation.get_proxy_sigma(mass, z)[0] mass_proxy_limits = (mean - 5 * sigma, mean + 5 * sigma) - def integrand(mass_proxy): + def integrand(mass_proxy) -> float: """Evaluate the unbinned distribution at fixed mass and redshift.""" # pylint: disable=cell-var-from-loop - return murata_unbinned_relation.distribution( - mass, z, mass_proxy, z, mass_proxy_limits, (0.0, 1.0) - ) + return murata_unbinned_relation.distribution(mass, z, mass_proxy)[0] result, _ = quad( integrand,