diff --git a/.gitignore b/.gitignore index f8eec899..a568c1ca 100644 --- a/.gitignore +++ b/.gitignore @@ -197,4 +197,6 @@ papers/catalog/plots/*.pdf papers/cosmo_val/logs/ # Ignore scratch notebooks -scratch/*/*.ipynb \ No newline at end of file +scratch/*/*.ipynb +scratch/guerrini/work_notebooks +scratch/guerrini/launch_scripts \ No newline at end of file diff --git a/cosmo_val/cat_config.yaml b/cosmo_val/cat_config.yaml index cc1326df..ea36a81d 100644 --- a/cosmo_val/cat_config.yaml +++ b/cosmo_val/cat_config.yaml @@ -1259,3 +1259,49 @@ SP_v1.6.6: e1_col: e1 e2_col: e2 path: unions_shapepipe_star_2024_v1.6.a.fits + +GLASS_mock_validation: + subdir: /n09data/guerrini/glass_mock_test/results/ + pipeline: SP + colour: violet + getdist_colour: 0.0, 0.5, 1.0 + ls: dashed + marker: '*' + cov_th: + A: 2405.3892055695346 + n_e: 6.128201234871523 + n_psf: 0.752316232272063 + sigma_e: 0.379587601488189 + mask: /home/guerrini/sp_validation/cosmo_inference/data/mask/mask_map_footprint_nside_4096.fits + psf: + PSF_flag: FLAG_PSF_HSM + PSF_size: SIGMA_PSF_HSM + square_size: true + star_flag: FLAG_STAR_HSM + star_size: SIGMA_STAR_HSM + hdu: 1 + path: unions_shapepipe_psf_2024_v1.6.a.fits + ra_col: RA + dec_col: Dec + e1_PSF_col: E1_PSF_HSM + e1_star_col: E1_STAR_HSM + e2_PSF_col: E2_PSF_HSM + e2_star_col: E2_STAR_HSM + shear: + R: 1.0 + path: tomo_test_2_glass_sim_00001_1024.fits + redshift_path: /home/guerrini/sp_validation_cosmostat/config/glass_mock/test_data/redshift_distribution_tomo.txt + w_col: w + e1_col: e1 + e1_PSF_col: e1_PSF + e2_col: e2 + e2_PSF_col: e2_PSF + cols: RA,Dec + tomo_bin_col: tom_bin_id + star: + ra_col: RA + dec_col: Dec + e1_col: e1 + e2_col: e2 + path: unions_shapepipe_star_2024_v1.6.a.fits + diff --git a/src/sp_validation/cosmo_val/core.py b/src/sp_validation/cosmo_val/core.py index 48251bbb..c6204c65 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -1,5 +1,6 @@ # %% import copy +import itertools import os import re from pathlib import Path @@ -7,6 +8,7 @@ import colorama import numpy as np import yaml +from astropy.io import fits from cs_util.cosmo import get_cosmo from shear_psf_leakage import run_object, run_scale @@ -88,7 +90,7 @@ class CosmologyValidation( Number of ell bins for pseudo-C_ell analysis (used with binning='powspace'). ell_step : int, default 10 Bin width in ell for linear binning (used with binning='linear'). - pol_factor : bool, default True + pol_factor : int, default -1 Apply polarization correction factor in pseudo-C_ell calculations. nrandom_cell : int, default 10 Number of random realizations for C_ell error estimation. @@ -97,6 +99,10 @@ class CosmologyValidation( noise debiasing, making those realizations reproducible run-to-run. cosmo_params : dict, optional Cosmological parameters to pass to get_cosmo(). If None, uses Planck 2018. + compute_tomography : bool, default False + Whether to compute tomographic correlation functions and pseudo-C_ell. + force_run : bool, default False + If True, forces re-computation of results even if cached outputs exist. Attributes ---------- @@ -214,7 +220,7 @@ def __init__( power=1 / 2, n_ell_bins=32, ell_step=10, - pol_factor=True, + pol_factor=-1, cell_method="map", noise_bias_method="analytic", fiducial_input_inka="coupled", @@ -223,6 +229,8 @@ def __init__( path_onecovariance=None, cosmo_params=None, blind=None, + compute_tomography=False, + force_run=False, ): self.rho_tau_method = rho_tau_method self.cov_estimate_method = cov_estimate_method @@ -243,7 +251,10 @@ def __init__( self.power = power self.n_ell_bins = n_ell_bins self.ell_step = ell_step + + assert pol_factor in (-1, 1), "The polarisatio factor must be -1 or 1." self.pol_factor = pol_factor + self.nrandom_cell = nrandom_cell self.cell_seed = cell_seed self.cell_method = cell_method @@ -252,6 +263,8 @@ def __init__( self.nside_mask = nside_mask self.path_onecovariance = path_onecovariance self.blind = blind + self.compute_tomography = compute_tomography + self.force_run = force_run assert self.cell_method in ["map", "catalog"], ( "cell_method must be 'map' or 'catalog'" @@ -636,3 +649,36 @@ def summarize_bmodes(self, fiducial_scale_cut=(12, 83), versions=None): print() return summary + + def _get_tomo_bins(self, version): + """ + Return the tomo_bin_ids for a given version. If the version does not have tomography, return None. + + Returns + ------- + tomo_bin_ids : list or None + List of unique tomographic bin IDs for the version, or None if no tomography is available + tomo_bin_pairs : list of tuples or None + List of unique pairs of tomographic bin IDs (including self-pairs) for the version, or None if no tomography is available + """ + if "tomo_bin_col" in self.cc[version]["shear"]: + self.print_cyan( + f"Extracting tomography information from version {version}." + ) + cat_gal = fits.getdata(self.cc[version]["shear"]["path"]) + tomo_bin = cat_gal[self.cc[version]["shear"]["tomo_bin_col"]] + tomo_bin_ids = np.unique(tomo_bin) + tomo_bin_ids = tomo_bin_ids[ + tomo_bin_ids > 0 + ] # Exclude zero or negative bins + self.print_cyan( + f"Found {len(tomo_bin_ids)} tomographic bins for version {version}: {tomo_bin_ids}." + ) + + tomo_bin_pairs = list( + itertools.combinations_with_replacement(tomo_bin_ids, 2) + ) + return tomo_bin_ids, tomo_bin_pairs + else: + self.print_cyan(f"Version {version} does not have tomography information.") + return None, None diff --git a/src/sp_validation/cosmo_val/pseudo_cl.py b/src/sp_validation/cosmo_val/pseudo_cl.py index 514049a6..13f747ad 100644 --- a/src/sp_validation/cosmo_val/pseudo_cl.py +++ b/src/sp_validation/cosmo_val/pseudo_cl.py @@ -7,33 +7,37 @@ on pymaster (NaMaster), healpy, and OneCovariance. """ +import colorsys import configparser +import itertools import os import healpy as hp import matplotlib.pyplot as plt import numpy as np -import pymaster as nmt from astropy.io import fits -from cs_util.cosmo import get_theo_c_ell - -from ..pseudo_cl import ( - apply_random_rotation, - get_n_gal_map, - get_pseudo_cls_catalog, - get_pseudo_cls_map, - make_namaster_bin, -) +from matplotlib.colors import to_rgb + +import sp_validation.pseudo_cl as spv_pseudo_cl + from ..rho_tau import get_params_rho_tau -from ..statistics import chi2_and_pte, cov_from_one_covariance +from ..statistics import cov_from_one_covariance class PseudoClMixin: + # ---------------- Pseudo-Cl properties ---------------- # @property def pseudo_cls(self): if not hasattr(self, "_pseudo_cls"): - self.calculate_pseudo_cl() - self.calculate_pseudo_cl_eb_cov() + self.calculate_pseudo_cl(compute_tomography=False) + self.calculate_pseudo_cl_inka_cov( + compute_tomography=False, load_all_block=True + ) + if self.compute_tomography: + self.calculate_pseudo_cl(compute_tomography=True) + self.calculate_pseudo_cl_inka_cov( + compute_tomography=True, load_all_block=True + ) return self._pseudo_cls @property @@ -42,51 +46,11 @@ def pseudo_cls_onecov(self): self.calculate_pseudo_cl_onecovariance() return self._pseudo_cls_onecov - def get_namaster_bin(self, lmin, lmax, b_lmax): - """Build NaMaster binning object (thin wrapper, state -> primitive).""" - return make_namaster_bin( - lmin, - lmax, - b_lmax, - self.binning, - ell_step=self.ell_step, - n_ell_bins=self.n_ell_bins, - power=self.power, - ) - - def get_variance_map(self, nside, e1, e2, w, unique_pix, idx_rep): - """ - Create a variance map from the input catalog. - """ - - variance_map = np.zeros(hp.nside2npix(nside)) - - variance_map[unique_pix] = np.bincount( - idx_rep, weights=(e1**2 + e2**2) / 2 * w**2 - ) - - return variance_map - - def get_field_and_workspace_from_map(self, mask, lmax, b): - """ - Create a NaMaster field and workspace from the input map. - """ - - nside = hp.npix2nside(len(mask)) - - # Create NaMaster field - f = nmt.NmtField( - mask=mask, - maps=[np.zeros(hp.nside2npix(nside)), np.zeros(hp.nside2npix(nside))], - lmax=lmax, - ) - - # Create NaMaster workspace - wsp = nmt.NmtWorkspace.from_fields(f, f, b) - - return f, wsp - - def calculate_pseudo_cl_eb_cov(self): + # ---------------- Pseudo-Cl calculation methods ---------------- # + # TODO: some cleaning to clearly separate DV, covariance, and utility functions. + def calculate_pseudo_cl_inka_cov( + self, compute_tomography=True, load_all_block=False + ): """ Compute a theoretical Gaussian covariance of the Pseudo-Cl for EE, EB and BB. """ @@ -94,182 +58,305 @@ def calculate_pseudo_cl_eb_cov(self): nside = self.nside - try: - self._pseudo_cls - except AttributeError: - self._pseudo_cls = {} + self._pseudo_cls = getattr(self, "_pseudo_cls", {}) for ver in self.versions: self.print_magenta(ver) + out_dir_block = self._output_path(f"pseudo_cl/iNKA_block_{ver}/") + os.makedirs(out_dir_block, exist_ok=True) + if ver not in self._pseudo_cls.keys(): self._pseudo_cls[ver] = {} - out_path = self._output_path(f"pseudo_cl_cov_{ver}.fits") - if os.path.exists(out_path): - self.print_done( - f"Skipping Pseudo-Cl covariance calculation, {out_path} exists" - ) - self._pseudo_cls[ver]["cov"] = fits.open(out_path) - else: - params = get_params_rho_tau(self.cc[ver], survey=ver) + out_path_merged = self._output_path_pseudo_cl_cov( + ver, "iNKA", tomography=compute_tomography + ) - self.print_cyan(f"Extracting the fiducial power spectrum for {ver}") + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - lmax = 2 * self.nside - ell = np.arange(1, lmax + 1) - pw = hp.pixwin(nside, lmax=lmax) - if pw.shape[0] != len(ell) + 1: + if tomo_bin_ids is None or tomo_bin_pairs is None: raise ValueError( - "Unexpected pixwin length for lmax=" - f"{lmax}: got {pw.shape[0]}, expected {len(ell) + 1}" - ) - pw = pw[1 : len(ell) + 1] - - # Load redshift distribution and calculate theory C_ell - path_redshift_distr = self.cc[ver]["shear"]["redshift_path"] - z, dndz = np.loadtxt(path_redshift_distr, unpack=True) - fiducial_cl = ( - get_theo_c_ell( - ell=ell, - z=z, - nz=dndz, - backend="ccl", - cosmo=self.cosmo, + f"Version {ver} does not have tomography information." ) - * pw**2 - ) - self.print_cyan("Getting a binning, n_gal_map, field and workspace.") + else: + tomo_bin_pairs = [("all", "all")] - lmin = 8 - lmax = 2 * self.nside - b_lmax = lmax - 1 + if os.path.exists(out_path_merged) and not self.force_run: + self.print_done( + f"Skipping Pseudo-Cl iNKA covariance calculation, {out_path_merged} exists" + ) + self._pseudo_cls[ver]["cov_iNKA"] = fits.open(out_path_merged) - b = self.get_namaster_bin(lmin, lmax, b_lmax) + if load_all_block: + self.print_done("Loading all the iNKA covariance blocks") + n_spectra = len(tomo_bin_pairs) - # Load data and create shear and noise maps - cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) + # indices into tomo_bin_pairs for all unique covariance blocks + block_indices = list( + itertools.combinations_with_replacement(range(n_spectra), 2) + ) - n_gal, unique_pix, _idx, idx_rep = self.get_n_gal_map( - params, nside, cat_gal - ) + for index_a, index_b in block_indices: + bin_key_a1, bin_key_a2 = tomo_bin_pairs[index_a] + bin_key_b1, bin_key_b2 = tomo_bin_pairs[index_b] + + load_path = self._output_path_iNKA_block_cov( + ver, + tomo_bin_quad=( + bin_key_a1, + bin_key_a2, + bin_key_b1, + bin_key_b2, + ), + ) - f, wsp = self.get_field_and_workspace_from_map(n_gal, b_lmax, b) - - if self.noise_bias_method == "randoms": - self.print_cyan("Getting a sample of Cls with noise bias.") - - cl_noise, f, wsp = self.get_sample( - params, - self.nside, - b_lmax, - b, - cat_gal, - n_gal, - n_gal, - unique_pix, - idx_rep, - np.random.default_rng(self.cell_seed), - ) + if os.path.exists(load_path): + self._pseudo_cls[ver][ + f"spectra_{bin_key_a1}{bin_key_a2}_spectra_{bin_key_b1}{bin_key_b2}" + ] = fits.open(load_path) + + if bin_key_a1 == bin_key_b1 and bin_key_a2 == bin_key_b2: + self._pseudo_cls[ver][ + f"tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}" + ]["cov"] = fits.open(load_path) + else: + raise FileNotFoundError( + "The file does not exist, please run `cosmo_val` with `force_run` to True." + ) - noise_bias_cl = np.mean(cl_noise, axis=0) + continue - elif self.noise_bias_method == "analytic": - self.print_cyan("Getting analytic noise bias.") + # Initialise dictionnary to store field and workspace + n_gal_map_dict = {} + field_dict = {} + wsp_dict = {} - e1, e2, w = ( - cat_gal[self.cc[ver]["shear"]["e1_col"]], - cat_gal[self.cc[ver]["shear"]["e2_col"]], - cat_gal[self.cc[ver]["shear"]["w_col"]], - ) - variance_map = self.get_variance_map( - self.nside, e1, e2, w, unique_pix, idx_rep - ) + self.print_cyan(f"Extracting the fiducial power spectrum for {ver}") - noise_bias = hp.nside2pixarea(self.nside) * np.mean(variance_map) + fiducial_cl = self.get_fiducial_cl(ver, compute_tomography) - noise_bias_cl = np.zeros((4, lmax)) - noise_bias_cl[0, :] = noise_bias - noise_bias_cl[3, :] = noise_bias + self.print_cyan( + "Estimating and adding the noise bias to the fiducial power spectra" + ) + self.print_cyan(f"Method used: {self.noise_bias_method}") - noise_bias_cl = wsp.decouple_cell(noise_bias_cl) # Decouple + params = get_params_rho_tau(self.cc[ver]) + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) - else: - raise ValueError( - f"Noise bias method {self.noise_bias_method} not recognized. It should be 'randoms' or 'analytic'." - ) + for bin_key1, bin_key2 in tomo_bin_pairs: + if bin_key1 == bin_key2: + cat_gal_ = self._get_tomographic_bin(params, cat_gal, bin_key1) - # Unbin, then fill the data vector below lmin with the lowest-ell value - noise_bias_cl = b.unbin_cell(noise_bias_cl) - lowest_ell = b.get_ell_list(0)[0] - noise_bias_cl[:, :lowest_ell] = noise_bias_cl[:, [lowest_ell]] + noise_bias_cl = self.get_noise_bias(params, nside, cat_gal_) - self.print_cyan("Adding noise bias to the fiducial Cls.") + else: + noise_bias_cl = np.zeros((4, 2 * nside)) - fiducial_cl = ( + # Update the fiducial_cl dictionnary + fiducial_cl[f"W{bin_key1}xW{bin_key2}"] = ( np.array( [ - fiducial_cl, - 0.0 * fiducial_cl, - 0.0 * fiducial_cl, - 0.0 * fiducial_cl, + fiducial_cl[f"W{bin_key1}xW{bin_key2}"], + 0.0 * fiducial_cl[f"W{bin_key1}xW{bin_key2}"], + 0.0 * fiducial_cl[f"W{bin_key1}xW{bin_key2}"], + 0.0 * fiducial_cl[f"W{bin_key1}xW{bin_key2}"], ] ) + noise_bias_cl ) - if self.fiducial_input_inka == "coupled": - self.print_cyan("Coupling the fiducial Cls.") + # Compute the fields and workspaces + for bin_key1, bin_key2 in tomo_bin_pairs: + self.print_cyan( + f"Computing fields and workspaces for {bin_key1}, {bin_key2}" + ) + lmin, lmax, b_lmax = spv_pseudo_cl.pseudo_cl_geometry(self.nside) + b = self.get_namaster_bin(lmin, lmax, b_lmax) + + # Get the tomographic bins + cat_gal_a = self._get_tomographic_bin(params, cat_gal, bin_key1) + cat_gal_b = self._get_tomographic_bin(params, cat_gal, bin_key2) + + # Compute the n_gal_maps and the wsp object + unique_pix_a, idx_a, idx_rep_a = self.get_pixels( + params, nside, cat_gal_a + ) + unique_pix_b, idx_b, idx_rep_b = self.get_pixels( + params, nside, cat_gal_b + ) + + # Compute the number density maps + n_gal_map_a = self.get_n_gal_map(params, nside, cat_gal_a) + n_gal_map_b = self.get_n_gal_map(params, nside, cat_gal_b) + + # Get the shear maps + shear_map_a_e1, shear_map_a_e2 = self.get_shear_map( + params, + self.nside, + cat_gal_a, + unique_pix=unique_pix_a, + idx=idx_a, + idx_rep=idx_rep_a, + ) + shear_map_b_e1, shear_map_b_e2 = self.get_shear_map( + params, + self.nside, + cat_gal_b, + unique_pix=unique_pix_b, + idx=idx_b, + idx_rep=idx_rep_b, + ) + + # Get the fields and workspaces + field_a, field_b, wsp = spv_pseudo_cl.get_field_and_workspace_from_map( + b, + mask_a=n_gal_map_a, + e1_map_a=shear_map_a_e1, + e2_map_a=shear_map_a_e2, + mask_b=n_gal_map_b, + e1_map_b=shear_map_b_e1, + e2_map_b=shear_map_b_e2, + pol_factor=self.pol_factor, + return_wsp=True, + ) + + # Save in the dictionnaries + if f"W{bin_key1}" not in n_gal_map_dict: + n_gal_map_dict[f"W{bin_key1}"] = n_gal_map_a + if f"W{bin_key2}" not in n_gal_map_dict: + n_gal_map_dict[f"W{bin_key2}"] = n_gal_map_b + if f"W{bin_key1}" not in field_dict: + field_dict[f"W{bin_key1}"] = field_a + if f"W{bin_key2}" not in field_dict: + field_dict[f"W{bin_key2}"] = field_b + if bin_key1 <= bin_key2 and f"W{bin_key1}xW{bin_key2}" not in wsp_dict: + wsp_dict[f"W{bin_key1}xW{bin_key2}"] = wsp + + if self.fiducial_input_inka == "coupled": + # Couple the cell if required + self.print_cyan("Coupling the fiducial Cls.") + for bin_key1, bin_key2 in tomo_bin_pairs: + # Get the wsp object + n_gal_map_a = n_gal_map_dict[f"W{bin_key1}"] + n_gal_map_b = n_gal_map_dict[f"W{bin_key2}"] + wsp = wsp_dict[f"W{bin_key1}xW{bin_key2}"] coupling_mat = wsp.get_coupling_matrix() coupling_mat_re = np.reshape( coupling_mat, (4, lmax, 4, lmax), order="F" ) - fiducial_cl = np.tensordot(coupling_mat_re, fiducial_cl) / np.mean( - n_gal**2 - ) # couple and divide by the mean of the mask squared + fiducial_cl[f"W{bin_key1}xW{bin_key2}"] = np.tensordot( + coupling_mat_re, fiducial_cl[f"W{bin_key1}xW{bin_key2}"] + ) / np.mean( + n_gal_map_a * n_gal_map_b + ) # couple and divide by the product of the mask + + # To compute all the blocks in the covariance, + # we must compute all the (i, j, k, l) bin combinations + # Or equivalently, compute the covariance for all spectra pairs + n_spectra = len(tomo_bin_pairs) + + # indices into tomo_bin_pairs for all unique covariance blocks + block_indices = list( + itertools.combinations_with_replacement(range(n_spectra), 2) + ) + # Loop on the different tomographic bin pairs to compute the covariance + for index_a, index_b in block_indices: + bin_key_a1, bin_key_a2 = tomo_bin_pairs[index_a] + bin_key_b1, bin_key_b2 = tomo_bin_pairs[index_b] + self.print_cyan( + f"Tomo Bin Quad: ({bin_key_a1}, {bin_key_a2}, {bin_key_b1}, {bin_key_b2})" + ) + + if ( + bin_key_a1 == bin_key_b1 + and bin_key_a2 == bin_key_b2 + and ( + f"tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}" + not in self._pseudo_cls[ver].keys() + ) + ): + self._pseudo_cls[ver][ + f"tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}" + ] = {} + + out_path = self._output_path_iNKA_block_cov( + ver, tomo_bin_quad=(bin_key_a1, bin_key_a2, bin_key_b1, bin_key_b2) + ) + + if os.path.exists(out_path) and not self.force_run: + self.print_done( + f"Skipping Pseudo-Cl covariance block Cl ({bin_key_a1, bin_key_a2}), Cl ({bin_key_b1, bin_key_b2}) calculation, {out_path} exists" + ) + + self._pseudo_cls[ver][ + f"spectra_{bin_key_a1}{bin_key_a2}_spectra_{bin_key_b1}{bin_key_b2}" + ] = fits.open(out_path) + + if bin_key_a1 == bin_key_b1 and bin_key_a2 == bin_key_b2: + self._pseudo_cls[ver][ + f"tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}" + ]["cov"] = fits.open(out_path) + + continue self.print_cyan("Computing the Pseudo-Cl covariance") - cw = nmt.NmtCovarianceWorkspace.from_fields(f, f, f, f) - - # Get actual number of ell bins from binning scheme - n_ell_actual = b.get_n_bands() - - covar_22_22 = nmt.gaussian_covariance( - cw, - 2, - 2, - 2, - 2, - fiducial_cl, - fiducial_cl, - fiducial_cl, - fiducial_cl, - wsp, - wb=wsp, - ).reshape([n_ell_actual, 4, n_ell_actual, 4]) + input_cl_a1_b1 = ( + fiducial_cl[f"W{bin_key_a1}xW{bin_key_b1}"] + if bin_key_a1 <= bin_key_b1 + else fiducial_cl[f"W{bin_key_b1}xW{bin_key_a1}"] + ) + input_cl_a1_b2 = ( + fiducial_cl[f"W{bin_key_a1}xW{bin_key_b2}"] + if bin_key_a1 <= bin_key_b2 + else fiducial_cl[f"W{bin_key_b2}xW{bin_key_a1}"] + ) + input_cl_a2_b1 = ( + fiducial_cl[f"W{bin_key_a2}xW{bin_key_b1}"] + if bin_key_a2 <= bin_key_b1 + else fiducial_cl[f"W{bin_key_b1}xW{bin_key_a2}"] + ) + input_cl_a2_b2 = ( + fiducial_cl[f"W{bin_key_a2}xW{bin_key_b2}"] + if bin_key_a2 <= bin_key_b2 + else fiducial_cl[f"W{bin_key_b2}xW{bin_key_a2}"] + ) - self.print_cyan("Saving Pseudo-Cl covariance") + covar_22_22 = spv_pseudo_cl.get_pseudo_cl_iNKA_covariance( + input_cl_a1_b1, + input_cl_a1_b2, + input_cl_a2_b1, + input_cl_a2_b2, + field_dict[f"W{bin_key_a1}"], + field_dict[f"W{bin_key_a2}"], + field_dict[f"W{bin_key_b1}"], + field_dict[f"W{bin_key_b2}"], + wsp_a=wsp_dict[f"W{bin_key_a1}xW{bin_key_a2}"], + wsp_b=wsp_dict[f"W{bin_key_b1}xW{bin_key_b2}"], + b=b, + ) - # covar_22_22 is indexed [ell, pol_a, ell, pol_b]; store each of the - # 16 EE/EB/BE/BB cross-blocks as a named HDU (row-major pol order). - # Append rather than construct from a list so astropy promotes the - # first HDU to a PrimaryHDU on write. - pols = ["EE", "EB", "BE", "BB"] - hdu = fits.HDUList() - for i, pa in enumerate(pols): - for j, pb in enumerate(pols): - hdu.append( - fits.ImageHDU( - covar_22_22[:, i, :, j], name=f"COVAR_{pa}_{pb}" - ) - ) + self.print_cyan("Saving Pseudo-Cl covariance") - hdu.writeto(out_path, overwrite=True) + self._pseudo_cls[ver][ + f"spectra_{bin_key_a1}{bin_key_a2}_spectra_{bin_key_b1}{bin_key_b2}" + ] = self._save_iNKA_covariance(covar_22_22, out_path) - self._pseudo_cls[ver]["cov"] = hdu + if bin_key_a1 == bin_key_b1 and bin_key_a2 == bin_key_b2: + self._pseudo_cls[ver][ + f"tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}" + ]["cov"] = self._pseudo_cls[ver][ + f"spectra_{bin_key_a1}{bin_key_a2}_spectra_{bin_key_b1}{bin_key_b2}" + ] + # Merge the covariance blocks + self._pseudo_cls[ver]["cov_iNKA"] = self._merge_iNKA_covariance( + ver, tomography=compute_tomography + ) + self.print_done(f"Done Pseudo-Cl covariance calculation for {ver}") self.print_done("Done Pseudo-Cl covariance") def calculate_pseudo_cl_onecovariance(self): @@ -299,8 +386,11 @@ def calculate_pseudo_cl_onecovariance(self): out_dir = self._output_path(f"pseudo_cl_cov_onecov_{ver}/") os.makedirs(out_dir, exist_ok=True) - if os.path.exists( - os.path.join(out_dir, "covariance_list_3x2pt_pure_Cell.dat") + if ( + os.path.exists( + os.path.join(out_dir, "covariance_list_3x2pt_pure_Cell.dat") + ) + and not self.force_run ): self.print_done(f"Skipping OneCovariance calculation, {out_dir} exists") self._load_onecovariance_cov(out_dir, ver) @@ -417,7 +507,7 @@ def calculate_pseudo_cl_g_ng_cov(self, gaussian_part="iNKA"): out_file = self._output_path( f"pseudo_cl_cov_g_ng_{gaussian_part}_{ver}.fits" ) - if os.path.exists(out_file): + if os.path.exists(out_file) and not self.force_run: self.print_done( f"Skipping Gaussian and Non-Gaussian covariance calculation, {out_file} exists" ) @@ -451,182 +541,356 @@ def calculate_pseudo_cl_g_ng_cov(self, gaussian_part="iNKA"): f"Done Gaussian and Non-Gaussian covariance of the Pseudo-Cl's using {gaussian_part} for the Gaussian part" ) - def calculate_pseudo_cl(self): + def calculate_pseudo_cl(self, compute_tomography=True): """ - Compute the pseudo-Cl of given catalogs. + Compute the pseudo-Cl of a `CosmologyValidation` inputs with tomography. """ - self.print_start("Computing pseudo-Cl's") + out_dir = self._output_path("pseudo_cl") + os.makedirs(out_dir, exist_ok=True) - nside = self.nside + if compute_tomography: + self.print_start("Computing tomographic pseudo-Cl's") + else: + self.print_start("Computing non-tomographic pseudo-Cl's") + + self._pseudo_cls = getattr(self, "_pseudo_cls", {}) - try: - self._pseudo_cls - except AttributeError: - self._pseudo_cls = {} for ver in self.versions: self.print_magenta(ver) - self._pseudo_cls[ver] = {} - - out_path = self._output_path(f"pseudo_cl_{ver}.fits") - if os.path.exists(out_path): - self.print_done(f"Skipping Pseudo-Cl's calculation, {out_path} exists") - cl_shear = fits.getdata(out_path) - self._pseudo_cls[ver]["pseudo_cl"] = cl_shear - elif self.cell_method == "map": - self.calculate_pseudo_cl_map(ver, nside, out_path) - elif self.cell_method == "catalog": - self.calculate_pseudo_cl_catalog(ver, out_path) - else: - raise ValueError(f"Unknown cell method: {self.cell_method}") - - self.print_done("Done pseudo-Cl's") - - def calculate_pseudo_cl_map(self, ver, nside, out_path): - params = get_params_rho_tau(self.cc[ver], survey=ver) - - # Load data and create shear and noise maps - cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) - - w = cat_gal[params["w_col"]] - self.print_cyan("Creating maps and computing Cl's...") - n_gal_map, unique_pix, _idx, idx_rep = self.get_n_gal_map( - params, nside, cat_gal - ) - mask = n_gal_map != 0 + if ver not in self.pseudo_cls.keys(): + self._pseudo_cls[ver] = {} - shear_map_e1 = np.zeros(hp.nside2npix(nside)) - shear_map_e2 = np.zeros(hp.nside2npix(nside)) + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - e1 = cat_gal[params["e1_col"]] - e2 = cat_gal[params["e2_col"]] + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) - del cat_gal + else: + tomo_bin_pairs = [("all", "all")] + + # Loop on the different tomographic bin pairs + for bin_key1, bin_key2 in tomo_bin_pairs: + self.print_cyan(f"Tomo Bin Pair: ({bin_key1}, {bin_key2})") + + if ( + f"tomo_bin_{bin_key1}_tomo_bin_{bin_key2}" + not in self._pseudo_cls[ver].keys() + ): + self._pseudo_cls[ver][ + f"tomo_bin_{bin_key1}_tomo_bin_{bin_key2}" + ] = {} + + out_path = self._output_path_pseudo_cl( + ver, tomo_bin_pair=(bin_key1, bin_key2) + ) + if os.path.exists(out_path) and not self.force_run: + self.print_done( + f"Skipping Pseudo-Cl's calculation, {out_path} exists" + ) + cl_shear = fits.getdata(out_path) + self._pseudo_cls[ver][f"tomo_bin_{bin_key1}_tomo_bin_{bin_key2}"][ + "pseudo_cl" + ] = cl_shear + continue + + if self.cell_method == "map": + self.calculate_pseudo_cl_map( + ver, self.nside, out_path, bin_key1, bin_key2 + ) + elif self.cell_method == "catalog": + self.calculate_pseudo_cl_catalog(ver, out_path, bin_key1, bin_key2) + else: + raise ValueError(f"Unknown cell method: {self.cell_method}") - shear_map_e1[unique_pix] += np.bincount(idx_rep, weights=e1 * w) - shear_map_e2[unique_pix] += np.bincount(idx_rep, weights=e2 * w) - shear_map_e1[mask] /= n_gal_map[mask] - shear_map_e2[mask] /= n_gal_map[mask] + def calculate_pseudo_cl_map(self, ver, nside, out_path, tomo_bin_a, tomo_bin_b): + assert (tomo_bin_a == "all" and tomo_bin_b == "all") or ( + isinstance(tomo_bin_a, (int, np.integer)) + and isinstance(tomo_bin_b, (int, np.integer)) + ), "tomo_bin_a and tomo_bin_b must be either both 'all' or both integers." - shear_map = shear_map_e1 + 1j * shear_map_e2 + params = get_params_rho_tau(self.cc[ver]) - del shear_map_e1, shear_map_e2 + self.print_cyan( + f"Computing pseudo-Cl's for tomographic bins {tomo_bin_a} and {tomo_bin_b}..." + ) - ell_eff, cl_shear, wsp = self.get_pseudo_cls_map(shear_map, n_gal_map) + # Load data and create shear and noise maps + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) - cl_noise = np.zeros_like(cl_shear) - rng = np.random.default_rng(self.cell_seed) + # Get the tomographic bin + cat_gal_a = self._get_tomographic_bin(params, cat_gal, tomo_bin_a) + cat_gal_b = self._get_tomographic_bin(params, cat_gal, tomo_bin_b) - for i in range(self.nrandom_cell): - noise_map_e1 = np.zeros(hp.nside2npix(nside)) - noise_map_e2 = np.zeros(hp.nside2npix(nside)) + del cat_gal - e1_rot, e2_rot = self.apply_random_rotation(e1, e2, rng) + self.print_cyan("Creating maps and computing Cl's...") + # Get the pixels and indices for the catalogs + unique_pix_a, idx_a, idx_rep_a = self.get_pixels(params, nside, cat_gal_a) + unique_pix_b, idx_b, idx_rep_b = self.get_pixels(params, nside, cat_gal_b) - noise_map_e1[unique_pix] += np.bincount(idx_rep, weights=e1_rot * w) - noise_map_e2[unique_pix] += np.bincount(idx_rep, weights=e2_rot * w) + # Create number density maps for each tomographic bin + n_gal_map_a = self.get_n_gal_map( + params, + nside, + cat_gal_a, + unique_pix=unique_pix_a, + idx=idx_a, + idx_rep=idx_rep_a, + ) + n_gal_map_b = self.get_n_gal_map( + params, + nside, + cat_gal_b, + unique_pix=unique_pix_b, + idx=idx_b, + idx_rep=idx_rep_b, + ) - noise_map_e1[mask] /= n_gal_map[mask] - noise_map_e2[mask] /= n_gal_map[mask] + # Create shear maps for each tomographic bin + shear_map_a_e1, shear_map_a_e2 = self.get_shear_map( + params, + nside, + cat_gal_a, + unique_pix=unique_pix_a, + idx=idx_a, + idx_rep=idx_rep_a, + n_gal_map=n_gal_map_a, + ) + shear_map_a = shear_map_a_e1 + 1j * shear_map_a_e2 + del shear_map_a_e1, shear_map_a_e2 - noise_map = noise_map_e1 + 1j * noise_map_e2 - del noise_map_e1, noise_map_e2 + shear_map_b_e1, shear_map_b_e2 = self.get_shear_map( + params, + nside, + cat_gal_b, + unique_pix=unique_pix_b, + idx=idx_b, + idx_rep=idx_rep_b, + n_gal_map=n_gal_map_b, + ) + shear_map_b = shear_map_b_e1 + 1j * shear_map_b_e2 + del shear_map_b_e1, shear_map_b_e2 - _, cl_noise_, _ = self.get_pseudo_cls_map(noise_map, n_gal_map, wsp) - cl_noise += cl_noise_ + # Compute the pseudo-Cl's + ell_eff, cl_shear, wsp = self.get_pseudo_cls_map( + shear_map_a, n_gal_map_a, shear_map_b=shear_map_b, mask_b=n_gal_map_b + ) - cl_noise /= self.nrandom_cell - del e1, e2, w - try: - del e1_rot, e2_rot - except NameError: # Continue if the random generation has been skipped. - pass - del n_gal_map + # Remove the noise bias for auto-correlations. + if tomo_bin_a == tomo_bin_b: + # Compute the noise bias using noise_bias_method + cl_noise = self.get_noise_bias_from_gaussian_real( + params, + nside, + cat_gal_a, + unique_pix=unique_pix_a, + idx=idx_a, + idx_rep=idx_rep_a, + n_gal_map=n_gal_map_a, + wsp=wsp, + ) - # Noise realizations are now reproducible (seeded rng from self.cell_seed). - cl_shear = cl_shear - cl_noise + # Subtract the noise bias from the pseudo-Cl's + cl_shear = cl_shear - cl_noise self.print_cyan("Saving pseudo-Cl's...") self.save_pseudo_cl(ell_eff, cl_shear, out_path) cl_shear = fits.getdata(out_path) - self._pseudo_cls[ver]["pseudo_cl"] = cl_shear + self._pseudo_cls[ver][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ] = cl_shear + + def calculate_pseudo_cl_catalog(self, ver, out_path, tomo_bin_a, tomo_bin_b): + assert (tomo_bin_a == "all" and tomo_bin_b == "all") or ( + isinstance(tomo_bin_a, (int, np.integer)) + and isinstance(tomo_bin_b, (int, np.integer)) + ), "tomo_bin_a and tomo_bin_b must be either both 'all' or both integers." - def calculate_pseudo_cl_catalog(self, ver, out_path): - params = get_params_rho_tau(self.cc[ver], survey=ver) + params = get_params_rho_tau(self.cc[ver]) # Load data and create shear and noise maps cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) ell_eff, cl_shear, wsp = self.get_pseudo_cls_catalog( - catalog=cat_gal, params=params + catalog=cat_gal, params=params, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b ) self.print_cyan("Saving pseudo-Cl's...") self.save_pseudo_cl(ell_eff, cl_shear, out_path) cl_shear = fits.getdata(out_path) - self._pseudo_cls[ver]["pseudo_cl"] = cl_shear + self._pseudo_cls[ver][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ] = cl_shear - def get_n_gal_map(self, params, nside, cat_gal): + # ---------------- Utility functions for pseudo-Cl calculations ---------------- # + def get_namaster_bin(self, lmin, lmax, b_lmax): + """Build NaMaster binning object (thin wrapper, state -> primitive).""" + return spv_pseudo_cl.make_namaster_bin( + lmin, + lmax, + b_lmax, + self.binning, + ell_step=self.ell_step, + n_ell_bins=self.n_ell_bins, + power=self.power, + ) + + def get_pixels(self, params, nside, cat_gal): + """Get unique pixels and indices for a catalog (thin wrapper -> primitive).""" + return spv_pseudo_cl.get_pixels( + cat_gal[params["ra_col"]], cat_gal[params["dec_col"]], nside + ) + + def get_n_gal_map( + self, params, nside, cat_gal, unique_pix=None, idx=None, idx_rep=None + ): """Weighted galaxy number-density map (thin wrapper -> primitive).""" - return get_n_gal_map( + return spv_pseudo_cl.get_n_gal_map( nside, cat_gal[params["ra_col"]], cat_gal[params["dec_col"]], weights=cat_gal[params["w_col"]], + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, ) - def get_gaussian_real( - self, params, nside, lmax, cat_gal, n_gal, mask, unique_pix, idx_rep, rng=None + def get_shear_map( + self, + params, + nside, + cat_gal, + unique_pix=None, + idx=None, + idx_rep=None, + n_gal_map=None, ): - e1_rot, e2_rot = self.apply_random_rotation( - cat_gal[params["e1_col"]], cat_gal[params["e2_col"]], rng + """Weighted shear map (thin wrapper -> primitive).""" + return spv_pseudo_cl.get_shear_map( + cat_gal[params["ra_col"]], + cat_gal[params["dec_col"]], + cat_gal[params["e1_col"]], + cat_gal[params["e2_col"]], + cat_gal[params["w_col"]], + nside, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + n_gal_map=n_gal_map, ) - noise_map_e1 = np.zeros(hp.nside2npix(nside)) - noise_map_e2 = np.zeros(hp.nside2npix(nside)) - w = cat_gal[params["w_col"]] - noise_map_e1[unique_pix] += np.bincount(idx_rep, weights=e1_rot * w) - noise_map_e2[unique_pix] += np.bincount(idx_rep, weights=e2_rot * w) - noise_map_e1[mask] /= n_gal[mask] - noise_map_e2[mask] /= n_gal[mask] - - return noise_map_e1 + 1j * noise_map_e2 - - def get_sample( + def get_noise_realisation( self, params, nside, - lmax, - b, cat_gal, - n_gal, - mask, - unique_pix, - idx_rep, + n_gal=None, + unique_pix=None, + idx=None, + idx_rep=None, rng=None, ): - noise_map = self.get_gaussian_real( - params, nside, lmax, cat_gal, n_gal, mask, unique_pix, idx_rep, rng + """ + Get a single Gaussian noise realization (thin wrapper -> primitive). + """ + return spv_pseudo_cl.get_noise_realisation( + cat_gal[params["ra_col"]], + cat_gal[params["dec_col"]], + cat_gal[params["e1_col"]], + cat_gal[params["e2_col"]], + cat_gal[params["w_col"]], + nside, + n_gal_map=n_gal, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + rng=rng, ) - f = nmt.NmtField(mask=mask, maps=[noise_map.real, noise_map.imag], lmax=lmax) - - wsp = nmt.NmtWorkspace.from_fields(f, f, b) + def get_noise_bias_from_gaussian_real( + self, + params, + nside, + cat_gal, + unique_pix=None, + idx=None, + idx_rep=None, + n_gal_map=None, + wsp=None, + ): + """Noise-bias from Gaussian realisations (thin wrapper, state -> primitive)""" + return spv_pseudo_cl.get_noise_bias_from_gaussian_real( + cat_gal[params["ra_col"]], + cat_gal[params["dec_col"]], + cat_gal[params["e1_col"]], + cat_gal[params["e2_col"]], + cat_gal[params["w_col"]], + nside, + nrandom_cell=self.nrandom_cell, + binning=self.binning, + ell_step=self.ell_step, + n_ell_bins=self.n_ell_bins, + power=self.power, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + n_gal_map=n_gal_map, + wsp=wsp, + seed=self.cell_seed, + ) - cl_noise = nmt.compute_coupled_cell(f, f) - cl_noise = wsp.decouple_cell(cl_noise) + def get_noise_bias_analytical( + self, params, nside, cat_gal, unique_pix=None, idx=None, idx_rep=None + ): + """Noise-bias from analytical prescription (thin wrapper, state -> primitive)""" + return spv_pseudo_cl.get_noise_bias_analytical( + cat_gal[params["ra_col"]], + cat_gal[params["dec_col"]], + cat_gal[params["e1_col"]], + cat_gal[params["e2_col"]], + cat_gal[params["w_col"]], + lmax=2 * nside, + nside=nside, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + ) - return cl_noise, f, wsp + def get_noise_bias(self, params, nside, cat_gal): + """Noise-bias estimation (thin wrapper, state -> primitive)""" + return spv_pseudo_cl.get_noise_bias( + cat_gal[params["ra_col"]], + cat_gal[params["dec_col"]], + cat_gal[params["e1_col"]], + cat_gal[params["e2_col"]], + cat_gal[params["w_col"]], + nside, + noise_bias_method=self.noise_bias_method, + binning=self.binning, + ell_step=self.ell_step, + n_ell_bins=self.n_ell_bins, + power=self.power, + nrandom_cell=self.nrandom_cell, + seed=self.cell_seed, + ) - def get_pseudo_cls_map(self, map, mask, wsp=None): + def get_pseudo_cls_map( + self, map_a, mask_a, wsp=None, shear_map_b=None, mask_b=None + ): """Map-based pseudo-cl (thin wrapper, state -> primitive).""" - return get_pseudo_cls_map( - map, - mask, + return spv_pseudo_cl.get_pseudo_cls_map( + map_a, + mask_a, self.nside, self.binning, + shear_map_b=shear_map_b, + mask_b=mask_b, pol_factor=self.pol_factor, wsp=wsp, ell_step=self.ell_step, @@ -634,13 +898,17 @@ def get_pseudo_cls_map(self, map, mask, wsp=None): power=self.power, ) - def get_pseudo_cls_catalog(self, catalog, params, wsp=None): + def get_pseudo_cls_catalog( + self, catalog, params, wsp=None, tomo_bin_a=None, tomo_bin_b=None + ): """Catalog-based pseudo-cl (thin wrapper, state -> primitive).""" - return get_pseudo_cls_catalog( + return spv_pseudo_cl.get_pseudo_cls_catalog( catalog, params, self.nside, self.binning, + tomo_bin_a=tomo_bin_a, + tomo_bin_b=tomo_bin_b, pol_factor=self.pol_factor, wsp=wsp, ell_step=self.ell_step, @@ -648,12 +916,68 @@ def get_pseudo_cls_catalog(self, catalog, params, wsp=None): power=self.power, ) - def apply_random_rotation(self, e1, e2, rng=None): - """Random ellipticity rotation (thin wrapper -> primitive). + def read_redshift_distribution(self, ver, is_tomography): + path_redshift_distr = self.cc[ver]["shear"]["redshift_path"] + redshift_distribution = np.loadtxt(path_redshift_distr) + z = redshift_distribution[:, 0] + dndz = redshift_distribution[:, 1:] - Pass a seeded ``rng`` for reproducible noise realizations. - """ - return apply_random_rotation(e1, e2, rng) + # Here it is assumed that the tomographic redshift distribution sum to the non-tomographic one and that the latter is normalised + if not is_tomography: + dndz = np.sum(dndz, axis=1) + + return z, dndz + + def get_fiducial_cl(self, ver, is_tomography): + """Get a theory prediction for the angular power spectra (thin wrapper, state -> primitive).""" + lmax = 2 * self.nside + + z, dndz = self.read_redshift_distribution(ver, is_tomography) + + fiducial_cl = spv_pseudo_cl.get_fiducial_cl(z, dndz, lmax, self.cosmo) + + # If non-tomographic, change the key to 'WallxWall' + if not is_tomography: + fiducial_cl = {"WallxWall": fiducial_cl["W1xW1"]} + + return fiducial_cl + + def _get_tomographic_bin(self, params, cat_gal, tomo_bin): + """Extract tomographic bin from a given catalogue""" + if tomo_bin == "all": + return cat_gal + else: + tomo_bin_id = cat_gal[params["tomo_bin_col"]] + mask = tomo_bin_id == tomo_bin + return cat_gal[mask] + + def _output_path_pseudo_cl(self, ver, tomo_bin_pair=None): + if tomo_bin_pair is None: + return self._output_path( + "pseudo_cl", + f"pseudo_cl_from_{self.cell_method}_non_tomo_{ver}_binning_{self.binning}_nbins_{self.n_ell_bins}.fits", + ) + else: + bin_key1, bin_key2 = tomo_bin_pair + return self._output_path( + "pseudo_cl", + f"pseudo_cl_from_{self.cell_method}_tomo_bin_{bin_key1}_tomo_bin_{bin_key2}_{ver}_binning_{self.binning}_nbins_{self.n_ell_bins}.fits", + ) + + def _output_path_pseudo_cl_cov(self, ver, method, tomography): + is_tomo = "tomo" if tomography else "non_tomo" + return self._output_path( + "pseudo_cl", + f"pseudo_cl_cov_{is_tomo}_{ver}_from_{method}_binning_{self.binning}_nbins_{self.n_ell_bins}.fits", + ) + + def _output_path_iNKA_block_cov(self, ver, tomo_bin_quad): + bin_key_a1, bin_key_a2, bin_key_b1, bin_key_b2 = tomo_bin_quad + return self._output_path( + "pseudo_cl", + f"iNKA_block_{ver}", + f"pseudo_cl_cov_from_iNKA_tomo_bin_{bin_key_a1}_tomo_bin_{bin_key_a2}_tomo_bin_{bin_key_b1}_tomo_bin_{bin_key_b2}_{ver}_binning_{self.binning}_nbins_{self.n_ell_bins}.fits", + ) def save_pseudo_cl(self, ell_eff, pseudo_cl, out_path): """ @@ -670,208 +994,438 @@ def save_pseudo_cl(self, ell_eff, pseudo_cl, out_path): col1 = fits.Column(name="ELL", format="D", array=ell_eff) col2 = fits.Column(name="EE", format="D", array=pseudo_cl[0]) col3 = fits.Column(name="EB", format="D", array=pseudo_cl[1]) - col4 = fits.Column(name="BB", format="D", array=pseudo_cl[3]) - coldefs = fits.ColDefs([col1, col2, col3, col4]) + col4 = fits.Column(name="BE", format="D", array=pseudo_cl[2]) + col5 = fits.Column(name="BB", format="D", array=pseudo_cl[3]) + coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) cell_hdu = fits.BinTableHDU.from_columns(coldefs, name="PSEUDO_CELL") cell_hdu.writeto(out_path, overwrite=True) - def plot_pseudo_cl(self): + def _save_iNKA_covariance(self, covar, out_path): + # covar_22_22 is indexed [ell, pol_a, ell, pol_b]; store each of the + # 16 EE/EB/BE/BB cross-blocks as a named HDU (row-major pol order). + # Append rather than construct from a list so astropy promotes the + # first HDU to a PrimaryHDU on write. + pols = ["EE", "EB", "BE", "BB"] + hdu = fits.HDUList() + for i, pa in enumerate(pols): + for j, pb in enumerate(pols): + hdu.append(fits.ImageHDU(covar[:, i, :, j], name=f"COVAR_{pa}_{pb}")) + + hdu.writeto(out_path, overwrite=True) + + return hdu + + def _merge_iNKA_covariance(self, ver, tomography): """ - Plot pseudo-Cl's for given catalogs. + Merge the iNKA covariance matrices for a given version to get the data vector covariance. """ - self.print_cyan("Plotting pseudo-Cl's") + out_path = self._output_path_pseudo_cl_cov( + ver, method="iNKA", tomography=tomography + ) - # Plotting EE - out_path = self._output_path("cell_ee.png") - fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) + if tomography: + # Merge the tomographic covariance matrices + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_EE_EE"].data - ax[0].errorbar( - ell, - ell * self.pseudo_cls[ver]["pseudo_cl"]["EE"], - yerr=ell * np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " EE", - color=self.cc[ver]["colour"], - capsize=2, - ) + # Get the number of bins from the pseudo_cls attribute + # The non-tomographic pseudo-cl are computed from the call + # to this attribute if not already computed. + n_ell = self._pseudo_cls[ver]["tomo_bin_all_tomo_bin_all"]["pseudo_cl"][ + "ELL" + ].shape[0] - ax[0].set_ylabel(r"$\ell C_\ell$") - - ax[0].set_xlim(ell.min() - 10, ell.max() + 100) - ax[0].set_xscale("squareroot") - ax[0].set_xticks(np.array([100, 400, 900, 1600])) - ax[0].minorticks_on() - ax[0].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[0].xaxis.set_ticks(minor_ticks, minor=True) + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise AssertionError( + f"Tomographic bin IDs of version {ver} is not available." + ) - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_EE_EE"].data - ax[1].errorbar( - ell, - self.pseudo_cls[ver]["pseudo_cl"]["EE"], - yerr=np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " EE", - color=self.cc[ver]["colour"], + n_spectra = len(tomo_bin_pairs) + block_indices = list( + itertools.combinations_with_replacement(range(n_spectra), 2) ) - ax[1].set_xlabel(r"$\ell$") - ax[1].set_ylabel(r"$C_\ell$") + pols = ["EE", "EB", "BE", "BB"] + covar = fits.HDUList() + for pa in pols: + for pb in pols: + full_cov = np.zeros((n_spectra * n_ell, n_spectra * n_ell)) + for index_a, index_b in block_indices: + bin_key_a1, bin_key_a2 = tomo_bin_pairs[index_a] + bin_key_b1, bin_key_b2 = tomo_bin_pairs[index_b] + + block_path = self._output_path_iNKA_block_cov( + ver, + tomo_bin_quad=( + bin_key_a1, + bin_key_a2, + bin_key_b1, + bin_key_b2, + ), + ) + block = fits.open(block_path)[f"COVAR_{pa}_{pb}"].data - ax[1].set_xlim(ell.min() - 10, ell.max() + 100) - ax[1].set_xscale("squareroot") - ax[1].set_yscale("log") - ax[1].set_xticks(np.array([100, 400, 900, 1600])) - ax[1].minorticks_on() - ax[1].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[1].xaxis.set_ticks(minor_ticks, minor=True) + sl_a = slice(index_a * n_ell, (index_a + 1) * n_ell) + sl_b = slice(index_b * n_ell, (index_b + 1) * n_ell) - plt.suptitle("Pseudo-Cl EE (Gaussian covariance)") - plt.legend() - plt.savefig(out_path) + full_cov[sl_a, sl_b] = block - # Plotting EB - out_path = self._output_path("cell_eb.png") + if index_a != index_b: + full_cov[sl_b, sl_a] = block.T - fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) + covar.append(fits.ImageHDU(full_cov, name=f"COVAR_{pa}_{pb}")) - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_EB_EB"].data - ax[0].errorbar( - ell, - ell * self.pseudo_cls[ver]["pseudo_cl"]["EB"], - yerr=ell * np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " EB", - color=self.cc[ver]["colour"], - capsize=2, + else: + block_path = self._output_path_iNKA_block_cov( + ver, tomo_bin_quad=("all", "all", "all", "all") ) + covar = fits.open(block_path) - ax[0].axhline(0, color="black", linestyle="--") - ax[0].set_ylabel(r"$\ell C_\ell$") + covar.writeto(out_path, overwrite=True) - ax[0].set_xlim(ell.min() - 10, ell.max() + 100) - ax[0].set_xscale("squareroot") - ax[0].set_xticks(np.array([100, 400, 900, 1600])) - ax[0].minorticks_on() - ax[0].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[0].xaxis.set_ticks(minor_ticks, minor=True) + return covar - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_EB_EB"].data - ax[1].errorbar( - ell, - self.pseudo_cls[ver]["pseudo_cl"]["EB"], - yerr=np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " EB", - color=self.cc[ver]["colour"], + # ---------------- Plotting functions for pseudo-Cl's ---------------- # + def plot_pseudo_cl( + self, + pol_list, + versions=None, + ell_factor="ell", + cov_type="iNKA", + offset=0.15, + tomography=True, + savefig=None, + show=True, + ): + """ + Plot the pseudo-Cl for EE power spectrum. + + Parameters + ---------- + pol_list : list + List of polarization types to plot (e.g., ["EE", "BB"]). + ell_factor : {"None", "ell", "ell(ell+1)"} + Factor to multiply the ell values by. + cov_type : {"iNKA"} + Type of covariance to use. + tomography : bool, optional + Whether to plot tomographic power spectra. + savefig : str, optional + Path to save the figure. + show : bool, optional + Whether to show the figure. + + Returns + ------- + fig, ax : matplotlib.figure.Figure, matplotlib.axes.Axes + Figure and axes objects for the plot. + """ + if versions is None: + versions = self.versions + else: + for ver in versions: + if ver not in self.versions: + raise ValueError( + f"Version {ver} is not available. Available versions: {self.versions}" + ) + # Check that the method for the covariance is valid + if cov_type not in ["iNKA"]: + raise ValueError( + f"Invalid covariance type: {cov_type}. Valid options are: ['iNKA']" ) - ax[1].set_xlabel(r"$\ell$") - ax[1].set_ylabel(r"$C_\ell$") + # Check that the ell_factor is valid + if ell_factor not in ["None", "ell", "ell(ell+1)"]: + raise ValueError( + f"Invalid ell_factor: {ell_factor}. Valid options are: ['None', 'ell', 'ell(ell+1)']" + ) - ax[1].set_xlim(ell.min() - 10, ell.max() + 100) - ax[1].set_xscale("squareroot") - ax[1].set_yscale("log") - ax[1].set_xticks(np.array([100, 400, 900, 1600])) - ax[1].minorticks_on() - ax[1].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[1].xaxis.set_ticks(minor_ticks, minor=True) + def get_ell_factor(ell): + """Given some ell, return the ell_factor for the plot.""" + if ell_factor == "None": + return 1 + elif ell_factor == "ell": + return ell + elif ell_factor == "ell(ell+1)": + return ell * (ell + 1) + + # Check that all items in the list are valid polarisation + valid_pols = ["EE", "BB", "EB", "BE"] + for pol in pol_list: + if pol not in valid_pols: + raise ValueError( + f"Invalid polarization type: {pol}. Valid options are: {valid_pols}" + ) - plt.suptitle("Pseudo-Cl EB (Gaussian covariance)") - plt.legend() - plt.savefig(out_path) + fmt_dict = {"EE": "o", "BB": "s", "EB": "^", "BE": "v"} + # From all the versions, get the maximum number of tomo_bin_ids + tomo_bins = {} + for ver in versions: + if tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + else: + tomo_bin_ids, tomo_bin_pairs = ["all"], [("all", "all")] - # Plotting BB - out_path = self._output_path("cell_bb.png") + tomo_bins[ver] = {"ids": tomo_bin_ids, "pairs": tomo_bin_pairs} - fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) + n_tomo_bins_plot = max(len(bins["ids"]) for bins in tomo_bins.values()) - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_BB_BB"].data - ax[0].errorbar( - ell, - ell * self.pseudo_cls[ver]["pseudo_cl"]["BB"], - yerr=ell * np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " BB", - color=self.cc[ver]["colour"], - capsize=2, - ) + fig, axs = plt.subplots( + n_tomo_bins_plot, + n_tomo_bins_plot, + figsize=(12, 12), + sharex=True, + sharey=True, + ) + + for j, ver in enumerate(versions): + # Plot the pseudo-cl for each tomo bin of the considered version + for tomo_bin_a, tomo_bin_b in tomo_bins[ver]["pairs"]: + if tomography: + ax = axs[tomo_bin_b - 1, tomo_bin_a - 1] + else: + ax = axs + ver_tomo_info = self._pseudo_cls[ver][ + f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}" + ] + ver_label = self.cc[ver]["label"] if "label" in self.cc[ver] else ver + ver_color = ( + self.cc[ver]["colour"] if "colour" in self.cc[ver] else "black" + ) + + pseudo_cls = ver_tomo_info["pseudo_cl"] + cov = ver_tomo_info["cov"] + + ell = pseudo_cls["ell"] + + ell_widths = np.diff(ell) + ell_widths = np.append( + ell_widths, ell_widths[-1] + ) # Assume last bin width is same as second last + + # Better jittering: symmetric around original ell values + jitter_fraction = (j - (len(versions) - 1) / 2) * offset + jittered_ell = ell + jitter_fraction * ell_widths + ell_factor_ = get_ell_factor(jittered_ell) + + for pol in pol_list: + pol_color = self.get_pol_color(ver_color, pol, pol_list) + ax.errorbar( + jittered_ell, + ell_factor_ * pseudo_cls[pol], + yerr=np.sqrt(np.diag(cov[f"COVAR_{pol}_{pol}"].data)) + * ell_factor_, + fmt=fmt_dict[pol], + label=ver_label + f" {pol}", + color=pol_color, + capsize=2, + ) - ax[0].axhline(0, color="black", linestyle="--") - ax[0].set_ylabel(r"$\ell C_\ell$") + # Draw to extract the yaxis text offset + fig.canvas.draw() - ax[0].set_xlim(ell.min() - 10, ell.max() + 100) - ax[0].set_xscale("squareroot") - ax[0].set_xticks(np.array([100, 400, 900, 1600])) - ax[0].minorticks_on() - ax[0].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[0].xaxis.set_ticks(minor_ticks, minor=True) + if ell_factor == "None": + ell_label = r"$C_\ell$" + elif ell_factor == "ell": + ell_label = r"$\ell C_\ell$" + else: + ell_label = r"$\ell(\ell+1) C_\ell$" - for ver in self.versions: - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - cov = self.pseudo_cls[ver]["cov"]["COVAR_BB_BB"].data - ax[1].errorbar( - ell, - self.pseudo_cls[ver]["pseudo_cl"]["BB"], - yerr=np.sqrt(np.diag(cov)), - fmt=self.cc[ver]["marker"], - label=ver + " BB", - color=self.cc[ver]["colour"], + for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + if tomography: + ax = axs[tomo_bin_b - 1, tomo_bin_a - 1] + else: + ax = axs + ax.text( + 0.8, + 0.95, + f"{tomo_bin_a}-{tomo_bin_b}", + transform=ax.transAxes, + verticalalignment="top", ) + ax.axhline(0, color="black", ls="--") + ax.set_xlim(ell.min(), ell.max()) + ax.set_xscale("squareroot") + ax.set_xticks(np.array([100, 400, 900, 1600])) + ax.minorticks_on() + minor_ticks = [i * 10 for i in range(1, 10)] + [ + i * 100 for i in range(1, 21) + ] + ax.xaxis.set_ticks(minor_ticks, minor=True) + ax.tick_params(axis="both", which="both", direction="in") + if tomo_bin_b == 6 or tomo_bin_b == "all": + ax.set_xlabel(r"$\ell$") + if tomo_bin_a == 1 or tomo_bin_a == "all": + text_offset = ax.yaxis.get_offset_text().get_text() + ax.yaxis.get_offset_text().set_visible(False) + ax.set_ylabel(f"{ell_label}{text_offset}") + else: + ax.yaxis.get_offset_text().set_visible(False) + if tomo_bin_a != tomo_bin_b: + ax = axs[tomo_bin_a - 1, tomo_bin_b - 1] + ax.set_visible(False) - ax[1].set_xlabel(r"$\ell$") - ax[1].set_ylabel(r"$C_\ell$") + # Setup the legend + plt.subplots_adjust(hspace=0.0, wspace=0.0) # Remove space between subplots - ax[1].set_xlim(ell.min() - 10, ell.max() + 100) - ax[1].set_xscale("squareroot") - ax[1].set_yscale("log") - ax[1].set_xticks(np.array([100, 400, 900, 1600])) - ax[1].minorticks_on() - ax[1].tick_params(axis="x", which="minor", length=2, width=0.8) - minor_ticks = [i * 10 for i in range(1, 10)] + [i * 100 for i in range(1, 21)] - ax[1].xaxis.set_ticks(minor_ticks, minor=True) + legend_ax = self._add_grouped_legend(fig, versions, self.cc, pol_list, fmt_dict) - plt.suptitle("Pseudo-Cl BB (Gaussian covariance)") - plt.legend() - plt.savefig(out_path) + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") + if show: + plt.show() - # Print C_l^BB PTE for each version and save BB data - print("\nC_l^BB PTE summary:") - for ver in self.versions: - cl_bb = self.pseudo_cls[ver]["pseudo_cl"]["BB"] - cov_bb = self.pseudo_cls[ver]["cov"]["COVAR_BB_BB"].data - chi2_bb, _, pte_bb = chi2_and_pte(cl_bb, cov_bb) - chi2_bb = float(chi2_bb) - print( - f" {ver}: C_l^BB PTE = {pte_bb:.4f} " - f"(chi2/dof = {chi2_bb:.1f}/{len(cl_bb)})" + return fig, axs, legend_ax + + def _add_grouped_legend( + self, + fig, + versions, + cc, + pol_list, + fmt_dict, + row_height=0.35, + label_width=2, + col_width=0.9, + gap_below=0.10, + capsize=2, + fontsize=10, + ): + """ + Add a custom legend below the figure with one row per version: + ... + + The box size (in inches) scales with the number of versions (rows) + and polarizations (columns), then gets converted to figure-fraction + coordinates so it works for any figsize. + + Parameters + ---------- + row_height : float + Height per row (per version), in inches. + label_width : float + Width reserved for the version label column, in inches. + col_width : float + Width per polarization column (marker + pol label), in inches. + gap_below : float + Vertical gap between the bottom of the subplot grid and the + top of the legend box, in inches. + """ + n_rows = len(versions) + n_pol = len(pol_list) + + # Desired box size in inches + box_width_in = label_width + n_pol * col_width + box_height_in = n_rows * row_height + + fig_w_in, fig_h_in = fig.get_size_inches() + + # Convert to figure-fraction + width_frac = box_width_in / fig_w_in + height_frac = box_height_in / fig_h_in + gap_frac = gap_below / fig_h_in + + left = 0.5 - width_frac / 2 # centered horizontally + bottom = -gap_frac - height_frac # just below the subplot grid (y=0) + + legend_ax = fig.add_axes([left, bottom, width_frac, height_frac]) + legend_ax.axis("off") + legend_ax.set_xlim(0, 1) + legend_ax.set_ylim(0, 1) + + # Fractions *within* legend_ax's own 0-1 coordinate system, derived + # from the same inch-based proportions so columns stay consistent + label_frac = label_width / box_width_in + col_frac = col_width / box_width_in + + dummy_yerr = 0.15 / n_rows + + for i, ver in enumerate(versions): + y = 1.0 - (i + 0.5) / n_rows + + ver_label = cc[ver]["label"] if "label" in cc[ver] else ver + ver_color = cc[ver]["colour"] if "colour" in cc[ver] else "black" + + legend_ax.text( + 0.02 * label_frac, + y, + ver_label, + ha="left", + va="center", + fontweight="bold", + fontsize=fontsize, ) - # Save BB data + covariance to .npz - ell = self.pseudo_cls[ver]["pseudo_cl"]["ELL"] - bb_out = self._output_path(f"{ver}_cell_bb_data.npz") - np.savez( - bb_out, - ell=ell, - cl_bb=cl_bb, - cov_bb=cov_bb, - chi2_bb=np.array(chi2_bb), - pte_bb=np.array(pte_bb), + for k, pol in enumerate(pol_list): + pol_color = self.get_pol_color(ver_color, pol, pol_list) + x_marker = label_frac + k * col_frac + 0.15 * col_frac + x_text = x_marker + 0.15 * col_frac + + legend_ax.errorbar( + [x_marker], + [y], + yerr=dummy_yerr, + fmt=fmt_dict[pol], + color=pol_color, + markersize=6, + capsize=capsize, + clip_on=False, + ) + legend_ax.text( + x_text, + y, + pol, + ha="left", + va="center", + fontsize=fontsize - 1, + ) + + return legend_ax + + def get_pol_color(self, base_color, pol, pol_list, lightness_range=(-0.25, 0.25)): + """ + Given a base version color, return a shade variant for a given + polarization, by adjusting lightness in HLS space while keeping + hue and saturation fixed. + + Parameters + ---------- + base_color : str or tuple + Any matplotlib-recognized color (hex, name, RGB tuple, etc.) + pol : str + The polarization this color is for (e.g. "EE"). + pol_list : list + Full list of polarizations being plotted, used to compute + this pol's relative position in the lightness range. + lightness_range : tuple of float + (min_offset, max_offset) added to the base lightness, spread + evenly across pol_list. Negative = darker, positive = lighter. + Values are in HLS lightness units (0-1 scale), so keep these + modest (e.g. +/-0.25) to avoid washing out to white or black. + + Returns + ------- + tuple + RGB color tuple in [0, 1] range, usable directly in matplotlib. + """ + r, g, b = to_rgb(base_color) + h, lightness, s = colorsys.rgb_to_hls(r, g, b) + + n_pol = len(pol_list) + idx = pol_list.index(pol) + + if n_pol == 1: + l_offset = 0.0 + else: + # spread idx evenly across [lightness_range[0], lightness_range[1]] + frac = idx / (n_pol - 1) # 0 to 1 + l_offset = lightness_range[0] + frac * ( + lightness_range[1] - lightness_range[0] ) - print(f" Saved BB data to {bb_out}") + + lightness_new = min( + max(lightness + l_offset, 0.05), 0.95 + ) # clamp to avoid pure black/white + + r_new, g_new, b_new = colorsys.hls_to_rgb(h, lightness_new, s) + return (r_new, g_new, b_new) diff --git a/src/sp_validation/pseudo_cl.py b/src/sp_validation/pseudo_cl.py index c9355ec9..3508e804 100644 --- a/src/sp_validation/pseudo_cl.py +++ b/src/sp_validation/pseudo_cl.py @@ -16,11 +16,13 @@ import healpy as hp import numpy as np import pymaster as nmt +from cs_util.cosmo import get_theo_c_ell # Lowest multipole retained by the pseudo-Cl estimators. LMIN = 8 +# ---------------------- Binning utility functions ---------------------- def pseudo_cl_geometry(nside): """Return ``(lmin, lmax, b_lmax)`` for the pseudo-Cl estimator at ``nside``. @@ -87,7 +89,38 @@ def make_namaster_bin( return b -def get_n_gal_map(nside, ra, dec, weights=None): +# ---------------------- Map computation utility functions ---------------------- +def get_pixels(ra, dec, nside): + """ + Get the HEALPix pixel indices for given RA and Dec. + + Parameters + ---------- + ra : np.ndarray + Right ascension in degrees. + dec : np.ndarray + Declination in degrees. + nside : int + HEALPix nside parameter. + + Returns + ------- + unique_pix : np.ndarray + Sorted unique pixel indices. + idx : np.ndarray + First-occurrence indices into the input from ``np.unique``. + idx_rep : np.ndarray + Inverse map: pixel-group index for each input object. + """ + pixels = hp.ang2pix(nside, theta=np.radians(90 - dec), phi=np.radians(ra)) + + unique_pix, idx, idx_rep = np.unique(pixels, return_index=True, return_inverse=True) + return unique_pix, idx, idx_rep + + +def get_n_gal_map( + nside, ra, dec, weights=None, unique_pix=None, idx=None, idx_rep=None +): """Weighted galaxy number-density HEALPix map plus pixel bookkeeping. Bins ``(ra, dec)`` (degrees) onto an ``nside`` HEALPix grid. With @@ -98,21 +131,106 @@ def get_n_gal_map(nside, ra, dec, weights=None): ------- n_gal : np.ndarray Map of summed weights (or counts) per pixel, shape ``(npix,)``. - unique_pix : np.ndarray - Sorted unique occupied pixel indices. - idx : np.ndarray - First-occurrence indices into the input from ``np.unique``. - idx_rep : np.ndarray - Inverse map: pixel-group index for each input object. """ - theta = (90.0 - dec) * np.pi / 180.0 - phi = ra * np.pi / 180.0 - pix = hp.ang2pix(nside, theta, phi) + if unique_pix is None or idx is None or idx_rep is None: + unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) - unique_pix, idx, idx_rep = np.unique(pix, return_index=True, return_inverse=True) n_gal = np.zeros(hp.nside2npix(nside)) n_gal[unique_pix] = np.bincount(idx_rep, weights=weights) - return n_gal, unique_pix, idx, idx_rep + return n_gal + + +def get_shear_map( + ra, dec, e1, e2, w, nside, unique_pix=None, idx=None, idx_rep=None, n_gal_map=None +): + """Weighted shear HEALPix maps plus pixel bookkeeping. + + Bins ``(ra, dec)`` (degrees) onto an ``nside`` HEALPix grid. The shear + components ``(e1, e2)`` are weighted by ``w`` and summed per pixel. If + ``unique_pix``, ``idx``, and ``idx_rep`` are provided, they are used to + avoid recomputing the pixel indices. + If ``n_gal_map`` is provided, it is used to normalize the shear maps by the galaxy density. + + Returns + ------- + e1_map : np.ndarray + Weighted sum of E1 per pixel, shape ``(npix,)``. + e2_map : np.ndarray + Weighted sum of E2 per pixel, shape ``(npix,)``. + """ + if unique_pix is None or idx is None or idx_rep is None: + unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) + + if n_gal_map is None: + n_gal_map = get_n_gal_map( + nside, ra, dec, weights=w, unique_pix=unique_pix, idx=idx, idx_rep=idx_rep + ) + + npix = hp.nside2npix(nside) + e1_map = np.zeros(npix) + e2_map = np.zeros(npix) + + e1_map[unique_pix] = np.bincount(idx_rep, weights=e1 * w) + e2_map[unique_pix] = np.bincount(idx_rep, weights=e2 * w) + + non_zero = n_gal_map > 0 + e1_map[non_zero] /= n_gal_map[non_zero] + e2_map[non_zero] /= n_gal_map[non_zero] + + return e1_map, e2_map + + +def get_variance_map( + nside, ra, dec, e1, e2, w, unique_pix=None, idx=None, idx_rep=None +): + """Compute the variance map of the shear components. + + The variance is computed as the weighted variance of the shear components in each pixel. + + Returns + ------- + variance_map : np.ndarray + Variance map of the shear components, shape ``(npix,)``. + """ + if unique_pix is None or idx is None or idx_rep is None: + unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) + + npix = hp.nside2npix(nside) + variance_map = np.zeros(npix) + + variance_map[unique_pix] = np.bincount( + idx_rep, weights=0.5 * (e1**2 + e2**2) * w**2 + ) + + return variance_map + + +def get_noise_bias_analytical( + ra, dec, e1, e2, w, lmax, nside=1024, unique_pix=None, idx=None, idx_rep=None +): + """ + Compute the analytical noise bias for shear power spectrum. + """ + variance_map = get_variance_map( + nside=nside, + ra=ra, + dec=dec, + e1=e1, + e2=e2, + w=w, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + ) + + noise_bias = hp.nside2pixarea(nside) * np.mean(variance_map) + + noise_bias_cl = np.zeros((4, lmax)) + + noise_bias_cl[0, :] = noise_bias # EE + noise_bias_cl[3, :] = noise_bias # BB + + return noise_bias_cl def apply_random_rotation(e1, e2, rng=None): @@ -140,13 +258,493 @@ def apply_random_rotation(e1, e2, rng=None): return e1_out, e2_out +def get_noise_realisation( + ra, + dec, + e1, + e2, + w, + nside, + unique_pix=None, + idx=None, + idx_rep=None, + n_gal_map=None, + rng=None, +): + """ + Generate a random noise realisation of the shear maps by applying a random rotation to the ellipticity components. + + Parameters + ---------- + ra, dec : np.ndarray + Right ascension and declination of the sources. + e1, e2 : np.ndarray + Ellipticity components. + w : np.ndarray + Weights of the sources. + nside : int + HEALPix resolution. + unique_pix, idx, idx_rep : np.ndarray, optional + Pixel indices and bookkeeping arrays. If not provided, they will be computed. + n_gal_map : np.ndarray, optional + Galaxy number density map. If not provided, it will be computed. + + Returns + ------- + noise_map_e1, noise_map_e2 : np.ndarray + Noise map for ellipticity components. + """ + # Apply random rotation to the ellipticity components + e1_rot, e2_rot = apply_random_rotation(e1, e2, rng=rng) + + # Compute the noise maps using the rotated ellipticity components + noise_map_e1, noise_map_e2 = get_shear_map( + ra=ra, + dec=dec, + e1=e1_rot, + e2=e2_rot, + w=w, + nside=nside, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + n_gal_map=n_gal_map, + ) + + return noise_map_e1, noise_map_e2 + + +def get_noise_bias_from_gaussian_real( + ra, + dec, + e1, + e2, + w, + nside, + nrandom_cell, + binning, + ell_step=10, + n_ell_bins=32, + power=0.5, + unique_pix=None, + idx=None, + idx_rep=None, + n_gal_map=None, + wsp=None, + seed=42, +): + """ + Compute the power spectrum of the noise bias from random realisations. + + Parameters + ---------- + ra, dec : np.ndarray + Right ascension and declination of the sources. + e1, e2 : np.ndarray + Ellipticity components. + w : np.ndarray + Weights of the sources. + nside : int + HEALPix resolution. + nrandom_cell : int + Number of random cells to use for the noise estimation. + binning, ell_step, n_ell_bins, power : str, int, int, float + Binning scheme and parameters. + unique_pix, idx, idx_rep : np.ndarray, optional + Pixel indices and bookkeeping arrays. If not provided, they will be computed. + n_gal_map : np.ndarray, optional + Galaxy number density map. If not provided, it will be computed. + + Returns + ------- + noise_bias_cl : np.ndarray + Power spectrum of the noise bias. + """ + lmin, lmax, b_lmax = pseudo_cl_geometry(nside) + + b = make_namaster_bin( + lmin, + lmax, + b_lmax, + binning, + ell_step=ell_step, + n_ell_bins=n_ell_bins, + power=power, + ) + + ell_eff = b.get_effective_ells() + noise_bias_cl = np.zeros((4, ell_eff.size)) + + rng = np.random.default_rng(seed) + + if unique_pix is None or idx is None or idx_rep is None: + unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) + + if n_gal_map is None: + n_gal_map = get_n_gal_map( + nside, ra, dec, weights=w, unique_pix=unique_pix, idx=idx, idx_rep=idx_rep + ) + + if wsp is None: + _, _, wsp = get_field_and_workspace_from_map(b, mask_a=n_gal_map) + + for _ in range(nrandom_cell): + noise_map_e1, noise_map_e2 = get_noise_realisation( + ra, + dec, + e1, + e2, + w, + nside, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + n_gal_map=n_gal_map, + rng=rng, + ) + + noise_map = noise_map_e1 + 1j * noise_map_e2 + del noise_map_e1, noise_map_e2 + + _, cl_noise_, _ = get_pseudo_cls_map( + noise_map, + n_gal_map, + nside, + binning, + ell_step=ell_step, + n_ell_bins=n_ell_bins, + power=power, + wsp=wsp, + ) + + noise_bias_cl += cl_noise_ + + noise_bias_cl /= nrandom_cell + + return noise_bias_cl + + +def get_noise_bias( + ra, + dec, + e1, + e2, + w, + nside, + noise_bias_method, + binning, + *, + ell_step=10, + n_ell_bins=32, + power=0.5, + nrandom_cell=100, + seed=42, +): + """ + Compute the noise bias from object positions and ellipticities. + + Parameters + ---------- + ra, dec : np.ndarray + Right ascension and declination of the sources. + e1, e2 : np.ndarray + Ellipticity components. + w : np.ndarray + Weights of the sources. + nside : int + HEALPix resolution. + noise_bias_method : {'randoms', 'analytic'} + Method used to estimate the noise bias. + binning : {'linear', 'logspace', 'powspace'} + Binning scheme. + ell_step : int, optional + Bin width in ell for ``'linear'`` binning. + n_ell_bins : int, optional + Number of ell bins for ``'logspace'`` / ``'powspace'`` binning. + power : float, optional + Exponent for ``'powspace'`` binning. + nrandom_cell : int, optional + Number of random cells to use for the noise estimation. (only for the `randoms` method) + seed : int, optional + Random seed for reproducibility. (only for the `randoms` method) + + Returns + ------- + noise_bias_cl : np.ndarray + Power spectrum of the noise bias. + """ + if noise_bias_method not in ["randoms", "analytic"]: + raise ValueError("noise_bias_method must be 'randoms' or 'analytic'") + + lmin, lmax, b_lmax = pseudo_cl_geometry(nside) + + b = make_namaster_bin( + lmin, + lmax, + b_lmax, + binning, + ell_step=ell_step, + n_ell_bins=n_ell_bins, + power=power, + ) + + unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) + + if noise_bias_method == "analytic": + noise_bias_cl = get_noise_bias_analytical( + ra, + dec, + e1, + e2, + w, + lmax, + nside, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + ) + + elif noise_bias_method == "randoms": + noise_bias_cl = get_noise_bias_from_gaussian_real( + ra, + dec, + e1, + e2, + w, + nside, + nrandom_cell, + binning, + ell_step=ell_step, + n_ell_bins=n_ell_bins, + power=power, + unique_pix=unique_pix, + idx=idx, + idx_rep=idx_rep, + seed=seed, + ) + + noise_bias_cl = b.unbin_cell(noise_bias_cl) + else: + raise ValueError( + f"Invalid noise bias method `{noise_bias_method}`. Must be 'analytic' or 'randoms'." + ) + + return noise_bias_cl + + +# ---------------------- Cl computation functions ---------------------- +def get_field_and_workspace_from_map( + b, + mask_a, + e1_map_a=None, + e2_map_a=None, + mask_b=None, + e1_map_b=None, + e2_map_b=None, + pol_factor=-1, + return_wsp=True, +): + """Compute a NaMaster field and workspace object from the input maps. + + If the shear maps are None, returns field objects but only the workspace objects is relevant and contains the mixing matrix. + If the second mask and shear maps (indexed b) are provided, the mixing matrix is computed between the two fields. + + Parameters + ---------- + b : nmt.NmtBin + NaMaster binning object. + mask_a : np.ndarray + Field mask for the first map. + e1_map_a : np.ndarray, optional + E1 map for the first field. + e2_map_a : np.ndarray, optional + E2 map for the first field. + mask_b : np.ndarray, optional + Field mask for the second map. + e1_map_b : np.ndarray, optional + E1 map for the second field. + e2_map_b : np.ndarray, optional + E2 map for the second field. + pol_factor : float, optional + Polarization factor to apply to the E2 map. + return_wsp : bool, optional + If True, return the NaMaster workspace object containing the mixing matrix. + + Returns + ------- + field_a : nmt.NmtField + NaMaster field object for the first map. + field_b : nmt.NmtField + NaMaster field object for the second map (if provided, same than the first map otherwise). + wsp : nmt.NmtWorkspace + NaMaster workspace object containing the mixing matrix. + + """ + nside = hp.npix2nside(len(mask_a)) + lmax = b.lmax + if e1_map_a is None or e2_map_a is None: + e1_map_a = np.zeros(hp.nside2npix(nside)) + e2_map_a = np.zeros(hp.nside2npix(nside)) + + # Create NaMaster field + field_a = nmt.NmtField( + mask=mask_a, maps=[e1_map_a, pol_factor * e2_map_a], lmax=lmax + ) + + if mask_b is not None: + if e1_map_b is None or e2_map_b is None: + e1_map_b = np.zeros(hp.nside2npix(nside)) + e2_map_b = np.zeros(hp.nside2npix(nside)) + + field_b = nmt.NmtField( + mask=mask_b, maps=[e1_map_b, pol_factor * e2_map_b], lmax=lmax + ) + else: + field_b = field_a + + if return_wsp: + # Create NaMaster workspace + wsp = nmt.NmtWorkspace.from_fields(field_a, field_b, b) + + return field_a, field_b, wsp + else: + return field_a, field_b, None + + +def get_field_and_workspace_from_catalog( + b, + ra_a, + dec_a, + e1_a, + e2_a, + w_a, + ra_b=None, + dec_b=None, + e1_b=None, + e2_b=None, + w_b=None, + pol_factor=-1, + return_wsp=True, + same_bin=False, +): + """Create a NaMaster field and workspace from the input catalog. + + If the second catalog is provided, the mixing matrix is computed between the two fields. + + Parameters + ---------- + b : nmt.NmtBin + NaMaster binning object. + ra_a : np.ndarray + Right ascension of sources in the first catalog. + dec_a : np.ndarray + Declination of sources in the first catalog. + e1_a : np.ndarray + E1 shear component of sources in the first catalog. + e2_a : np.ndarray + E2 shear component of sources in the first catalog. + w_a : np.ndarray + Weights of sources in the first catalog. + ra_b : np.ndarray, optional + Right ascension of sources in the second catalog. + dec_b : np.ndarray, optional + Declination of sources in the second catalog. + e1_b : np.ndarray, optional + E1 shear component of sources in the second catalog. + e2_b : np.ndarray, optional + E2 shear component of sources in the second catalog. + w_b : np.ndarray, optional + Weights of sources in the second catalog. + pol_factor : float, optional + Polarization factor to apply to the E2 component. + return_wsp : bool, optional + If True, return the NaMaster workspace object containing the mixing matrix. + + Returns + ------- + field_a : nmt.NmtFieldCatalog + NaMaster field object for the first catalog. + field_b : nmt.NmtFieldCatalog + NaMaster field object for the second catalog (if provided, same as the first catalog otherwise). + wsp : nmt.NmtWorkspace + NaMaster workspace object containing the mixing matrix. + + """ + lmax = b.lmax + # Get field for input catalog a + field_a = nmt.NmtFieldCatalog( + positions=[ra_a, dec_a], + weights=w_a, + field=[e1_a, pol_factor * e2_a], + lmax=lmax, + lmax_mask=lmax, + spin=2, + lonlat=True, + ) + + if ( + ra_b is not None + and dec_b is not None + and e1_b is not None + and e2_b is not None + and w_b is not None + and not same_bin + ): + field_b = nmt.NmtFieldCatalog( + positions=[ra_b, dec_b], + weights=w_b, + field=[e1_b, pol_factor * e2_b], + lmax=lmax, + lmax_mask=lmax, + spin=2, + lonlat=True, + ) + else: + field_b = field_a + + if return_wsp: + wsp = nmt.NmtWorkspace.from_fields(field_a, field_b, b) + return field_a, field_b, wsp + else: + return field_a, field_b, None + + +def compute_cl_from_field_and_workspace(field_a, field_b, wsp, b): + """Compute the angular power spectrum from the input NaMaster field and workspace + + Parameters + ---------- + field_a : nmt.NmtField + NaMaster field object for the first catalog. + field_b : nmt.NmtField + NaMaster field object for the second catalog. + wsp : nmt.NmtWorkspace + NaMaster workspace object containing the mixing matrix. + b : nmt.NmtBin + NaMaster binning object. + + Returns + ------- + cl_coupled : np.ndarray + Coupled angular power spectrum. + cl_decoupled : np.ndarray + Decoupled angular power spectrum. + """ + cl_coupled = nmt.compute_coupled_cell(field_a, field_b) + cl_decoupled = wsp.decouple_cell(cl_coupled) + + return cl_coupled, cl_decoupled + + def get_pseudo_cls_map( - shear_map, - mask, + shear_map_a, + mask_a, nside, binning, *, - pol_factor=True, + shear_map_b=None, + mask_b=None, + pol_factor=-1, wsp=None, ell_step=10, n_ell_bins=32, @@ -156,16 +754,20 @@ def get_pseudo_cls_map( Parameters ---------- - shear_map : np.ndarray + shear_map_a : np.ndarray Complex shear map (``e1 + 1j * e2``). - mask : np.ndarray + mask_a : np.ndarray Field mask (the galaxy number-density map). nside : int HEALPix resolution; fixes the harmonic geometry. binning : str Binning scheme passed to :func:`make_namaster_bin`. - pol_factor : bool, optional - If ``True`` flip the sign of the imaginary (e2) component. + shear_map_b : np.ndarray, optional + Complex shear map for the second field (``e1 + 1j * e2``). + mask_b : np.ndarray, optional + Field mask for the second field (the galaxy number-density map). + pol_factor : float, optional + Polarization factor to apply to the E2 component. wsp : nmt.NmtWorkspace, optional Reuse a coupling workspace; built from the field if ``None``. ell_step, n_ell_bins, power : optional @@ -180,6 +782,27 @@ def get_pseudo_cls_map( wsp : nmt.NmtWorkspace The coupling workspace (newly built or the one passed in). """ + # First do some assertion checks + if shear_map_b is not None: + assert mask_b is not None, "mask_b must be provided if shear_map_b is provided" + assert shear_map_a.shape == shear_map_b.shape, ( + "shear_map_a and shear_map_b must have the same shape" + ) + assert mask_a.shape == mask_b.shape, ( + "mask_a and mask_b must have the same shape" + ) + + if mask_b is not None: + assert shear_map_b is not None, ( + "shear_map_b must be provided if mask_b is provided" + ) + assert shear_map_a.shape == shear_map_b.shape, ( + "shear_map_a and shear_map_b must have the same shape" + ) + assert mask_a.shape == mask_b.shape, ( + "mask_a and mask_b must have the same shape" + ) + lmin, lmax, b_lmax = pseudo_cl_geometry(nside) b = make_namaster_bin( @@ -193,18 +816,36 @@ def get_pseudo_cls_map( ) ell_eff = b.get_effective_ells() - factor = -1 if pol_factor else 1 - - f_all = nmt.NmtField( - mask=mask, maps=[shear_map.real, factor * shear_map.imag], lmax=b_lmax - ) if wsp is None: - wsp = nmt.NmtWorkspace.from_fields(f_all, f_all, b) + field_a, field_b, wsp = get_field_and_workspace_from_map( + b, + mask_a, + e1_map_a=shear_map_a.real, + e2_map_a=shear_map_a.imag, + mask_b=mask_b, + e1_map_b=shear_map_b.real if shear_map_b is not None else None, + e2_map_b=shear_map_b.imag if shear_map_b is not None else None, + pol_factor=pol_factor, + return_wsp=True, + ) + else: + field_a, field_b, _ = get_field_and_workspace_from_map( + b, + mask_a, + e1_map_a=shear_map_a.real, + e2_map_a=shear_map_a.imag, + mask_b=mask_b, + e1_map_b=shear_map_b.real if shear_map_b is not None else None, + e2_map_b=shear_map_b.imag if shear_map_b is not None else None, + pol_factor=pol_factor, + return_wsp=False, + ) - cl_coupled = nmt.compute_coupled_cell(f_all, f_all) - cl_all = wsp.decouple_cell(cl_coupled) + cl_coupled, cl_decoupled = compute_cl_from_field_and_workspace( + field_a, field_b, wsp, b + ) - return ell_eff, cl_all, wsp + return ell_eff, cl_decoupled, wsp def get_pseudo_cls_catalog( @@ -213,7 +854,9 @@ def get_pseudo_cls_catalog( nside, binning, *, - pol_factor=True, + tomo_bin_a="all", + tomo_bin_b="all", + pol_factor=-1, wsp=None, ell_step=10, n_ell_bins=32, @@ -232,8 +875,10 @@ def get_pseudo_cls_catalog( HEALPix resolution; fixes the harmonic geometry. binning : str Binning scheme passed to :func:`make_namaster_bin`. - pol_factor : bool, optional - If ``True`` flip the sign of the e2 component. + tomo_bin_a, tomo_bin_b : str or int or None, optional + Tomographic bin IDs for the two fields. + pol_factor : int, optional + Polarization factor to apply to the E2 component. wsp : nmt.NmtWorkspace, optional Reuse a coupling workspace; built from the field if ``None``. ell_step, n_ell_bins, power : optional @@ -248,6 +893,11 @@ def get_pseudo_cls_catalog( wsp : nmt.NmtWorkspace The coupling workspace (newly built or the one passed in). """ + # First make some assertion checks reagarding the run mode + assert (tomo_bin_a == "all" and tomo_bin_b == "all") or ( + tomo_bin_a != "all" and tomo_bin_b != "all" + ), "Both tomo_bin_a and tomo_bin_b must be provided or both must be 'all'" + lmin, lmax, b_lmax = pseudo_cl_geometry(nside) b = make_namaster_bin( @@ -261,22 +911,140 @@ def get_pseudo_cls_catalog( ) ell_eff = b.get_effective_ells() - factor = -1 if pol_factor else 1 + is_tomography = tomo_bin_a != "all" and tomo_bin_b != "all" + if is_tomography: + assert params["tomo_bin_col"] is not None, ( + "The column of tomographic bin ids is not specified." + ) + mask_tomo_a = catalog[params["tomo_bin_col"]] == tomo_bin_a + mask_tomo_b = catalog[params["tomo_bin_col"]] == tomo_bin_b + catalog_a = catalog[mask_tomo_a] + catalog_b = catalog[mask_tomo_b] + same_bin = tomo_bin_a == tomo_bin_b + else: + catalog_a = catalog + catalog_b = catalog + same_bin = True - f_all = nmt.NmtFieldCatalog( - positions=[catalog[params["ra_col"]], catalog[params["dec_col"]]], - weights=catalog[params["w_col"]], - field=[catalog[params["e1_col"]], factor * catalog[params["e2_col"]]], - lmax=b_lmax, - lmax_mask=b_lmax, - spin=2, - lonlat=True, + if wsp is None: + field_a, field_b, wsp = get_field_and_workspace_from_catalog( + b, + ra_a=catalog_a[params["ra_col"]], + dec_a=catalog_a[params["dec_col"]], + e1_a=catalog_a[params["e1_col"]], + e2_a=catalog_a[params["e2_col"]], + w_a=catalog_a[params["w_col"]], + ra_b=catalog_b[params["ra_col"]], + dec_b=catalog_b[params["dec_col"]], + e1_b=catalog_b[params["e1_col"]], + e2_b=catalog_b[params["e2_col"]], + w_b=catalog_b[params["w_col"]], + pol_factor=pol_factor, + return_wsp=True, + same_bin=same_bin, + ) + else: + field_a, field_b, _ = get_field_and_workspace_from_catalog( + b, + ra_a=catalog_a[params["ra_col"]], + dec_a=catalog_a[params["dec_col"]], + e1_a=catalog_a[params["e1_col"]], + e2_a=catalog_a[params["e2_col"]], + w_a=catalog_a[params["w_col"]], + ra_b=catalog_b[params["ra_col"]], + dec_b=catalog_b[params["dec_col"]], + e1_b=catalog_b[params["e1_col"]], + e2_b=catalog_b[params["e2_col"]], + w_b=catalog_b[params["w_col"]], + pol_factor=pol_factor, + return_wsp=False, + same_bin=same_bin, + ) + + cl_coupled, cl_decoupled = compute_cl_from_field_and_workspace( + field_a, field_b, wsp, b ) - if wsp is None: - wsp = nmt.NmtWorkspace.from_fields(f_all, f_all, b) + return ell_eff, cl_decoupled, wsp + + +# ---------------------- Covariance computation functions ---------------------- +def get_fiducial_cl(z, dndz, lmax, cosmo, backend="camb"): + """ + Get the fiducial Cl's using the redshift distribution. + Cosmology is determined by the input cosmo object. + """ + ell = np.arange(1, lmax + 1) + + fiducial_cl = get_theo_c_ell(ell=ell, z=z, nz=dndz, backend=backend, cosmo=cosmo) + + return fiducial_cl + + +def get_pseudo_cl_iNKA_covariance( + input_cl_a1_b1, + input_cl_a1_b2, + input_cl_a2_b1, + input_cl_a2_b2, + field_a1, + field_a2, + field_b1, + field_b2, + wsp_a, + wsp_b, + b, +): + """Compute the iNKA covariance for pseudo-Cl. + + Parameters + ---------- + input_cl_a1_b1 : np.ndarray + Input Cl for field a1 and b1. + input_cl_a1_b2 : np.ndarray + Input Cl for field a1 and b2. + input_cl_a2_b1 : np.ndarray + Input Cl for field a2 and b1. + input_cl_a2_b2 : np.ndarray + Input Cl for field a2 and b2. + field_a1 : nmt.NmtField + NaMaster field object for the first catalog (a1). + field_a2 : nmt.NmtField + NaMaster field object for the second catalog (a2). + field_b1 : nmt.NmtField + NaMaster field object for the first catalog (b1). + field_b2 : nmt.NmtField + NaMaster field object for the second catalog (b2). + wsp_a : nmt.NmtWorkspace + NaMaster workspace object containing the mixing matrix for fields a. + wsp_b : nmt.NmtWorkspace + NaMaster workspace object containing the mixing matrix for fields b. + b : nmt.NmtBin + NaMaster binning object. + + Returns + ------- + cov_matrix : np.ndarray + Covariance matrix of the pseudo-Cl, shape ``(n_bins, n_bins)``. + """ + # Compute the coupling coefficients for the covariance + cw = nmt.NmtCovarianceWorkspace.from_fields(field_a1, field_a2, field_b1, field_b2) + + # Get actual number of ell bins from binning scheme + n_ell_actual = b.get_n_bands() - cl_coupled = nmt.compute_coupled_cell(f_all, f_all) - cl_all = wsp.decouple_cell(cl_coupled) + # Compute the covariance using NaMaster's built-in function + cov_matrix = nmt.gaussian_covariance( + cw, + 2, + 2, + 2, + 2, + input_cl_a1_b1, + input_cl_a1_b2, + input_cl_a2_b1, + input_cl_a2_b2, + wsp_a, + wb=wsp_b, + ).reshape([n_ell_actual, 4, n_ell_actual, 4]) - return ell_eff, cl_all, wsp + return cov_matrix diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 6874a08c..95a720f9 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -51,6 +51,10 @@ def get_params_rho_tau(cat, survey="other"): params["w_col"] = cat["shear"]["w_col"] params["e1_col"] = cat["shear"]["e1_col"] params["e2_col"] = cat["shear"]["e2_col"] + try: + params["tomo_bin_col"] = cat["shear"]["tomo_bin_col"] + except KeyError: + params["tomo_bin_col"] = None params["R11"] = cat["shear"].get("R11") params["R22"] = cat["shear"].get("R22") diff --git a/src/sp_validation/tests/data/generate_test_cl_catalog_reference.py b/src/sp_validation/tests/data/generate_test_cl_catalog_reference.py new file mode 100644 index 00000000..7c41d564 --- /dev/null +++ b/src/sp_validation/tests/data/generate_test_cl_catalog_reference.py @@ -0,0 +1,137 @@ +# %% +from pathlib import Path + +import IPython +import numpy as np +import yaml +from astropy.io import fits +from astropy.table import Table + +from sp_validation.cosmo_val import CosmologyValidation +from sp_validation.rho_tau import get_params_rho_tau + +ipython = IPython.get_ipython() + +# %% +NSIDE = 64 +SEED = 1234 +N_ELL_BINS = 8 + +rng = np.random.default_rng(SEED) +version = "TestCatalog" + +cat_dir = Path("./catalog") +nz_dir = Path("./nz") +output_dir = Path("./output_test") + +for d in (cat_dir, nz_dir, output_dir): + d.mkdir(exist_ok=True) + +n_gal = 5000 +ra = rng.uniform(10.0, 30.0, n_gal) +dec = rng.uniform(10.0, 30.0, n_gal) +e1 = rng.normal(0, 0.25, n_gal) +e2 = rng.normal(0, 0.25, n_gal) +w = rng.uniform(0.5, 1.0, n_gal) +tomo_bin_id = rng.integers(1, 3, n_gal) # Create a two bin catalogue +Table( + {"RA": ra, "Dec": dec, "e1": e1, "e2": e2, "w": w, "tomo_bin_id": tomo_bin_id} +).write(cat_dir / "shear.fits", overwrite=True) + +shear_cfg = { + "path": "shear.fits", + "w_col": "w", + "e1_col": "e1", + "e2_col": "e2", + "R": 1.0, + "e1_col_corrected": "e1", + "e2_col_corrected": "e2", + "ra_col": "RA", + "dec_col": "Dec", + "tomo_bin_col": "tomo_bin_id", +} + +psf_cfg = { + "path": "shear.fits", + "ra_col": "RA", + "dec_col": "Dec", + "e1_PSF_col": "e1", + "e2_PSF_col": "e2", + "e1_star_col": "e1", + "e2_star_col": "e2", + "PSF_size": "w", + "star_size": "w", + "PSF_flag": "w", + "star_flag": "w", +} +config_data = { + "nz": {"subdir": str(nz_dir), "dndz": {"blind": "A", "path": "dndz"}}, + "paths": {"output": str(output_dir)}, + version: { + "subdir": str(cat_dir), + "pipeline": "SP", + "shear": shear_cfg, + "star": {**psf_cfg}, + "psf": psf_cfg, + }, +} + +config_path = Path("./config.yaml") +config_path.write_text(yaml.dump(config_data, sort_keys=False)) + +# %% +cv = CosmologyValidation( + versions=[version], + catalog_config=config_path, + output_dir=str(output_dir), + nside=NSIDE, + binning="powspace", + power=0.5, + n_ell_bins=N_ELL_BINS, + pol_factor=-1, +) +cv._test_version = version + +ver = cv._test_version +params = get_params_rho_tau(cv.cc[ver], survey=ver) +cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) + +cat_gal_tomo_bin_1 = cat_gal[cat_gal[params["tomo_bin_col"]] == 1] +cat_gal_tomo_bin_2 = cat_gal[cat_gal[params["tomo_bin_col"]] == 2] +n_gal_map_a = cv.get_n_gal_map(params, NSIDE, cat_gal_tomo_bin_1) +shear_map_a_e1, shear_map_a_e2 = cv.get_shear_map(params, NSIDE, cat_gal_tomo_bin_1) +n_gal_map_b = cv.get_n_gal_map(params, NSIDE, cat_gal_tomo_bin_2) +shear_map_b_e1, shear_map_b_e2 = cv.get_shear_map(params, NSIDE, cat_gal_tomo_bin_2) + +shear_map_a = shear_map_a_e1 + 1j * shear_map_a_e2 +shear_map_b = shear_map_b_e1 + 1j * shear_map_b_e2 + + +# %% +ell_eff, cl_all, wsp = cv.get_pseudo_cls_map( + shear_map_a, n_gal_map_a, shear_map_b=shear_map_b, mask_b=n_gal_map_b +) + +# %% +# Get the catalogue output +ver = cv._test_version +cv._pseudo_cls = { + ver: { + "tomo_bin_1_tomo_bin_1": {}, + "tomo_bin_1_tomo_bin_2": {}, + "tomo_bin_2_tomo_bin_2": {}, + } +} +out_path = cv._output_path(f"pseudo_cl_cat_{ver}.fits") +tomo_bin_ids, tomo_bin_pairs = cv._get_tomo_bins(ver) +for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + out_path = cv._output_path(f"pseudo_cl_cat_{ver}_{tomo_bin_a}_{tomo_bin_b}.fits") + cv.calculate_pseudo_cl_catalog( + ver, out_path, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b + ) + + +# %% +np.savez("./test_cl_catalog", cv._pseudo_cls[ver]) +# %% +test_result = np.load("./test_cl_catalog.npz", allow_pickle=True) diff --git a/src/sp_validation/tests/data/test_cl_catalog.npz b/src/sp_validation/tests/data/test_cl_catalog.npz new file mode 100644 index 00000000..fdba323b Binary files /dev/null and b/src/sp_validation/tests/data/test_cl_catalog.npz differ diff --git a/src/sp_validation/tests/test_pseudo_cl.py b/src/sp_validation/tests/test_pseudo_cl.py index 45a2366d..c07afd00 100644 --- a/src/sp_validation/tests/test_pseudo_cl.py +++ b/src/sp_validation/tests/test_pseudo_cl.py @@ -51,6 +51,7 @@ """ import os +from pathlib import Path import numpy as np import numpy.testing as npt @@ -58,6 +59,7 @@ import yaml from sp_validation.cosmo_val import CosmologyValidation +from sp_validation.pseudo_cl import apply_random_rotation from sp_validation.rho_tau import get_params_rho_tau # These tests need the full harmonic-space stack (pymaster/NaMaster + healpy), @@ -73,6 +75,8 @@ SEED = 1234 N_ELL_BINS = 8 +REFERENCE_TOMO = Path(__file__).parent / "data" / "test_cl_catalog.npz" + # --------------------------------------------------------------------------- # Synthetic-catalog fixture (deterministic; no cluster data) @@ -101,9 +105,10 @@ def _write_synthetic_config(tmp_path): e1 = rng.normal(0, 0.25, n_gal) e2 = rng.normal(0, 0.25, n_gal) w = rng.uniform(0.5, 1.0, n_gal) - Table({"RA": ra, "Dec": dec, "e1": e1, "e2": e2, "w": w}).write( - cat_dir / "shear.fits", overwrite=True - ) + tomo_bin_id = rng.integers(1, 3, n_gal) # Create a two bin catalogue + Table( + {"RA": ra, "Dec": dec, "e1": e1, "e2": e2, "w": w, "tomo_bin_id": tomo_bin_id} + ).write(cat_dir / "shear.fits", overwrite=True) z_edges = np.linspace(0.05, 3.0, 31) dndz = np.exp(-(((z_edges - 0.7) / 0.3) ** 2)) @@ -120,6 +125,7 @@ def _write_synthetic_config(tmp_path): "e2_col_corrected": "e2", "ra_col": "RA", "dec_col": "Dec", + "tomo_bin_col": "tomo_bin_id", } # Minimal psf block so get_params_rho_tau() succeeds; the pseudo-Cl # primitives only read the shear-side keys, but the production call path @@ -165,7 +171,7 @@ def cv(tmp_path): binning="powspace", power=0.5, n_ell_bins=N_ELL_BINS, - pol_factor=True, + pol_factor=-1, ) cv._test_version = version return cv @@ -180,6 +186,11 @@ def cat_and_params(cv): return cat_gal, params +def test_reference_exists(): + """Guard the guard: a missing reference must fail loudly, not skip.""" + assert REFERENCE_TOMO.exists(), f"committed reference missing: {REFERENCE_TOMO}" + + # Tolerances ---------------------------------------------------------------- # Bitwise-stable primitives (binning math, n_gal map, map-based pseudo-Cl). RTOL_DET = 1e-9 @@ -265,18 +276,38 @@ def test_unknown_binning_raises(self, cv): cv.get_namaster_bin(LMIN, LMAX, B_LMAX) +# ========================================================================== +# get_pixels -- fetch the pixels from ra, dec, and nside (HEALPix) +# ========================================================================== +def test_get_pixels(cv, cat_and_params): + """Pin the HEALPix pixelization of the synthetic catalog.""" + cat_gal, params = cat_and_params + unique_pix, idx, idx_rep = cv.get_pixels(params, NSIDE, cat_gal) + + # Structural invariants of the HEALPix occupancy map. + assert unique_pix.size == 485 + assert idx_rep.size == 5000 # one entry per galaxy + + # Pinned scalar summaries: total weight is conserved (sum of weights), + # peak occupancy, and the pixel index bookkeeping. + npt.assert_array_equal(unique_pix[:5], [12167, 12168, 12169, 12170, 12171]) + assert int(unique_pix.sum()) == 7849989 + + # =========================================================================== # get_n_gal_map -- weighted galaxy number-density map # =========================================================================== def test_get_n_gal_map(cv, cat_and_params): cat_gal, params = cat_and_params - n_gal, unique_pix, idx, idx_rep = cv.get_n_gal_map(params, NSIDE, cat_gal) + + unique_pix, idx, idx_rep = cv.get_pixels(params, NSIDE, cat_gal) + n_gal = cv.get_n_gal_map( + params, NSIDE, cat_gal, unique_pix=unique_pix, idx=None, idx_rep=idx_rep + ) # Structural invariants of the HEALPix occupancy map. assert n_gal.shape == (healpy.nside2npix(NSIDE),) assert int(np.count_nonzero(n_gal)) == 485 - assert unique_pix.size == 485 - assert idx_rep.size == 5000 # one entry per galaxy # The map is supported exactly on the occupied pixels. npt.assert_array_equal(np.nonzero(n_gal)[0], np.sort(unique_pix)) @@ -284,8 +315,6 @@ def test_get_n_gal_map(cv, cat_and_params): # peak occupancy, and the pixel index bookkeeping. npt.assert_allclose(n_gal.sum(), 3760.657282591494, rtol=RTOL_DET) npt.assert_allclose(n_gal.max(), 16.063410433711447, rtol=RTOL_DET) - npt.assert_array_equal(unique_pix[:5], [12167, 12168, 12169, 12170, 12171]) - assert int(unique_pix.sum()) == 7849989 npt.assert_allclose( n_gal[unique_pix][:5], np.array( @@ -307,7 +336,10 @@ def test_get_n_gal_map(cv, cat_and_params): # =========================================================================== def _build_shear_map(cv, cat_gal, params): """Replicate calculate_pseudo_cl_map's weighted shear-map construction.""" - n_gal, unique_pix, _idx, idx_rep = cv.get_n_gal_map(params, NSIDE, cat_gal) + unique_pix, _idx, idx_rep = cv.get_pixels(params, NSIDE, cat_gal) + n_gal = cv.get_n_gal_map( + params, NSIDE, cat_gal, unique_pix=unique_pix, idx=_idx, idx_rep=idx_rep + ) w = cat_gal[params["w_col"]] e1 = cat_gal[params["e1_col"]] e2 = cat_gal[params["e2_col"]] @@ -321,10 +353,32 @@ def _build_shear_map(cv, cat_gal, params): return m1 + 1j * m2, n_gal +def test_build_shear_map(cv, cat_and_params): + """Pin the weighted shear map construction from the synthetic catalog.""" + cat_gal, params = cat_and_params + unique_pix, _idx, idx_rep = cv.get_pixels(params, NSIDE, cat_gal) + n_gal = cv.get_n_gal_map( + params, NSIDE, cat_gal, unique_pix=unique_pix, idx=_idx, idx_rep=idx_rep + ) + m1, m2 = cv.get_shear_map( + params, NSIDE, cat_gal, unique_pix=unique_pix, idx=_idx, idx_rep=idx_rep + ) + + m = m1 + 1j * m2 + + m_replicate, n_gal_replicate = _build_shear_map(cv, cat_gal, params) + + npt.assert_allclose(m, m_replicate, rtol=RTOL_DET, atol=ATOL_DET) + npt.assert_allclose(n_gal, n_gal_replicate, rtol=RTOL_DET, atol=ATOL_DET) + + def test_get_pseudo_cls_map(cv, cat_and_params): cat_gal, params = cat_and_params shear_map, n_gal = _build_shear_map(cv, cat_gal, params) ell_eff, cl_all, wsp = cv.get_pseudo_cls_map(shear_map, n_gal) + ell_eff_2, cl_all_2, wsp_2 = cv.get_pseudo_cls_map( + shear_map, n_gal, shear_map_b=shear_map, mask_b=n_gal + ) assert cl_all.shape == (4, N_ELL_BINS) npt.assert_allclose( @@ -387,13 +441,106 @@ def test_get_pseudo_cls_map(cv, cat_and_params): # BE (index 2) is the transpose-symmetric partner of EB for an auto-spectrum. npt.assert_allclose(cl_all[2], cl_all[1], rtol=RTOL_DET, atol=1e-18) + # Assert that running the same map against itself and against itself as a second map gives the same result. + npt.assert_allclose(cl_all, cl_all_2, rtol=RTOL_DET, atol=ATOL_DET) + npt.assert_allclose(ell_eff, ell_eff_2, rtol=RTOL_DET, atol=ATOL_DET) + + +def test_get_pseudo_cls_map_with_tomo(cv, cat_and_params): + cat_gal, params = cat_and_params + cat_gal_tomo1 = cat_gal[cat_gal[params["tomo_bin_col"]] == 1] + cat_gal_tomo2 = cat_gal[cat_gal[params["tomo_bin_col"]] == 2] + shear_map_1, n_gal_1 = _build_shear_map(cv, cat_gal_tomo1, params) + shear_map_2, n_gal_2 = _build_shear_map(cv, cat_gal_tomo2, params) + ell_eff, cl_all, wsp = cv.get_pseudo_cls_map( + shear_map_1, n_gal_1, shear_map_b=shear_map_2, mask_b=n_gal_2 + ) + + assert cl_all.shape == (4, N_ELL_BINS) + npt.assert_allclose( + ell_eff, + np.array([11.5, 20.0, 30.5, 43.5, 58.5, 75.5, 95.0, 116.5]), + rtol=RTOL_DET, + atol=ATOL_DET, + ) + npt.assert_allclose( + cl_all[0], # EE + np.array( + [ + -8.648947250571888e-06, + 4.386882223005456e-06, + -2.3193526251107696e-06, + 1.0651397314115168e-06, + -2.722591256637468e-07, + -3.862270431501105e-07, + -1.40420405782524e-06, + 2.1531429566476774e-07, + ] + ), + rtol=RTOL_DET, + atol=ATOL_DET, + ) + npt.assert_allclose( + cl_all[1], # EB + np.array( + [ + -2.8032134416262087e-07, + -1.271298363785575e-06, + 1.7112863041724342e-08, + -2.899403854283714e-07, + 1.360217254538336e-06, + 7.682614612656511e-08, + 6.620365648170535e-07, + -3.164479179586416e-07, + ] + ), + rtol=RTOL_DET, + atol=ATOL_DET, + ) + npt.assert_allclose( + cl_all[2], # BE + np.array( + [ + -3.83934537153214e-06, + 6.218772676611559e-06, + -5.542920986484232e-06, + 8.810792449470091e-07, + -3.2973293732012423e-07, + -1.3155110757359602e-08, + -1.420039280183438e-07, + 3.09464700816748e-07, + ] + ), + rtol=RTOL_DET, + atol=ATOL_DET, + ) + npt.assert_allclose( + cl_all[3], # BB + np.array( + [ + 9.387700364892905e-06, + -2.9045253374872504e-06, + -1.7006849005134948e-06, + -3.6869001010423035e-07, + 4.4986878855942184e-07, + 6.021221048891767e-07, + 8.413198154662088e-08, + -5.782957119801217e-07, + ] + ), + rtol=RTOL_DET, + atol=ATOL_DET, + ) + # =========================================================================== # get_pseudo_cls_catalog -- NmtFieldCatalog path (drifts ~2e-12) # =========================================================================== def test_get_pseudo_cls_catalog(cv, cat_and_params): cat_gal, params = cat_and_params - ell_eff, cl_all, wsp = cv.get_pseudo_cls_catalog(catalog=cat_gal, params=params) + ell_eff, cl_all, wsp = cv.get_pseudo_cls_catalog( + catalog=cat_gal, params=params, tomo_bin_a="all", tomo_bin_b="all" + ) assert cl_all.shape == (4, N_ELL_BINS) # Effective ells share the binning math with the map path: bitwise-stable. @@ -471,7 +618,7 @@ def test_apply_random_rotation_preserves_magnitude(cv, cat_and_params): e1 = np.asarray(cat_gal[params["e1_col"]], dtype=np.float64) e2 = np.asarray(cat_gal[params["e2_col"]], dtype=np.float64) - e1_rot, e2_rot = cv.apply_random_rotation(e1, e2) + e1_rot, e2_rot = apply_random_rotation(e1, e2) assert e1_rot.shape == e1.shape assert e2_rot.shape == e2.shape @@ -491,16 +638,16 @@ def test_apply_random_rotation_reproducible_with_seed(cv, cat_and_params): e1 = np.asarray(cat_gal[params["e1_col"]], dtype=np.float64) e2 = np.asarray(cat_gal[params["e2_col"]], dtype=np.float64) - a1, a2 = cv.apply_random_rotation(e1, e2, np.random.default_rng(42)) - b1, b2 = cv.apply_random_rotation(e1, e2, np.random.default_rng(42)) + a1, a2 = apply_random_rotation(e1, e2, np.random.default_rng(42)) + b1, b2 = apply_random_rotation(e1, e2, np.random.default_rng(42)) npt.assert_array_equal(a1, b1) npt.assert_array_equal(a2, b2) - c1, _ = cv.apply_random_rotation(e1, e2, np.random.default_rng(7)) + c1, _ = apply_random_rotation(e1, e2, np.random.default_rng(7)) assert not np.allclose(a1, c1) - d1, _ = cv.apply_random_rotation(e1, e2) - f1, _ = cv.apply_random_rotation(e1, e2) + d1, _ = apply_random_rotation(e1, e2) + f1, _ = apply_random_rotation(e1, e2) assert not np.allclose(d1, f1) @@ -515,9 +662,9 @@ def test_calculate_pseudo_cl_catalog_end_to_end(cv, tmp_path): (it drops the BE row); we pin the round-tripped table. """ ver = cv._test_version - cv._pseudo_cls = {ver: {}} + cv._pseudo_cls = {ver: {"tomo_bin_all_tomo_bin_all": {}}} out_path = cv._output_path(f"pseudo_cl_cat_{ver}.fits") - cv.calculate_pseudo_cl_catalog(ver, out_path) + cv.calculate_pseudo_cl_catalog(ver, out_path, tomo_bin_a="all", tomo_bin_b="all") assert os.path.exists(out_path) d = fits.getdata(out_path) @@ -588,5 +735,94 @@ def test_calculate_pseudo_cl_catalog_end_to_end(cv, tmp_path): # (same computation, FITS round-trip) -- consistency, not an independent pin. cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) params = get_params_rho_tau(cv.cc[ver], survey=ver) - _, cl_prim, _ = cv.get_pseudo_cls_catalog(catalog=cat_gal, params=params) + _, cl_prim, _ = cv.get_pseudo_cls_catalog( + catalog=cat_gal, params=params, tomo_bin_a="all", tomo_bin_b="all" + ) npt.assert_allclose(ee, cl_prim[0], rtol=RTOL_CAT, atol=ATOL_CAT) + + +def test_calculate_pseudo_cl_catalog_end_to_end_tomo(cv, tmp_path): + """End-to-end catalog path: FITS round-trip of ell + EE/EB/BB. + + The catalog method has no random noise debiasing, so it is reproducible to + the same ~2e-12 catalog-path float noise. save_pseudo_cl stores ELL/EE/EB/BB + (it drops the BE row); we pin the round-tripped table. + """ + ver = cv._test_version + cv._pseudo_cls = { + ver: { + "tomo_bin_1_tomo_bin_1": {}, + "tomo_bin_1_tomo_bin_2": {}, + "tomo_bin_2_tomo_bin_2": {}, + } + } + out_path = cv._output_path(f"pseudo_cl_cat_{ver}.fits") + tomo_bin_ids, tomo_bin_pairs = cv._get_tomo_bins(ver) + + result_to_compare = np.load(REFERENCE_TOMO, allow_pickle=True)["arr_0"].item() + for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + out_path = cv._output_path( + f"pseudo_cl_cat_{ver}_{tomo_bin_a}_{tomo_bin_b}.fits" + ) + cv.calculate_pseudo_cl_catalog( + ver, out_path, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b + ) + + assert os.path.exists(out_path) + d = fits.getdata(out_path) + # FITS gives big-endian f8; normalize for value comparison. + ell = np.asarray(d["ELL"], dtype=np.float64) + ee = np.asarray(d["EE"], dtype=np.float64) + eb = np.asarray(d["EB"], dtype=np.float64) + be = np.asarray(d["BE"], dtype=np.float64) + bb = np.asarray(d["BB"], dtype=np.float64) + + npt.assert_allclose( + ell, + result_to_compare[f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ]["ELL"], + rtol=RTOL_DET, + atol=ATOL_DET, + ) + npt.assert_allclose( + ee, + result_to_compare[f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ]["EE"], + rtol=RTOL_CAT, + atol=ATOL_CAT, + ) + npt.assert_allclose( + eb, + result_to_compare[f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ]["EB"], + rtol=RTOL_CAT, + atol=ATOL_CAT, + ) + npt.assert_allclose( + be, + result_to_compare[f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ]["BE"], + rtol=RTOL_CAT, + atol=ATOL_CAT, + ) + npt.assert_allclose( + bb, + result_to_compare[f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ + "pseudo_cl" + ]["BB"], + rtol=RTOL_CAT, + atol=ATOL_CAT, + ) + + # The end-to-end catalog EE matches the primitive get_pseudo_cls_catalog EE + # (same computation, FITS round-trip) -- consistency, not an independent pin. + cat_gal = fits.getdata(cv.cc[ver]["shear"]["path"]) + params = get_params_rho_tau(cv.cc[ver], survey=ver) + _, cl_prim, _ = cv.get_pseudo_cls_catalog( + catalog=cat_gal, params=params, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b + ) + npt.assert_allclose(ee, cl_prim[0], rtol=RTOL_CAT, atol=ATOL_CAT)