diff --git a/examples/cluster_number_counts/cluster_redshift_richness.py b/examples/cluster_number_counts/cluster_redshift_richness.py index 0077050ef..a7acf0b3c 100644 --- a/examples/cluster_number_counts/cluster_redshift_richness.py +++ b/examples/cluster_number_counts/cluster_redshift_richness.py @@ -1,75 +1,73 @@ """Likelihood factory function for cluster number counts.""" import os -from typing import Any, Dict import pyccl as ccl import sacc -from firecrown.likelihood.gauss_family.statistic.cluster_number_counts import ( - ClusterNumberCounts, -) +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.modeling_tools import ModelingTools -from firecrown.models.cluster_abundance import ClusterAbundance -from firecrown.models.cluster_mass_rich_proxy import ClusterMassRich -from firecrown.models.cluster_redshift_spec import ClusterRedshiftSpec - +from firecrown.models.cluster.abundance import ClusterAbundance +from firecrown.models.cluster.kernel import ( + SpectroscopicRedshift, +) +from firecrown.models.cluster.mass_proxy import MurataBinned +from firecrown.likelihood.likelihood import NamedParameters, Likelihood +from typing import Tuple -def build_likelihood(build_parameters): - """ - Here we instantiate the number density (or mass function) object. - """ - pivot_mass = 14.625862906 - pivot_redshift = 0.6 +def get_cluster_abundance() -> ClusterAbundance: + hmf = ccl.halos.MassFuncBocquet16() + min_mass, max_mass = 13.0, 16.0 + min_z, max_z = 0.2, 0.8 + cluster_abundance = ClusterAbundance(min_mass, max_mass, min_z, max_z, hmf) - cluster_mass_r = ClusterMassRich(pivot_mass, pivot_redshift) - cluster_z = ClusterRedshiftSpec() + # 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) - # TODO: remove try/except when pyccl 3.0 is released - try: - hmd_200 = ccl.halos.MassDef200m() - except TypeError: - hmd_200 = ccl.halos.MassDef200m + redshift_proxy_kernel = SpectroscopicRedshift() + cluster_abundance.add_kernel(redshift_proxy_kernel) - hmf_args: Dict[str, Any] = {} - hmf_name = "Tinker08" + return cluster_abundance - cluster_abundance = ClusterAbundance(hmd_200, hmf_name, hmf_args) - stats = ClusterNumberCounts( - "numcosmo_simulated_redshift_richness", - cluster_abundance, - cluster_mass_r, - cluster_z, - use_cluster_counts=build_parameters.get_bool("use_cluster_counts", True), - use_mean_log_mass=build_parameters.get_bool("use_mean_log_mass", False), +def build_likelihood( + build_parameters: NamedParameters, +) -> Tuple[Likelihood, ModelingTools]: + """ + 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 + if build_parameters.get_bool("use_cluster_counts", True): + average_properties |= ClusterProperty.COUNTS + if build_parameters.get_bool("use_mean_log_mass", False): + average_properties |= ClusterProperty.MASS + + survey_name = "numcosmo_simulated_redshift_richness" + likelihood = ConstGaussian( + [BinnedClusterNumberCounts(average_properties, survey_name, integrator)] ) - stats_list = [stats] - # Here we instantiate the actual likelihood. The statistics argument carry - # the order of the data/theory vector. - - lk = ConstGaussian(stats_list) - - # We load the correct SACC file. - saccfile = os.path.expanduser( - os.path.expandvars( - "${FIRECROWN_DIR}/examples/cluster_number_counts/" - "cluster_redshift_richness_sacc_data.fits" - ) + # Read in sacc data + sacc_file_nm = "cluster_redshift_richness_sacc_data.fits" + sacc_path = os.path.expanduser( + os.path.expandvars("${FIRECROWN_DIR}/examples/cluster_number_counts/") ) - sacc_data = sacc.Sacc.load_fits(saccfile) - - # The read likelihood method is called passing the loaded SACC file, the - # cluster number count functions will receive the appropriated sections - # of the SACC file. - lk.read(sacc_data) + sacc_data = sacc.Sacc.load_fits(os.path.join(sacc_path, sacc_file_nm)) + likelihood.read(sacc_data) - # This script will be loaded by the appropriated connector. The framework - # then looks for the `likelihood` variable to find the instance that will - # be used to compute the likelihood. - modeling = ModelingTools() + cluster_abundance = get_cluster_abundance() + modeling_tools = ModelingTools(cluster_abundance=cluster_abundance) - return lk, modeling + return likelihood, modeling_tools diff --git a/examples/cluster_number_counts/compare_integration_methods.py b/examples/cluster_number_counts/compare_integration_methods.py index 62bc92adf..fb5563408 100644 --- a/examples/cluster_number_counts/compare_integration_methods.py +++ b/examples/cluster_number_counts/compare_integration_methods.py @@ -1,26 +1,20 @@ """Test integral methods for cluster abundance.""" - -from typing import Any, Dict import time import itertools -import pyccl as ccl +import pyccl import numpy as np -from numcosmo_py import Ncm - -from firecrown.models.cluster_abundance import ClusterAbundance -from firecrown.models.cluster_mass_rich_proxy import ( - ClusterMassRich, - ClusterMassRichBinArgument, -) -from firecrown.models.cluster_redshift_spec import ( - ClusterRedshiftSpecArgument, +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 -def compare_integration(): - """Compare integration methods.""" - +def get_cosmology() -> pyccl.Cosmology: Omega_c = 0.262 Omega_b = 0.049 Omega_k = 0.0 @@ -32,7 +26,8 @@ def compare_integration(): Neff = 3.046 w0 = -1.0 wa = 0.0 - cosmo_ccl = ccl.Cosmology( + + cosmo_ccl = pyccl.Cosmology( Omega_c=Omega_c, Omega_b=Omega_b, Neff=Neff, @@ -45,56 +40,75 @@ def compare_integration(): T_CMB=Tcmb0, m_nu=[0.00, 0.0, 0.0], ) + + return cosmo_ccl + + +def get_mass_richness() -> Kernel: pivot_mass = 14.625862906 pivot_redshift = 0.6 - cluster_mass_r = ClusterMassRich(pivot_mass, pivot_redshift) - - cluster_mass_r.mu_p0 = 3.0 - cluster_mass_r.mu_p1 = 0.86 - cluster_mass_r.mu_p2 = 0.0 - cluster_mass_r.sigma_p0 = 3.0 - cluster_mass_r.sigma_p1 = 0.7 - cluster_mass_r.sigma_p2 = 0.0 - - # TODO: remove try/except when pyccl 3.0 is released - try: - hmd_200 = ccl.halos.MassDef200m() - except TypeError: - hmd_200 = ccl.halos.MassDef200m - - hmf_args: Dict[str, Any] = {} - hmf_name = "Tinker08" - z_bins = np.linspace(0.0, 1.0, 4) - r_bins = np.linspace(1.0, 2.5, 5) - - integ_options = [ - Ncm.IntegralNDMethod.P, - Ncm.IntegralNDMethod.P_V, - Ncm.IntegralNDMethod.H, - Ncm.IntegralNDMethod.H_V, - ] - for integ_method in integ_options: - abundance_test = ClusterAbundance( - hmd_200, hmf_name, hmf_args, integ_method=integ_method - ) - test_m_list = [] - t1 = time.time() + 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 + - for i, j in itertools.product(range(3), range(4)): - cluster_mass_bin = ClusterMassRichBinArgument( - cluster_mass_r, 13, 15, r_bins[j], r_bins[j + 1] +def compare_integration() -> None: + """Compare integration methods.""" + hmf = pyccl.halos.MassFuncTinker08() + abundance = ClusterAbundance(13, 15, 0, 4, hmf) + + mass_richness = get_mass_richness() + abundance.add_kernel(mass_richness) + + redshift_proxy_kernel = SpectroscopicRedshift() + abundance.add_kernel(redshift_proxy_kernel) + + 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)) + + 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_proxy_bins[mass_proxy_idx], + mass_proxy_bins[mass_proxy_idx + 1], ) - cluster_z_bin = ClusterRedshiftSpecArgument(z_bins[i], z_bins[i + 1]) - cluster_counts = abundance_test.compute( - cosmo_ccl, cluster_mass_bin, cluster_z_bin + + nc_integrator.set_integration_bounds( + abundance, + sky_area, + z_proxy_limits, + mass_proxy_limits, ) - test_m_list.append(cluster_counts) - t2 = time.time() + + counts = nc_integrator.integrate(integrand) + counts_list.append(counts) + + t_stop = time.time() print( - f"The time for {integ_method} is {t2-t1}\n\n" - f"The counts value is {test_m_list}\n\n" + f"The time for NumCosmo integration method {method} is {t_stop-t_start}\n\n" + f"The counts value is {counts_list}\n\n" ) diff --git a/examples/cluster_number_counts/generate_rich_mean_mass_sacc_data.py b/examples/cluster_number_counts/generate_rich_mean_mass_sacc_data.py old mode 100644 new mode 100755 index 43c2bc78b..5666faf4f --- a/examples/cluster_number_counts/generate_rich_mean_mass_sacc_data.py +++ b/examples/cluster_number_counts/generate_rich_mean_mass_sacc_data.py @@ -19,11 +19,11 @@ from astropy.io import fits from scipy import stats - +from typing import Any import sacc -def generate_sacc_file(): +def generate_sacc_file() -> Any: """Generate a SACC file for cluster number counts.""" H0 = 71.0 @@ -40,8 +40,7 @@ def generate_sacc_file(): cosmo.add_submodel(prim) dist = Nc.Distance.new(2.0) - - tf = Nc.TransferFunc.new_from_name("NcTransferFuncEH") + tf = Nc.TransferFuncEH.new() psml = Nc.PowspecMLTransfer.new(tf) @@ -74,13 +73,9 @@ def generate_sacc_file(): zu = 0.65 # NumCosmo proxy model based on arxiv 1904.07524v2 - cluster_z = Nc.ClusterRedshift.new_from_name( - f"NcClusterRedshiftNodist{{'z-min': <{zl:22.15e}>, 'z-max':<{zu:22.15e}>}}" - ) - - cluster_m = Nc.ClusterMass.new_from_name( - f"NcClusterMassAscaso{{'M0':<{3.0e14 / 0.71:22.15e}>,'z0':<0.6>, " - f"'lnRichness-min':<{lnRl:22.15e}>, 'lnRichness-max':<{lnRu:22.15e}>}}" + cluster_z = Nc.ClusterRedshiftNodist(z_max=zu, z_min=zl) + cluster_m = Nc.ClusterMassAscaso( + M0=3.0e14 / 0.71, z0=0.6, lnRichness_min=lnRl, lnRichness_max=lnRu ) cluster_m.param_set_by_name("mup0", 3.19) cluster_m.param_set_by_name("mup1", 2 / np.log(10)) diff --git a/examples/des_y1_3x2pt/numcosmo_run.py b/examples/des_y1_3x2pt/numcosmo_run.py index 1dd3e6098..ef6d029e4 100755 --- a/examples/des_y1_3x2pt/numcosmo_run.py +++ b/examples/des_y1_3x2pt/numcosmo_run.py @@ -129,7 +129,7 @@ def setup_likelihood( def setup_fit(likelihood: Ncm.Likelihood, mset: Ncm.MSet) -> Ncm.Fit: """Setup the fit object.""" - fit = Ncm.Fit.new( + fit = Ncm.Fit.factory( Ncm.FitType.NLOPT, "ln-neldermead", likelihood, @@ -151,7 +151,7 @@ def setup_fit_all() -> Ncm.Fit: return fit -def run_test(): +def run_test() -> None: """Run the fit.""" fit = setup_fit_all() @@ -162,7 +162,7 @@ def run_test(): fit.run_restart(Ncm.FitRunMsgs.FULL, 1.0e-3, 0.0, None, None) -def run_compute_best_fit(): +def run_compute_best_fit() -> None: """Run the fit.""" fit = setup_fit_all() @@ -178,7 +178,7 @@ def run_compute_best_fit(): fit.run_restart(Ncm.FitRunMsgs.FULL, 1.0e-3, 0.0, None, None) -def run_apes_sampler(ssize: int): +def run_apes_sampler(ssize: int) -> None: """Run the fit.""" fit = setup_fit_all() @@ -192,7 +192,7 @@ def run_apes_sampler(ssize: int): nwalkers = mset.fparam_len * 100 esmcmc = create_esmcmc( - fit.lh, mset, "des_y1_3x2pt_apes", nwalkers=nwalkers, nthreads=1 + fit.props.likelihood, mset, "des_y1_3x2pt_apes", nwalkers=nwalkers, nthreads=1 ) esmcmc.start_run() diff --git a/examples/des_y1_3x2pt/numcosmo_run_PT.py b/examples/des_y1_3x2pt/numcosmo_run_PT.py index 4935b00c0..a73821539 100755 --- a/examples/des_y1_3x2pt/numcosmo_run_PT.py +++ b/examples/des_y1_3x2pt/numcosmo_run_PT.py @@ -77,7 +77,7 @@ lh.priors_add_gauss_param_name(mset, "NcFirecrownPT:src0_delta_z", -0.001, 0.016) lh.priors_add_gauss_param_name(mset, "NcFirecrownPT:lens0_delta_z", +0.001, 0.008) -fit = Ncm.Fit.new( +fit = Ncm.Fit.factory( Ncm.FitType.NLOPT, "ln-neldermead", lh, diff --git a/examples/srd_sn/numcosmo_run.py b/examples/srd_sn/numcosmo_run.py index 6577f0c34..2375618bd 100755 --- a/examples/srd_sn/numcosmo_run.py +++ b/examples/srd_sn/numcosmo_run.py @@ -81,7 +81,7 @@ lh = Ncm.Likelihood(dataset=dset) -fit = Ncm.Fit.new( +fit = Ncm.Fit.factory( Ncm.FitType.NLOPT, "ln-neldermead", lh, diff --git a/firecrown/connector/cosmosis/likelihood.py b/firecrown/connector/cosmosis/likelihood.py index 8edaee08c..c77a6cbee 100644 --- a/firecrown/connector/cosmosis/likelihood.py +++ b/firecrown/connector/cosmosis/likelihood.py @@ -7,8 +7,6 @@ likelihood abstract base class; it the implementation of a CosmoSIS module, not a specific likelihood. """ - - import cosmosis.datablock from cosmosis.datablock import option_section from cosmosis.datablock import names as section_names @@ -175,7 +173,7 @@ def execute(self, sample: cosmosis.datablock) -> int: return 0 - def form_error_message(self, exc): + def form_error_message(self, exc: MissingSamplerParameterError) -> str: """Form the error message that will be used to report a missing parameter, when that parameter should have been supplied by the sampler.""" @@ -228,6 +226,6 @@ def execute(sample: cosmosis.datablock, instance: FirecrownLikelihood) -> int: return instance.execute(sample) -def cleanup(_): +def cleanup(_) -> int: """Cleanup hook for a CosmoSIS module. This one has nothing to do.""" return 0 diff --git a/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py b/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py new file mode 100644 index 000000000..8d5e9e01f --- /dev/null +++ b/firecrown/likelihood/gauss_family/statistic/binned_cluster_number_counts.py @@ -0,0 +1,152 @@ +"""This module holds classes needed to predict the binned cluster number counts + +The binned cluster number counts statistic predicts the number of galaxy +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 +from firecrown.likelihood.gauss_family.statistic.statistic import ( + Statistic, + DataVector, + TheoryVector, +) +from firecrown.likelihood.gauss_family.statistic.source.source import SourceSystematic +from firecrown.modeling_tools import ModelingTools + + +class BinnedClusterNumberCounts(Statistic): + """The Binned Cluster Number Counts statistic + + This class will make a prediction for the number of clusters in a z, mass bin + and compare that prediction to the data provided in the sacc file. + """ + + def __init__( + self, + cluster_properties: ClusterProperty, + survey_name: str, + integrator: Integrator, + systematics: Optional[List[SourceSystematic]] = None, + ): + super().__init__() + self.systematics = systematics or [] + self.theory_vector: Optional[TheoryVector] = None + self.cluster_properties = cluster_properties + self.survey_name = survey_name + self.integrator = integrator + self.data_vector = DataVector.from_list([]) + self.sky_area = 0.0 + self.bins: List[SaccBin] = [] + + def read(self, sacc_data: sacc.Sacc) -> None: + # Build the data vector and indices needed for the likelihood + if self.cluster_properties == ClusterProperty.NONE: + raise ValueError("You must specify at least one cluster property.") + + sacc_adapter = AbundanceData(sacc_data) + self.sky_area = sacc_adapter.get_survey_tracer(self.survey_name).sky_area + + data, indices = sacc_adapter.get_observed_data_and_indices_by_survey( + self.survey_name, self.cluster_properties + ) + self.data_vector = DataVector.from_list(data) + self.sacc_indices = np.array(indices) + + self.bins = sacc_adapter.get_bin_edges( + self.survey_name, self.cluster_properties + ) + for bin_edge in self.bins: + if bin_edge.dimension != self.bins[0].dimension: + raise ValueError( + "The cluster number counts statistic requires all bins to be the " + "same dimension." + ) + + super().read(sacc_data) + + def get_data_vector(self) -> DataVector: + assert self.data_vector is not None + return self.data_vector + + def compute_theory_vector(self, tools: ModelingTools) -> TheoryVector: + assert tools.cluster_abundance is not None + + theory_vector_list: List[float] = [] + cluster_counts = [] + + cluster_counts = self.get_binned_cluster_counts(tools) + + for cl_property in ClusterProperty: + include_prop = cl_property & self.cluster_properties + if not include_prop: + continue + + if cl_property == ClusterProperty.COUNTS: + theory_vector_list += cluster_counts + continue + + theory_vector_list += self.get_binned_cluster_property( + tools, cluster_counts, cl_property + ) + + return TheoryVector.from_list(theory_vector_list) + + def get_binned_cluster_property( + self, + tools: ModelingTools, + cluster_counts: List[float], + cluster_properties: ClusterProperty, + ) -> List[float]: + """Computes the mean mass of clusters in each bin + + Using the data from the sacc file, this function evaluates the likelihood for + a single point of the parameter space, and returns the predicted mean mass of + 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, + ) + + total_mass = self.integrator.integrate(integrand) + mean_mass = total_mass / counts + cluster_masses.append(mean_mass) + + return cluster_masses + + def get_binned_cluster_counts(self, tools: ModelingTools) -> List[float]: + """Computes the number of clusters in each bin + + Using the data from the sacc file, this function evaluates the likelihood for + a single point of the parameter space, and returns the predicted number of + clusters in each bin.""" + 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, + ) + + integrand = tools.cluster_abundance.get_integrand() + counts = self.integrator.integrate(integrand) + cluster_counts.append(counts) + + return cluster_counts diff --git a/firecrown/likelihood/gauss_family/statistic/cluster_number_counts.py b/firecrown/likelihood/gauss_family/statistic/cluster_number_counts.py deleted file mode 100644 index a7d41b3fc..000000000 --- a/firecrown/likelihood/gauss_family/statistic/cluster_number_counts.py +++ /dev/null @@ -1,216 +0,0 @@ -"""Cluster Number Count statistic support. -This module reads the necessary data from a SACC file to compute the -theoretical prediction of cluster number counts inside bins of redshift -and a mass proxy. -""" - -from __future__ import annotations -from typing import List, Dict, Tuple, Optional - -import numpy as np - -import sacc -from sacc.tracers import SurveyTracer - -from .statistic import Statistic, DataVector, TheoryVector -from .source.source import SourceSystematic -from ....models.cluster_abundance import ClusterAbundance -from ....models.cluster_mass import ClusterMass, ClusterMassArgument -from ....models.cluster_redshift import ClusterRedshift, ClusterRedshiftArgument -from ....modeling_tools import ModelingTools - - -class ClusterNumberCounts(Statistic): - """A Cluster Number Count statistic (e.g., halo mass function, - multiplicity functions, volume element, etc.). - This subclass implements the read and computes method for - the Statistic class. It is used to compute the theoretical prediction of - cluster number counts. - """ - - def __init__( - self, - survey_tracer: str, - cluster_abundance: ClusterAbundance, - cluster_mass: ClusterMass, - cluster_redshift: ClusterRedshift, - systematics: Optional[List[SourceSystematic]] = None, - use_cluster_counts: bool = True, - use_mean_log_mass: bool = False, - ): - """Initialize the ClusterNumberCounts object. - Parameters - - :param survey_tracer: name of the survey tracer in the SACC data. - :param cluster_abundance: The cluster abundance model to use. - :param systematics: A list of the statistics-level systematics to apply to - the statistic. The default of `None` implies no systematics. - """ - super().__init__(parameter_prefix=survey_tracer) - self.survey_tracer = survey_tracer - self.systematics = systematics or [] - self.data_vector: Optional[DataVector] = None - self.theory_vector: Optional[TheoryVector] = None - self.cluster_abundance: ClusterAbundance = cluster_abundance - self.cluster_mass: ClusterMass = cluster_mass - self.cluster_redshift: ClusterRedshift = cluster_redshift - self.tracer_args: List[Tuple[ClusterRedshiftArgument, ClusterMassArgument]] = [] - self.use_cluster_counts: bool = use_cluster_counts - self.use_mean_log_mass: bool = use_mean_log_mass - - if not self.use_cluster_counts and not self.use_mean_log_mass: - raise ValueError( - "At least one of use_cluster_counts and use_mean_log_mass must be True." - ) - - def _read_data_type(self, sacc_data, data_type): - """Internal function to read the data from the SACC file.""" - tracers_combinations = np.array( - sacc_data.get_tracer_combinations(data_type=data_type) - ) - - if len(tracers_combinations) == 0: - raise ValueError( - f"The SACC file does not contain any tracers for the " - f"{data_type} data type." - ) - - if tracers_combinations.shape[1] != 3: - raise ValueError( - "The SACC file must contain 3 tracers for the " - "cluster_counts data type: cluster_survey, " - "redshift argument and mass argument tracers." - ) - - cluster_survey_tracers = tracers_combinations[:, 0] - - if self.survey_tracer not in cluster_survey_tracers: - raise ValueError( - f"The SACC tracer {self.survey_tracer} is not " - f"present in the SACC file." - ) - - survey_selection = cluster_survey_tracers == self.survey_tracer - - z_tracers = np.unique(tracers_combinations[survey_selection, 1]) - logM_tracers = np.unique(tracers_combinations[survey_selection, 2]) - - z_tracer_bins: Dict[str, ClusterRedshiftArgument] = { - z_tracer: self.cluster_redshift.gen_bin_from_tracer( - sacc_data.get_tracer(z_tracer) - ) - for z_tracer in z_tracers - } - logM_tracer_bins: Dict[str, ClusterMassArgument] = { - logM_tracer: self.cluster_mass.gen_bin_from_tracer( - sacc_data.get_tracer(logM_tracer) - ) - for logM_tracer in logM_tracers - } - - self.tracer_args = [ - (z_tracer_bins[z_tracer], logM_tracer_bins[logM_tracer]) - for _, z_tracer, logM_tracer in tracers_combinations[survey_selection] - ] - - self.cluster_abundance.read(sacc_data) - self.cluster_mass.read(sacc_data) - self.cluster_redshift.read(sacc_data) - - data_vector_list = list( - sacc_data.get_mean(data_type=data_type)[survey_selection] - ) - sacc_indices_list = list( - sacc_data.indices(data_type=data_type)[survey_selection] - ) - - return data_vector_list, sacc_indices_list - - def read(self, sacc_data: sacc.Sacc): - """Read the data for this statistic from the SACC file. - - :param sacc_data: The data in the SACC format. - """ - - try: - survey_tracer: SurveyTracer = sacc_data.get_tracer(self.survey_tracer) - except KeyError as exc: - raise ValueError( - f"The SACC file does not contain the SurveyTracer " - f"{self.survey_tracer}." - ) from exc - if not isinstance(survey_tracer, SurveyTracer): - raise ValueError( - f"The SACC tracer {self.survey_tracer} is not a SurveyTracer." - ) - - self.cluster_abundance.sky_area = survey_tracer.sky_area - - data_vector_list = [] - sacc_indices_list = [] - - if self.use_cluster_counts: - # pylint: disable-next=no-member - cluster_counts = sacc.standard_types.cluster_counts - ( - cluster_counts_data_vector_list, - cluster_counts_sacc_indices_list, - ) = self._read_data_type(sacc_data, cluster_counts) - data_vector_list += cluster_counts_data_vector_list - sacc_indices_list += cluster_counts_sacc_indices_list - - if self.use_mean_log_mass: - # pylint: disable-next=no-member - cluster_mean_log_mass = sacc.standard_types.cluster_mean_log_mass - ( - mean_log_mass_data_vector_list, - mean_log_mass_sacc_indices_list, - ) = self._read_data_type(sacc_data, cluster_mean_log_mass) - - data_vector_list += mean_log_mass_data_vector_list - sacc_indices_list += mean_log_mass_sacc_indices_list - - self.data_vector = DataVector.from_list(data_vector_list) - self.sacc_indices = np.array(sacc_indices_list) - super().read(sacc_data) - - def get_data_vector(self) -> DataVector: - """Return the data vector; raise exception if there is none.""" - assert self.data_vector is not None - return self.data_vector - - def compute_theory_vector(self, tools: ModelingTools) -> TheoryVector: - """Compute a Number Count statistic using the data from the - Read method, the cosmology object, and the Bocquet16 halo mass function. - - :param tools: ModelingTools firecrown object - used to load the required cosmology. - - :return: Numpy Array of floats - An array with the theoretical prediction of the number of clusters - in each bin of redsfhit and mass. - """ - ccl_cosmo = tools.get_ccl_cosmology() - theory_vector_list = [] - cluster_counts_list = [] - - if self.use_cluster_counts or self.use_mean_log_mass: - cluster_counts_list = [ - self.cluster_abundance.compute(ccl_cosmo, logM_tracer_arg, z_tracer_arg) - for z_tracer_arg, logM_tracer_arg in self.tracer_args - ] - if self.use_cluster_counts: - theory_vector_list += cluster_counts_list - - if self.use_mean_log_mass: - mean_log_mass_list = [ - self.cluster_abundance.compute_unormalized_mean_logM( - ccl_cosmo, logM_tracer_arg, z_tracer_arg - ) - / counts - for (z_tracer_arg, logM_tracer_arg), counts in zip( - self.tracer_args, cluster_counts_list - ) - ] - theory_vector_list += mean_log_mass_list - return TheoryVector.from_list(theory_vector_list) diff --git a/firecrown/likelihood/gauss_family/statistic/supernova.py b/firecrown/likelihood/gauss_family/statistic/supernova.py index 986de4840..0c1a9e8aa 100644 --- a/firecrown/likelihood/gauss_family/statistic/supernova.py +++ b/firecrown/likelihood/gauss_family/statistic/supernova.py @@ -20,7 +20,7 @@ class Supernova(Statistic): """A statistic that applies an additive shift M to a supernova's distance modulus.""" - def __init__(self, sacc_tracer) -> None: + def __init__(self, sacc_tracer: str) -> None: """Initialize this statistic.""" super().__init__(parameter_prefix=sacc_tracer) @@ -29,7 +29,7 @@ def __init__(self, sacc_tracer) -> None: self.a: Optional[npt.NDArray[np.float64]] = None self.M = parameters.register_new_updatable_parameter() - def read(self, sacc_data: sacc.Sacc): + def read(self, sacc_data: sacc.Sacc) -> None: """Read the data for this statistic from the SACC file.""" # We do not actually need the tracer, but we want to make sure the SACC diff --git a/firecrown/likelihood/likelihood.py b/firecrown/likelihood/likelihood.py index 14b63c1f5..f185ce9a9 100644 --- a/firecrown/likelihood/likelihood.py +++ b/firecrown/likelihood/likelihood.py @@ -52,7 +52,7 @@ """ from __future__ import annotations -from typing import Mapping, Tuple, Union, Optional +from typing import Mapping, Tuple, Union, Optional, Set from abc import abstractmethod import types import warnings @@ -89,7 +89,7 @@ def __init__(self, parameter_prefix: Optional[str] = None) -> None: self.statistics: UpdatableCollection = UpdatableCollection() @abstractmethod - def read(self, sacc_data: sacc.Sacc): + def read(self, sacc_data: sacc.Sacc) -> None: """Read the covariance matrix for this likelihood from the SACC file.""" @abstractmethod @@ -184,7 +184,18 @@ def get_float_array(self, name: str) -> npt.NDArray[np.float64]: assert val.dtype == np.float64 return val - def to_set(self): + def to_set( + self, + ) -> Set[ + Union[ + str, + int, + bool, + float, + npt.NDArray[np.int64], + npt.NDArray[np.float64], + ] + ]: """Return the contained data as a set.""" return set(self.data) diff --git a/firecrown/modeling_tools.py b/firecrown/modeling_tools.py index 8873bf6a6..62521e675 100644 --- a/firecrown/modeling_tools.py +++ b/firecrown/modeling_tools.py @@ -9,6 +9,8 @@ from abc import ABC, abstractmethod import pyccl.nl_pt +from firecrown.models.cluster.abundance import ClusterAbundance + from .updatable import Updatable, UpdatableCollection @@ -21,6 +23,7 @@ def __init__( *, pt_calculator: Optional[pyccl.nl_pt.EulerianPTCalculator] = None, pk_modifiers: Optional[Collection[PowerspectrumModifier]] = None, + cluster_abundance: Optional[ClusterAbundance] = None, ): super().__init__() self.ccl_cosmo: Optional[pyccl.Cosmology] = None @@ -29,9 +32,10 @@ def __init__( self.pk_modifiers: UpdatableCollection = UpdatableCollection(pk_modifiers) self.powerspectra: Dict[str, pyccl.Pk2D] = {} self._prepared: bool = False + self.cluster_abundance = cluster_abundance - def add_pk(self, name: str, powerspectrum: pyccl.Pk2D): - """Add a :class:`pyccl.Pk2D` to the table of power spectra.""" + def add_pk(self, name: str, powerspectrum: pyccl.Pk2D) -> None: + """Add a :python:`pyccl.Pk2D` to the table of power spectra.""" if name in self.powerspectra: raise KeyError(f"Power spectrum {name} already exists") @@ -87,6 +91,9 @@ def prepare(self, ccl_cosmo: pyccl.Cosmology) -> None: for pkm in self.pk_modifiers: self.add_pk(name=pkm.name, powerspectrum=pkm.compute_p_of_k_z(tools=self)) + if self.cluster_abundance is not None: + self.cluster_abundance.update_ingredients(ccl_cosmo) + self._prepared = True def _reset(self) -> None: diff --git a/firecrown/models/cluster/__init__.py b/firecrown/models/cluster/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/firecrown/models/cluster/abundance.py b/firecrown/models/cluster/abundance.py new file mode 100644 index 000000000..6d836ff1c --- /dev/null +++ b/firecrown/models/cluster/abundance.py @@ -0,0 +1,168 @@ +"""The module responsible for building the cluster abundance calculation. + +The galaxy cluster abundance integral is a combination of both theoretical +and phenomenological predictions. This module contains the classes and +functions that produce those predictions. +""" +from typing import List, Callable, Optional, 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): + """The class that calculates the predicted number counts of galaxy clusters + + The abundance is a function of a specific cosmology, a mass and redshift range, + an area on the sky, a halo mass function, as well as multiple kernels, where + each kernel represents a different distribution involved in the final cluster + abundance integrand. + """ + + @property + 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, + max_mass: float, + min_z: float, + max_z: float, + halo_mass_function: pyccl.halos.MassFunc, + ) -> None: + super().__init__() + self.kernels: UpdatableCollection = UpdatableCollection() + self.halo_mass_function = halo_mass_function + self.min_mass = min_mass + self.max_mass = max_mass + self.min_z = min_z + self.max_z = max_z + 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 + self._hmf_cache = {} + + def comoving_volume( + self, z: npt.NDArray[np.float64], sky_area: float = 0 + ) -> npt.NDArray[np.float64]: + """The differential comoving volume given area sky_area at redshift z. + + :param sky_area: The area of the survey on the sky in square degrees.""" + scale_factor = 1.0 / (1.0 + z) + angular_diam_dist = bkg.angular_diameter_distance(self.cosmo, scale_factor) + h_over_h0 = bkg.h_over_h0(self.cosmo, scale_factor) + + dV = ( + pyccl.physical_constants.CLIGHT_HMPC + * (angular_diam_dist**2) + * ((1.0 + z) ** 2) + / (self.cosmo["h"] * h_over_h0) + ) + assert isinstance(dV, np.ndarray) + + sky_area_rad = sky_area * (np.pi / 180.0) ** 2 + + return dV * sky_area_rad + + def mass_function( + self, + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + """The mass function at z and mass.""" + scale_factor = 1.0 / (1.0 + z) + return_vals = [] + + for m, a in zip(mass, scale_factor): + val = self._hmf_cache.get((m, a)) + if val is None: + val = self.halo_mass_function(self.cosmo, 10**m, a) + self._hmf_cache[(m, a)] = val + 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/abundance_data.py b/firecrown/models/cluster/abundance_data.py new file mode 100644 index 000000000..a37975952 --- /dev/null +++ b/firecrown/models/cluster/abundance_data.py @@ -0,0 +1,136 @@ +"""The module responsible for extracting cluster data from a sacc file. +""" +from typing import Tuple, List +import numpy as np +import numpy.typing as npt +import sacc +from sacc.tracers import SurveyTracer +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.models.cluster.binning import SaccBin + + +class AbundanceData: + """The class used to wrap a sacc file and return the cluster abundance data. + + The sacc file is a complicated set of tracers (bins) and surveys. This class + manipulates that data and returns only the data relevant for the cluster + number count statistic. The data in this class is specific to a single + survey name.""" + + _survey_index = 0 + _redshift_index = 1 + _mass_index = 2 + + def __init__(self, sacc_data: sacc.Sacc): + self.sacc_data = sacc_data + + def get_survey_tracer(self, survey_nm: str) -> SurveyTracer: + """Returns the SurveyTracer for the specified survey name.""" + try: + survey_tracer: SurveyTracer = self.sacc_data.get_tracer(survey_nm) + except KeyError as exc: + raise ValueError( + f"The SACC file does not contain the SurveyTracer " f"{survey_nm}." + ) from exc + + if not isinstance(survey_tracer, SurveyTracer): + raise ValueError(f"The SACC tracer {survey_nm} is not a SurveyTracer.") + + return survey_tracer + + def get_observed_data_and_indices_by_survey( + self, + survey_nm: str, + properties: ClusterProperty, + ) -> Tuple[List[float], List[int]]: + """Will return observed data (and sacc indices) for a specified survey based on + the properties requested by the caller. + + For example if the caller has enabled COUNTS then the observed cluster counts + within each N dimensional bin will be returned.""" + + data_vectors = [] + sacc_indices = [] + + for cluster_property in ClusterProperty: + include_prop = cluster_property & properties + if not include_prop: + continue + + if cluster_property == ClusterProperty.COUNTS: + # pylint: disable=no-member + data_type = sacc.standard_types.cluster_counts + elif cluster_property == ClusterProperty.MASS: + # pylint: disable=no-member + data_type = sacc.standard_types.cluster_mean_log_mass + else: + raise NotImplementedError(cluster_property) + + bin_combinations = self._all_bin_combinations_for_data_type(data_type) + + my_survey_mask = bin_combinations[:, self._survey_index] == survey_nm + + data_vectors += list( + self.sacc_data.get_mean(data_type=data_type)[my_survey_mask] + ) + + sacc_indices += list( + self.sacc_data.indices(data_type=data_type)[my_survey_mask] + ) + + return data_vectors, sacc_indices + + def _all_bin_combinations_for_data_type(self, data_type: str) -> npt.NDArray: + bins_combos_for_type = np.array( + self.sacc_data.get_tracer_combinations(data_type=data_type) + ) + + if len(bins_combos_for_type) == 0: + raise ValueError( + f"The SACC file does not contain any tracers for the " + f"{data_type} data type." + ) + + if bins_combos_for_type.shape[1] != 3: + raise ValueError( + "The SACC file must contain 3 tracers for the " + "cluster_counts data type: cluster_survey, " + "redshift argument and mass argument tracers." + ) + + return bins_combos_for_type + + def get_bin_edges( + self, survey_nm: str, properties: ClusterProperty + ) -> List[SaccBin]: + """Returns the limits for all z, mass bins for the requested data type.""" + bins = [] + + for cluster_property in ClusterProperty: + if not cluster_property & properties: + continue + + if cluster_property == ClusterProperty.COUNTS: + # pylint: disable=no-member + data_type = sacc.standard_types.cluster_counts + elif cluster_property == ClusterProperty.MASS: + # pylint: disable=no-member + data_type = sacc.standard_types.cluster_mean_log_mass + else: + raise NotImplementedError(cluster_property) + + bin_combinations = self._all_bin_combinations_for_data_type(data_type) + my_survey_mask = bin_combinations[:, self._survey_index] == survey_nm + bin_combinations_for_survey = bin_combinations[my_survey_mask] + + for _, z_tracer, mass_tracer in bin_combinations_for_survey: + z_data: sacc.tracers.BinZTracer = self.sacc_data.get_tracer(z_tracer) + mass_data: sacc.tracers.BinRichnessTracer = self.sacc_data.get_tracer( + mass_tracer + ) + sacc_bin = SaccBin([z_data, mass_data]) + bins.append(sacc_bin) + + # Remove duplicates while preserving order (i.e. dont use set()) + unique_bins = list(dict.fromkeys(bins)) + return unique_bins diff --git a/firecrown/models/cluster/binning.py b/firecrown/models/cluster/binning.py new file mode 100644 index 000000000..f4183b718 --- /dev/null +++ b/firecrown/models/cluster/binning.py @@ -0,0 +1,86 @@ +"""This module contains the classes that define the bins and binning +used for cluster theoretical predictions within Firecrown.""" + +from typing import Tuple, List, TypeVar, Generic +from abc import ABC, abstractmethod +import sacc + +T = TypeVar("T") + + +class NDimensionalBin(Generic[T], ABC): + """Class which defines the interface for an N dimensional bin used in + the cluster likelihood.""" + + def __init__(self, coordinate_bins: List[T]): + """_summary_ + + Args: + coordinate_bins (List[T]): _description_ + dimension (int): _description_ + """ + self.coordinate_bins = coordinate_bins + self.dimension = len(coordinate_bins) + + @property + @abstractmethod + def z_edges(self) -> Tuple[float, float]: + """Redshift bin edges""" + + @property + @abstractmethod + def mass_proxy_edges(self) -> Tuple[float, float]: + """Mass proxy bin edges""" + + 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.""" + + def __init__(self, bins: List[sacc.BaseTracer]): + super().__init__(bins) + + @property + def z_edges(self) -> Tuple[float, float]: + z_bin = [x for x in self.coordinate_bins if x.tracer_type == "bin_z"] + if len(z_bin) != 1: + raise ValueError("SaccBin must have exactly one z bin") + return z_bin[0].lower, z_bin[0].upper + + @property + def mass_proxy_edges(self) -> Tuple[float, float]: + mass_bin = [x for x in self.coordinate_bins if x.tracer_type == "bin_richness"] + if len(mass_bin) != 1: + raise ValueError("SaccBin must have exactly one richness bin") + return mass_bin[0].lower, mass_bin[0].upper + + def __eq__(self, other: object) -> bool: + """Two bins are equal if they have the same lower/upper bound.""" + if not isinstance(other, SaccBin): + return False + + if self.dimension != other.dimension: + return False + + for my_bin in self.coordinate_bins: + 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: + return False + + return True + + 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))) diff --git a/firecrown/models/cluster/integrator/__init__.py b/firecrown/models/cluster/integrator/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/firecrown/models/cluster/integrator/integrator.py b/firecrown/models/cluster/integrator/integrator.py new file mode 100644 index 000000000..4503753bd --- /dev/null +++ b/firecrown/models/cluster/integrator/integrator.py @@ -0,0 +1,64 @@ +"""The integrator module + +This module holds the classes that define the interface required to +integrate an assembled cluster abundance. +""" +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 + + +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.""" + + 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() + + @abstractmethod + def integrate( + self, + integrand: AbundanceIntegrand, + ) -> 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 new file mode 100644 index 000000000..d3105307d --- /dev/null +++ b/firecrown/models/cluster/integrator/numcosmo_integrator.py @@ -0,0 +1,161 @@ +"""The NumCosmo integrator module + +This module holds the NumCosmo implementation of the integrator classes +""" +from typing import Tuple, Callable, Sequence, Optional +from enum import Enum +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 + + +class NumCosmoIntegralMethod(Enum): + """The available NumCosmo integration methods.""" + + P = Ncm.IntegralNDMethod.P + P_V = Ncm.IntegralNDMethod.P_V + H = Ncm.IntegralNDMethod.H + H_V = Ncm.IntegralNDMethod.H_V + + +class NumCosmoIntegrator(Integrator): + """The NumCosmo implementation of the Integrator base class.""" + + def __init__( + self, + method: Optional[NumCosmoIntegralMethod] = None, + relative_tolerance: float = 1e-4, + absolute_tolerance: float = 1e-12, + ) -> None: + super().__init__() + 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, + ) -> 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.set_reltol(self._relative_tolerance) + int_nd.set_abstol(self._absolute_tolerance) + res = Ncm.Vector.new(1) + err = Ncm.Vector.new(1) + + bl, bu = zip(*self.integral_bounds) + 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.""" + + def __init__( + self, + dim: int, + fun: Callable[[npt.NDArray], Sequence[float]], + ) -> None: + super().__init__() + self.dim = dim + self.fun = fun + + # pylint: disable-next=arguments-differ + def do_get_dimensions(self) -> Tuple[int, int]: + """Returns the dimensionality of the integral.""" + return self.dim, 1 + + # pylint: disable-next=arguments-differ + def do_integrand( + self, + x_vec: Ncm.Vector, + dim: int, + npoints: int, + _fdim: int, + fval_vec: Ncm.Vector, + ) -> 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)) diff --git a/firecrown/models/cluster/integrator/scipy_integrator.py b/firecrown/models/cluster/integrator/scipy_integrator.py new file mode 100644 index 000000000..87bc8420e --- /dev/null +++ b/firecrown/models/cluster/integrator/scipy_integrator.py @@ -0,0 +1,115 @@ +"""The SciPy integrator module + +This module holds the scipy implementation of the integrator classes +""" +from typing import Callable, Dict, Tuple +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): + """The scipy implementation of the Integrator base class using nquad.""" + + def __init__( + self, relative_tolerance: float = 1e-4, absolute_tolerance: float = 1e-12 + ) -> None: + super().__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, + ) -> float: + scipy_integrand = self._integral_wrapper(integrand) + val = nquad( + scipy_integrand, + ranges=self.integral_bounds, + opts={ + "epsabs": self._absolute_tolerance, + "epsrel": self._relative_tolerance, + }, + )[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 new file mode 100644 index 000000000..3e094269c --- /dev/null +++ b/firecrown/models/cluster/kernel.py @@ -0,0 +1,167 @@ +"""The cluster kernel module + +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 +import numpy.typing as npt +import numpy as np +from firecrown.updatable import Updatable + + +class KernelType(Enum): + """The kernels that can be included in the cluster abundance integrand""" + + MASS = 1 + Z = 2 + MASS_PROXY = 3 + Z_PROXY = 4 + COMPLETENESS = 5 + 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 + + 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]: + a_nc = 1.1321 + b_nc = 0.7751 + a_mc = 13.31 + b_mc = 0.2025 + log_mc = a_mc + b_mc * (1.0 + z) + nc = a_nc + b_nc * (1.0 + z) + completeness = (mass / log_mc) ** nc / ((mass / log_mc) ** nc + 1.0) + assert isinstance(completeness, np.ndarray) + return completeness + + +class Purity(Kernel): + """The purity kernel + + 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 + ln_rc = a_rc + b_rc * (1.0 + z) + return ln_rc + + def _nc(self, z: npt.NDArray[np.float64]) -> npt.NDArray[np.float64]: + b_nc = np.log(10) * 0.3527 + a_nc = np.log(10) * 0.8612 + nc = a_nc + b_nc * (1.0 + z) + assert isinstance(nc, np.ndarray) + return nc + + 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]: + if all(mass_proxy == -1.0): + mean_mass = (mass_proxy_limits[0] + mass_proxy_limits[1]) / 2 + ln_r = np.log(10**mean_mass) + else: + ln_r = np.log(10**mass_proxy) + + r_over_rc = ln_r / self._ln_rc(z) + + purity = (r_over_rc) ** self._nc(z) / (r_over_rc ** self._nc(z) + 1.0) + assert isinstance(purity, np.ndarray) + return purity + + +class TrueMass(Kernel): + """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]: + return np.atleast_1d(1.0) + + +class SpectroscopicRedshift(Kernel): + """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]: + return np.atleast_1d(1.0) diff --git a/firecrown/models/cluster/mass_proxy.py b/firecrown/models/cluster/mass_proxy.py new file mode 100644 index 000000000..f64a4f6a9 --- /dev/null +++ b/firecrown/models/cluster/mass_proxy.py @@ -0,0 +1,234 @@ +"""The mass richness kernel module + +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 +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 + + +class MassRichnessGaussian(Kernel): + """The representation of mass richness relations that are of a gaussian form.""" + + @staticmethod + def observed_value( + p: Tuple[float, float, float], + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + pivot_mass: float, + log1p_pivot_redshift: float, + ) -> npt.NDArray[np.float64]: + """Return observed quantity corrected by redshift and mass.""" + + ln_mass = mass * np.log(10) + delta_ln_mass = ln_mass - pivot_mass + delta_z = np.log1p(z) - log1p_pivot_redshift + + result = p[0] + p[1] * delta_ln_mass + p[2] * delta_z + assert isinstance(result, np.ndarray) + return result + + @abstractmethod + def get_proxy_mean( + self, + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + """Return observed quantity corrected by redshift and mass.""" + + @abstractmethod + def get_proxy_sigma( + self, + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + """Return observed scatter corrected by redshift and mass.""" + + 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) + + x_min = (proxy_mean - mass_proxy_limits[0] * np.log(10.0)) / ( + np.sqrt(2.0) * proxy_sigma + ) + x_max = (proxy_mean - mass_proxy_limits[1] * np.log(10.0)) / ( + np.sqrt(2.0) * proxy_sigma + ) + + return_vals = np.empty_like(x_min) + mask1 = (x_max > 3.0) | (x_min < -3.0) + mask2 = ~mask1 + + # pylint: disable=no-member + return_vals[mask1] = ( + -(special.erfc(x_min[mask1]) - special.erfc(x_max[mask1])) / 2.0 + ) + # pylint: disable=no-member + return_vals[mask2] = ( + special.erf(x_min[mask2]) - special.erf(x_max[mask2]) + ) / 2.0 + assert isinstance(return_vals, np.ndarray) + return return_vals + + def _distribution_unbinned( + 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) + + normalization = 1 / np.sqrt(2 * np.pi * proxy_sigma**2) + result = normalization * np.exp( + -0.5 * ((mass_proxy - proxy_mean) / proxy_sigma) ** 2 + ) + + assert isinstance(result, np.ndarray) + return result + + +class MurataBinned(MassRichnessGaussian): + """The mass richness relation defined in Murata 19 for a binned data vector.""" + + 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) + + 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) + + # Updatable parameters + self.mu_p0 = parameters.register_new_updatable_parameter() + self.mu_p1 = parameters.register_new_updatable_parameter() + self.mu_p2 = parameters.register_new_updatable_parameter() + self.sigma_p0 = parameters.register_new_updatable_parameter() + 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], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + return MassRichnessGaussian.observed_value( + (self.mu_p0, self.mu_p1, self.mu_p2), + mass, + z, + self.pivot_mass, + self.log1p_pivot_redshift, + ) + + def get_proxy_sigma( + self, + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + return MassRichnessGaussian.observed_value( + (self.sigma_p0, self.sigma_p1, self.sigma_p2), + mass, + z, + self.pivot_mass, + self.log1p_pivot_redshift, + ) + + 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 + ) + + +class MurataUnbinned(MassRichnessGaussian): + """The mass richness relation defined in Murata 19 for a unbinned data vector.""" + + 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) + + 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) + + # Updatable parameters + self.mu_p0 = parameters.register_new_updatable_parameter() + self.mu_p1 = parameters.register_new_updatable_parameter() + self.mu_p2 = parameters.register_new_updatable_parameter() + self.sigma_p0 = parameters.register_new_updatable_parameter() + 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], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + return MassRichnessGaussian.observed_value( + (self.mu_p0, self.mu_p1, self.mu_p2), + mass, + z, + self.pivot_mass, + self.log1p_pivot_redshift, + ) + + def get_proxy_sigma( + self, + mass: npt.NDArray[np.float64], + z: npt.NDArray[np.float64], + ) -> npt.NDArray[np.float64]: + return MassRichnessGaussian.observed_value( + (self.sigma_p0, self.sigma_p1, self.sigma_p2), + mass, + z, + self.pivot_mass, + self.log1p_pivot_redshift, + ) + + 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_unbinned( + mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits + ) diff --git a/firecrown/models/cluster/properties.py b/firecrown/models/cluster/properties.py new file mode 100644 index 000000000..e141b13b9 --- /dev/null +++ b/firecrown/models/cluster/properties.py @@ -0,0 +1,13 @@ +"""Module containing classes relevant to defining cluster properties.""" +from enum import Flag, auto + + +class ClusterProperty(Flag): + """Flag containing the possible cluster properties we can make a theoretical + prediction for.""" + + NONE = 0 + COUNTS = auto() + MASS = auto() + REDSHIFT = auto() + SHEAR = auto() diff --git a/firecrown/models/cluster_abundance.py b/firecrown/models/cluster_abundance.py deleted file mode 100644 index 387273085..000000000 --- a/firecrown/models/cluster_abundance.py +++ /dev/null @@ -1,358 +0,0 @@ -r"""Cluster Abundance Module -abstract class to compute cluster abundance. -======================================== -The implemented functions use PyCCL library as backend. -""" -from __future__ import annotations -from typing import Optional, Any, Dict, List, Tuple, final - -import numpy as np -import pyccl as ccl -from numcosmo_py import Ncm -import scipy.integrate - -from ..updatable import Updatable -from .cluster_mass import ClusterMassArgument -from .cluster_redshift import ClusterRedshiftArgument - - -class CountsIntegralND(Ncm.IntegralND): - """Integral subclass used by the ClusterAbundance - to compute the integrals using numcosmo.""" - - def __init__(self, dim, fun, *args): - super().__init__() - self.dim = dim - self.fun = fun - self.args = args - - # pylint: disable-next=arguments-differ - def do_get_dimensions(self) -> Tuple[int, int]: - """Get number of dimensions.""" - return self.dim, 1 - - # pylint: disable-next=arguments-differ - def do_integrand( - self, - x_vec: Ncm.Vector, - dim: int, - npoints: int, - _fdim: int, - fval_vec: Ncm.Vector, - ) -> None: - """Integrand function.""" - x = np.array(x_vec.dup_array()).reshape(npoints, dim) - fval_vec.set_array([self.fun(x_i, *self.args) for x_i in x]) - - -class ClusterAbundance(Updatable): - r"""Cluster Abundance class""" - - def __init__( - self, - halo_mass_definition: ccl.halos.MassDef, - halo_mass_function_name: str, - halo_mass_function_args: Dict[str, Any], - sky_area: float = 100.0, - use_completness: bool = False, - use_purity: bool = False, - integ_method: Ncm.IntegralNDMethod = Ncm.IntegralNDMethod.P_V, - prefer_scipy_integration: bool = False, - reltol: float = 1.0e-4, - abstol: float = 1.0e-12, - ): - """Initialize the ClusterAbundance class. - - :param halo_mass_definition: Halo mass definition. - :param halo_mass_function_name: Halo mass function name. - :param halo_mass_function_args: Halo mass function arguments. - :param sky_area: Sky area in square degrees, defaults to 100 sq deg. - :param use_completness: Use completeness function, defaults to False. - :param use_purity: Use purity function, defaults to False. - - :return: ClusterAbundance object. - """ - super().__init__() - self.sky_area = sky_area - self.halo_mass_definition = halo_mass_definition - self.halo_mass_function_name = halo_mass_function_name - self.halo_mass_function_args = halo_mass_function_args - self.halo_mass_function: Optional[ccl.halos.MassFunc] = None - self.use_purity = use_purity - self.integ_method = integ_method - self.reltol = reltol - self.abstol = abstol - self.prefer_scipy_integration = prefer_scipy_integration - if use_completness: - self.base_mf_d2N_dz_dlnM = self.mf_d2N_dz_dlnM_completeness - else: - self.base_mf_d2N_dz_dlnM = self.mf_d2N_dz_dlnM - - @property - def sky_area(self) -> float: - """Return the sky area.""" - return self.sky_area_rad * (180.0 / np.pi) ** 2 - - @sky_area.setter - def sky_area(self, sky_area: float) -> None: - """Set the sky area.""" - self.sky_area_rad = sky_area * (np.pi / 180.0) ** 2 - - @final - def _reset(self) -> None: - """Implementation of the Updatable interface method `_reset`.""" - self.halo_mass_function = None - - def read(self, sacc_data): - """Read the data for this statistic from the SACC file. - - :param sacc_data: The data in the sacc format. - """ - - def dV_dz(self, ccl_cosmo: ccl.Cosmology, z) -> float: - """Differential Comoving Volume at z. - - parameters - :param ccl_cosmo: pyccl Cosmology - :param z: Cluster Redshift. - - :return: Differential Comoving Volume at z in units of Mpc^3 (comoving). - """ - a = 1.0 / (1.0 + z) - da = ccl.background.angular_diameter_distance(ccl_cosmo, a) - E = ccl.background.h_over_h0(ccl_cosmo, a) - dV = ( - ((1.0 + z) ** 2) - * (da**2) - * ccl.physical_constants.CLIGHT_HMPC - / ccl_cosmo["h"] - / E - ) - return dV * self.sky_area_rad - - def mf_d2N_dV_dlnM(self, ccl_cosmo: ccl.Cosmology, logM: float, z: float) -> float: - """ - Computes the mass function at z and logM. - - :param ccl_cosmo: pyccl Cosmology - :param logM: Cluster mass given by log10(M) where M is in units of - M_sun (comoving). - :param z: Cluster Redshift. - :return: Number Density pdf at z and logm in units of Mpc^-3 (comoving). - """ - a = 1.0 / (1.0 + z) - mass = 10 ** (logM) - if self.halo_mass_function is None: - self.halo_mass_function = ccl.halos.MassFunc.from_name( - self.halo_mass_function_name - )(**self.halo_mass_function_args) - nm = self.halo_mass_function(ccl_cosmo, mass, a) - return nm - - def mf_d2N_dz_dlnM(self, ccl_cosmo: ccl.Cosmology, logM: float, z: float) -> float: - """ - Computes the mass function at z and logM. - - :param ccl_cosmo: pyccl Cosmology - :param logM: Cluster mass given by log10(M) where M is in units of - M_sun (comoving). - :param z: Cluster Redshift. - :return: Number Density pdf at z and logm in units of Mpc^-3 (comoving). - """ - d2N_dV_dlnM = self.mf_d2N_dV_dlnM(ccl_cosmo, logM, z) - dV_dz = self.dV_dz(ccl_cosmo, z) - - return d2N_dV_dlnM * dV_dz - - def mf_d2N_dz_dlnM_completeness( - self, ccl_cosmo: ccl.Cosmology, logM: float, z: float - ) -> float: - """ - Computes the mass function at z and logM. - - :param ccl_cosmo: pyccl Cosmology - :param logM: Cluster mass given by log10(M) where M is in units of - M_sun (comoving). - :param z: Cluster Redshift. - :return: Number Density pdf at z and logm in units of Mpc^-3 (comoving). - """ - d2N_dV_dlnM = self.mf_d2N_dV_dlnM(ccl_cosmo, logM, z) - dV_dz = self.dV_dz(ccl_cosmo, z) - completeness = self._cluster_abundance_compute_completeness(logM, z) - - return d2N_dV_dlnM * dV_dz * completeness - - def _cluster_abundance_compute_purity(self, logM_obs, z): - ln_r = np.log(10**logM_obs) - a_nc = np.log(10) * 0.8612 - b_nc = np.log(10) * 0.3527 - a_rc = 2.2183 - b_rc = -0.6592 - nc = a_nc + b_nc * (1.0 + z) - ln_rc = a_rc + b_rc * (1.0 + z) - purity = (ln_r / ln_rc) ** nc / ((ln_r / ln_rc) ** nc + 1.0) - return purity - - def _cluster_abundance_compute_completeness(self, logM, z): - a_nc = 1.1321 - b_nc = 0.7751 - a_mc = 13.31 - b_mc = 0.2025 - log_mc = a_mc + b_mc * (1.0 + z) - nc = a_nc + b_nc * (1.0 + z) - C = (logM / log_mc) ** nc / ((logM / log_mc) ** nc + 1.0) - return C - - def _process_args(self, args): - x = np.array(args[0:-5]) - index_map, arg, ccl_cosmo, mass_arg, redshift_arg = args[-5:] - arg[index_map] = x - redshift_start_index = 2 + redshift_arg.dim - - logM, z = arg[0:2] - proxy_z = arg[2:redshift_start_index] - proxy_m = arg[redshift_start_index:] - - return logM, z, proxy_z, proxy_m, ccl_cosmo, mass_arg, redshift_arg - - # Generic integrand for the cluster abundance - # The arg array always has the following structure: - # [logM, z, proxy_z, proxy_m] - # where proxy_z and proxy_m are the proxy parameters for the redshift and - # mass arguments respectively. - # The index_map array is used to map the proxy parameters to the correct - # position in the arg array. - def _compute_integrand(self, *args): - ( - logM, - z, - proxy_z, - proxy_m, - ccl_cosmo, - mass_arg, - redshift_arg, - ) = self._process_args(args) - - return ( - self.base_mf_d2N_dz_dlnM(ccl_cosmo, logM, z) - * mass_arg.p(logM, z, *proxy_m) - * redshift_arg.p(logM, z, *proxy_z) - ) - - # As above but for the mean mass - def _compute_integrand_mean_logM(self, *args) -> float: - ( - logM, - z, - proxy_z, - proxy_m, - ccl_cosmo, - mass_arg, - redshift_arg, - ) = self._process_args(args) - - return ( - self.base_mf_d2N_dz_dlnM(ccl_cosmo, logM, z) - * mass_arg.p(logM, z, *proxy_m) - * redshift_arg.p(logM, z, *proxy_z) - * logM - ) - - def _compute_any_from_args( - self, - integrand, - ccl_cosmo: ccl.Cosmology, - mass_arg: ClusterMassArgument, - redshift_arg: ClusterRedshiftArgument, - ) -> float: - last_index = 0 - - arg = np.zeros(2 + mass_arg.dim + redshift_arg.dim) - index_map: List[int] = [] - bounds_list: List[Tuple[float, float]] = [] - - if mass_arg.is_dirac_delta(): - arg[0] = mass_arg.get_logM() - else: - index_map.append(last_index) - bounds_list.append(mass_arg.get_logM_bounds()) - last_index += 1 - - if redshift_arg.is_dirac_delta(): - arg[1] = redshift_arg.get_z() - else: - index_map.append(last_index) - bounds_list.append(redshift_arg.get_z_bounds()) - last_index += 1 - - if mass_arg.dim > 0: - index_map += list(range(last_index, last_index + mass_arg.dim)) - bounds_list += mass_arg.get_proxy_bounds() - - last_index += mass_arg.dim - - if redshift_arg.dim > 0: - index_map += list(range(last_index, last_index + redshift_arg.dim)) - bounds_list += redshift_arg.get_proxy_bounds() - last_index += redshift_arg.dim - - if len(index_map) == 0: - # No proxy bins - return ( - self.mf_d2N_dz_dlnM(ccl_cosmo, arg[0], arg[1]) - * mass_arg.p(arg[0], arg[1]) - * redshift_arg.p(arg[0], arg[1]) - ) - - if self.prefer_scipy_integration: - return scipy.integrate.nquad( - integrand, - args=(index_map, arg, ccl_cosmo, mass_arg, redshift_arg), - ranges=bounds_list, - opts={"epsabs": self.abstol, "epsrel": self.reltol}, - )[0] - - Ncm.cfg_init() - int_nd = CountsIntegralND( - len(index_map), - integrand, - index_map, - arg, - ccl_cosmo, - mass_arg, - redshift_arg, - ) - int_nd.set_method(self.integ_method) - int_nd.set_reltol(self.reltol) - int_nd.set_abstol(self.abstol) - res = Ncm.Vector.new(1) - err = Ncm.Vector.new(1) - - bl, bu = zip(*bounds_list) - int_nd.eval(Ncm.Vector.new_array(bl), Ncm.Vector.new_array(bu), res, err) - return res.get(0) - - def compute( - self, - ccl_cosmo: ccl.Cosmology, - mass_arg: ClusterMassArgument, - redshift_arg: ClusterRedshiftArgument, - ) -> float: - """Compute the integrand for the given cosmology at the given mass - and redshift.""" - return self._compute_any_from_args( - self._compute_integrand, ccl_cosmo, mass_arg, redshift_arg - ) - - def compute_unormalized_mean_logM( - self, - ccl_cosmo: ccl.Cosmology, - mass_arg: ClusterMassArgument, - redshift_arg: ClusterRedshiftArgument, - ): - """Compute the mean log(M) * integrand for the given cosmology at the - given mass and redshift. - """ - return self._compute_any_from_args( - self._compute_integrand_mean_logM, ccl_cosmo, mass_arg, redshift_arg - ) diff --git a/firecrown/models/cluster_mass.py b/firecrown/models/cluster_mass.py deleted file mode 100644 index 567e05944..000000000 --- a/firecrown/models/cluster_mass.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Cluster Mass Module -Abstract class to compute cluster mass function. - - -The implemented functions use the :mod:`pyccl` library. -""" - -from __future__ import annotations -from typing import final, List, Tuple, Optional -from abc import abstractmethod - -import numpy as np -import sacc - -from ..updatable import Updatable -from ..parameters import ParamsMap - - -class ClusterMassArgument: - """Cluster Mass argument class.""" - - def __init__(self, logMl: float, logMu: float): - self.logMl: float = logMl - self.logMu: float = logMu - self.logM: Optional[float] = None - self.dirac_delta: bool = False - - if logMl > logMu: - raise ValueError("logMl must be smaller than logMu") - if logMl == logMu: - self.dirac_delta = True - self.logM = logMl - - def is_dirac_delta(self) -> bool: - """Check if the argument is a dirac delta.""" - - return self.dirac_delta - - def get_logM(self) -> float: - """Return the logM value if the argument is a dirac delta.""" - - if self.logM is not None: - return self.logM - raise ValueError("Argument is not a Dirac delta") - - @property - @abstractmethod - def dim(self) -> int: - """Return the dimension of the argument.""" - - @abstractmethod - def get_logM_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster mass argument.""" - - @abstractmethod - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster mass proxy argument.""" - - @abstractmethod - def p(self, logM: float, z: float, *proxy_args) -> float: - """Return the probability of the argument.""" - - -class ClusterMass(Updatable): - """Cluster Mass module.""" - - @abstractmethod - def read(self, sacc_data: sacc.Sacc): - """Abstract method to read the data for this source from the SACC file.""" - - def _update_cluster_mass(self, params: ParamsMap): - """Method to update the ClusterMass from the given ParamsMap. - Subclasses that need to do more than update their contained - :class:`Updatable` instance variables should implement this method.""" - - @final - def _update(self, params: ParamsMap): - """Implementation of Updatable interface method `_update`.""" - - self._update_cluster_mass(params) - - @abstractmethod - def gen_bins_by_array(self, logM_obs_bins: np.ndarray) -> List[ClusterMassArgument]: - """Generate bins by an array of bin edges.""" - - @abstractmethod - def gen_bin_from_tracer(self, tracer: sacc.BaseTracer) -> ClusterMassArgument: - """Return the bin for the given tracer.""" diff --git a/firecrown/models/cluster_mass_rich_proxy.py b/firecrown/models/cluster_mass_rich_proxy.py deleted file mode 100644 index 3dfb5d0cf..000000000 --- a/firecrown/models/cluster_mass_rich_proxy.py +++ /dev/null @@ -1,215 +0,0 @@ -"""Cluster Mass Richness proxy module - -Define the Cluster Mass Richness proxy module and its arguments. -""" -from typing import List, Tuple, final - -import numpy as np -from scipy import special -import sacc - -from ..parameters import ( - ParamsMap, -) -from .cluster_mass import ClusterMass, ClusterMassArgument -from .. import parameters - - -class ClusterMassRich(ClusterMass): - """Cluster Mass Richness proxy. - - The following parameters are special Updatable parameters, which means that - they can be updated by the sampler: - - :ivar mu_p0: mu parameter 0 - :ivar mu_p1: mu parameter 1 - :ivar mu_p2: mu parameter 2 - :ivar sigma_p0: sigma parameter 0 - :ivar sigma_p1: sigma parameter 1 - :ivar sigma_p2: sigma parameter 2 - """ - - def __init__( - self, pivot_mass, pivot_redshift, logMl: float = 13.0, logMu: float = 16.0 - ): - """Initialize the ClusterMassRich object.""" - - super().__init__() - self.pivot_mass = pivot_mass - self.pivot_redshift = pivot_redshift - self.log_pivot_mass = pivot_mass * np.log(10.0) - self.log1p_pivot_redshift = np.log1p(self.pivot_redshift) - self.logMl = logMl - self.logMu = logMu - - # Updatable parameters - self.mu_p0 = parameters.register_new_updatable_parameter() - self.mu_p1 = parameters.register_new_updatable_parameter() - self.mu_p2 = parameters.register_new_updatable_parameter() - self.sigma_p0 = parameters.register_new_updatable_parameter() - self.sigma_p1 = parameters.register_new_updatable_parameter() - self.sigma_p2 = parameters.register_new_updatable_parameter() - - self.logM_obs_min = 0.0 - self.logM_obs_max = np.inf - - @final - def _update_cluster_mass(self, params: ParamsMap): - """Perform any updates necessary after the parameters have being updated. - - This implementation has nothing to do.""" - - def read(self, _: sacc.Sacc): - """Method to read the data for this source from the SACC file.""" - - @staticmethod - def cluster_mass_parameters_function( - log_pivot_mass, log1p_pivot_redshift, p: Tuple[float, float, float], logM, z - ): - """Return observed quantity corrected by redshift and mass.""" - - lnM = logM * np.log(10) - Delta_lnM = lnM - log_pivot_mass - Delta_z = np.log1p(z) - log1p_pivot_redshift - - return p[0] + p[1] * Delta_lnM + p[2] * Delta_z - - def cluster_mass_lnM_obs_mu_sigma(self, logM, z): - """Return the mean and standard deviation of the observed mass.""" - - return [ - ClusterMassRich.cluster_mass_parameters_function( - self.log_pivot_mass, - self.log1p_pivot_redshift, - (self.mu_p0, self.mu_p1, self.mu_p2), - logM, - z, - ), - ClusterMassRich.cluster_mass_parameters_function( - self.log_pivot_mass, - self.log1p_pivot_redshift, - (self.sigma_p0, self.sigma_p1, self.sigma_p2), - logM, - z, - ), - ] - - def gen_bins_by_array(self, logM_obs_bins: np.ndarray) -> List[ClusterMassArgument]: - """Generate bins by an array of bin edges.""" - - if len(logM_obs_bins) < 2: - raise ValueError("logM_obs_bins must have at least two elements") - - # itertools.pairwise is only available in Python 3.10 - # using zip instead - return [ - ClusterMassRichBinArgument( - self, self.logMl, self.logMu, logM_obs_lower, logM_obs_upper - ) - for logM_obs_lower, logM_obs_upper in zip( - logM_obs_bins[:-1], logM_obs_bins[1:] - ) - ] - - def point_arg(self, logM_obs: float) -> ClusterMassArgument: - """Return the argument generator of the cluster mass function.""" - - return ClusterMassRichPointArgument(self, self.logMl, self.logMu, logM_obs) - - def gen_bin_from_tracer(self, tracer: sacc.BaseTracer) -> ClusterMassArgument: - """Return the argument for the given tracer.""" - - if not isinstance(tracer, sacc.tracers.BinRichnessTracer): - raise ValueError("Tracer must be a BinRichnessTracer") - - return ClusterMassRichBinArgument( - self, self.logMl, self.logMu, tracer.lower, tracer.upper - ) - - -class ClusterMassRichPointArgument(ClusterMassArgument): - """Argument for the Cluster Mass Richness proxy.""" - - def __init__( - self, - richness: ClusterMassRich, - logMl: float, - logMu: float, - logM_obs: float, - ): - super().__init__(logMl, logMu) - self.richness: ClusterMassRich = richness - self.logM_obs: float = logM_obs - - @property - def dim(self) -> int: - """Return the dimension of the argument.""" - return 0 - - def get_logM_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster mass argument.""" - return (self.logMl, self.logMu) - - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster mass proxy argument.""" - return [] - - def p(self, logM: float, z: float, *_) -> float: - """Return the probability of the point argument.""" - - lnM_obs = self.logM_obs * np.log(10.0) - - lnM_mu, sigma = self.richness.cluster_mass_lnM_obs_mu_sigma(logM, z) - x = lnM_obs - lnM_mu - chisq = np.dot(x, x) / (2.0 * sigma**2) - likelihood = np.exp(-chisq) / (np.sqrt(2.0 * np.pi * sigma**2)) - return likelihood * np.log(10.0) - - -class ClusterMassRichBinArgument(ClusterMassArgument): - """Argument for the Cluster Mass Richness proxy.""" - - def __init__( - self, - richness: ClusterMassRich, - logMl: float, - logMu: float, - logM_obs_lower: float, - logM_obs_upper: float, - ): - super().__init__(logMl, logMu) - self.richness: ClusterMassRich = richness - self.logM_obs_lower: float = logM_obs_lower - self.logM_obs_upper: float = logM_obs_upper - if logM_obs_lower >= logM_obs_upper: - raise ValueError("logM_obs_lower must be less than logM_obs_upper") - - @property - def dim(self) -> int: - """Return the dimension of the argument.""" - return 0 - - def get_logM_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster mass argument.""" - return (self.logMl, self.logMu) - - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster mass proxy argument.""" - return [] - - def p(self, logM: float, z: float, *_) -> float: - """Return the probability of the binned argument.""" - - lnM_obs_mu, sigma = self.richness.cluster_mass_lnM_obs_mu_sigma(logM, z) - x_min = (lnM_obs_mu - self.logM_obs_lower * np.log(10.0)) / ( - np.sqrt(2.0) * sigma - ) - x_max = (lnM_obs_mu - self.logM_obs_upper * np.log(10.0)) / ( - np.sqrt(2.0) * sigma - ) - - if x_max > 3.0 or x_min < -3.0: - # pylint: disable-next=no-member - return -(special.erfc(x_min) - special.erfc(x_max)) / 2.0 - # pylint: disable-next=no-member - return (special.erf(x_min) - special.erf(x_max)) / 2.0 diff --git a/firecrown/models/cluster_mass_true.py b/firecrown/models/cluster_mass_true.py deleted file mode 100644 index 33d3adc38..000000000 --- a/firecrown/models/cluster_mass_true.py +++ /dev/null @@ -1,75 +0,0 @@ -"""Cluster Mass True Module - -Class to compute cluster mass functions with no proxy, -i.e., assuming we have the true masses of the clusters. - -""" - -from typing import final, List, Tuple - -import numpy as np -import sacc - - -from ..parameters import ( - ParamsMap, -) -from .cluster_mass import ClusterMass, ClusterMassArgument - - -class ClusterMassTrue(ClusterMass): - """Cluster Mass class.""" - - @final - def _update_cluster_mass(self, params: ParamsMap): - """Method to update the ClusterMassTrue from the given ParamsMap.""" - - def read(self, sacc_data: sacc.Sacc): - """Method to read the data for this source from the SACC file.""" - - def gen_bins_by_array(self, logM_obs_bins: np.ndarray) -> List[ClusterMassArgument]: - """Generate the bins by an array of bin edges.""" - - if len(logM_obs_bins) < 2: - raise ValueError("logM_bins must have at least two elements") - - # itertools.pairwise is only available in Python 3.10 - # using zip instead - return [ - ClusterMassTrueArgument(lower, upper) - for lower, upper in zip(logM_obs_bins[:-1], logM_obs_bins[1:]) - ] - - def point_arg(self, logM: float) -> ClusterMassArgument: - """Return the argument for the given mass.""" - - return ClusterMassTrueArgument(logM, logM) - - def gen_bin_from_tracer(self, tracer: sacc.BaseTracer) -> ClusterMassArgument: - """Return the argument for the given tracer.""" - - if not isinstance(tracer, sacc.tracers.BinLogMTracer): - raise ValueError("Tracer must be a BinLogMTracer") - - return ClusterMassTrueArgument(tracer.lower, tracer.upper) - - -class ClusterMassTrueArgument(ClusterMassArgument): - """Cluster mass true argument class.""" - - @property - def dim(self) -> int: - """Return the dimension of the argument.""" - return 0 - - def get_logM_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster mass argument.""" - return (self.logMl, self.logMu) - - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster mass proxy argument.""" - return [] - - def p(self, logM: float, z: float, *_) -> float: - """Return the probability of the argument.""" - return 1.0 diff --git a/firecrown/models/cluster_redshift.py b/firecrown/models/cluster_redshift.py deleted file mode 100644 index 74444acd5..000000000 --- a/firecrown/models/cluster_redshift.py +++ /dev/null @@ -1,88 +0,0 @@ -"""Cluster Redshift Module -Abstract class to compute cluster redshift functions. - - -The implemented functions use the :mod:`pyccl` library. -""" - -from typing import final, List, Tuple, Optional -from abc import abstractmethod - -import numpy as np - -import sacc - -from ..updatable import Updatable -from ..parameters import ParamsMap - - -class ClusterRedshiftArgument: - """Cluster Redshift argument class.""" - - def __init__(self, zl: float, zu: float): - self.zl: float = zl - self.zu: float = zu - self.z: Optional[float] = None - self.dirac_delta: bool = False - - if zl > zu: - raise ValueError("zl must be smaller than zu") - if zl == zu: - self.dirac_delta = True - self.z = zl - - def is_dirac_delta(self) -> bool: - """Check if the argument is a dirac delta.""" - - return self.dirac_delta - - def get_z(self) -> float: - """Return the z value if the argument is a dirac delta.""" - - if self.z is not None: - return self.z - raise ValueError("Argument is not a Dirac delta") - - @property - @abstractmethod - def dim(self) -> int: - """Return the dimension of the argument.""" - - @abstractmethod - def get_z_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster redshift argument.""" - - @abstractmethod - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster redshift proxy argument.""" - - @abstractmethod - def p(self, logM: float, z: float, *args) -> float: - """Return the probability of the argument.""" - - -class ClusterRedshift(Updatable): - """Cluster Redshift class.""" - - @abstractmethod - def read(self, sacc_data: sacc.Sacc): - """Abstract method to read the data for this source from the SACC file.""" - - def _update_cluster_redshift(self, params: ParamsMap): - """Method to update the ClusterRedshift from the given ParamsMap. - Subclasses that need to do more than update their contained - :class:`Updatable` instance variables should implement this method.""" - - @final - def _update(self, params: ParamsMap): - """Implementation of Updatable interface method `_update`.""" - - self._update_cluster_redshift(params) - - @abstractmethod - def gen_bins_by_array(self, z_bins: np.ndarray) -> List[ClusterRedshiftArgument]: - """Generate the bins by an array of bin edges.""" - - @abstractmethod - def gen_bin_from_tracer(self, tracer: sacc.BaseTracer) -> ClusterRedshiftArgument: - """Return the bin for the given tracer.""" diff --git a/firecrown/models/cluster_redshift_spec.py b/firecrown/models/cluster_redshift_spec.py deleted file mode 100644 index 32de42946..000000000 --- a/firecrown/models/cluster_redshift_spec.py +++ /dev/null @@ -1,73 +0,0 @@ -"""Cluster Redshift Spectroscopy Module - -Class to compute cluster redshift spectroscopy functions. - -""" - -from typing import final, List, Tuple - -import numpy as np -import sacc - -from ..parameters import ( - ParamsMap, -) -from .cluster_redshift import ClusterRedshift, ClusterRedshiftArgument - - -class ClusterRedshiftSpec(ClusterRedshift): - """Cluster Redshift class.""" - - @final - def _update_cluster_redshift(self, params: ParamsMap): - """Method to update the ClusterRedshiftSpec from the given ParamsMap.""" - - def read(self, sacc_data: sacc.Sacc): - """Method to read the data for this source from the SACC file.""" - - def gen_bins_by_array(self, z_bins: np.ndarray) -> List[ClusterRedshiftArgument]: - """Generate the bins by an array of bin edges.""" - - if len(z_bins) < 2: - raise ValueError("z_bins must have at least two elements") - - # itertools.pairwise is only available in Python 3.10 - # using zip instead - return [ - ClusterRedshiftSpecArgument(lower, upper) - for lower, upper in zip(z_bins[:-1], z_bins[1:]) - ] - - def point_arg(self, z: float) -> ClusterRedshiftArgument: - """Return the argument for the given redshift.""" - - return ClusterRedshiftSpecArgument(z, z) - - def gen_bin_from_tracer(self, tracer: sacc.BaseTracer) -> ClusterRedshiftArgument: - """Return the argument for the given tracer.""" - - if not isinstance(tracer, sacc.tracers.BinZTracer): - raise ValueError("Tracer must be a BinZTracer") - - return ClusterRedshiftSpecArgument(tracer.lower, tracer.upper) - - -class ClusterRedshiftSpecArgument(ClusterRedshiftArgument): - """Cluster Redshift spectroscopy argument class.""" - - @property - def dim(self) -> int: - """Return the dimension of the argument.""" - return 0 - - def get_z_bounds(self) -> Tuple[float, float]: - """Return the bounds of the cluster redshift argument.""" - return (self.zl, self.zu) - - def get_proxy_bounds(self) -> List[Tuple[float, float]]: - """Return the bounds of the cluster redshift proxy argument.""" - return [] - - def p(self, logM: float, z: float, *args) -> float: - """Return the probability of the argument.""" - return 1.0 diff --git a/firecrown/parameters.py b/firecrown/parameters.py index ced6f8e0f..25da404f7 100644 --- a/firecrown/parameters.py +++ b/firecrown/parameters.py @@ -91,7 +91,7 @@ def __len__(self): """Return the number of parameters contained.""" return len(self.params_names) - def __add__(self, other: RequiredParameters): + def __add__(self, other: RequiredParameters) -> RequiredParameters: """Return a new RequiredParameters with the concatenated names. Note that this function returns a new object that does not share state diff --git a/firecrown/updatable.py b/firecrown/updatable.py index 769553cb5..ecb68c28c 100644 --- a/firecrown/updatable.py +++ b/firecrown/updatable.py @@ -14,7 +14,7 @@ """ from __future__ import annotations -from typing import final, Dict, Optional, Any, List, Union +from typing import final, Dict, Optional, Any, List, Union, Iterable from abc import ABC from collections import UserList from .parameters import ( @@ -299,7 +299,7 @@ def _get_derived_parameters(self) -> DerivedParameterCollection: return DerivedParameterCollection([]) -class UpdatableCollection(UserList): +class UpdatableCollection(UserList[Any]): """UpdatableCollection is a list of Updatable objects and is itself supports :meth:`update` and :meth:`reset` (although it does not inherit @@ -310,7 +310,7 @@ class UpdatableCollection(UserList): updated. """ - def __init__(self, iterable=None): + def __init__(self, iterable: Optional[Iterable[Any]] = None) -> None: """Initialize the UpdatableCollection from the supplied iterable. If the iterable contains any object that is not Updatable, a TypeError @@ -327,7 +327,7 @@ def __init__(self, iterable=None): ) @final - def update(self, params: ParamsMap): + def update(self, params: ParamsMap) -> None: """Update self by calling update() on each contained item. :param params: new parameter values @@ -348,7 +348,7 @@ def is_updated(self) -> bool: return self._updated @final - def reset(self): + def reset(self) -> None: """Resets self by calling reset() on each contained item.""" self._updated = False for updatable in self: @@ -392,7 +392,7 @@ def append(self, item: Updatable) -> None: ) super().append(item) - def __setitem__(self, key, value): + def __setitem__(self, key, value) -> None: """Set self[key] to value; raise TypeError if Value is not Updatable.""" if not isinstance(value, Updatable): raise TypeError( diff --git a/pytest.ini b/pytest.ini index c130f5b83..758c873a0 100644 --- a/pytest.ini +++ b/pytest.ini @@ -6,7 +6,8 @@ addopts = testpaths = tests markers = slow: Mark slow tests to ignore them unless they are requested - + integration: Tests to determine if independent units of code work together. Typically carried out after significant changes. + precision_sensitive: Tests that are sensitive to the numerical precision used in calculations. filterwarnings = ignore::DeprecationWarning:pkg_resources.*: ignore::DeprecationWarning:cobaya.*: diff --git a/setup.cfg b/setup.cfg index 09da03503..7115eb579 100644 --- a/setup.cfg +++ b/setup.cfg @@ -55,4 +55,3 @@ explicit_package_bases = True [mypy-firecrown.connector.cobaya.*] disallow_subclassing_any = False - diff --git a/tests/conftest.py b/tests/conftest.py index 07ff19d62..ea5b5c33c 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -15,6 +15,7 @@ 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): @@ -128,3 +129,43 @@ def fixture_tools_with_vanilla_cosmology(): result = ModelingTools() result.update(ParamsMap()) result.prepare(pyccl.CosmologyVanillaLCDM()) + + +@pytest.fixture(name="cluster_sacc_data") +def fixture_cluster_sacc_data() -> sacc.Sacc: + # pylint: disable=no-member + cc = sacc.standard_types.cluster_counts + # pylint: disable=no-member + mlm = sacc.standard_types.cluster_mean_log_mass + + s = sacc.Sacc() + s.add_tracer("survey", "my_survey", 4000) + s.add_tracer("survey", "not_my_survey", 4000) + s.add_tracer("bin_z", "z_bin_tracer_1", 0, 2) + s.add_tracer("bin_z", "z_bin_tracer_2", 2, 4) + s.add_tracer("bin_richness", "mass_bin_tracer_1", 0, 2) + s.add_tracer("bin_richness", "mass_bin_tracer_2", 2, 4) + + s.add_data_point(cc, ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_1"), 1) + s.add_data_point(cc, ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_2"), 1) + s.add_data_point(cc, ("not_my_survey", "z_bin_tracer_2", "mass_bin_tracer_1"), 1) + s.add_data_point(cc, ("not_my_survey", "z_bin_tracer_2", "mass_bin_tracer_2"), 1) + + s.add_data_point(mlm, ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_1"), 1) + s.add_data_point(mlm, ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_2"), 1) + s.add_data_point(mlm, ("not_my_survey", "z_bin_tracer_2", "mass_bin_tracer_1"), 1) + 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 4cbf49f42..ff34e4d88 100644 --- a/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py +++ b/tests/likelihood/gauss_family/statistic/test_cluster_number_counts.py @@ -1,61 +1,119 @@ -"""Tests for the ClusterNumberCounts statistic. -""" +"""Tests for the binned cluster number counts module.""" +from unittest.mock import Mock +import sacc import pytest +import pyccl +from firecrown.models.cluster.integrator.integrator import Integrator +from firecrown.likelihood.gauss_family.statistic.source.source import SourceSystematic +from firecrown.modeling_tools import ModelingTools +from firecrown.models.cluster.properties import ClusterProperty +from firecrown.parameters import ParamsMap +from firecrown.models.cluster.abundance import ClusterAbundance +from firecrown.likelihood.gauss_family.statistic.binned_cluster_number_counts import ( + BinnedClusterNumberCounts, +) -import pyccl.halos -import sacc -from firecrown.likelihood.gauss_family.statistic.cluster_number_counts import ( - ClusterNumberCounts, - ClusterAbundance, -) -from firecrown.models.cluster_mass_rich_proxy import ClusterMassRich -from firecrown.models.cluster_redshift_spec import ClusterRedshiftSpec - - -@pytest.fixture(name="minimal_stat") -def fixture_minimal_stat() -> ClusterNumberCounts: - """Return a correctly initialized :class:`ClusterNumberCounts` object.""" - stat = ClusterNumberCounts( - survey_tracer="SDSS", - cluster_abundance=ClusterAbundance( - halo_mass_definition=pyccl.halos.MassDef(0.5, "matter"), - halo_mass_function_name="200m", - halo_mass_function_args={}, - ), - cluster_mass=ClusterMassRich(pivot_mass=10.0, pivot_redshift=1.25), - cluster_redshift=ClusterRedshiftSpec(), +def test_create_binned_number_counts(): + integrator = Mock(spec=Integrator) + integrator.integrate.return_value = 1.0 + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", integrator) + assert bnc is not None + assert bnc.cluster_properties == ClusterProperty.NONE + assert bnc.survey_name == "Test" + assert bnc.systematics == [] + assert bnc.theory_vector is None + assert len(bnc.data_vector) == 0 + + bnc = BinnedClusterNumberCounts( + (ClusterProperty.COUNTS | ClusterProperty.MASS), "Test", integrator ) - return stat + 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 + ) + assert bnc.systematics == systematics -@pytest.fixture(name="missing_survey_tracer") -def fixture_missing_survey_tracer() -> sacc.Sacc: - """Return a sacc.Sacc object that lacks a survey_tracer.""" - return sacc.Sacc() +def test_get_data_vector(): + integrator = Mock(spec=Integrator) + integrator.integrate.return_value = 1.0 + bnc = BinnedClusterNumberCounts(ClusterProperty.NONE, "Test", integrator) + dv = bnc.get_data_vector() + assert dv is not None + assert len(dv) == 0 -@pytest.fixture(name="good_sacc_data") -def fixture_sacc_data(): - """Return a sacc.Sacc object sufficient to correctly set a - :class:`ClusterNumberCounts` object. - """ - data = sacc.Sacc() - return data +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_missing_survey_tracer( - minimal_stat: ClusterNumberCounts, missing_survey_tracer: sacc.Sacc -): with pytest.raises( - ValueError, match="The SACC file does not contain the SurveyTracer SDSS." + ValueError, + match="You must specify at least one cluster property", ): - minimal_stat.read(missing_survey_tracer) + bnc.read(cluster_sacc_data) + bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", integrator) + bnc.read(cluster_sacc_data) + assert bnc.sky_area == 4000 + assert len(bnc.bins) == 2 + assert len(bnc.data_vector) == 2 + assert bnc.sacc_indices is not None + assert len(bnc.sacc_indices) == 2 -def test_read_works(): - """After read() is called, we should be able to get the statistic's + bnc = BinnedClusterNumberCounts(ClusterProperty.MASS, "my_survey", integrator) + bnc.read(cluster_sacc_data) + assert bnc.sky_area == 4000 + assert len(bnc.bins) == 2 + assert len(bnc.data_vector) == 2 + assert bnc.sacc_indices is not None + assert len(bnc.sacc_indices) == 2 + + bnc = BinnedClusterNumberCounts( + (ClusterProperty.COUNTS | ClusterProperty.MASS), "my_survey", integrator + ) + bnc.read(cluster_sacc_data) + assert bnc.sky_area == 4000 + assert len(bnc.bins) == 2 + assert len(bnc.data_vector) == 4 + assert bnc.sacc_indices is not None + assert len(bnc.sacc_indices) == 4 - :class:`DataVector` and also should be able to call - :meth:`compute_theory_vector`. - """ + +def test_compute_theory_vector(cluster_sacc_data: sacc.Sacc): + integrator = Mock(spec=Integrator) + integrator.integrate.return_value = 1.0 + tools = ModelingTools() + + hmf = pyccl.halos.MassFuncBocquet16() + cosmo = pyccl.cosmology.CosmologyVanillaLCDM() + params = ParamsMap() + + tools.cluster_abundance = ClusterAbundance(13, 17, 0, 2, hmf) + tools.update(params) + tools.prepare(cosmo) + + bnc = BinnedClusterNumberCounts(ClusterProperty.COUNTS, "my_survey", integrator) + 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.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 + ) + bnc.read(cluster_sacc_data) + tv = bnc.compute_theory_vector(tools) + assert tv is not None + assert len(tv) == 4 diff --git a/tests/test_cluster.py b/tests/test_cluster.py deleted file mode 100644 index 18a10bd99..000000000 --- a/tests/test_cluster.py +++ /dev/null @@ -1,167 +0,0 @@ -"""Tests for the cluster module.""" - -from typing import Any, Dict, List -import itertools -import math - -import pytest -import pyccl as ccl -import numpy as np - -from firecrown.models.cluster_mass import ClusterMass, ClusterMassArgument -from firecrown.models.cluster_redshift import ClusterRedshift, ClusterRedshiftArgument -from firecrown.models.cluster_mass_true import ClusterMassTrue -from firecrown.models.cluster_abundance import ClusterAbundance -from firecrown.models.cluster_mass_rich_proxy import ClusterMassRich -from firecrown.models.cluster_redshift_spec import ClusterRedshiftSpec -from firecrown.parameters import ParamsMap - - -@pytest.fixture(name="ccl_cosmo") -def fixture_ccl_cosmo(): - """Fixture for a CCL cosmology object.""" - - return ccl.Cosmology( - Omega_c=0.22, Omega_b=0.0448, h=0.71, sigma8=0.8, n_s=0.963, Neff=3.44 - ) - - -@pytest.fixture(name="parameters") -def fixture_parameters(): - """Fixture for a parameter map.""" - - parameters = ParamsMap( - { - "mu_p0": 3.19, - "mu_p1": 0.8, - "mu_p2": 0.0, - "sigma_p0": 0.3, - "sigma_p1": 0.08, - "sigma_p2": 0.0, - } - ) - return parameters - - -@pytest.fixture(name="z_args") -def fixture_cluster_z_args(parameters): - """Fixture for cluster redshifts.""" - z_bins = np.array([0.2000146, 0.31251036, 0.42500611, 0.53750187, 0.64999763]) - cluster_z = ClusterRedshiftSpec() - assert isinstance(cluster_z, ClusterRedshift) - - z_args = cluster_z.gen_bins_by_array(z_bins) - - cluster_z.update(parameters) - - return z_args - - -@pytest.fixture(name="logM_args") -def fixture_cluster_mass_logM_args(parameters): - """Fixture for cluster masses.""" - logM_bins = np.array([13.0, 13.5, 14.0, 14.5, 15.0]) - cluster_mass_t = ClusterMassTrue() - assert isinstance(cluster_mass_t, ClusterMass) - - logM_args = cluster_mass_t.gen_bins_by_array(logM_bins) - cluster_mass_t.update(parameters) - - return logM_args - - -@pytest.fixture(name="rich_args") -def fixture_cluster_mass_rich_args(parameters): - """Fixture for cluster masses.""" - pivot_mass = 14.0 - pivot_redshift = 0.6 - proxy_bins = np.array([0.45805137, 0.81610273, 1.1741541, 1.53220547, 1.89025684]) - - cluster_mass_r = ClusterMassRich(pivot_mass, pivot_redshift) - assert isinstance(cluster_mass_r, ClusterMass) - - rich_args = cluster_mass_r.gen_bins_by_array(proxy_bins) - cluster_mass_r.update(parameters) - - return rich_args - - -@pytest.fixture(name="cluster_abundance") -def fixture_cluster_abundance(parameters): - """Fixture for cluster objects.""" - - # TODO: remove try/except when pyccl 3.0 is released - try: - hmd_200 = ccl.halos.MassDef200c() - except TypeError: - hmd_200 = ccl.halos.MassDef200c - - hmf_args: Dict[str, Any] = {} - hmf_name = "Bocquet16" - sky_area = 489 - - cluster_abundance = ClusterAbundance(hmd_200, hmf_name, hmf_args, sky_area) - assert isinstance(cluster_abundance, ClusterAbundance) - - cluster_abundance.update(parameters) - - return cluster_abundance - - -def test_initialize_objects( - ccl_cosmo: ccl.Cosmology, cluster_abundance, z_args, logM_args, rich_args -): - """Test initialization of cluster objects.""" - - for z_arg in z_args: - assert isinstance(z_arg, ClusterRedshiftArgument) - - for logM_arg in logM_args: - assert isinstance(logM_arg, ClusterMassArgument) - - for rich_arg in rich_args: - assert isinstance(rich_arg, ClusterMassArgument) - - assert isinstance(ccl_cosmo, ccl.Cosmology) - assert isinstance(cluster_abundance, ClusterAbundance) - - -def test_cluster_mass_function_compute( - ccl_cosmo: ccl.Cosmology, - cluster_abundance: ClusterAbundance, - z_args: List[ClusterRedshiftArgument], - logM_args: List[ClusterMassArgument], - rich_args: List[ClusterMassArgument], -): - """Test cluster mass function computations.""" - - for redshift_arg, logM_arg, rich_arg in itertools.product( - z_args, logM_args, rich_args - ): - count_rich = cluster_abundance.compute(ccl_cosmo, rich_arg, redshift_arg) - assert math.isfinite(count_rich) - assert count_rich > 0.0 - count_mass = cluster_abundance.compute(ccl_cosmo, logM_arg, redshift_arg) - assert math.isfinite(count_mass) - assert count_mass > 0.0 - - -def test_cluster_mass_function_compute_scipy( - ccl_cosmo: ccl.Cosmology, - cluster_abundance: ClusterAbundance, - z_args: List[ClusterRedshiftArgument], - logM_args: List[ClusterMassArgument], - rich_args: List[ClusterMassArgument], -): - """Test cluster mass function computations.""" - - for redshift_arg, logM_arg, rich_arg in itertools.product( - z_args, logM_args, rich_args - ): - cluster_abundance.prefer_scipy_integration = True - count_rich = cluster_abundance.compute(ccl_cosmo, rich_arg, redshift_arg) - assert math.isfinite(count_rich) - assert count_rich > 0.0 - count_mass = cluster_abundance.compute(ccl_cosmo, logM_arg, redshift_arg) - assert math.isfinite(count_mass) - assert count_mass > 0.0 diff --git a/tests/test_cluster_abundance.py b/tests/test_cluster_abundance.py new file mode 100644 index 000000000..f811d8738 --- /dev/null +++ b/tests/test_cluster_abundance.py @@ -0,0 +1,227 @@ +"""Tests for the cluster abundance module.""" +from unittest.mock import Mock +import pytest +import pyccl +import numpy as np +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") +def fixture_cluster_abundance(): + """Test fixture that represents an assembled cluster abundance class.""" + hmf = pyccl.halos.MassFuncBocquet16() + ca = ClusterAbundance(13, 17, 0, 2, hmf) + 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 + # pylint: disable=protected-access + assert cluster_abundance._hmf_cache == {} + assert isinstance( + cluster_abundance.halo_mass_function, pyccl.halos.MassFuncBocquet16 + ) + assert cluster_abundance.min_mass == 13.0 + assert cluster_abundance.max_mass == 17.0 + assert cluster_abundance.min_z == 0.0 + assert cluster_abundance.max_z == 2.0 + assert len(cluster_abundance.kernels) == 0 + + +def test_cluster_update_ingredients(cluster_abundance: ClusterAbundance): + cosmo = pyccl.CosmologyVanillaLCDM() + cluster_abundance.update_ingredients(cosmo) + assert cluster_abundance.cosmo is not None + assert cluster_abundance.cosmo == cosmo + # pylint: disable=protected-access + 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) + + result = cluster_abundance.comoving_volume(np.linspace(0.1, 1, 10), 360**2) + assert isinstance(result, np.ndarray) + assert np.issubdtype(result.dtype, np.float64) + assert len(result) == 10 + assert np.all(result > 0) + + +@pytest.mark.slow +def test_abundance_massfunc_returns_value(cluster_abundance: ClusterAbundance): + cosmo = pyccl.CosmologyVanillaLCDM() + cluster_abundance.update_ingredients(cosmo) + + result = cluster_abundance.mass_function( + np.linspace(13, 17, 5), np.linspace(0.1, 1, 5) + ) + assert isinstance(result, np.ndarray) + 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 new file mode 100644 index 000000000..2304e1a65 --- /dev/null +++ b/tests/test_cluster_binning.py @@ -0,0 +1,97 @@ +"""Tests for the cluster binning module""" +from unittest.mock import Mock +import sacc +from firecrown.models.cluster.binning import SaccBin, NDimensionalBin + + +def test_create_sacc_bin_with_correct_dimension(): + tracer = sacc.tracers.SurveyTracer("", 1) + for i in range(10): + tracers = [tracer for _ in range(i)] + sb = SaccBin(tracers) + assert sb is not None + assert sb.dimension == i + + +def test_sacc_bin_z_edges(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_z, tracer_lambda]) + assert sb.z_edges == (0, 1) + + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_lambda, tracer_z]) + assert sb.z_edges == (0, 1) + + +def test_sacc_bin_richness_edges(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_z, tracer_lambda]) + assert sb.mass_proxy_edges == (4, 5) + + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_lambda, tracer_z]) + assert sb.mass_proxy_edges == (4, 5) + + +def test_equal_sacc_bins_are_equal(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb1 = SaccBin([tracer_z, tracer_lambda]) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 == sb2 + assert sb1 is not sb2 + + sb1 = SaccBin([tracer_lambda, tracer_z]) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 == sb2 + assert sb1 is not sb2 + + +def test_sacc_bin_different_edges_not_equal(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb1 = SaccBin([tracer_z, tracer_lambda]) + + tracer_z = sacc.tracers.BinZTracer("", 0, 2) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 != sb2 + + tracer_z = sacc.tracers.BinZTracer("", -1, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 != sb2 + + tracer_z = sacc.tracers.BinZTracer("", -1, 2) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 != sb2 + + +def test_sacc_bin_different_dimensions_not_equal(): + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_z2 = sacc.tracers.BinZTracer("", 1, 2) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb1 = SaccBin([tracer_z, tracer_z2, tracer_lambda]) + sb2 = SaccBin([tracer_z, tracer_lambda]) + + assert sb1 != sb2 + + +def test_sacc_bin_must_be_equal_type(): + other_bin = Mock(spec=NDimensionalBin) + + tracer_z = sacc.tracers.BinZTracer("", 0, 1) + tracer_lambda = sacc.tracers.BinRichnessTracer("", 4, 5) + sb = SaccBin([tracer_z, tracer_lambda]) + + assert sb != other_bin diff --git a/tests/test_cluster_data.py b/tests/test_cluster_data.py new file mode 100644 index 000000000..4c6084902 --- /dev/null +++ b/tests/test_cluster_data.py @@ -0,0 +1,141 @@ +"""Tests for the cluster abundance data module.""" +import pytest +import sacc +from firecrown.models.cluster.abundance_data import AbundanceData +from firecrown.models.cluster.properties import ClusterProperty + + +def test_create_abundance_data(): + s = sacc.Sacc() + ad = AbundanceData(s) + + # pylint: disable=protected-access + assert ad._mass_index == 2 + # pylint: disable=protected-access + assert ad._redshift_index == 1 + # pylint: disable=protected-access + assert ad._survey_index == 0 + + +def test_get_survey_tracer_missing_survey_name(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + with pytest.raises( + ValueError, + match="The SACC file does not contain the SurveyTracer the_black_lodge.", + ): + ad.get_survey_tracer("the_black_lodge") + + +def test_get_survey_tracer_wrong_tracer_type(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + with pytest.raises( + ValueError, + match="The SACC tracer z_bin_tracer_1 is not a SurveyTracer.", + ): + ad.get_survey_tracer("z_bin_tracer_1") + + +def test_get_survey_tracer_returns_survey_tracer(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + survey = ad.get_survey_tracer("my_survey") + assert survey is not None + assert isinstance(survey, sacc.tracers.SurveyTracer) + + +def test_get_bin_edges_cluster_counts(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + bins = ad.get_bin_edges("my_survey", ClusterProperty.COUNTS) + assert len(bins) == 2 + + +def test_get_bin_edges_cluster_mass(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + bins = ad.get_bin_edges("my_survey", ClusterProperty.MASS) + assert len(bins) == 2 + + +def test_get_bin_edges_counts_and_mass(cluster_sacc_data: sacc.Sacc): + ad = AbundanceData(cluster_sacc_data) + bins = ad.get_bin_edges( + "my_survey", (ClusterProperty.MASS | ClusterProperty.COUNTS) + ) + assert len(bins) == 2 + + +def test_get_bin_edges_not_implemented_cluster_property_throws( + cluster_sacc_data: sacc.Sacc, +): + ad = AbundanceData(cluster_sacc_data) + with pytest.raises(NotImplementedError): + ad.get_bin_edges("my_survey", ClusterProperty.SHEAR) + + +def test_observed_data_and_indices_by_survey_cluster_counts( + cluster_sacc_data: sacc.Sacc, +): + ad = AbundanceData(cluster_sacc_data) + data, indices = ad.get_observed_data_and_indices_by_survey( + "my_survey", ClusterProperty.COUNTS + ) + assert len(data) == 2 + assert len(indices) == 2 + + +def test_observed_data_and_indices_by_survey_cluster_mass( + cluster_sacc_data: sacc.Sacc, +): + ad = AbundanceData(cluster_sacc_data) + data, indices = ad.get_observed_data_and_indices_by_survey( + "my_survey", ClusterProperty.MASS + ) + assert len(data) == 2 + assert len(indices) == 2 + + +def test_observed_data_and_indices_by_survey_cluster_counts_and_mass( + cluster_sacc_data: sacc.Sacc, +): + ad = AbundanceData(cluster_sacc_data) + data, indices = ad.get_observed_data_and_indices_by_survey( + "my_survey", ClusterProperty.MASS | ClusterProperty.COUNTS + ) + assert len(data) == 4 + assert len(indices) == 4 + + +def test_observed_data_and_indices_by_survey_not_implemented_throws( + cluster_sacc_data: sacc.Sacc, +): + ad = AbundanceData(cluster_sacc_data) + with pytest.raises(NotImplementedError): + ad.get_observed_data_and_indices_by_survey("my_survey", ClusterProperty.SHEAR) + + +def test_observed_data_and_indices_no_data_throws(): + # pylint: disable=no-member + cc = sacc.standard_types.cluster_counts + + s = sacc.Sacc() + s.add_tracer("survey", "my_survey", 4000) + s.add_tracer("bin_z", "z_bin_tracer_1", 0, 2) + s.add_tracer("bin_z", "z_bin_tracer_2", 2, 4) + s.add_tracer("bin_richness", "mass_bin_tracer_1", 0, 2) + s.add_tracer("bin_richness", "mass_bin_tracer_2", 2, 4) + + s.add_data_point( + cc, + ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_1"), + 1, + ) + s.add_data_point( + cc, + ("my_survey", "z_bin_tracer_1", "mass_bin_tracer_2"), + 1, + ) + + ad = AbundanceData(s) + + with pytest.raises( + ValueError, match="The SACC file does not contain any tracers for the" + ): + ad.get_observed_data_and_indices_by_survey("my_survey", ClusterProperty.MASS) diff --git a/tests/test_cluster_integrators.py b/tests/test_cluster_integrators.py new file mode 100644 index 000000000..acaf31a32 --- /dev/null +++ b/tests/test_cluster_integrators.py @@ -0,0 +1,291 @@ +"""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 + ) + + 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) + + 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 + ) + + 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_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 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], + ) -> npt.NDArray[np.float64]: + # xy + result = a * b + return result + + integrator.set_integration_bounds( + empty_cluster_abundance, + sky_area, + (0, 1), + (0, 1), + ) + 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 new file mode 100644 index 000000000..1719bdbf7 --- /dev/null +++ b/tests/test_cluster_kernels.py @@ -0,0 +1,185 @@ +"""Tests for the cluster kernel module.""" +import pytest +import numpy as np +from firecrown.models.cluster.kernel import ( + Completeness, + Purity, + KernelType, + Kernel, + SpectroscopicRedshift, + TrueMass, +) + + +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 + + +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 + + +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 + + +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 + + +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 + ) + + +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 + ) + + +@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( + [ + 0.77657274, + 0.96966127, + 0.99286409, + 0.99780586, + 0.999224, + 0.99970302, + 0.99988111, + 0.99995125, + 0.99997982, + 0.99999166, + ] + ) + + purity = pk.distribution( + mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_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) + + +@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( + [ + 0.9978693724040568, + 0.9984319673134954, + 0.9988620014089232, + 0.9991864843696077, + 0.9994279315032029, + 0.999604893383804, + 0.9997324678841709, + 0.9998227843987537, + 0.9998854531462606, + 0.9999279749997235, + ] + ) + + purity = pk.distribution( + mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_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) + + +@pytest.mark.precision_sensitive +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( + [ + 0.0056502277493542, + 0.01896566878380423, + 0.03805597500308377, + 0.06224888967250564, + 0.09124569979282898, + 0.12486247682690908, + 0.16290218589569144, + 0.20507815091349266, + 0.2509673905442634, + 0.2999886170051561, + ] + ) + + comp = ck.distribution( + mass, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits + ) + 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 new file mode 100644 index 000000000..7914cdcc9 --- /dev/null +++ b/tests/test_cluster_mass_richness.py @@ -0,0 +1,233 @@ +"""Tests for the cluster mass richness module.""" +import pytest +import numpy as np +from scipy.integrate import quad +from firecrown.models.cluster.mass_proxy import ( + MurataBinned, + MurataUnbinned, + MassRichnessGaussian, +) +from firecrown.models.cluster.kernel import ( + KernelType, + Kernel, +) + +PIVOT_Z = 0.6 +PIVOT_MASS = 14.625862906 + + +@pytest.fixture(name="murata_binned_relation") +def fixture_murata_binned() -> MurataBinned: + """Initialize cluster object.""" + + mr = MurataBinned(PIVOT_MASS, PIVOT_Z) + + # Set the parameters to the values used in the test + # they should be such that the variance is always positive. + mr.mu_p0 = 3.00 + mr.mu_p1 = 0.086 + mr.mu_p2 = 0.01 + mr.sigma_p0 = 3.0 + mr.sigma_p1 = 0.07 + mr.sigma_p2 = 0.01 + + return mr + + +@pytest.fixture(name="murata_unbinned_relation") +def fixture_murata_unbinned() -> MurataUnbinned: + """Initialize cluster object.""" + + mr = MurataUnbinned(PIVOT_MASS, PIVOT_Z) + + # Set the parameters to the values used in the test + # they should be such that the variance is always positive. + mr.mu_p0 = 3.00 + mr.mu_p1 = 0.086 + mr.mu_p2 = 0.01 + mr.sigma_p0 = 3.0 + mr.sigma_p1 = 0.07 + mr.sigma_p2 = 0.01 + + return mr + + +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) + + assert mb.mu_p0 is None + assert mb.mu_p1 is None + assert mb.mu_p2 is None + assert mb.sigma_p0 is None + assert mb.sigma_p1 is None + assert mb.sigma_p2 is None + + +def test_cluster_observed_z(): + for z in np.geomspace(1.0e-18, 2.0, 20): + z = np.atleast_1d(z) + mass = np.atleast_1d(0) + f_z = MassRichnessGaussian.observed_value((0.0, 0.0, 1.0), mass, z, 0, 0) + assert f_z == pytest.approx(np.log1p(z), 1.0e-7, 0.0) + + +def test_cluster_observed_mass(): + for mass in np.linspace(10.0, 16.0, 20): + z = np.atleast_1d(0) + mass = np.atleast_1d(mass) + f_logM = MassRichnessGaussian.observed_value((0.0, 1.0, 0.0), mass, z, 0, 0) + + assert f_logM == pytest.approx(mass * np.log(10.0), 1.0e-7, 0.0) + + +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 + for mass1, mass2 in zip(mass_array[:-1], mass_array[1:]): + 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 + ) + probability_1 = murata_binned_relation.distribution( + mass2, z, mass_proxy, z_proxy, mass_proxy_limits, z_proxy_limits + ) + + assert probability_0 >= 0 + assert probability_1 >= 0 + + # Probability should be initially monotonically increasing + # and then monotonically decreasing. It should flip only once. + + # Test for the flip + if (not flip) and (probability_1 < probability_0): + flip = True + + # Test for the second flip + if flip and (probability_1 > probability_0): + raise ValueError("Probability flipped twice") + + if flip: + assert probability_1 <= probability_0 + else: + assert probability_1 >= probability_0 + + +def test_cluster_murata_binned_mean(murata_binned_relation: MurataBinned): + for mass in np.linspace(7.0, 26.0, 20): + for z in np.geomspace(1.0e-18, 2.0, 20): + mass = np.atleast_1d(mass) + z = np.atleast_1d(z) + test = murata_binned_relation.get_proxy_mean(mass, z) + + true = MassRichnessGaussian.observed_value( + (3.00, 0.086, 0.01), + mass, + z, + PIVOT_MASS * np.log(10.0), + np.log1p(PIVOT_Z), + ) + + assert test == pytest.approx(true, rel=1e-7, abs=0.0) + + +def test_cluster_murata_binned_variance(murata_binned_relation: MurataBinned): + for mass in np.linspace(7.0, 26.0, 20): + for z in np.geomspace(1.0e-18, 2.0, 20): + mass = np.atleast_1d(mass) + z = np.atleast_1d(z) + test = murata_binned_relation.get_proxy_sigma(mass, z) + + true = MassRichnessGaussian.observed_value( + (3.00, 0.07, 0.01), + mass, + z, + PIVOT_MASS * np.log(10.0), + np.log1p(PIVOT_Z), + ) + + assert test == pytest.approx(true, rel=1e-7, abs=0.0) + + +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 + for mass1, mass2 in zip(mass_array[:-1], mass_array[1:]): + 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 density should be initially monotonically increasing + # and then monotonically decreasing. It should flip only once. + + # Test for the flip + if (not flip) and (probability_1 < probability_0): + flip = True + + # Test for the second flip + if flip and (probability_1 > probability_0): + raise ValueError("Probability flipped twice") + + if flip: + assert probability_1 <= probability_0 + else: + assert probability_1 >= probability_0 + + +@pytest.mark.precision_sensitive +def test_cluster_murata_unbinned_distribution_is_normalized( + murata_unbinned_relation: MurataUnbinned, +): + for mass_i, z_i in zip(np.linspace(7.0, 26.0, 20), np.geomspace(1.0e-18, 2.0, 20)): + mass = np.atleast_1d(mass_i) + z = np.atleast_1d(z_i) + + mean = murata_unbinned_relation.get_proxy_mean(mass, z)[0] + sigma = murata_unbinned_relation.get_proxy_sigma(mass, z)[0] + mass_proxy_limits = (mean - 5 * sigma, mean + 5 * sigma) + + def integrand(mass_proxy): + """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) + ) + + result, _ = quad( + integrand, + mass_proxy_limits[0], + mass_proxy_limits[1], + epsabs=1e-12, + epsrel=1e-12, + ) + + assert result == pytest.approx(1.0, rel=1.0e-6, abs=0.0) diff --git a/tests/test_clustermodel.py b/tests/test_clustermodel.py deleted file mode 100644 index 2da4cc0fe..000000000 --- a/tests/test_clustermodel.py +++ /dev/null @@ -1,168 +0,0 @@ -"""Test cluster mass function computations.""" - -import itertools -import numpy as np -import pytest - -from firecrown.models.cluster_mass_rich_proxy import ( - ClusterMassRich, - ClusterMassRichBinArgument, - ClusterMassRichPointArgument, -) - - -@pytest.fixture(name="cluster_mass_rich") -def fixture_cluster_mass_rich() -> ClusterMassRich: - """Initialize cluster object.""" - pivot_redshift = 0.6 - pivot_mass = 14.625862906 - - cluster_mass_rich = ClusterMassRich(pivot_mass, pivot_redshift) - - # Set the parameters to the values used in the test - # they should be such that the variance is always positive. - cluster_mass_rich.mu_p0 = 3.00 - cluster_mass_rich.mu_p1 = 0.086 - cluster_mass_rich.mu_p2 = 0.01 - cluster_mass_rich.sigma_p0 = 3.0 - cluster_mass_rich.sigma_p1 = 0.07 - cluster_mass_rich.sigma_p2 = 0.01 - - return cluster_mass_rich - - -def test_cluster_mass_parameters_function_z(): - for z in np.geomspace(1.0e-18, 2.0, 20): - f_z = ClusterMassRich.cluster_mass_parameters_function( - 0.0, 0.0, (0.0, 0.0, 1.0), 0.0, z - ) - assert f_z == pytest.approx(np.log1p(z), 1.0e-7, 0.0) - - -def test_cluster_mass_parameters_function_M(): - for logM in np.linspace(10.0, 16.0, 20): - f_logM = ClusterMassRich.cluster_mass_parameters_function( - 0.0, 0.0, (0.0, 1.0, 0.0), logM, 0.0 - ) - - assert f_logM == pytest.approx(logM * np.log(10.0), 1.0e-7, 0.0) - - -def test_cluster_mass_lnM_obs_mu_sigma(cluster_mass_rich: ClusterMassRich): - for logM, z in itertools.product( - np.linspace(10.0, 16.0, 20), np.geomspace(1.0e-18, 2.0, 20) - ): - lnM_obs_mu_direct = ClusterMassRich.cluster_mass_parameters_function( - cluster_mass_rich.log_pivot_mass, - cluster_mass_rich.log1p_pivot_redshift, - (cluster_mass_rich.mu_p0, cluster_mass_rich.mu_p1, cluster_mass_rich.mu_p2), - logM, - z, - ) - sigma_direct = ClusterMassRich.cluster_mass_parameters_function( - cluster_mass_rich.log_pivot_mass, - cluster_mass_rich.log1p_pivot_redshift, - ( - cluster_mass_rich.sigma_p0, - cluster_mass_rich.sigma_p1, - cluster_mass_rich.sigma_p2, - ), - logM, - z, - ) - lnM_obs_mu, sigma = cluster_mass_rich.cluster_mass_lnM_obs_mu_sigma(logM, z) - - assert lnM_obs_mu == pytest.approx(lnM_obs_mu_direct, 1.0e-7, 0.0) - assert sigma == pytest.approx(sigma_direct, 1.0e-7, 0.0) - - -def test_cluster_richness_bins_invalid_edges(cluster_mass_rich: ClusterMassRich): - rich_bin_edges = np.array([20]) - - with pytest.raises(ValueError): - _ = cluster_mass_rich.gen_bins_by_array(rich_bin_edges) - - rich_bin_edges = np.array([20, 30, 10]) - with pytest.raises(ValueError): - _ = cluster_mass_rich.gen_bins_by_array(rich_bin_edges) - - rich_bin_edges = np.array([20, 30, 30]) - with pytest.raises(ValueError): - _ = cluster_mass_rich.gen_bins_by_array(rich_bin_edges) - - -def test_cluster_richness_bins(cluster_mass_rich: ClusterMassRich): - rich_bin_edges = np.array([20, 40, 60, 80]) - - richness_bins = cluster_mass_rich.gen_bins_by_array(rich_bin_edges) - - for Rl, Ru, richness_bin in zip( - rich_bin_edges[:-1], rich_bin_edges[1:], richness_bins - ): - assert isinstance(richness_bin, ClusterMassRichBinArgument) - assert Rl == richness_bin.logM_obs_lower - assert Ru == richness_bin.logM_obs_upper - - -def test_cluster_bin_probability(cluster_mass_rich: ClusterMassRich): - cluser_mass_rich_bin_argument = ClusterMassRichBinArgument( - cluster_mass_rich, 13.0, 17.0, 1.0, 5.0 - ) - - logM_array = np.linspace(7.0, 26.0, 2000) - for z in np.geomspace(1.0e-18, 2.0, 20): - flip = False - for logM_0, logM_1 in zip(logM_array[:-1], logM_array[1:]): - probability_0 = cluser_mass_rich_bin_argument.p(logM_0, z) - probability_1 = cluser_mass_rich_bin_argument.p(logM_1, z) - - assert probability_0 >= 0 - assert probability_1 >= 0 - - # Probability should be initially monotonically increasing - # and then monotonically decreasing. It should flip only once. - - # Test for the flip - if (not flip) and (probability_1 < probability_0): - flip = True - - # Test for the second flip - if flip and (probability_1 > probability_0): - raise ValueError("Probability flipped twice") - - if flip: - assert probability_1 <= probability_0 - else: - assert probability_1 >= probability_0 - - -def test_cluster_point_probability(cluster_mass_rich: ClusterMassRich): - cluser_mass_rich_bin_argument = ClusterMassRichPointArgument( - cluster_mass_rich, 13.0, 17.0, 2.5 - ) - - logM_array = np.linspace(7.0, 26.0, 2000) - for z in np.geomspace(1.0e-18, 2.0, 20): - flip = False - for logM_0, logM_1 in zip(logM_array[:-1], logM_array[1:]): - probability_0 = cluser_mass_rich_bin_argument.p(logM_0, z) - probability_1 = cluser_mass_rich_bin_argument.p(logM_1, z) - - assert probability_0 >= 0 - assert probability_1 >= 0 - - # Probability density should be initially monotonically increasing - # and then monotonically decreasing. It should flip only once. - - # Test for the flip - if (not flip) and (probability_1 < probability_0): - flip = True - - # Test for the second flip - if flip and (probability_1 > probability_0): - raise ValueError("Probability flipped twice") - - if flip: - assert probability_1 <= probability_0 - else: - assert probability_1 >= probability_0 diff --git a/tests/test_modeling_tools.py b/tests/test_modeling_tools.py index 493408d28..081669135 100644 --- a/tests/test_modeling_tools.py +++ b/tests/test_modeling_tools.py @@ -19,6 +19,7 @@ def test_default_constructed_state(): # Default constructed state is pretty barren... assert tools.ccl_cosmo is None assert tools.pt_calculator is None + assert tools.cluster_abundance is None assert len(tools.powerspectra) == 0