diff --git a/scripts/xip_xim.py b/scripts/xip_xim.py index 1345f01a..763752ce 100755 --- a/scripts/xip_xim.py +++ b/scripts/xip_xim.py @@ -20,6 +20,8 @@ def params_default(): params = { "input_path": "shape_catalog_ngmix.fits", + "tomo_bin1": None, + "tomo_bin2": None, "key_ra": "RA", "key_dec": "DEC", "key_e1": "e1", @@ -30,20 +32,27 @@ def params_default(): "theta_max": 200, "n_theta": 20, "output_path": "./xip_xim.txt", + "key_tomo_bin_col": "tomo_bin_id", } short_options = { "input_path": "-i", "output_path": "-o", + "tomo_bin1": "-b1", + "tomo_bin2": "-b2", } types = { "sign_e1": "int", "sign_e2": "int", + "tomo_bin1": "int", + "tomo_bin2": "int", } help_strings = { "input_path": "shear catalogue input path, default={}", + "tomo_bin1": "First tomographic bin, default={} (non-tomographic)", + "tomo_bin2": "Second tomographic bin, default={} (non-tomographic)", "key_ra": "column name for right ascension, default={}", "key_dec": "column name for declination, default={}", "key_e1": "column name for ellipticity component 1, default={}", @@ -54,6 +63,7 @@ def params_default(): "theta_max": "maximum angular scale [arcmin], default={}", "n_theta": "number of angular scales, default={}", "output_path": "output path, default={}", + "key_tomo_bin_col": "column name for tomography bin, default={}", } return params, short_options, types, help_strings @@ -132,18 +142,52 @@ def main(argv=None): "Signs for ellipticity components =" + f" ({params['sign_e1']:+d}, {params['sign_e2']:+d})" ) - g1 = data[params["key_e1"]] * params["sign_e1"] - g2 = data[params["key_e2"]] * params["sign_e2"] - cat = treecorr.Catalog( - ra=data[params["key_ra"]], - dec=data[params["key_dec"]], - g1=g1, - g2=g2, - w=data["w"], - ra_units=coord_units, - dec_units=coord_units, - ) - + tomographic = params["tomo_bin1"] is not None and params["tomo_bin2"] is not None + cat2 = None + if tomographic: + if params["verbose"]: + print( + f"Calculating 2PCF for tomographic bins {params['tomo_bin1']} and {params['tomo_bin2']}" + ) + tomo_bin1_idx = data[params["key_tomo_bin_col"]] == params["tomo_bin1"] + tomo_bin2_idx = data[params["key_tomo_bin_col"]] == params["tomo_bin2"] + + g1 = data[params["key_e1"]] * params["sign_e1"] + g2 = data[params["key_e2"]] * params["sign_e2"] + + cat1 = treecorr.Catalog( + ra=data[params["key_ra"]][tomo_bin1_idx], + dec=data[params["key_dec"]][tomo_bin1_idx], + g1=g1[tomo_bin1_idx], + g2=g2[tomo_bin1_idx], + w=data["w"][tomo_bin1_idx], + ra_units=coord_units, + dec_units=coord_units, + ) + if params["tomo_bin1"] != params["tomo_bin2"]: + cat2 = treecorr.Catalog( + ra=data[params["key_ra"]][tomo_bin2_idx], + dec=data[params["key_dec"]][tomo_bin2_idx], + g1=g1[tomo_bin2_idx], + g2=g2[tomo_bin2_idx], + w=data["w"][tomo_bin2_idx], + ra_units=coord_units, + dec_units=coord_units, + ) + else: + if params["verbose"]: + print("Calculating non-tomographic 2PCF") + g1 = data[params["key_e1"]] * params["sign_e1"] + g2 = data[params["key_e2"]] * params["sign_e2"] + cat1 = treecorr.Catalog( + ra=data[params["key_ra"]], + dec=data[params["key_dec"]], + g1=g1, + g2=g2, + w=data["w"], + ra_units=coord_units, + dec_units=coord_units, + ) # Set treecorr config info for correlation sep_units = "arcmin" TreeCorrConfig = { @@ -159,7 +203,7 @@ def main(argv=None): # Compute correlation if params["verbose"]: print("Correlating...") - gg.process(cat, cat) + gg.process(cat1, cat2=cat2) # Write to file if params["verbose"]: diff --git a/src/sp_validation/cosmo_val/cosebis.py b/src/sp_validation/cosmo_val/cosebis.py index aa4657ba..93832f2e 100644 --- a/src/sp_validation/cosmo_val/cosebis.py +++ b/src/sp_validation/cosmo_val/cosebis.py @@ -29,6 +29,7 @@ def calculate_cosebis( cov_path=None, scale_cuts=None, evaluate_all_scale_cuts=False, + compute_tomography=False, min_sep=None, max_sep=None, nbins=None, @@ -68,6 +69,8 @@ def calculate_cosebis( If True, evaluates COSEBIs for all possible scale cut combinations using the reporting binning parameters. Ignored when scale_cuts is provided. Defaults to False. + compute_tomography : bool, optional + If True, computes COSEBIs for all tomographic bin combinations. Defaults to False.s min_sep : float, optional Minimum separation for reporting binning (only used when evaluate_all_scale_cuts=True). Defaults to self.treecorr_config["min_sep"]. @@ -99,44 +102,63 @@ def calculate_cosebis( f"Computing fine-binned 2PCF with {nbins_int} bins from {min_sep_int} to " f"{max_sep_int} arcmin" ) - gg = self.calculate_2pcf(version, npatch=npatch, **treecorr_config) - if scale_cuts is not None: - # Explicit scale cuts provided - print(f"Evaluating {len(scale_cuts)} explicit scale cuts") - results = calculate_cosebis( - gg=gg, nmodes=nmodes, scale_cuts=scale_cuts, cov_path=cov_path - ) - elif evaluate_all_scale_cuts: - # Use reporting binning parameters or inherit from class config - binning = self._binning(min_sep, max_sep, nbins) - min_sep, max_sep, nbins = ( - binning["min_sep"], - binning["max_sep"], - binning["nbins"], - ) + ggs = self.calculate_2pcf( + version, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config, + ) + results = { + bin_key: None for bin_key in ggs.keys() + } # Initialize results dictionary - # Generate scale cuts using np.geomspace (no TreeCorr needed) - bin_edges = np.geomspace(min_sep, max_sep, nbins + 1) - generated_cuts = [ - (bin_edges[start], bin_edges[stop]) - for start in range(nbins) - for stop in range(start + 1, nbins + 1) - ] + for bin_key in ggs.keys(): + # LG TO-DO: account for tomographic scale-cuts + if scale_cuts is not None: + # Explicit scale cuts provided + print(f"Evaluating {len(scale_cuts)} explicit scale cuts") + results[bin_key] = calculate_cosebis( + gg=ggs[bin_key], + nmodes=nmodes, + scale_cuts=scale_cuts, + cov_path=cov_path, + ) + elif evaluate_all_scale_cuts: + # Use reporting binning parameters or inherit from class config + binning = self._binning(min_sep, max_sep, nbins) + min_sep, max_sep, nbins = ( + binning["min_sep"], + binning["max_sep"], + binning["nbins"], + ) - print(f"Evaluating {len(generated_cuts)} scale cut combinations") + # Generate scale cuts using np.geomspace (no TreeCorr needed) + bin_edges = np.geomspace(min_sep, max_sep, nbins + 1) + generated_cuts = [ + (bin_edges[start], bin_edges[stop]) + for start in range(nbins) + for stop in range(start + 1, nbins + 1) + ] - # Call b_modes function with scale cuts list - results = calculate_cosebis( - gg=gg, nmodes=nmodes, scale_cuts=generated_cuts, cov_path=cov_path - ) - else: - # Single scale cut behavior: use full range - results = calculate_cosebis( - gg=gg, nmodes=nmodes, scale_cuts=None, cov_path=cov_path - ) - # Extract single results dict from scale_cuts dictionary - results = next(iter(results.values())) + print(f"Evaluating {len(generated_cuts)} scale cut combinations") + + # Call b_modes function with scale cuts list + results = calculate_cosebis( + gg=ggs[bin_key], + nmodes=nmodes, + scale_cuts=generated_cuts, + cov_path=cov_path, + ) + else: + # Single scale cut behavior: use full range + results = calculate_cosebis( + gg=ggs[bin_key], nmodes=nmodes, scale_cuts=None, cov_path=cov_path + ) + # Extract single results dict from scale_cuts dictionary + results = next(iter(results.values())) + + results[bin_key] = results # Store results for this bin_key return results @@ -152,6 +174,7 @@ def plot_cosebis( cov_path=None, scale_cuts=None, # Explicit scale cuts evaluate_all_scale_cuts=False, # Grid-based scale cuts + compute_tomography=False, min_sep=None, max_sep=None, nbins=None, # Reporting binning @@ -228,6 +251,7 @@ def plot_cosebis( cov_path=cov_path, scale_cuts=scale_cuts, evaluate_all_scale_cuts=evaluate_all_scale_cuts, + compute_tomography=False, # LG: Hardcoded to False for now min_sep=min_sep, max_sep=max_sep, nbins=nbins, diff --git a/src/sp_validation/cosmo_val/pure_eb.py b/src/sp_validation/cosmo_val/pure_eb.py index 7524074e..be41e8d8 100644 --- a/src/sp_validation/cosmo_val/pure_eb.py +++ b/src/sp_validation/cosmo_val/pure_eb.py @@ -29,6 +29,7 @@ def calculate_pure_eb( max_sep_int=300, nbins_int=100, npatch=256, + compute_tomography=False, var_method="jackknife", cov_path_int=None, cosmo_cov=None, @@ -61,6 +62,8 @@ def calculate_pure_eb( npatch : int, optional Number of patches for the jackknife or bootstrap resampling. Defaults to the value in self.npatch if not provided. + compute_tomography : bool, optional + Whether to compute tomographic (cross-bin) pure E/B modes. Defaults to False. var_method : str, optional Variance estimation method. Defaults to "jackknife". cov_path_int : str, optional @@ -109,10 +112,21 @@ def calculate_pure_eb( treecorr_config_int = self._binning(min_sep_int, max_sep_int, nbins_int) # Calculate correlation functions - gg = self.calculate_2pcf(version, npatch=npatch, **treecorr_config) - gg_int = self.calculate_2pcf(version, npatch=npatch, **treecorr_config_int) + ggs = self.calculate_2pcf( + version, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config, + ) + ggs_int = self.calculate_2pcf( + version, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config_int, + ) # Get redshift distribution if using analytic covariance + # LG TO-DO: check tomographic redshift distribution handling z_dist = ( np.column_stack(self.get_redshift(version)) if cov_path_int is not None @@ -120,15 +134,20 @@ def calculate_pure_eb( ) # Delegate to b_modes module - results = calculate_pure_eb_correlation( - gg=gg, - gg_int=gg_int, - var_method=var_method, - cov_path_int=cov_path_int, - cosmo_cov=cosmo_cov, - n_samples=n_samples, - z_dist=z_dist, - ) + results = { + bin_key: None for bin_key in ggs.keys() + } # Initialize results dictionary + + for bin_key in ggs.keys(): + results[bin_key] = calculate_pure_eb_correlation( + gg=ggs[bin_key], + gg_int=ggs_int[bin_key], + var_method=var_method, + cov_path_int=cov_path_int, + cosmo_cov=cosmo_cov, + n_samples=n_samples, + z_dist=z_dist, + ) return results diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 76d05d0a..49f32424 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -12,12 +12,18 @@ import matplotlib.ticker as mticker import numpy as np import treecorr -from astropy.io import fits from cs_util import plots as cs_plots class RealSpaceMixin: - def calculate_2pcf(self, ver, npatch=None, save_fits=False, **treecorr_config): + def calculate_2pcf( + self, + ver, + npatch=None, + compute_tomography=False, + # save_fits=False, + **treecorr_config, + ): """ Calculate the two-point correlation function (2PCF) ξ± for a given catalog version with TreeCorr. @@ -34,6 +40,9 @@ def calculate_2pcf(self, ver, npatch=None, save_fits=False, **treecorr_config): npatch (int, optional): The number of patches to use for the calculation. Defaults to the instance's `npatch` attribute. + compute_tomography (bool, optional): Whether to compute tomographic correlations. + Defaults to False. + save_fits (bool, optional): Whether to save the ξ± results to FITS files. Defaults to False. @@ -53,139 +62,185 @@ def calculate_2pcf(self, ver, npatch=None, save_fits=False, **treecorr_config): headers if `save_fits` is True. """ - self.print_magenta(f"Computing {ver} ξ±") - npatch = npatch or self.npatch treecorr_config = { **self._binning(**treecorr_config), "var_method": "jackknife" if int(npatch) > 1 else "shot", } - gg = treecorr.GGCorrelation(treecorr_config) + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + self.print_magenta( + f"Computing tomographic ξ± for {ver} with {len(tomo_bin_pairs)} bins." + ) - # If the output file already exists, skip the calculation + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError(f"Version {ver} does not have tomography information.") + else: + self.print_magenta(f"Computing non-tomographic ξ± for {ver}.") + tomo_bin_pairs = [("all", "all")] + + ggs = {f"tomo_bin_{b1}_tomo_bin_{b2}": None for b1, b2 in tomo_bin_pairs} + + # LG TO-DO: Change to sacc_io method out_fname = self._output_path( f"{ver}_xi_minsep={treecorr_config['min_sep']}_maxsep={treecorr_config['max_sep']}_nbins={treecorr_config['nbins']}_npatch={npatch}.txt" ) - if os.path.exists(out_fname): - self.print_done(f"Skipping 2PCF calculation, {out_fname} exists") - gg.read(out_fname) + self.print_done(f"Skipping 2PCF calculation, {out_fname} exists.") + for bin1, bin2 in tomo_bin_pairs: + ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = treecorr.GGCorrelation( + treecorr_config + ) + ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"].read(out_fname) else: - # Load data and create a catalog - with self.results[ver].temporarily_read_data(): - g1, g2 = self._calibrated_g(ver) - w = self._read_shear_cols(ver, "w_col") - - # Use patch file if it exists - patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=g2, - w=w, - ra_units=self.treecorr_config["ra_units"], - dec_units=self.treecorr_config["dec_units"], - npatch=npatch, - patch_centers=patch_file if os.path.exists(patch_file) else None, - ) + patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - # If no patch file exists, save the current patches - if not os.path.exists(patch_file): - cat_gal.write_patch_centers(patch_file) + for bin1, bin2 in tomo_bin_pairs: + gg = treecorr.GGCorrelation(treecorr_config) - # Process the catalog & write the correlation functions - gg.process(cat_gal) - gg.write(out_fname, write_patch_results=True, write_cov=True) + # Load data and create a catalog + with self.results[ver].temporarily_read_data(): + if bin1 == "all" and bin2 == "all": + mask_bin1 = np.ones( + len(self.results[ver].dat_shear), dtype=bool + ) + else: + mask_bin1 = self.results[ver].dat_shear["tom_bin_id"] == bin1 + + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") + + cat_gal1 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][mask_bin1], + dec=self.results[ver].dat_shear["Dec"][mask_bin1], + g1=g1[mask_bin1], + g2=g2[mask_bin1], + w=w[mask_bin1], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + patch_centers=patch_file + if os.path.exists(patch_file) + else None, + ) + cat_gal2 = None + + if bin1 != bin2: + mask_bin2 = self.results[ver].dat_shear["tom_bin_id"] == bin2 + cat_gal2 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][mask_bin2], + dec=self.results[ver].dat_shear["Dec"][mask_bin2], + g1=g1[mask_bin2], + g2=g2[mask_bin2], + w=w[mask_bin2], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + patch_centers=patch_file + if os.path.exists(patch_file) + else None, + ) + + # If no patch file exists, save the current patches + if not os.path.exists(patch_file): + cat_gal1.write_patch_centers(patch_file) + + # Process the catalog & write the correlation functions + gg.process(cat_gal1, cat2=cat_gal2) + ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = gg + # LG TO-DO: No longer writing out text file, change to sacc_io method + # gg.write(out_fname, write_patch_results=True, write_cov=True) + + # LG: FITS writeout function deprecated, now writing to SACC format # Save xi_p and xi_m results to fits file # (moved outside so it runs even if txt exists) - if save_fits: - lst = np.arange(1, treecorr_config["nbins"] + 1) - - col1 = fits.Column(name="BIN1", format="K", array=np.ones(len(lst))) - col2 = fits.Column(name="BIN2", format="K", array=np.ones(len(lst))) - col3 = fits.Column(name="ANGBIN", format="K", array=lst) - col4 = fits.Column(name="VALUE", format="D", array=gg.xip) - col5 = fits.Column(name="ANG", format="D", unit="arcmin", array=gg.meanr) - coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) - xiplus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_PLUS") - - col4 = fits.Column(name="VALUE", format="D", array=gg.xim) - coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) - ximinus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_MINUS") - - # append xi_plus header info - xiplus_dict = { - "2PTDATA": "T", - "QUANT1": "G+R", - "QUANT2": "G+R", - "KERNEL_1": "NZ_SOURCE", - "KERNEL_2": "NZ_SOURCE", - "WINDOWS": "SAMPLE", - } - for key in xiplus_dict: - xiplus_hdu.header[key] = xiplus_dict[key] - - col1 = fits.Column(name="BIN1", format="K", array=np.ones(len(lst))) - col2 = fits.Column(name="BIN2", format="K", array=np.ones(len(lst))) - col3 = fits.Column(name="ANGBIN", format="K", array=lst) - col4 = fits.Column(name="VALUE", format="D", array=gg.xip) - col5 = fits.Column(name="ANG", format="D", unit="arcmin", array=gg.rnom) - coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) - xiplus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_PLUS") - - col4 = fits.Column(name="VALUE", format="D", array=gg.xim) - coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) - ximinus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_MINUS") - - # append xi_plus header info - xiplus_dict = { - "2PTDATA": "T", - "QUANT1": "G+R", - "QUANT2": "G+R", - "KERNEL_1": "NZ_SOURCE", - "KERNEL_2": "NZ_SOURCE", - "WINDOWS": "SAMPLE", - } - for key in xiplus_dict: - xiplus_hdu.header[key] = xiplus_dict[key] - # Use same naming format as txt output - fits_base = out_fname.replace(".txt", "").replace("_xi_", "_") - xiplus_hdu.writeto( - f"{fits_base.replace(ver, f'xi_plus_{ver}')}.fits", - overwrite=True, - ) - - # append xi_minus header info - ximinus_dict = {**xiplus_dict, "QUANT1": "G-R", "QUANT2": "G-R"} - for key in ximinus_dict: - ximinus_hdu.header[key] = ximinus_dict[key] - ximinus_hdu.writeto( - f"{fits_base.replace(ver, f'xi_minus_{ver}')}.fits", - overwrite=True, - ) + # if save_fits: + # lst = np.arange(1, treecorr_config["nbins"] + 1) + + # col1 = fits.Column(name="BIN1", format="K", array=np.ones(len(lst))) + # col2 = fits.Column(name="BIN2", format="K", array=np.ones(len(lst))) + # col3 = fits.Column(name="ANGBIN", format="K", array=lst) + # col4 = fits.Column(name="VALUE", format="D", array=gg.xip) + # col5 = fits.Column(name="ANG", format="D", unit="arcmin", array=gg.meanr) + # coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) + # xiplus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_PLUS") + + # col4 = fits.Column(name="VALUE", format="D", array=gg.xim) + # coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) + # ximinus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_MINUS") + + # # append xi_plus header info + # xiplus_dict = { + # "2PTDATA": "T", + # "QUANT1": "G+R", + # "QUANT2": "G+R", + # "KERNEL_1": "NZ_SOURCE", + # "KERNEL_2": "NZ_SOURCE", + # "WINDOWS": "SAMPLE", + # } + # for key in xiplus_dict: + # xiplus_hdu.header[key] = xiplus_dict[key] + + # col1 = fits.Column(name="BIN1", format="K", array=np.ones(len(lst))) + # col2 = fits.Column(name="BIN2", format="K", array=np.ones(len(lst))) + # col3 = fits.Column(name="ANGBIN", format="K", array=lst) + # col4 = fits.Column(name="VALUE", format="D", array=gg.xip) + # col5 = fits.Column(name="ANG", format="D", unit="arcmin", array=gg.rnom) + # coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) + # xiplus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_PLUS") + + # col4 = fits.Column(name="VALUE", format="D", array=gg.xim) + # coldefs = fits.ColDefs([col1, col2, col3, col4, col5]) + # ximinus_hdu = fits.BinTableHDU.from_columns(coldefs, name="XI_MINUS") + + # # append xi_plus header info + # xiplus_dict = { + # "2PTDATA": "T", + # "QUANT1": "G+R", + # "QUANT2": "G+R", + # "KERNEL_1": "NZ_SOURCE", + # "KERNEL_2": "NZ_SOURCE", + # "WINDOWS": "SAMPLE", + # } + # for key in xiplus_dict: + # xiplus_hdu.header[key] = xiplus_dict[key] + # # Use same naming format as txt output + # fits_base = out_fname.replace(".txt", "").replace("_xi_", "_") + # xiplus_hdu.writeto( + # f"{fits_base.replace(ver, f'xi_plus_{ver}')}.fits", + # overwrite=True, + # ) + + # # append xi_minus header info + # ximinus_dict = {**xiplus_dict, "QUANT1": "G-R", "QUANT2": "G-R"} + # for key in ximinus_dict: + # ximinus_hdu.header[key] = ximinus_dict[key] + # ximinus_hdu.writeto( + # f"{fits_base.replace(ver, f'xi_minus_{ver}')}.fits", + # overwrite=True, + # ) # Add correlation object to class if not hasattr(self, "cat_ggs"): self.cat_ggs = {} - self.cat_ggs[ver] = gg + self.cat_ggs[ver] = ggs - self.print_done("Done 2PCF") + self.print_done(f"Done 2PCF for {ver}.") - return gg + return ggs def plot_2pcf(self): # Plot of n_pairs plt.subplots(ncols=1, nrows=1) for ver in self.versions: + # FORCE non-tomography for now self.calculate_2pcf(ver) plt.plot( - self.cat_ggs[ver].meanr, - self.cat_ggs[ver].npairs, + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].npairs, label=ver, ls=self.cc[ver]["ls"], color=self.cc[ver]["colour"], @@ -202,9 +257,10 @@ def plot_2pcf(self): plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) for idx, ver in enumerate(self.versions): plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, fx=1.05, nx=len(ver)), - self.cat_ggs[ver].xip, - yerr=np.sqrt(self.cat_ggs[ver].varxip), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, fx=1.05, nx=len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xip, + yerr=np.sqrt(self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxip), label=ver, ls=self.cc[ver]["ls"], color=self.cc[ver]["colour"], @@ -225,9 +281,10 @@ def plot_2pcf(self): plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) for idx, ver in enumerate(self.versions): plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, fx=1.05, nx=len(ver)), - self.cat_ggs[ver].xim, - yerr=np.sqrt(self.cat_ggs[ver].varxim), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, fx=1.05, nx=len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xim, + yerr=np.sqrt(self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxim), label=ver, ls=self.cc[ver]["ls"], color=self.cc[ver]["colour"], @@ -248,9 +305,11 @@ def plot_2pcf(self): plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) for idx, ver in enumerate(self.versions): plt.errorbar( - self.cat_ggs[ver].meanr, - self.cat_ggs[ver].xip * self.cat_ggs[ver].meanr, - yerr=np.sqrt(self.cat_ggs[ver].varxip) * self.cat_ggs[ver].meanr, + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xip + * self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, + yerr=np.sqrt(self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxip) + * self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, label=ver, ls=self.cc[ver]["ls"], color=self.cc[ver]["colour"], @@ -270,9 +329,12 @@ def plot_2pcf(self): plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) for idx, ver in enumerate(self.versions): plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, len(ver)), - self.cat_ggs[ver].xim * self.cat_ggs[ver].meanr, - yerr=np.sqrt(self.cat_ggs[ver].varxim) * self.cat_ggs[ver].meanr, + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xim + * self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, + yerr=np.sqrt(self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxim) + * self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr, label=ver, ls=self.cc[ver]["ls"], color=self.cc[ver]["colour"], @@ -294,15 +356,17 @@ def plot_2pcf(self): for idx, ver in enumerate(self.versions): plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, len(ver)), - self.cat_ggs[ver].xip, - yerr=np.sqrt(self.cat_ggs[ver].varxim), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xip, + yerr=np.sqrt(self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxip), label=r"$\xi_+$", ls="solid", color="green", ) plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, len(ver)), self.xi_psf_sys[ver]["mean"], yerr=np.sqrt(self.xi_psf_sys[ver]["var"]), label=r"$\xi^{\rm psf}_{+, {\rm sys}}$", @@ -310,16 +374,18 @@ def plot_2pcf(self): color="red", ) plt.errorbar( - self.cat_ggs[ver].meanr * cs_plots.dx(idx, len(ver)), - self.cat_ggs[ver].xip + self.xi_psf_sys[ver]["mean"], + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].meanr + * cs_plots.dx(idx, len(ver)), + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].xip + + self.xi_psf_sys[ver]["mean"], yerr=np.sqrt( - self.cat_ggs[ver].varxip + self.xi_psf_sys[ver]["var"] + self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"].varxip + + self.xi_psf_sys[ver]["var"] ), label=r"$\xi_+ + \xi^{\rm psf}_{+, {\rm sys}}$", ls="dashdot", color="magenta", ) - plt.xscale("log") plt.yscale("log") plt.legend() @@ -340,7 +406,7 @@ def plot_ratio_xi_sys_xi(self, threshold=0.1, offset=0.02): for idx, ver in enumerate(self.versions): self.calculate_2pcf(ver) xi_psf_sys = self.xi_psf_sys[ver] - gg = self.cat_ggs[ver] + gg = self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"] ratio = xi_psf_sys["mean"] / gg.xip ratio_err = np.sqrt( @@ -394,9 +460,14 @@ def plot_ratio_xi_sys_xi(self, threshold=0.1, offset=0.02): print(f"Ratio of xi_psf_sys to xi plot saved to {out_path}") def calculate_aperture_mass_dispersion( - self, theta_min=0.3, theta_max=200, nbins=500, nbins_map=15, npatch=25 + self, + theta_min=0.3, + theta_max=200, + nbins=500, + nbins_map=15, + npatch=25, + compute_tomography=False, ): - self.print_start("Computing aperture-mass dispersion") self._map2 = {} theta_map = np.geomspace(theta_min * 5, theta_max / 2, nbins_map) @@ -405,53 +476,91 @@ def calculate_aperture_mass_dispersion( treecorr_config = self._binning(theta_min, theta_max, nbins) for ver in self.versions: - self.print_magenta(ver) + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + self.print_magenta( + f"Computing MAP for {ver} with {len(tomo_bin_pairs)} bins." + ) + + if tomo_bin_ids is None or tomo_bin_pairs is None: + raise ValueError( + f"Version {ver} does not have tomography information." + ) + else: + self.print_magenta(f"Computing non-tomographic MAP for {ver}.") + + tomo_bin_pairs = [("all", "all")] - gg = treecorr.GGCorrelation(treecorr_config) + self._map2.setdefault(ver, {}) + # LG TO-DO: Change to sacc_io method out_fname = self._output_path(f"xi_for_map2_{ver}.txt") if os.path.exists(out_fname): self.print_green(f"Skipping xi for Map2, {out_fname} exists") - gg.read(out_fname) - else: - with self.results[ver].temporarily_read_data(): - g1, g2 = self._calibrated_g(ver) - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=g2, - w=self._read_shear_cols(ver, "w_col"), - ra_units=self.treecorr_config["ra_units"], - dec_units=self.treecorr_config["dec_units"], - npatch=npatch, + for bin1, bin2 in tomo_bin_pairs: + self._map2[ver][f"tomo_bin_{bin1}_{bin2}"] = treecorr.GGCorrelation( + treecorr_config ) + self._map2[ver][f"tomo_bin_{bin1}_{bin2}"].read(out_fname) - gg.process(cat_gal) - gg.write(out_fname) - del cat_gal - del g1 - del g2 - - mapsq, mapsq_im, mxsq, mxsq_im, varmapsq = gg.calculateMapSq( - R=theta_map, - m2_uform="Schneider", - ) - out_fname_map2 = self._output_path(f"map2_{ver}.txt") - if os.path.exists(out_fname_map2): - self.print_green(f"Skipping Map2, {out_fname_map2} exists") else: - print(f"Writing Map2 to output file {out_fname_map2} ") - gg.writeMapSq(out_fname_map2, R=theta_map, m2_uform="Schneider") - self._map2[ver] = { - "mapsq": mapsq, - "mapsq_im": mapsq_im, - "mxsq": mxsq, - "mxsq_im": mxsq_im, - "varmapsq": varmapsq, - } - - self.print_done("Done aperture-mass dispersion") + for bin1, bin2 in tomo_bin_pairs: + gg = treecorr.GGCorrelation(treecorr_config) + + # Load data and create a catalog + with self.results[ver].temporarily_read_data(): + if bin1 == "all" and bin2 == "all": + mask_bin1 = np.ones( + len(self.results[ver].dat_shear), dtype=bool + ) + else: + mask_bin1 = ( + self.results[ver].dat_shear["tom_bin_id"] == bin1 + ) + + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") + + cat_gal1 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][mask_bin1], + dec=self.results[ver].dat_shear["Dec"][mask_bin1], + g1=g1[mask_bin1], + g2=g2[mask_bin1], + w=w[mask_bin1], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + ) + cat_gal2 = None + + if bin1 != bin2: + mask_bin2 = ( + self.results[ver].dat_shear["tom_bin_id"] == bin2 + ) + cat_gal2 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][mask_bin2], + dec=self.results[ver].dat_shear["Dec"][mask_bin2], + g1=g1[mask_bin2], + g2=g2[mask_bin2], + w=w[mask_bin2], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + ) + gg.process(cat_gal1, cat2=cat_gal2) + + mapsq, mapsq_im, mxsq, mxsq_im, varmapsq = gg.calculateMapSq( + R=theta_map, + m2_uform="Schneider", + ) + self._map2[ver][f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = { + "mapsq": mapsq, + "mapsq_im": mapsq_im, + "mxsq": mxsq, + "mxsq_im": mxsq_im, + "varmapsq": varmapsq, + } + self.print_done(f"Done aperture-mass dispersion for {ver}.") @property def map2(self): @@ -462,8 +571,15 @@ def map2(self): def plot_aperture_mass_dispersion(self): for mode in ["mapsq", "mapsq_im", "mxsq", "mxsq_im"]: x = [self.map2["theta_map"] for ver in self.versions] - y = [self.map2[ver][mode] for ver in self.versions] - yerr = [np.sqrt(self.map2[ver]["varmapsq"]) for ver in self.versions] + # FORCE non-tomography for now + y = [ + self.map2[ver]["tomo_bin_all_tomo_bin_all"][mode] + for ver in self.versions + ] + yerr = [ + np.sqrt(self.map2[ver]["tomo_bin_all_tomo_bin_all"]["varmapsq"]) + for ver in self.versions + ] labels = list(self.versions) colors = [self.cc[ver]["colour"] for ver in self.versions] linestyles = [self.cc[ver]["ls"] for ver in self.versions] @@ -494,8 +610,15 @@ def plot_aperture_mass_dispersion(self): for mode in ["mapsq", "mapsq_im", "mxsq", "mxsq_im"]: x = [self.map2["theta_map"] for ver in self.versions] - y = [np.abs(self.map2[ver][mode]) for ver in self.versions] - yerr = [np.sqrt(self.map2[ver]["varmapsq"]) for ver in self.versions] + # FORCE non-tomography for now + y = [ + np.abs(self.map2[ver]["tomo_bin_all_tomo_bin_all"][mode]) + for ver in self.versions + ] + yerr = [ + np.sqrt(self.map2[ver]["tomo_bin_all_tomo_bin_all"]["varmapsq"]) + for ver in self.versions + ] xlabel = r"$\theta$ [arcmin]" ylabel = "dispersion" title = f"Aperture-mass dispersion mode {mode}" diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index f50992d4..2e34c6c9 100644 --- a/src/sp_validation/tests/test_cosmo_val.py +++ b/src/sp_validation/tests/test_cosmo_val.py @@ -512,13 +512,14 @@ def test_calculate_2pcf_runs_on_synthetic_catalog(self, tmp_path): **params, ) - gg = cv.calculate_2pcf(version) + ggs = cv.calculate_2pcf(version) # treecorr GGCorrelation with xi+/- on the configured angular grid - assert gg.xip.shape == (nbins,) - assert gg.xim.shape == (nbins,) - assert np.all(np.isfinite(gg.xip)) - assert np.all(np.isfinite(gg.xim)) + # LG:Hardcoded to be non-tomographic for now + assert ggs["tomo_bin_all_tomo_bin_all"].xip.shape == (nbins,) + assert ggs["tomo_bin_all_tomo_bin_all"].xim.shape == (nbins,) + assert np.all(np.isfinite(ggs["tomo_bin_all_tomo_bin_all"].xip)) + assert np.all(np.isfinite(ggs["tomo_bin_all_tomo_bin_all"].xim)) # The additive-bias subtraction in the pipeline must have run. assert version in cv.c1 and version in cv.c2