From 122bc97cbf96dd76b1816c060e0e71ab5c46319f Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Wed, 22 Jul 2026 17:51:11 +0200 Subject: [PATCH 1/6] added tomographic 2pcf implementation --- scripts/xip_xim.py | 72 +++- src/sp_validation/cosmo_val/real_space.py | 463 +++++++++++++++------- 2 files changed, 374 insertions(+), 161 deletions(-) diff --git a/scripts/xip_xim.py b/scripts/xip_xim.py index 1345f01a..88261266 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,56 @@ 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( + f"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 +207,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/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 76d05d0a..0dbc51ee 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -17,7 +17,12 @@ class RealSpaceMixin: - def calculate_2pcf(self, ver, npatch=None, save_fits=False, **treecorr_config): + def calculate_2pcf(self, ver, + npatch=None, + tomo_bin1 = None, + tomo_bin2 = None, + # save_fits=False, + **treecorr_config): """ Calculate the two-point correlation function (2PCF) ξ± for a given catalog version with TreeCorr. @@ -33,6 +38,12 @@ 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. + + tomo_bin1 (int, optional): The first tomographic bin to use for the calculation. + If None, the calculation is non-tomographic. Defaults to None. + + tomo_bin2 (int, optional): The second tomographic bin to use for the calculation. + If None, the calculation is non-tomographic. Defaults to None. save_fits (bool, optional): Whether to save the ξ± results to FITS files. Defaults to False. @@ -52,129 +63,197 @@ def calculate_2pcf(self, ver, npatch=None, save_fits=False, **treecorr_config): - FITS files for ξ+ and ξ− are saved with additional metadata in their 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", } - + pol_factor = self.pol_factor gg = treecorr.GGCorrelation(treecorr_config) - - # If the output file already exists, skip the calculation - out_fname = self._output_path( + + if tomo_bin1 is None and tomo_bin2 is None: + self.print_magenta(f"Computing non-tomographic ξ± for {ver}.") + # 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) - if os.path.exists(out_fname): - self.print_done(f"Skipping 2PCF calculation, {out_fname} exists") - gg.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") - 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, - ) + #LG: need tomographic patch file? + 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) + cat_gal = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"], + dec=self.results[ver].dat_shear["Dec"], + g1=g1, + g2=pol_factor * 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, + ) - # Process the catalog & write the correlation functions - gg.process(cat_gal) - gg.write(out_fname, write_patch_results=True, write_cov=True) + # If no patch file exists, save the current patches + if not os.path.exists(patch_file): + cat_gal.write_patch_centers(patch_file) - # 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, - ) + # Process the catalog & write the correlation functions + gg.process(cat_gal) + # 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) + else: + self.print_magenta(f"Computing tomographic ξ± for {ver}, bin {tomo_bin1} and {tomo_bin2}.") - # 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, + # LG TO-DO: Change to sacc_io method + out_fname = self._output_path( + f"{ver}_xi_tomo_bins={tomo_bin1}_{tomo_bin2}_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) + + else: + with self.results[ver].temporarily_read_data(): + tomo_bin_idx_bin1 = self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") + + #LG: need tomographic patch file? + patch_file_bin1 = self._output_path(f"{ver}_patches_npatch={npatch}_bin={tomo_bin1}.dat") + + cat_gal_bin1 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin1], + dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin1], + g1=g1[tomo_bin_idx_bin1], + g2=pol_factor * g2[tomo_bin_idx_bin1], + w=w[tomo_bin_idx_bin1], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + patch_centers=patch_file_bin1 if os.path.exists(patch_file_bin1) else None, + ) + + if tomo_bin1 == tomo_bin2: + gg.process(cat_gal_bin1) + # 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) + + if not os.path.exists(patch_file_bin1): + cat_gal_bin1.write_patch_centers(patch_file_bin1) + + else: + tomo_bin_idx_bin2 = self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 + + #LG: need tomographic patch file? + patch_file_bin2 = self._output_path(f"{ver}_patches_npatch={npatch}_bin={tomo_bin2}.dat") + + cat_gal_bin2 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], + dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin2], + g1=g1[tomo_bin_idx_bin2], + g2=pol_factor * g2[tomo_bin_idx_bin2], + w=w[tomo_bin_idx_bin2], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + patch_centers=patch_file_bin2 if os.path.exists(patch_file_bin2) else None, + ) + + gg.process(cat_gal_bin1, cat2=cat_gal_bin2) + # 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) + + if not os.path.exists(patch_file_bin1) or not os.path.exists(patch_file_bin2): + cat_gal_bin1.write_patch_centers(patch_file_bin1) + cat_gal_bin2.write_patch_centers(patch_file_bin2) + + # 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, + # ) + # Add correlation object to class if not hasattr(self, "cat_ggs"): self.cat_ggs = {} self.cat_ggs[ver] = gg - self.print_done("Done 2PCF") + self.print_done(f"Done 2PCF for bins {tomo_bin1} and {tomo_bin2} for {ver}.") return gg @@ -182,6 +261,7 @@ 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, @@ -394,7 +474,8 @@ 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, + tomo_bin1=None, tomo_bin2=None ): self.print_start("Computing aperture-mass dispersion") @@ -403,55 +484,137 @@ def calculate_aperture_mass_dispersion( self._map2["theta_map"] = theta_map treecorr_config = self._binning(theta_min, theta_max, nbins) + pol_factor = self.pol_factor for ver in self.versions: self.print_magenta(ver) gg = treecorr.GGCorrelation(treecorr_config) + self._map2.setdefault(ver, {}) + + if tomo_bin1 is None and tomo_bin2 is None: + self.print_magenta("Computing non-tomographic aperture-mass dispersion.") + # 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=pol_factor * 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, + ) + + 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]["1_1"] = { + "mapsq": mapsq, + "mapsq_im": mapsq_im, + "mxsq": mxsq, + "mxsq_im": mxsq_im, + "varmapsq": varmapsq, + } - 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, - ) - - gg.process(cat_gal) - gg.write(out_fname) - del cat_gal - del g1 - del g2 + self.print_magenta( + f"Computing tomographic aperture-mass dispersion for bins {tomo_bin1} and {tomo_bin2}." + ) + # LG TO-DO: Change to sacc_io method + out_fname = self._output_path( + f"xi_for_map2_{ver}_tomo_bins={tomo_bin1}_{tomo_bin2}.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(): + tomo_bin_idx_bin1 = ( + self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 + ) + g1, g2 = self._calibrated_g(ver) + cat_gal_bin1 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin1], + dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin1], + g1=g1[tomo_bin_idx_bin1], + g2=pol_factor * g2[tomo_bin_idx_bin1], + w=self._read_shear_cols(ver, "w_col")[tomo_bin_idx_bin1], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + ) + + if tomo_bin1 == tomo_bin2: + gg.process(cat_gal_bin1) + # gg.write(out_fname) + del cat_gal_bin1 + del g1 + del g2 + + else: + tomo_bin_idx_bin2 = ( + self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 + ) + cat_gal_bin2 = treecorr.Catalog( + ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], + dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin2], + g1=g1[tomo_bin_idx_bin2], + g2=pol_factor * g2[tomo_bin_idx_bin2], + w=self._read_shear_cols(ver, "w_col")[tomo_bin_idx_bin2], + ra_units=self.treecorr_config["ra_units"], + dec_units=self.treecorr_config["dec_units"], + npatch=npatch, + ) + + gg.process(cat_gal_bin1, cat2=cat_gal_bin2) + # gg.write(out_fname) + del cat_gal_bin1 + del cat_gal_bin2 + 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}_tomo_bins={tomo_bin1}_{tomo_bin2}.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][f"{tomo_bin1}_{tomo_bin2}"] = { + "mapsq": mapsq, + "mapsq_im": mapsq_im, + "mxsq": mxsq, + "mxsq_im": mxsq_im, + "varmapsq": varmapsq, + } - 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") + self.print_done(f"Done aperture-mass dispersion for {ver}.") @property def map2(self): @@ -462,8 +625,9 @@ 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]["1_1"][mode] for ver in self.versions] + yerr = [np.sqrt(self.map2[ver]["1_1"]["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 +658,9 @@ 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]["1_1"][mode]) for ver in self.versions] + yerr = [np.sqrt(self.map2[ver]["1_1"]["varmapsq"]) for ver in self.versions] xlabel = r"$\theta$ [arcmin]" ylabel = "dispersion" title = f"Aperture-mass dispersion mode {mode}" From a63afa8349757444529eaf187cc44dacd410f73c Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Wed, 22 Jul 2026 17:51:57 +0200 Subject: [PATCH 2/6] ruff pass --- scripts/xip_xim.py | 12 +-- src/sp_validation/cosmo_val/real_space.py | 108 ++++++++++++++-------- 2 files changed, 74 insertions(+), 46 deletions(-) diff --git a/scripts/xip_xim.py b/scripts/xip_xim.py index 88261266..763752ce 100755 --- a/scripts/xip_xim.py +++ b/scripts/xip_xim.py @@ -142,9 +142,7 @@ def main(argv=None): "Signs for ellipticity components =" + f" ({params['sign_e1']:+d}, {params['sign_e2']:+d})" ) - tomographic = ( - params["tomo_bin1"] is not None and params["tomo_bin2"] is not None - ) + tomographic = params["tomo_bin1"] is not None and params["tomo_bin2"] is not None cat2 = None if tomographic: if params["verbose"]: @@ -153,10 +151,10 @@ def main(argv=None): ) 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], @@ -178,9 +176,7 @@ def main(argv=None): ) else: if params["verbose"]: - print( - f"Calculating non-tomographic 2PCF" - ) + print("Calculating non-tomographic 2PCF") g1 = data[params["key_e1"]] * params["sign_e1"] g2 = data[params["key_e2"]] * params["sign_e2"] cat1 = treecorr.Catalog( diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 0dbc51ee..3eaef4ba 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -12,17 +12,19 @@ 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, - tomo_bin1 = None, - tomo_bin2 = None, - # save_fits=False, - **treecorr_config): + def calculate_2pcf( + self, + ver, + npatch=None, + tomo_bin1=None, + tomo_bin2=None, + # save_fits=False, + **treecorr_config, + ): """ Calculate the two-point correlation function (2PCF) ξ± for a given catalog version with TreeCorr. @@ -38,10 +40,10 @@ def calculate_2pcf(self, ver, npatch (int, optional): The number of patches to use for the calculation. Defaults to the instance's `npatch` attribute. - + tomo_bin1 (int, optional): The first tomographic bin to use for the calculation. If None, the calculation is non-tomographic. Defaults to None. - + tomo_bin2 (int, optional): The second tomographic bin to use for the calculation. If None, the calculation is non-tomographic. Defaults to None. @@ -63,7 +65,7 @@ def calculate_2pcf(self, ver, - FITS files for ξ+ and ξ− are saved with additional metadata in their headers if `save_fits` is True. """ - + npatch = npatch or self.npatch treecorr_config = { **self._binning(**treecorr_config), @@ -71,13 +73,13 @@ def calculate_2pcf(self, ver, } pol_factor = self.pol_factor gg = treecorr.GGCorrelation(treecorr_config) - + if tomo_bin1 is None and tomo_bin2 is None: self.print_magenta(f"Computing non-tomographic ξ± for {ver}.") # 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" - ) + 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) @@ -88,7 +90,7 @@ def calculate_2pcf(self, ver, g1, g2 = self._calibrated_g(ver) w = self._read_shear_cols(ver, "w_col") - #LG: need tomographic patch file? + # LG: need tomographic patch file? patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") cat_gal = treecorr.Catalog( @@ -100,7 +102,9 @@ def calculate_2pcf(self, ver, 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_centers=patch_file + if os.path.exists(patch_file) + else None, ) # If no patch file exists, save the current patches @@ -112,7 +116,9 @@ def calculate_2pcf(self, ver, # 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) else: - self.print_magenta(f"Computing tomographic ξ± for {ver}, bin {tomo_bin1} and {tomo_bin2}.") + self.print_magenta( + f"Computing tomographic ξ± for {ver}, bin {tomo_bin1} and {tomo_bin2}." + ) # LG TO-DO: Change to sacc_io method out_fname = self._output_path( @@ -122,15 +128,19 @@ def calculate_2pcf(self, ver, if os.path.exists(out_fname): self.print_done(f"Skipping 2PCF calculation, {out_fname} exists.") gg.read(out_fname) - + else: with self.results[ver].temporarily_read_data(): - tomo_bin_idx_bin1 = self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 + tomo_bin_idx_bin1 = ( + self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 + ) g1, g2 = self._calibrated_g(ver) w = self._read_shear_cols(ver, "w_col") - - #LG: need tomographic patch file? - patch_file_bin1 = self._output_path(f"{ver}_patches_npatch={npatch}_bin={tomo_bin1}.dat") + + # LG: need tomographic patch file? + patch_file_bin1 = self._output_path( + f"{ver}_patches_npatch={npatch}_bin={tomo_bin1}.dat" + ) cat_gal_bin1 = treecorr.Catalog( ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin1], @@ -141,22 +151,28 @@ def calculate_2pcf(self, ver, ra_units=self.treecorr_config["ra_units"], dec_units=self.treecorr_config["dec_units"], npatch=npatch, - patch_centers=patch_file_bin1 if os.path.exists(patch_file_bin1) else None, + patch_centers=patch_file_bin1 + if os.path.exists(patch_file_bin1) + else None, ) - + if tomo_bin1 == tomo_bin2: gg.process(cat_gal_bin1) # 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) - + if not os.path.exists(patch_file_bin1): cat_gal_bin1.write_patch_centers(patch_file_bin1) else: - tomo_bin_idx_bin2 = self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 - - #LG: need tomographic patch file? - patch_file_bin2 = self._output_path(f"{ver}_patches_npatch={npatch}_bin={tomo_bin2}.dat") + tomo_bin_idx_bin2 = ( + self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 + ) + + # LG: need tomographic patch file? + patch_file_bin2 = self._output_path( + f"{ver}_patches_npatch={npatch}_bin={tomo_bin2}.dat" + ) cat_gal_bin2 = treecorr.Catalog( ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], @@ -167,19 +183,23 @@ def calculate_2pcf(self, ver, ra_units=self.treecorr_config["ra_units"], dec_units=self.treecorr_config["dec_units"], npatch=npatch, - patch_centers=patch_file_bin2 if os.path.exists(patch_file_bin2) else None, + patch_centers=patch_file_bin2 + if os.path.exists(patch_file_bin2) + else None, ) - + gg.process(cat_gal_bin1, cat2=cat_gal_bin2) # 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) - if not os.path.exists(patch_file_bin1) or not os.path.exists(patch_file_bin2): + if not os.path.exists(patch_file_bin1) or not os.path.exists( + patch_file_bin2 + ): cat_gal_bin1.write_patch_centers(patch_file_bin1) cat_gal_bin2.write_patch_centers(patch_file_bin2) - + # 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: @@ -474,8 +494,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, - tomo_bin1=None, tomo_bin2=None + self, + theta_min=0.3, + theta_max=200, + nbins=500, + nbins_map=15, + npatch=25, + tomo_bin1=None, + tomo_bin2=None, ): self.print_start("Computing aperture-mass dispersion") @@ -493,7 +519,9 @@ def calculate_aperture_mass_dispersion( self._map2.setdefault(ver, {}) if tomo_bin1 is None and tomo_bin2 is None: - self.print_magenta("Computing non-tomographic aperture-mass dispersion.") + self.print_magenta( + "Computing non-tomographic aperture-mass dispersion." + ) # 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): @@ -578,10 +606,14 @@ def calculate_aperture_mass_dispersion( ) cat_gal_bin2 = treecorr.Catalog( ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], - dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin2], + dec=self.results[ver].dat_shear["Dec"][ + tomo_bin_idx_bin2 + ], g1=g1[tomo_bin_idx_bin2], g2=pol_factor * g2[tomo_bin_idx_bin2], - w=self._read_shear_cols(ver, "w_col")[tomo_bin_idx_bin2], + w=self._read_shear_cols(ver, "w_col")[ + tomo_bin_idx_bin2 + ], ra_units=self.treecorr_config["ra_units"], dec_units=self.treecorr_config["dec_units"], npatch=npatch, From 65025089c8dad4f20a39cf64e317418294458a72 Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Thu, 23 Jul 2026 18:44:09 +0200 Subject: [PATCH 3/6] edits after review --- src/sp_validation/cosmo_val/cosebis.py | 92 +++-- src/sp_validation/cosmo_val/pure_eb.py | 41 +- src/sp_validation/cosmo_val/real_space.py | 438 +++++++++------------- src/sp_validation/tests/test_cosmo_val.py | 11 +- 4 files changed, 276 insertions(+), 306 deletions(-) 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 3eaef4ba..49f32424 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -20,8 +20,7 @@ def calculate_2pcf( self, ver, npatch=None, - tomo_bin1=None, - tomo_bin2=None, + compute_tomography=False, # save_fits=False, **treecorr_config, ): @@ -41,11 +40,8 @@ def calculate_2pcf( npatch (int, optional): The number of patches to use for the calculation. Defaults to the instance's `npatch` attribute. - tomo_bin1 (int, optional): The first tomographic bin to use for the calculation. - If None, the calculation is non-tomographic. Defaults to None. - - tomo_bin2 (int, optional): The second tomographic bin to use for the calculation. - If None, the calculation is non-tomographic. Defaults to None. + 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. @@ -71,132 +67,91 @@ def calculate_2pcf( **self._binning(**treecorr_config), "var_method": "jackknife" if int(npatch) > 1 else "shot", } - pol_factor = self.pol_factor - gg = treecorr.GGCorrelation(treecorr_config) - if tomo_bin1 is None and tomo_bin2 is None: - self.print_magenta(f"Computing non-tomographic ξ± for {ver}.") - # 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 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 os.path.exists(out_fname): - self.print_done(f"Skipping 2PCF calculation, {out_fname} exists.") - gg.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") - # LG: need tomographic patch file? - patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") + 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")] - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=pol_factor * 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, - ) + ggs = {f"tomo_bin_{b1}_tomo_bin_{b2}": None for b1, b2 in tomo_bin_pairs} - # If no patch file exists, save the current patches - if not os.path.exists(patch_file): - cat_gal.write_patch_centers(patch_file) + # 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.") + 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) - # Process the catalog & write the correlation functions - gg.process(cat_gal) - # 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) else: - self.print_magenta( - f"Computing tomographic ξ± for {ver}, bin {tomo_bin1} and {tomo_bin2}." - ) - - # LG TO-DO: Change to sacc_io method - out_fname = self._output_path( - f"{ver}_xi_tomo_bins={tomo_bin1}_{tomo_bin2}_minsep={treecorr_config['min_sep']}_maxsep={treecorr_config['max_sep']}_nbins={treecorr_config['nbins']}_npatch={npatch}.txt" - ) + patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - if os.path.exists(out_fname): - self.print_done(f"Skipping 2PCF calculation, {out_fname} exists.") - gg.read(out_fname) + for bin1, bin2 in tomo_bin_pairs: + gg = treecorr.GGCorrelation(treecorr_config) - else: + # Load data and create a catalog with self.results[ver].temporarily_read_data(): - tomo_bin_idx_bin1 = ( - self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 - ) + 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") - # LG: need tomographic patch file? - patch_file_bin1 = self._output_path( - f"{ver}_patches_npatch={npatch}_bin={tomo_bin1}.dat" - ) - - cat_gal_bin1 = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin1], - dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin1], - g1=g1[tomo_bin_idx_bin1], - g2=pol_factor * g2[tomo_bin_idx_bin1], - w=w[tomo_bin_idx_bin1], + 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_bin1 - if os.path.exists(patch_file_bin1) + patch_centers=patch_file + if os.path.exists(patch_file) else None, ) - - if tomo_bin1 == tomo_bin2: - gg.process(cat_gal_bin1) - # 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) - - if not os.path.exists(patch_file_bin1): - cat_gal_bin1.write_patch_centers(patch_file_bin1) - - else: - tomo_bin_idx_bin2 = ( - self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 - ) - - # LG: need tomographic patch file? - patch_file_bin2 = self._output_path( - f"{ver}_patches_npatch={npatch}_bin={tomo_bin2}.dat" - ) - - cat_gal_bin2 = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], - dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin2], - g1=g1[tomo_bin_idx_bin2], - g2=pol_factor * g2[tomo_bin_idx_bin2], - w=w[tomo_bin_idx_bin2], + 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_bin2 - if os.path.exists(patch_file_bin2) + patch_centers=patch_file + if os.path.exists(patch_file) else None, ) - gg.process(cat_gal_bin1, cat2=cat_gal_bin2) - # 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) + # If no patch file exists, save the current patches + if not os.path.exists(patch_file): + cat_gal1.write_patch_centers(patch_file) - if not os.path.exists(patch_file_bin1) or not os.path.exists( - patch_file_bin2 - ): - cat_gal_bin1.write_patch_centers(patch_file_bin1) - cat_gal_bin2.write_patch_centers(patch_file_bin2) + # 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 @@ -271,11 +226,11 @@ def calculate_2pcf( # 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(f"Done 2PCF for bins {tomo_bin1} and {tomo_bin2} for {ver}.") + self.print_done(f"Done 2PCF for {ver}.") - return gg + return ggs def plot_2pcf(self): # Plot of n_pairs @@ -284,8 +239,8 @@ def plot_2pcf(self): # 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"], @@ -302,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"], @@ -325,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"], @@ -348,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"], @@ -370,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"], @@ -394,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}}$", @@ -410,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() @@ -440,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( @@ -500,152 +466,100 @@ def calculate_aperture_mass_dispersion( nbins=500, nbins_map=15, npatch=25, - tomo_bin1=None, - tomo_bin2=None, + compute_tomography=False, ): - self.print_start("Computing aperture-mass dispersion") self._map2 = {} theta_map = np.geomspace(theta_min * 5, theta_max / 2, nbins_map) self._map2["theta_map"] = theta_map treecorr_config = self._binning(theta_min, theta_max, nbins) - pol_factor = self.pol_factor for ver in self.versions: - self.print_magenta(ver) - - gg = treecorr.GGCorrelation(treecorr_config) - self._map2.setdefault(ver, {}) - - if tomo_bin1 is None and tomo_bin2 is None: + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) self.print_magenta( - "Computing non-tomographic aperture-mass dispersion." + f"Computing MAP for {ver} with {len(tomo_bin_pairs)} bins." ) - # 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=pol_factor * 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, - ) - gg.process(cat_gal) - # gg.write(out_fname) - del cat_gal - del g1 - del g2 + 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")] - 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]["1_1"] = { - "mapsq": mapsq, - "mapsq_im": mapsq_im, - "mxsq": mxsq, - "mxsq_im": mxsq_im, - "varmapsq": varmapsq, - } + 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") + 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) else: - self.print_magenta( - f"Computing tomographic aperture-mass dispersion for bins {tomo_bin1} and {tomo_bin2}." - ) - # LG TO-DO: Change to sacc_io method - out_fname = self._output_path( - f"xi_for_map2_{ver}_tomo_bins={tomo_bin1}_{tomo_bin2}.txt" - ) - if os.path.exists(out_fname): - self.print_green(f"Skipping xi for Map2, {out_fname} exists") - gg.read(out_fname) - else: + 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(): - tomo_bin_idx_bin1 = ( - self.results[ver].dat_shear["tom_bin_id"] == tomo_bin1 - ) + 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) - cat_gal_bin1 = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin1], - dec=self.results[ver].dat_shear["Dec"][tomo_bin_idx_bin1], - g1=g1[tomo_bin_idx_bin1], - g2=pol_factor * g2[tomo_bin_idx_bin1], - w=self._read_shear_cols(ver, "w_col")[tomo_bin_idx_bin1], + 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 tomo_bin1 == tomo_bin2: - gg.process(cat_gal_bin1) - # gg.write(out_fname) - del cat_gal_bin1 - del g1 - del g2 - - else: - tomo_bin_idx_bin2 = ( - self.results[ver].dat_shear["tom_bin_id"] == tomo_bin2 + if bin1 != bin2: + mask_bin2 = ( + self.results[ver].dat_shear["tom_bin_id"] == bin2 ) - cat_gal_bin2 = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"][tomo_bin_idx_bin2], - dec=self.results[ver].dat_shear["Dec"][ - tomo_bin_idx_bin2 - ], - g1=g1[tomo_bin_idx_bin2], - g2=pol_factor * g2[tomo_bin_idx_bin2], - w=self._read_shear_cols(ver, "w_col")[ - tomo_bin_idx_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) - gg.process(cat_gal_bin1, cat2=cat_gal_bin2) - # gg.write(out_fname) - del cat_gal_bin1 - del cat_gal_bin2 - 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}_tomo_bins={tomo_bin1}_{tomo_bin2}.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][f"{tomo_bin1}_{tomo_bin2}"] = { - "mapsq": mapsq, - "mapsq_im": mapsq_im, - "mxsq": mxsq, - "mxsq_im": mxsq_im, - "varmapsq": varmapsq, - } - + 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 @@ -658,8 +572,14 @@ 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] # FORCE non-tomography for now - y = [self.map2[ver]["1_1"][mode] for ver in self.versions] - yerr = [np.sqrt(self.map2[ver]["1_1"]["varmapsq"]) for ver in self.versions] + 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] @@ -691,8 +611,14 @@ 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] # FORCE non-tomography for now - y = [np.abs(self.map2[ver]["1_1"][mode]) for ver in self.versions] - yerr = [np.sqrt(self.map2[ver]["1_1"]["varmapsq"]) for ver in self.versions] + 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 From 871bb2668567fc1f880a6e55f3e51da32934bca9 Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Fri, 24 Jul 2026 17:04:38 +0200 Subject: [PATCH 4/6] parent-child functio reconciliation --- src/sp_validation/cosmo_val/cosebis.py | 22 +- src/sp_validation/cosmo_val/pure_eb.py | 6 +- src/sp_validation/cosmo_val/real_space.py | 621 ++++++---------------- src/sp_validation/tests/test_cosmo_val.py | 2 +- 4 files changed, 176 insertions(+), 475 deletions(-) diff --git a/src/sp_validation/cosmo_val/cosebis.py b/src/sp_validation/cosmo_val/cosebis.py index 93832f2e..87c36ac2 100644 --- a/src/sp_validation/cosmo_val/cosebis.py +++ b/src/sp_validation/cosmo_val/cosebis.py @@ -103,7 +103,7 @@ def calculate_cosebis( f"{max_sep_int} arcmin" ) - ggs = self.calculate_2pcf( + ggs = self.calculate_2pcf_version( version, npatch=npatch, compute_tomography=compute_tomography, @@ -144,7 +144,7 @@ def calculate_cosebis( print(f"Evaluating {len(generated_cuts)} scale cut combinations") # Call b_modes function with scale cuts list - results = calculate_cosebis( + results[bin_key] = calculate_cosebis( gg=ggs[bin_key], nmodes=nmodes, scale_cuts=generated_cuts, @@ -152,13 +152,11 @@ def calculate_cosebis( ) else: # Single scale cut behavior: use full range - results = calculate_cosebis( + results[bin_key] = 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 + results[bin_key] = next(iter(results[bin_key].values())) return results @@ -241,7 +239,7 @@ def plot_cosebis( # Get or calculate results for this version if results is None: # Calculate COSEBIs using instance method - results = self.calculate_cosebis( + results_tomo = self.calculate_cosebis( version, min_sep_int=min_sep_int, max_sep_int=max_sep_int, @@ -256,6 +254,11 @@ def plot_cosebis( max_sep=max_sep, nbins=nbins, ) + results = results_tomo[ + "tomo_bin_all_tomo_bin_all" + ] # LG: Extract non-tomographic results + elif isinstance(results, dict) and "tomo_bin_all_tomo_bin_all" in results: + results = results["tomo_bin_all_tomo_bin_all"] # Generate plots using specialized plotting functions # Extract single result for plotting if multiple scale cuts were evaluated @@ -289,9 +292,12 @@ def plot_cosebis( if multiple_scale_cuts and len(results) > 1: # Create temporary gg object with correct binning for mapping treecorr_config_temp = self._binning(min_sep, max_sep, nbins) - gg_temp = self.calculate_2pcf( + ggs_temp = self.calculate_2pcf_version( version, npatch=npatch, **treecorr_config_temp ) + gg_temp = ggs_temp[ + "tomo_bin_all_tomo_bin_all" + ] # LG: Extract non-tomographic result plot_cosebis_scale_cut_heatmap( results, diff --git a/src/sp_validation/cosmo_val/pure_eb.py b/src/sp_validation/cosmo_val/pure_eb.py index be41e8d8..f585b0e3 100644 --- a/src/sp_validation/cosmo_val/pure_eb.py +++ b/src/sp_validation/cosmo_val/pure_eb.py @@ -70,7 +70,7 @@ def calculate_pure_eb( Path to the covariance matrix for the reporting binning. Replaces the treecorr covariance matrix if provided, meaning that var_method has no effect on the results although it is still passed to - CosmologyValidation.calculate_2pcf. + CosmologyValidation.calculate_2pcf_version. cosmo_cov : pyccl.Cosmology, optional Cosmology object to use for theoretical xi+/xi- predictions in the semi-analytical covariance calculation. Defaults to self.cosmo if not @@ -112,13 +112,13 @@ def calculate_pure_eb( treecorr_config_int = self._binning(min_sep_int, max_sep_int, nbins_int) # Calculate correlation functions - ggs = self.calculate_2pcf( + ggs = self.calculate_2pcf_version( version, npatch=npatch, compute_tomography=compute_tomography, **treecorr_config, ) - ggs_int = self.calculate_2pcf( + ggs_int = self.calculate_2pcf_version( version, npatch=npatch, compute_tomography=compute_tomography, diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 49f32424..d6d50e25 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -8,58 +8,48 @@ import os -import matplotlib.pyplot as plt -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( + def calculate_2pcf_version( self, ver, npatch=None, compute_tomography=False, - # save_fits=False, **treecorr_config, ): """ - Calculate the two-point correlation function (2PCF) ξ± for a given catalog + Calculate the two-point correlation function (2PCF) ξ± for a single catalog version with TreeCorr. + This is the per-version child function. Use :meth:`calculate_2pcf` to run over every + version in ``self.versions`` in one call. + By default the class instance's `npatch` and `treecorr_config` entries are - used to - initialize the TreeCorr Catalog and GGCorrelation objects, but may be - overridden - by passing keyword arguments. + used to initialize the TreeCorr Catalog and GGCorrelation objects, but may + be overridden by passing keyword arguments. Parameters: ver (str): The catalog version to process. - 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. + npatch (int, optional): The number of patches to use for the + calculation. Defaults to the instance's `npatch` attribute. - save_fits (bool, optional): Whether to save the ξ± results to FITS files. - Defaults to False. + compute_tomography (bool, optional): Whether to compute tomographic + correlations. Defaults to False. - **treecorr_config: Additional TreeCorr configuration parameters that will - override the instance's default `treecorr_config`. For example, `min_sep=1`. + **treecorr_config: Additional TreeCorr configuration parameters that + will override the instance's default `treecorr_config`. For example, + `min_sep=1`. Returns: - treecorr.GGCorrelation: The TreeCorr GGCorrelation object containing the - computed 2PCF results. - - Notes: - - If the output file for the given configuration already exists, the - calculation is skipped, and the results are loaded from the file. - - If a patch file for the given configuration does not exist, it is - created during the process. - - FITS files for ξ+ and ξ− are saved with additional metadata in their - headers if `save_fits` is True. + dict: Mapping of ``"tomo_bin_{b1}_tomo_bin_{b2}"`` to the corresponding + treecorr.GGCorrelation object. For the non-tomographic case the single + key is ``"tomo_bin_all_tomo_bin_all"``. """ npatch = npatch or self.npatch @@ -70,12 +60,11 @@ def calculate_2pcf( 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.print_magenta( f"Computing tomographic ξ± 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 ξ± for {ver}.") tomo_bin_pairs = [("all", "all")] @@ -83,381 +72,100 @@ def calculate_2pcf( 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.") - 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: - patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") + patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - for bin1, bin2 in tomo_bin_pairs: - gg = treecorr.GGCorrelation(treecorr_config) + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") - # 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, - # ) - - # Add correlation object to class - if not hasattr(self, "cat_ggs"): - self.cat_ggs = {} - self.cat_ggs[ver] = ggs + for bin1, bin2 in tomo_bin_pairs: + gg = treecorr.GGCorrelation(treecorr_config) + + # Load data and create a catalog + if bin1 == "all" and bin2 == "all": + mask_bin1 = np.ones(len(g1), dtype=bool) + else: + mask_bin1 = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == bin1 + + cat_gal1 = treecorr.Catalog( + ra=cat_gal["RA"][mask_bin1], + dec=cat_gal["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 = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == bin2 + cat_gal2 = treecorr.Catalog( + ra=cat_gal["RA"][mask_bin2], + dec=cat_gal["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 self.print_done(f"Done 2PCF for {ver}.") return ggs - def plot_2pcf(self): - # Plot of n_pairs - plt.subplots(ncols=1, nrows=1) + def calculate_2pcf( + self, + npatch=None, + compute_tomography=False, + **treecorr_config, + ): + """ + Calculate the 2PCF ξ± for every catalog version in ``self.versions``. + + Parent function that iterates over ``self.versions`` and delegates the + per-version computation to :meth:`calculate_2pcf_version`. Results are stored + in ``self.cat_ggs`` keyed by version. + + Parameters: + npatch (int, optional): Number of patches to use; defaults to the + instance's `npatch` attribute. + + compute_tomography (bool, optional): Whether to compute tomographic + correlations. Defaults to False. + + **treecorr_config: Additional TreeCorr configuration parameters passed + through to each per-version call. + + Returns: + dict: ``self.cat_ggs``, mapping each version to its + ``{"tomo_bin_{b1}_tomo_bin_{b2}": treecorr.GGCorrelation}`` dict. + """ + self.cat_ggs = {} for ver in self.versions: - # FORCE non-tomography for now - self.calculate_2pcf(ver) - plt.plot( - 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"], - ) - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.ylabel(r"$n_{\rm pair}$") - plt.legend() - out_path = self._output_path("n_pair.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"n_pair plot saved to {out_path}") - - # Plot of xi_+ - plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) - for idx, ver in enumerate(self.versions): - plt.errorbar( - 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"], - ) - plt.xscale("log") - plt.yscale("log") - plt.legend() - plt.ticklabel_format(axis="y") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.xlim([self.theta_min_plot, self.theta_max_plot]) - plt.ylabel(r"$\xi_+(\theta)$") - out_path = self._output_path("xi_p.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_plus plot saved to {out_path}") - - # Plot of xi_- - plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) - for idx, ver in enumerate(self.versions): - plt.errorbar( - 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"], - ) - plt.xscale("log") - plt.yscale("log") - plt.legend() - plt.ticklabel_format(axis="y") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.xlim([self.theta_min_plot, self.theta_max_plot]) - plt.ylabel(r"$\xi_-(\theta)$") - out_path = self._output_path("xi_m.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_minus plot saved to {out_path}") - - # Plot of xi_+(theta) * theta - plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) - for idx, ver in enumerate(self.versions): - plt.errorbar( - 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"], - ) - plt.xscale("log") - plt.legend() - plt.ticklabel_format(axis="y") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.xlim([self.theta_min_plot, self.theta_max_plot]) - plt.ylabel(r"$\theta \xi_+(\theta)$") - out_path = self._output_path("xi_p_theta.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_plus_theta plot saved to {out_path}") - - # Plot of xi_- * theta - plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) - for idx, ver in enumerate(self.versions): - plt.errorbar( - 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"], - ) - plt.xscale("log") - plt.legend() - plt.ticklabel_format(axis="y") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.xlim([self.theta_min_plot, self.theta_max_plot]) - plt.ylabel(r"$\theta \xi_-(\theta)$") - out_path = self._output_path("xi_m_theta.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_minus_theta plot saved to {out_path}") - - # Plot of xi_+ with and without xi_psf_sys - # but skip if xi_psf_sys is not calculated since that takes forever - if hasattr(self, "_xi_psf_sys"): - for idx, ver in enumerate(self.versions): - plt.subplots(ncols=1, nrows=1, figsize=(7, 7)) - plt.errorbar( - 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]["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}}$", - ls="dotted", - color="red", - ) - plt.errorbar( - 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]["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() - plt.ticklabel_format(axis="y") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.xlim([self.theta_min_plot, self.theta_max_plot]) - plt.ylim(1e-8, 5e-4) - plt.ylabel(r"$\xi_+(\theta)$") - out_path = self._output_path(f"xi_p_xi_psf_sys_{ver}.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"xi_plus_xi_psf_sys {ver} plot saved to {out_path}") - - def plot_ratio_xi_sys_xi(self, threshold=0.1, offset=0.02): - - plt.subplots(ncols=1, nrows=1, figsize=(10, 7)) - - for idx, ver in enumerate(self.versions): - self.calculate_2pcf(ver) - xi_psf_sys = self.xi_psf_sys[ver] - gg = self.cat_ggs[ver]["tomo_bin_all_tomo_bin_all"] - - ratio = xi_psf_sys["mean"] / gg.xip - ratio_err = np.sqrt( - (np.sqrt(xi_psf_sys["var"]) / gg.xip) ** 2 - + (xi_psf_sys["mean"] * np.sqrt(gg.varxip) / gg.xip**2) ** 2 + self.cat_ggs[ver] = self.calculate_2pcf_version( + ver, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config, ) - theta = gg.meanr - jittered_theta = theta * (1 + idx * offset) - - plt.errorbar( - jittered_theta, - ratio, - yerr=ratio_err, - label=ver, - ls=self.cc[ver]["ls"], - color=self.cc[ver]["colour"], - fmt=self.cc[ver].get("marker", None), - capsize=5, - ) + # LG TO-DO: No longer writing out text file, change to sacc_io method - plt.fill_between( - [self.theta_min_plot, self.theta_max_plot], - -threshold, - +threshold, - color="black", - alpha=0.1, - label=f"{threshold:.0%} threshold", - ) - plt.plot( - [self.theta_min_plot, self.theta_max_plot], - [threshold, threshold], - ls="dashed", - color="black", - ) - plt.plot( - [self.theta_min_plot, self.theta_max_plot], - [-threshold, -threshold], - ls="dashed", - color="black", - ) - plt.xscale("log") - plt.xlabel(rf"$\theta$ [{self.treecorr_config['sep_units']}]") - plt.ylabel(r"$\xi^{\rm psf}_{+, {\rm sys}} / \xi_+$") - plt.gca().yaxis.set_major_formatter(mticker.PercentFormatter(xmax=1)) - plt.legend() - plt.title("Ratio of PSF systematics to cosmic shear signal") - out_path = self._output_path("ratio_xi_sys_xi.png") - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - print(f"Ratio of xi_psf_sys to xi plot saved to {out_path}") + return self.cat_ggs def calculate_aperture_mass_dispersion( self, @@ -478,88 +186,73 @@ def calculate_aperture_mass_dispersion( for ver in self.versions: 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." ) + self.print_magenta( + f"Computing MAP for {ver} with {len(tomo_bin_pairs)} bins." + ) else: self.print_magenta(f"Computing non-tomographic MAP for {ver}.") tomo_bin_pairs = [("all", "all")] self._map2.setdefault(ver, {}) + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) # 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") - 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) - else: - 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", + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") + + for bin1, bin2 in tomo_bin_pairs: + gg = treecorr.GGCorrelation(treecorr_config) + + # Load data and create a catalog + if bin1 == "all" and bin2 == "all": + mask_bin1 = np.ones(len(cat_gal), dtype=bool) + else: + mask_bin1 = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == bin1 + + cat_gal1 = treecorr.Catalog( + ra=cat_gal["RA"][mask_bin1], + dec=cat_gal["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 = cat_gal[self.cc[ver]["shear"]["tomo_bin_col"]] == bin2 + cat_gal2 = treecorr.Catalog( + ra=cat_gal["RA"][mask_bin2], + dec=cat_gal["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, ) - self._map2[ver][f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = { - "mapsq": mapsq, - "mapsq_im": mapsq_im, - "mxsq": mxsq, - "mxsq_im": mxsq_im, - "varmapsq": varmapsq, - } + + 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 @@ -568,10 +261,12 @@ def map2(self): self.calculate_aperture_mass_dispersion() return self._map2 + # LG: plotting functions removed, perhaps can use Sacha's implementation in psf_systematics.py instead + 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] - # FORCE non-tomography for now + # LG: FORCE non-tomography for now y = [ self.map2[ver]["tomo_bin_all_tomo_bin_all"][mode] for ver in self.versions diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index 2e34c6c9..9536cbee 100644 --- a/src/sp_validation/tests/test_cosmo_val.py +++ b/src/sp_validation/tests/test_cosmo_val.py @@ -512,7 +512,7 @@ def test_calculate_2pcf_runs_on_synthetic_catalog(self, tmp_path): **params, ) - ggs = cv.calculate_2pcf(version) + ggs = cv.calculate_2pcf_version(version) # treecorr GGCorrelation with xi+/- on the configured angular grid # LG:Hardcoded to be non-tomographic for now From 54ee05459238300010682fe543a264386a483e1c Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Fri, 24 Jul 2026 17:52:17 +0200 Subject: [PATCH 5/6] pass CI tests --- src/sp_validation/cosmo_val/real_space.py | 11 ++++++----- 1 file changed, 6 insertions(+), 5 deletions(-) diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index d6d50e25..378c8863 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -76,8 +76,9 @@ def calculate_2pcf_version( patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) - g1, g2 = self._calibrated_g(ver) - w = self._read_shear_cols(ver, "w_col") + with self.results[ver].temporarily_read_data(): + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") for bin1, bin2 in tomo_bin_pairs: gg = treecorr.GGCorrelation(treecorr_config) @@ -202,9 +203,9 @@ def calculate_aperture_mass_dispersion( cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) # LG TO-DO: Change to sacc_io method - - g1, g2 = self._calibrated_g(ver) - w = self._read_shear_cols(ver, "w_col") + with self.results[ver].temporarily_read_data(): + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") for bin1, bin2 in tomo_bin_pairs: gg = treecorr.GGCorrelation(treecorr_config) From 584b8019ebaefa7738ca9dab261a36d43e9b5222 Mon Sep 17 00:00:00 2001 From: LisaGoh Date: Fri, 24 Jul 2026 18:11:55 +0200 Subject: [PATCH 6/6] updated 2pcf plotting functions --- .../cosmo_val/psf_systematics.py | 225 --------- src/sp_validation/cosmo_val/real_space.py | 451 +++++++++++++++--- 2 files changed, 377 insertions(+), 299 deletions(-) diff --git a/src/sp_validation/cosmo_val/psf_systematics.py b/src/sp_validation/cosmo_val/psf_systematics.py index 6765123e..465aa937 100644 --- a/src/sp_validation/cosmo_val/psf_systematics.py +++ b/src/sp_validation/cosmo_val/psf_systematics.py @@ -1113,231 +1113,6 @@ def plot_objectwise_leakage(self): self.print_done(f"Object-wise leakage coefficients plot saved to {out_path}") # --- utlility functions for plotting --- - def plot_2pcf_tomography( - self, - x_y_plot_function, - x_label, - y_label_plus, - y_label_minus, - tomo_bin_label_position, - extract_text_offset, - add_index_version_to_kwargs, - x_scale=None, - y_scale=None, - tomography=False, - versions=None, - colors=None, - savefig=None, - show=True, - close=True, - **kwargs, - ): - """ - Standard plot function for 2-point correlation functions with tomographic bins. - - Parameters - ---------- - x_y_plot_function : callable - Function to plot the x and y data. Should accept axes for `+' and `-' components, version, tomo_bin_indices, and kwargs. - x_label : str - Label for the x-axis. - y_label_plus : str - Label for the y-axis of the `+' component. - y_label_minus : str - Label for the y-axis of the `-' component. - tomo_bin_label_position : tuple - Position to place the tomographic bin labels in axes coordinates (x, y). - extract_text_offset : bool - If True, extract the y-axis offset text and include it in the y-label. - add_index_version_to_kwargs : bool - If True, add the index of the version to the kwargs for plotting. - x_scale : str, optional - Scale for the x-axis ('linear', 'log', etc.). If None, default scale is used. - y_scale : str, optional - Scale for the y-axis ('linear', 'log', etc.). If None, default scale is used. - tomography : bool, optional - If True, plot the tomographic bins. Default is False. - versions : list, optional - List of versions to plot. If None, all versions are plotted. - colors : list, optional - List of colors for each version. If None, default colors are used. - savefig : str, optional - If provided, save the figure to this file. - show : bool, optional - If True, show the figure. - close : bool, optional - If True, close the figure after saving or showing. - kwargs : dict - Additional keyword arguments to pass to the plotting function. - """ - versions = versions if versions is not None else self.versions - colors = ( - colors - if colors is not None - else [self.cc[ver]["colour"] for ver in versions] - ) - - # First get the max number of tomo bins among the versions. - tomo_bins = self._get_tomo_bins_for_versions(versions, tomography=tomography) - - max_key = max(tomo_bins, key=lambda k: len(tomo_bins[k]["ids"])) - n_tomo_bins_plot = len(tomo_bins[max_key]["ids"]) - reference_tomo_bin_pairs = tomo_bins[max_key]["pairs"] - - n_rows = n_tomo_bins_plot - n_cols = n_tomo_bins_plot + 2 - - # First start with the quantiles plot - fig, axs = plt.subplots( - n_rows, - n_cols, - figsize=(5 * n_cols, 5 * n_rows), - sharex=True, - sharey=True, - gridspec_kw={"wspace": 0, "hspace": 0}, - ) - - for idx, (ver, color) in enumerate(zip(versions, colors)): - tomo_bin_pairs = tomo_bins[ver]["pairs"] - - kwargs["color"] = color - if add_index_version_to_kwargs: - kwargs["idx"] = idx - kwargs["versions"] = versions - - for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: - # Plot the nsamples last samples - ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) - ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) - - # Apply the x_y_plot_function to plot the data - x_y_plot_function( - ax_plus, ax_minus, ver, tomo_bin_a, tomo_bin_b, **kwargs - ) - - # Draw to extract the y-axis text offset - fig.canvas.draw() - - # Set the visibility to false where necessary - self._set_ax_visibility_to_false(axs, n_tomo_bins_plot) - - # Set the labels and scales for the plots - for tomo_bin_a, tomo_bin_b in reference_tomo_bin_pairs: - ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) - ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) - - ax_plus.tick_params( - axis="both", - which="both", - direction="in", - bottom=True, - top=False, - labelbottom=tomo_bin_b == 1 or tomo_bin_b == "all", - left=True, - right=False, - labelleft=tomo_bin_a == 1 or tomo_bin_a == "all", - ) - if x_scale is not None: - ax_plus.set_xscale(x_scale) - - ax_plus.text( - tomo_bin_label_position[0], - tomo_bin_label_position[1], - f"{tomo_bin_a}-{tomo_bin_b}", - transform=ax_plus.transAxes, - verticalalignment="top", - bbox=dict( - boxstyle="square", - facecolor="white", - edgecolor="black", - alpha=0.8, - ), - ) - ax_plus.axhline(0, color="k", linestyle="--") - - if y_scale is not None: - ax_plus.set_yscale(y_scale) - if tomo_bin_b == 1 or tomo_bin_b == "all": - ax_plus.set_xlabel(r"$\theta$ [arcmin]") - if tomo_bin_a == 1 or tomo_bin_a == "all": - text_offset = ( - ax_plus.yaxis.get_offset_text().get_text() - if extract_text_offset - else "" - ) - ax_plus.set_ylabel(y_label_plus + text_offset) - ax_plus.yaxis.get_offset_text().set_visible( - False - ) # Hide the offset text for the plus ax - - # Move the ticks to the right for the minus ax - ax_minus.yaxis.tick_right() - ax_minus.yaxis.set_label_position("right") - ax_minus.tick_params( - axis="both", - which="both", - direction="in", - bottom=True, - top=False, - labelbottom=tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all", - left=False, - right=True, - labelleft=False, - labelright=tomo_bin_a == 1 or tomo_bin_a == "all", - ) - if x_scale is not None: - ax_minus.set_xscale(x_scale) - ax_minus.text( - tomo_bin_label_position[0], - tomo_bin_label_position[1], - f"{tomo_bin_a}-{tomo_bin_b}", - transform=ax_minus.transAxes, - verticalalignment="top", - bbox=dict( - boxstyle="square", - facecolor="white", - edgecolor="black", - alpha=0.8, - ), - ) - ax_minus.axhline(0, color="k", linestyle="--") - if y_scale is not None: - ax_minus.set_yscale(y_scale) - if tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all": - ax_minus.set_xlabel(x_label) - if tomo_bin_a == 1 or tomo_bin_a == "all": - text_offset = ( - ax_minus.yaxis.get_offset_text().get_text() - if extract_text_offset - else "" - ) - ax_minus.set_ylabel(y_label_minus + text_offset) - ax_minus.yaxis.get_offset_text().set_visible( - False - ) # Hide the offset text for the minus ax - - # Build the legend - handles = [] - for ver, color in zip(versions, colors): - label = self.cc[ver]["label"] if "label" in self.cc[ver] else ver - handles.append(plt.Line2D([0], [0], color=color, lw=2, label=label)) - fig.legend( - handles=handles, - loc="upper center", - ncol=3, - frameon=False, - bbox_to_anchor=(0.5, 0.0), - ) - - if savefig is not None: - plt.savefig(savefig, dpi=300, bbox_inches="tight") - self.print_done(f"Plot saved to {os.path.abspath(savefig)}") - - if show: - plt.show() - - if close: - plt.close() def _xi_psf_sys_sample_x_y_plot_function( self, diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 378c8863..96b7a077 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -8,10 +8,10 @@ import os +import matplotlib.pyplot as plt import numpy as np import treecorr from astropy.io import fits -from cs_util import plots as cs_plots class RealSpaceMixin: @@ -263,79 +263,382 @@ def map2(self): return self._map2 # LG: plotting functions removed, perhaps can use Sacha's implementation in psf_systematics.py instead + def plot_2pcf_tomography( + self, + x_y_plot_function, + x_label, + y_label_plus, + y_label_minus, + tomo_bin_label_position, + extract_text_offset, + add_index_version_to_kwargs, + x_scale=None, + y_scale=None, + tomography=False, + versions=None, + colors=None, + savefig=None, + show=True, + close=True, + **kwargs, + ): + """ + Standard plot function for 2-point correlation functions with tomographic bins. + + Parameters + ---------- + x_y_plot_function : callable + Function to plot the x and y data. Should accept axes for `+' and `-' components, version, tomo_bin_indices, and kwargs. + x_label : str + Label for the x-axis. + y_label_plus : str + Label for the y-axis of the `+' component. + y_label_minus : str + Label for the y-axis of the `-' component. + tomo_bin_label_position : tuple + Position to place the tomographic bin labels in axes coordinates (x, y). + extract_text_offset : bool + If True, extract the y-axis offset text and include it in the y-label. + add_index_version_to_kwargs : bool + If True, add the index of the version to the kwargs for plotting. + x_scale : str, optional + Scale for the x-axis ('linear', 'log', etc.). If None, default scale is used. + y_scale : str, optional + Scale for the y-axis ('linear', 'log', etc.). If None, default scale is used. + tomography : bool, optional + If True, plot the tomographic bins. Default is False. + versions : list, optional + List of versions to plot. If None, all versions are plotted. + colors : list, optional + List of colors for each version. If None, default colors are used. + savefig : str, optional + If provided, save the figure to this file. + show : bool, optional + If True, show the figure. + close : bool, optional + If True, close the figure after saving or showing. + kwargs : dict + Additional keyword arguments to pass to the plotting function. + """ + versions = versions if versions is not None else self.versions + colors = ( + colors + if colors is not None + else [self.cc[ver]["colour"] for ver in versions] + ) + + # First get the max number of tomo bins among the versions. + tomo_bins = self._get_tomo_bins_for_versions(versions, tomography=tomography) + + max_key = max(tomo_bins, key=lambda k: len(tomo_bins[k]["ids"])) + n_tomo_bins_plot = len(tomo_bins[max_key]["ids"]) + reference_tomo_bin_pairs = tomo_bins[max_key]["pairs"] + + n_rows = n_tomo_bins_plot + n_cols = n_tomo_bins_plot + 2 + + # First start with the quantiles plot + fig, axs = plt.subplots( + n_rows, + n_cols, + figsize=(5 * n_cols, 5 * n_rows), + sharex=True, + sharey=True, + gridspec_kw={"wspace": 0, "hspace": 0}, + ) + + for idx, (ver, color) in enumerate(zip(versions, colors)): + tomo_bin_pairs = tomo_bins[ver]["pairs"] + + kwargs["color"] = color + if add_index_version_to_kwargs: + kwargs["idx"] = idx + kwargs["versions"] = versions + + for tomo_bin_a, tomo_bin_b in tomo_bin_pairs: + # Plot the nsamples last samples + ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) + ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) + + # Apply the x_y_plot_function to plot the data + x_y_plot_function( + ax_plus, ax_minus, ver, tomo_bin_a, tomo_bin_b, **kwargs + ) - 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] - # LG: 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] - - xlabel = r"$\theta$ [arcmin]" - ylabel = "dispersion" - title = f"Aperture-mass dispersion {mode}" - out_path = self._output_path(f"{mode}.png") - cs_plots.plot_data_1d( - x, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - labels=labels, - xlog=True, - xlim=[self.theta_min_plot, self.theta_max_plot], - ylim=[-2e-6, 5e-6], - colors=colors, - linestyles=linestyles, - shift_x=True, + # Draw to extract the y-axis text offset + fig.canvas.draw() + + # Set the visibility to false where necessary + self._set_ax_visibility_to_false(axs, n_tomo_bins_plot) + + # Set the labels and scales for the plots + for tomo_bin_a, tomo_bin_b in reference_tomo_bin_pairs: + ax_plus = self._get_ax_plus(axs, tomo_bin_a, tomo_bin_b) + ax_minus = self._get_ax_minus(axs, tomo_bin_a, tomo_bin_b) + + ax_plus.tick_params( + axis="both", + which="both", + direction="in", + bottom=True, + top=False, + labelbottom=tomo_bin_b == 1 or tomo_bin_b == "all", + left=True, + right=False, + labelleft=tomo_bin_a == 1 or tomo_bin_a == "all", + ) + if x_scale is not None: + ax_plus.set_xscale(x_scale) + + ax_plus.text( + tomo_bin_label_position[0], + tomo_bin_label_position[1], + f"{tomo_bin_a}-{tomo_bin_b}", + transform=ax_plus.transAxes, + verticalalignment="top", + bbox=dict( + boxstyle="square", + facecolor="white", + edgecolor="black", + alpha=0.8, + ), + ) + ax_plus.axhline(0, color="k", linestyle="--") + + if y_scale is not None: + ax_plus.set_yscale(y_scale) + if tomo_bin_b == 1 or tomo_bin_b == "all": + ax_plus.set_xlabel(r"$\theta$ [arcmin]") + if tomo_bin_a == 1 or tomo_bin_a == "all": + text_offset = ( + ax_plus.yaxis.get_offset_text().get_text() + if extract_text_offset + else "" + ) + ax_plus.set_ylabel(y_label_plus + text_offset) + ax_plus.yaxis.get_offset_text().set_visible( + False + ) # Hide the offset text for the plus ax + + # Move the ticks to the right for the minus ax + ax_minus.yaxis.tick_right() + ax_minus.yaxis.set_label_position("right") + ax_minus.tick_params( + axis="both", + which="both", + direction="in", + bottom=True, + top=False, + labelbottom=tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all", + left=False, + right=True, + labelleft=False, + labelright=tomo_bin_a == 1 or tomo_bin_a == "all", ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"linear-scale {mode} plot saved to {out_path}") - - for mode in ["mapsq", "mapsq_im", "mxsq", "mxsq_im"]: - x = [self.map2["theta_map"] 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}" - out_path = self._output_path(f"{mode}_log.png") - cs_plots.plot_data_1d( - x, - y, - yerr, - title, - xlabel, - ylabel, - out_path=None, - labels=labels, - xlog=True, - ylog=True, - xlim=[self.theta_min_plot, self.theta_max_plot], - ylim=[1e-8, 1e-5], - colors=colors, - linestyles=linestyles, - shift_x=True, + if x_scale is not None: + ax_minus.set_xscale(x_scale) + ax_minus.text( + tomo_bin_label_position[0], + tomo_bin_label_position[1], + f"{tomo_bin_a}-{tomo_bin_b}", + transform=ax_minus.transAxes, + verticalalignment="top", + bbox=dict( + boxstyle="square", + facecolor="white", + edgecolor="black", + alpha=0.8, + ), ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"log-scale {mode} plot saved to {out_path}") + ax_minus.axhline(0, color="k", linestyle="--") + if y_scale is not None: + ax_minus.set_yscale(y_scale) + if tomo_bin_b == n_tomo_bins_plot or tomo_bin_b == "all": + ax_minus.set_xlabel(x_label) + if tomo_bin_a == 1 or tomo_bin_a == "all": + text_offset = ( + ax_minus.yaxis.get_offset_text().get_text() + if extract_text_offset + else "" + ) + ax_minus.set_ylabel(y_label_minus + text_offset) + ax_minus.yaxis.get_offset_text().set_visible( + False + ) # Hide the offset text for the minus ax + + # Build the legend + handles = [] + for ver, color in zip(versions, colors): + label = self.cc[ver]["label"] if "label" in self.cc[ver] else ver + handles.append(plt.Line2D([0], [0], color=color, lw=2, label=label)) + fig.legend( + handles=handles, + loc="upper center", + ncol=3, + frameon=False, + bbox_to_anchor=(0.5, 0.0), + ) + + if savefig is not None: + plt.savefig(savefig, dpi=300, bbox_inches="tight") + self.print_done(f"Plot saved to {os.path.abspath(savefig)}") + + if show: + plt.show() + + if close: + plt.close() + + def _xiplus_ximinus_sample_x_y_plot_function( + self, + ax_plus, + ax_minus, + version, + tomo_bin_a, + tomo_bin_b, + idx, + versions, + color, + offset, + times_theta, + alpha, + ): + """Plot the measured ξ± 2PCF for one version/tomographic-bin pair. + + Uses the jackknife covariance from TreeCorr to plot the error bars. This function is fed into :meth:`plot_2pcf_tomography` as the ``x_y_plot_function`` argument. + + Parameters + ---------- + ax_plus, ax_minus : matplotlib.axes.Axes + Axes for the ξ+ and ξ- components. + version : str + Catalog version to plot. + tomo_bin_a, tomo_bin_b : int or str + Tomographic bin pair (``"all"`` for the non-tomographic case). + idx : int + Index of ``version`` within ``versions`` (used for the x-jitter). + versions : list + Full list of versions being plotted. + color : str + Colour for this version. + offset : float + Fractional jitter applied to θ for readability. + times_theta : bool + If True, plot θ·ξ± rather than ξ±. + alpha : float + Opacity of the plotted points/error bars. + """ + # Get the measured 2PCF for this version and tomographic-bin pair. + gg = self.cat_ggs[version][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"] + + # Angular scales of the measurement. + theta = gg.meanr + + # Add the offset to the theta values for better visualisation. + jittered_theta = self._get_jittered_theta(theta, idx, len(versions), offset) + + scale = theta if times_theta else 1 + + y_plus = gg.xip * scale + y_minus = gg.xim * scale + yerr_plus = np.sqrt(gg.varxip) * scale + yerr_minus = np.sqrt(gg.varxim) * scale + + ax_plus.errorbar( + jittered_theta, + y_plus, + yerr=yerr_plus, + color=color, + alpha=alpha, + fmt="o", + markersize=3, + capsize=2, + ) + + ax_minus.errorbar( + jittered_theta, + y_minus, + yerr=yerr_minus, + color=color, + alpha=alpha, + fmt="o", + markersize=3, + capsize=2, + ) + + def _mapsq_mxsq_sample_x_y_plot_function( + self, + ax_plus, + ax_minus, + version, + tomo_bin_a, + tomo_bin_b, + idx, + versions, + color, + offset, + times_theta, + alpha, + ): + """Plot the aperture-mass dispersion ⟨M_ap²⟩ / ⟨M_ײ⟩ for one bin pair. + + Uses the jackknife covariance from TreeCorr to plot the error bars. This function is fed into :meth:`plot_2pcf_tomography` as the ``x_y_plot_function`` argument. The E-mode ⟨M_ap²⟩ is placed on the ``plus`` axis and the B-mode ⟨M_ײ⟩ on the ``minus`` axis. + + Parameters + ---------- + ax_plus, ax_minus : matplotlib.axes.Axes + Axes for the E-mode (⟨M_ap²⟩) and B-mode (⟨M_ײ⟩) components. + version : str + Catalog version to plot. + tomo_bin_a, tomo_bin_b : int or str + Tomographic bin pair (``"all"`` for the non-tomographic case). + idx : int + Index of ``version`` within ``versions`` (used for the x-jitter). + versions : list + Full list of versions being plotted. + color : str + Colour for this version. + offset : float + Fractional jitter applied to θ for readability. + times_theta : bool + If True, plot θ·⟨M²⟩ rather than ⟨M²⟩. + alpha : float + Opacity of the plotted points/error bars. + """ + # Get the aperture-mass dispersion for this version and bin pair. + map2 = self.map2[version][f"tomo_bin_{tomo_bin_a}_tomo_bin_{tomo_bin_b}"] + + # Angular scales of the measurement. + theta = self.map2["theta_map"] + + # Add the offset to the theta values for better visualisation. + jittered_theta = self._get_jittered_theta(theta, idx, len(versions), offset) + + scale = theta if times_theta else 1 + + y_plus = map2["mapsq"] * scale + y_minus = map2["mxsq"] * scale + # Both E- and B-mode share the same variance estimate. + yerr = np.sqrt(map2["varmapsq"]) * scale + + ax_plus.errorbar( + jittered_theta, + y_plus, + yerr=yerr, + color=color, + alpha=alpha, + fmt="o", + markersize=3, + capsize=2, + ) + + ax_minus.errorbar( + jittered_theta, + y_minus, + yerr=yerr, + color=color, + alpha=alpha, + fmt="o", + markersize=3, + capsize=2, + )