diff --git a/.gitignore b/.gitignore index a568c1ca..f8eec899 100644 --- a/.gitignore +++ b/.gitignore @@ -197,6 +197,4 @@ papers/catalog/plots/*.pdf papers/cosmo_val/logs/ # Ignore scratch notebooks -scratch/*/*.ipynb -scratch/guerrini/work_notebooks -scratch/guerrini/launch_scripts \ No newline at end of file +scratch/*/*.ipynb \ No newline at end of file diff --git a/cosmo_val/cat_config.yaml b/cosmo_val/cat_config.yaml index ea36a81d..cc1326df 100644 --- a/cosmo_val/cat_config.yaml +++ b/cosmo_val/cat_config.yaml @@ -1259,49 +1259,3 @@ 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 066bf4a1..48251bbb 100644 --- a/src/sp_validation/cosmo_val/core.py +++ b/src/sp_validation/cosmo_val/core.py @@ -1,6 +1,5 @@ # %% import copy -import itertools import os import re from pathlib import Path @@ -8,7 +7,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 from ..b_modes import ( @@ -89,7 +88,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 : int, default -1 + pol_factor : bool, default True Apply polarization correction factor in pseudo-C_ell calculations. nrandom_cell : int, default 10 Number of random realizations for C_ell error estimation. @@ -98,10 +97,6 @@ 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 ---------- @@ -219,7 +214,7 @@ def __init__( power=1 / 2, n_ell_bins=32, ell_step=10, - pol_factor=-1, + pol_factor=True, cell_method="map", noise_bias_method="analytic", fiducial_input_inka="coupled", @@ -228,8 +223,6 @@ 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 @@ -250,10 +243,7 @@ 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 @@ -262,8 +252,6 @@ 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'" @@ -648,36 +636,3 @@ 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 d8c1caf8..514049a6 100644 --- a/src/sp_validation/cosmo_val/pseudo_cl.py +++ b/src/sp_validation/cosmo_val/pseudo_cl.py @@ -7,37 +7,33 @@ 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 matplotlib.colors import to_rgb - -import sp_validation.pseudo_cl as spv_pseudo_cl - +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 ..rho_tau import get_params_rho_tau -from ..statistics import cov_from_one_covariance +from ..statistics import chi2_and_pte, cov_from_one_covariance class PseudoClMixin: - # ---------------- Pseudo-Cl properties ---------------- # @property def pseudo_cls(self): if not hasattr(self, "_pseudo_cls"): - 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 - ) + self.calculate_pseudo_cl() + self.calculate_pseudo_cl_eb_cov() return self._pseudo_cls @property @@ -46,11 +42,51 @@ def pseudo_cls_onecov(self): self.calculate_pseudo_cl_onecovariance() return self._pseudo_cls_onecov - # ---------------- 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 - ): + 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): """ Compute a theoretical Gaussian covariance of the Pseudo-Cl for EE, EB and BB. """ @@ -58,305 +94,182 @@ def calculate_pseudo_cl_inka_cov( nside = self.nside - 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) - 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_merged = self._output_path_pseudo_cl_cov( - ver, "iNKA", tomography=compute_tomography - ) + 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) - if compute_tomography: - tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + self.print_cyan(f"Extracting the fiducial power spectrum for {ver}") - if tomo_bin_ids is None or tomo_bin_pairs is None: + lmax = 2 * self.nside + ell = np.arange(1, lmax + 1) + pw = hp.pixwin(nside, lmax=lmax) + if pw.shape[0] != len(ell) + 1: raise ValueError( - f"Version {ver} does not have tomography information." + "Unexpected pixwin length for lmax=" + f"{lmax}: got {pw.shape[0]}, expected {len(ell) + 1}" ) - - else: - tomo_bin_pairs = [("all", "all")] - - 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" + 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, + ) + * pw**2 ) - self._pseudo_cls[ver]["cov_iNKA"] = fits.open(out_path_merged) - if load_all_block: - self.print_done("Loading all the iNKA covariance blocks") - n_spectra = len(tomo_bin_pairs) + self.print_cyan("Getting a binning, n_gal_map, field and workspace.") - # indices into tomo_bin_pairs for all unique covariance blocks - block_indices = list( - itertools.combinations_with_replacement(range(n_spectra), 2) - ) + lmin = 8 + lmax = 2 * self.nside + b_lmax = lmax - 1 - 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, - ), - ) + b = self.get_namaster_bin(lmin, lmax, b_lmax) - 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." - ) + # Load data and create shear and noise maps + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) - continue + n_gal, unique_pix, _idx, idx_rep = self.get_n_gal_map( + params, nside, cat_gal + ) - # Initialise dictionnary to store field and workspace - n_gal_map_dict = {} - field_dict = {} - wsp_dict = {} + 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), + ) - self.print_cyan(f"Extracting the fiducial power spectrum for {ver}") + noise_bias_cl = np.mean(cl_noise, axis=0) - fiducial_cl = self.get_fiducial_cl(ver, compute_tomography) + elif self.noise_bias_method == "analytic": + self.print_cyan("Getting analytic 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}") + 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 + ) - params = get_params_rho_tau(self.cc[ver]) - cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) + noise_bias = hp.nside2pixarea(self.nside) * np.mean(variance_map) - 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) + noise_bias_cl = np.zeros((4, lmax)) + noise_bias_cl[0, :] = noise_bias + noise_bias_cl[3, :] = noise_bias - noise_bias_cl = self.get_noise_bias(params, nside, cat_gal_) + noise_bias_cl = wsp.decouple_cell(noise_bias_cl) # Decouple else: - noise_bias_cl = np.zeros((4, 2 * nside)) + raise ValueError( + f"Noise bias method {self.noise_bias_method} not recognized. It should be 'randoms' or 'analytic'." + ) + + # 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]] + + self.print_cyan("Adding noise bias to the fiducial Cls.") - # Update the fiducial_cl dictionnary - fiducial_cl[f"W{bin_key1}xW{bin_key2}"] = ( + fiducial_cl = ( np.array( [ - 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}"], + fiducial_cl, + 0.0 * fiducial_cl, + 0.0 * fiducial_cl, + 0.0 * fiducial_cl, ] ) + noise_bias_cl ) - # 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 not 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}"] + if self.fiducial_input_inka == "coupled": + self.print_cyan("Coupling the fiducial Cls.") coupling_mat = wsp.get_coupling_matrix() coupling_mat_re = np.reshape( coupling_mat, (4, lmax, 4, lmax), order="F" ) - 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 + fiducial_cl = np.tensordot(coupling_mat_re, fiducial_cl) / np.mean( + n_gal**2 + ) # couple and divide by the mean of the mask squared self.print_cyan("Computing the Pseudo-Cl covariance") - 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_a2 - 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}"] - ) - - 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, - ) + 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]) self.print_cyan("Saving Pseudo-Cl covariance") - 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) + # 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}" + ) + ) - 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}" - ] + hdu.writeto(out_path, overwrite=True) + + self._pseudo_cls[ver]["cov"] = hdu - # 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): @@ -386,11 +299,8 @@ 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") - ) - and not self.force_run + if os.path.exists( + os.path.join(out_dir, "covariance_list_3x2pt_pure_Cell.dat") ): self.print_done(f"Skipping OneCovariance calculation, {out_dir} exists") self._load_onecovariance_cov(out_dir, ver) @@ -507,7 +417,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) and not self.force_run: + if os.path.exists(out_file): self.print_done( f"Skipping Gaussian and Non-Gaussian covariance calculation, {out_file} exists" ) @@ -541,356 +451,182 @@ 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, compute_tomography=True): + def calculate_pseudo_cl(self): """ - Compute the pseudo-Cl of a `CosmologyValidation` inputs with tomography. + Compute the pseudo-Cl of given catalogs. """ - out_dir = self._output_path("pseudo_cl") - os.makedirs(out_dir, exist_ok=True) + self.print_start("Computing pseudo-Cl's") - 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", {}) + nside = self.nside + try: + self._pseudo_cls + except AttributeError: + self._pseudo_cls = {} for ver in self.versions: self.print_magenta(ver) - if ver not in self.pseudo_cls.keys(): - self._pseudo_cls[ver] = {} - - if compute_tomography: - tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) - - if tomo_bin_ids is None or tomo_bin_pairs is None: - raise ValueError( - f"Version {ver} does not have tomography information." - ) - + 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: - 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}") + raise ValueError(f"Unknown cell method: {self.cell_method}") - 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." + self.print_done("Done pseudo-Cl's") - params = get_params_rho_tau(self.cc[ver]) - - self.print_cyan( - f"Computing pseudo-Cl's for tomographic bins {tomo_bin_a} and {tomo_bin_b}..." - ) + 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"]) - # 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) + 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 + + shear_map_e1 = np.zeros(hp.nside2npix(nside)) + shear_map_e2 = np.zeros(hp.nside2npix(nside)) + + e1 = cat_gal[params["e1_col"]] + e2 = cat_gal[params["e2_col"]] del cat_gal - 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) + 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] - # 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, - ) + shear_map = shear_map_e1 + 1j * shear_map_e2 - # 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 + del shear_map_e1, shear_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 + ell_eff, cl_shear, wsp = self.get_pseudo_cls_map(shear_map, n_gal_map) - # 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 = np.zeros_like(cl_shear) + rng = np.random.default_rng(self.cell_seed) - # 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, - ) + for i in range(self.nrandom_cell): + noise_map_e1 = np.zeros(hp.nside2npix(nside)) + noise_map_e2 = np.zeros(hp.nside2npix(nside)) + + e1_rot, e2_rot = self.apply_random_rotation(e1, e2, rng) + + 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) - # Subtract the noise bias from the pseudo-Cl's - cl_shear = cl_shear - cl_noise + noise_map_e1[mask] /= n_gal_map[mask] + noise_map_e2[mask] /= n_gal_map[mask] + + noise_map = noise_map_e1 + 1j * noise_map_e2 + del noise_map_e1, noise_map_e2 + + _, cl_noise_, _ = self.get_pseudo_cls_map(noise_map, n_gal_map, wsp) + cl_noise += cl_noise_ + + 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 + + # Noise realizations are now reproducible (seeded rng from self.cell_seed). + 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][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ - "pseudo_cl" - ] = cl_shear + self._pseudo_cls[ver]["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." - - params = get_params_rho_tau(self.cc[ver]) + def calculate_pseudo_cl_catalog(self, ver, 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"]) ell_eff, cl_shear, wsp = self.get_pseudo_cls_catalog( - catalog=cat_gal, params=params, tomo_bin_a=tomo_bin_a, tomo_bin_b=tomo_bin_b + catalog=cat_gal, params=params ) 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][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"][ - "pseudo_cl" - ] = cl_shear - - # ---------------- 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, - ) + self._pseudo_cls[ver]["pseudo_cl"] = cl_shear - 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 - ): + def get_n_gal_map(self, params, nside, cat_gal): """Weighted galaxy number-density map (thin wrapper -> primitive).""" - return spv_pseudo_cl.get_n_gal_map( + return 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_shear_map( - self, - params, - nside, - cat_gal, - unique_pix=None, - idx=None, - idx_rep=None, - n_gal_map=None, + def get_gaussian_real( + self, params, nside, lmax, cat_gal, n_gal, mask, unique_pix, idx_rep, rng=None ): - """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, + e1_rot, e2_rot = self.apply_random_rotation( + cat_gal[params["e1_col"]], cat_gal[params["e2_col"]], rng ) + 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_noise_realisation( + def get_sample( self, params, nside, + lmax, + b, cat_gal, - n_gal=None, - unique_pix=None, - idx=None, - idx_rep=None, + n_gal, + mask, + unique_pix, + idx_rep, rng=None, ): - """ - 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, + noise_map = self.get_gaussian_real( + params, nside, lmax, cat_gal, n_gal, mask, unique_pix, idx_rep, rng ) - 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, - ) + f = nmt.NmtField(mask=mask, maps=[noise_map.real, noise_map.imag], lmax=lmax) - 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, - ) + wsp = nmt.NmtWorkspace.from_fields(f, f, b) - 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, - ) + cl_noise = nmt.compute_coupled_cell(f, f) + cl_noise = wsp.decouple_cell(cl_noise) - def get_pseudo_cls_map( - self, map_a, mask_a, wsp=None, shear_map_b=None, mask_b=None - ): + return cl_noise, f, wsp + + def get_pseudo_cls_map(self, map, mask, wsp=None): """Map-based pseudo-cl (thin wrapper, state -> primitive).""" - return spv_pseudo_cl.get_pseudo_cls_map( - map_a, - mask_a, + return get_pseudo_cls_map( + map, + mask, 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, @@ -898,17 +634,13 @@ def get_pseudo_cls_map( power=self.power, ) - def get_pseudo_cls_catalog( - self, catalog, params, wsp=None, tomo_bin_a=None, tomo_bin_b=None - ): + def get_pseudo_cls_catalog(self, catalog, params, wsp=None): """Catalog-based pseudo-cl (thin wrapper, state -> primitive).""" - return spv_pseudo_cl.get_pseudo_cls_catalog( + return 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, @@ -916,68 +648,12 @@ def get_pseudo_cls_catalog( power=self.power, ) - 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:] - - # 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 apply_random_rotation(self, e1, e2, rng=None): + """Random ellipticity rotation (thin wrapper -> primitive). - 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", - ) + Pass a seeded ``rng`` for reproducible noise realizations. + """ + return apply_random_rotation(e1, e2, rng) def save_pseudo_cl(self, ell_eff, pseudo_cl, out_path): """ @@ -994,440 +670,208 @@ 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="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]) + col4 = fits.Column(name="BB", format="D", array=pseudo_cl[3]) + coldefs = fits.ColDefs([col1, col2, col3, col4]) cell_hdu = fits.BinTableHDU.from_columns(coldefs, name="PSEUDO_CELL") cell_hdu.writeto(out_path, overwrite=True) - 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): + def plot_pseudo_cl(self): """ - Merge the iNKA covariance matrices for a given version to get the data vector covariance. + Plot pseudo-Cl's for given catalogs. """ - out_path = self._output_path_pseudo_cl_cov( - ver, method="iNKA", tomography=tomography - ) + self.print_cyan("Plotting pseudo-Cl's") - print(self._pseudo_cls[ver]["tomo_bin_all_tomo_bin_all"]) + # 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. - 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$") - if tomo_bin_ids is None or tomo_bin_pairs is None: - raise AssertionError( - f"Tomographic bin IDs of version {ver} is not available." - ) + 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) - n_spectra = len(tomo_bin_pairs) - block_indices = list( - itertools.combinations_with_replacement(range(n_spectra), 2) + 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"], ) - 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_xlabel(r"$\ell$") + ax[1].set_ylabel(r"$C_\ell$") - sl_a = slice(index_a * n_ell, (index_a + 1) * n_ell) - sl_b = slice(index_b * n_ell, (index_b + 1) * n_ell) + 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) - full_cov[sl_a, sl_b] = block + plt.suptitle("Pseudo-Cl EE (Gaussian covariance)") + plt.legend() + plt.savefig(out_path) - if index_a != index_b: - full_cov[sl_b, sl_a] = block.T + # Plotting EB + out_path = self._output_path("cell_eb.png") - covar.append(fits.ImageHDU(full_cov, name=f"COVAR_{pa}_{pb}")) + fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) - else: - block_path = self._output_path_iNKA_block_cov( - ver, tomo_bin_quad=(bin_key_a1, bin_key_a2, bin_key_b1, bin_key_b2) + 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, ) - covar = fits.open(block_path) - - covar.writeto(out_path, overwrite=True) - return covar - - # ---------------- 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. + ax[0].axhline(0, color="black", linestyle="--") + ax[0].set_ylabel(r"$\ell C_\ell$") - 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[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) - # 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)']" + 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"], ) - 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}" - ) + ax[1].set_xlabel(r"$\ell$") + ax[1].set_ylabel(r"$C_\ell$") - 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")] + 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) - tomo_bins[ver] = {"ids": tomo_bin_ids, "pairs": tomo_bin_pairs} - - n_tomo_bins_plot = max(len(bins["ids"]) for bins in tomo_bins.values()) - - 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, - ) + plt.suptitle("Pseudo-Cl EB (Gaussian covariance)") + plt.legend() + plt.savefig(out_path) - # Draw to extract the yaxis text offset - fig.canvas.draw() + # Plotting BB + out_path = self._output_path("cell_bb.png") - 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$" + fig, ax = plt.subplots(nrows=2, ncols=1, figsize=(8, 8)) - 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", + 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, ) - 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) - - # Setup the legend - plt.subplots_adjust(hspace=0.0, wspace=0.0) # Remove space between subplots - - legend_ax = self._add_grouped_legend(fig, versions, self.cc, pol_list, fmt_dict) - - if savefig is not None: - plt.savefig(savefig, dpi=300, bbox_inches="tight") - if show: - plt.show() - - 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() + ax[0].axhline(0, color="black", linestyle="--") + ax[0].set_ylabel(r"$\ell C_\ell$") - # 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 + 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) - 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, + 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 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, - ) + ax[1].set_xlabel(r"$\ell$") + ax[1].set_ylabel(r"$C_\ell$") - return legend_ax + 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_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. + plt.suptitle("Pseudo-Cl BB (Gaussian covariance)") + plt.legend() + plt.savefig(out_path) - 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 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)})" ) - 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) + # 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), + ) + print(f" Saved BB data to {bb_out}") diff --git a/src/sp_validation/pseudo_cl.py b/src/sp_validation/pseudo_cl.py index 3508e804..c9355ec9 100644 --- a/src/sp_validation/pseudo_cl.py +++ b/src/sp_validation/pseudo_cl.py @@ -16,13 +16,11 @@ 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``. @@ -89,38 +87,7 @@ def make_namaster_bin( return b -# ---------------------- 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 -): +def get_n_gal_map(nside, ra, dec, weights=None): """Weighted galaxy number-density HEALPix map plus pixel bookkeeping. Bins ``(ra, dec)`` (degrees) onto an ``nside`` HEALPix grid. With @@ -131,106 +98,21 @@ def get_n_gal_map( ------- 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. """ - if unique_pix is None or idx is None or idx_rep is None: - unique_pix, idx, idx_rep = get_pixels(ra, dec, nside) + theta = (90.0 - dec) * np.pi / 180.0 + phi = ra * np.pi / 180.0 + pix = hp.ang2pix(nside, theta, phi) + 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 - - -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 + return n_gal, unique_pix, idx, idx_rep def apply_random_rotation(e1, e2, rng=None): @@ -258,493 +140,13 @@ 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_a, - mask_a, + shear_map, + mask, nside, binning, *, - shear_map_b=None, - mask_b=None, - pol_factor=-1, + pol_factor=True, wsp=None, ell_step=10, n_ell_bins=32, @@ -754,20 +156,16 @@ def get_pseudo_cls_map( Parameters ---------- - shear_map_a : np.ndarray + shear_map : np.ndarray Complex shear map (``e1 + 1j * e2``). - mask_a : np.ndarray + mask : 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`. - 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. + pol_factor : bool, optional + If ``True`` flip the sign of the imaginary (e2) component. wsp : nmt.NmtWorkspace, optional Reuse a coupling workspace; built from the field if ``None``. ell_step, n_ell_bins, power : optional @@ -782,27 +180,6 @@ 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( @@ -816,36 +193,18 @@ def get_pseudo_cls_map( ) ell_eff = b.get_effective_ells() - if wsp is None: - 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, - ) + factor = -1 if pol_factor else 1 - cl_coupled, cl_decoupled = compute_cl_from_field_and_workspace( - field_a, field_b, wsp, b + 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) + + cl_coupled = nmt.compute_coupled_cell(f_all, f_all) + cl_all = wsp.decouple_cell(cl_coupled) - return ell_eff, cl_decoupled, wsp + return ell_eff, cl_all, wsp def get_pseudo_cls_catalog( @@ -854,9 +213,7 @@ def get_pseudo_cls_catalog( nside, binning, *, - tomo_bin_a="all", - tomo_bin_b="all", - pol_factor=-1, + pol_factor=True, wsp=None, ell_step=10, n_ell_bins=32, @@ -875,10 +232,8 @@ def get_pseudo_cls_catalog( HEALPix resolution; fixes the harmonic geometry. binning : str Binning scheme passed to :func:`make_namaster_bin`. - 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. + pol_factor : bool, optional + If ``True`` flip the sign of the e2 component. wsp : nmt.NmtWorkspace, optional Reuse a coupling workspace; built from the field if ``None``. ell_step, n_ell_bins, power : optional @@ -893,11 +248,6 @@ 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( @@ -911,140 +261,22 @@ def get_pseudo_cls_catalog( ) ell_eff = b.get_effective_ells() - 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 + factor = -1 if pol_factor else 1 - 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 + 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, ) - 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() + if wsp is None: + wsp = nmt.NmtWorkspace.from_fields(f_all, f_all, b) - # 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]) + cl_coupled = nmt.compute_coupled_cell(f_all, f_all) + cl_all = wsp.decouple_cell(cl_coupled) - return cov_matrix + return ell_eff, cl_all, wsp diff --git a/src/sp_validation/rho_tau.py b/src/sp_validation/rho_tau.py index 95a720f9..6874a08c 100644 --- a/src/sp_validation/rho_tau.py +++ b/src/sp_validation/rho_tau.py @@ -51,10 +51,6 @@ 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/test_pseudo_cl.py b/src/sp_validation/tests/test_pseudo_cl.py index cdaef9d4..45a2366d 100644 --- a/src/sp_validation/tests/test_pseudo_cl.py +++ b/src/sp_validation/tests/test_pseudo_cl.py @@ -58,7 +58,6 @@ 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), @@ -166,7 +165,7 @@ def cv(tmp_path): binning="powspace", power=0.5, n_ell_bins=N_ELL_BINS, - pol_factor=-1, + pol_factor=True, ) cv._test_version = version return cv @@ -266,38 +265,18 @@ 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 - - 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 - ) + n_gal, unique_pix, idx, idx_rep = cv.get_n_gal_map(params, NSIDE, cat_gal) # 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)) @@ -305,6 +284,8 @@ 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( @@ -326,10 +307,7 @@ 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.""" - 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 - ) + n_gal, unique_pix, _idx, idx_rep = cv.get_n_gal_map(params, NSIDE, cat_gal) w = cat_gal[params["w_col"]] e1 = cat_gal[params["e1_col"]] e2 = cat_gal[params["e2_col"]] @@ -415,9 +393,7 @@ def test_get_pseudo_cls_map(cv, cat_and_params): # =========================================================================== 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, tomo_bin_a="all", tomo_bin_b="all" - ) + ell_eff, cl_all, wsp = cv.get_pseudo_cls_catalog(catalog=cat_gal, params=params) assert cl_all.shape == (4, N_ELL_BINS) # Effective ells share the binning math with the map path: bitwise-stable. @@ -495,7 +471,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 = apply_random_rotation(e1, e2) + e1_rot, e2_rot = cv.apply_random_rotation(e1, e2) assert e1_rot.shape == e1.shape assert e2_rot.shape == e2.shape @@ -515,16 +491,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 = apply_random_rotation(e1, e2, np.random.default_rng(42)) - b1, b2 = apply_random_rotation(e1, e2, np.random.default_rng(42)) + 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)) npt.assert_array_equal(a1, b1) npt.assert_array_equal(a2, b2) - c1, _ = apply_random_rotation(e1, e2, np.random.default_rng(7)) + c1, _ = cv.apply_random_rotation(e1, e2, np.random.default_rng(7)) assert not np.allclose(a1, c1) - d1, _ = apply_random_rotation(e1, e2) - f1, _ = apply_random_rotation(e1, e2) + d1, _ = cv.apply_random_rotation(e1, e2) + f1, _ = cv.apply_random_rotation(e1, e2) assert not np.allclose(d1, f1) @@ -539,9 +515,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: {"tomo_bin_all_tomo_bin_all": {}}} + cv._pseudo_cls = {ver: {}} out_path = cv._output_path(f"pseudo_cl_cat_{ver}.fits") - cv.calculate_pseudo_cl_catalog(ver, out_path, tomo_bin_a="all", tomo_bin_b="all") + cv.calculate_pseudo_cl_catalog(ver, out_path) assert os.path.exists(out_path) d = fits.getdata(out_path) @@ -612,7 +588,5 @@ 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, tomo_bin_a="all", tomo_bin_b="all" - ) + _, cl_prim, _ = cv.get_pseudo_cls_catalog(catalog=cat_gal, params=params) npt.assert_allclose(ee, cl_prim[0], rtol=RTOL_CAT, atol=ATOL_CAT)