diff --git a/scripts/xip_xim.py b/scripts/xip_xim.py index 1345f01a..763752ce 100755 --- a/scripts/xip_xim.py +++ b/scripts/xip_xim.py @@ -20,6 +20,8 @@ def params_default(): params = { "input_path": "shape_catalog_ngmix.fits", + "tomo_bin1": None, + "tomo_bin2": None, "key_ra": "RA", "key_dec": "DEC", "key_e1": "e1", @@ -30,20 +32,27 @@ def params_default(): "theta_max": 200, "n_theta": 20, "output_path": "./xip_xim.txt", + "key_tomo_bin_col": "tomo_bin_id", } short_options = { "input_path": "-i", "output_path": "-o", + "tomo_bin1": "-b1", + "tomo_bin2": "-b2", } types = { "sign_e1": "int", "sign_e2": "int", + "tomo_bin1": "int", + "tomo_bin2": "int", } help_strings = { "input_path": "shear catalogue input path, default={}", + "tomo_bin1": "First tomographic bin, default={} (non-tomographic)", + "tomo_bin2": "Second tomographic bin, default={} (non-tomographic)", "key_ra": "column name for right ascension, default={}", "key_dec": "column name for declination, default={}", "key_e1": "column name for ellipticity component 1, default={}", @@ -54,6 +63,7 @@ def params_default(): "theta_max": "maximum angular scale [arcmin], default={}", "n_theta": "number of angular scales, default={}", "output_path": "output path, default={}", + "key_tomo_bin_col": "column name for tomography bin, default={}", } return params, short_options, types, help_strings @@ -132,18 +142,52 @@ def main(argv=None): "Signs for ellipticity components =" + f" ({params['sign_e1']:+d}, {params['sign_e2']:+d})" ) - g1 = data[params["key_e1"]] * params["sign_e1"] - g2 = data[params["key_e2"]] * params["sign_e2"] - cat = treecorr.Catalog( - ra=data[params["key_ra"]], - dec=data[params["key_dec"]], - g1=g1, - g2=g2, - w=data["w"], - ra_units=coord_units, - dec_units=coord_units, - ) - + tomographic = params["tomo_bin1"] is not None and params["tomo_bin2"] is not None + cat2 = None + if tomographic: + if params["verbose"]: + print( + f"Calculating 2PCF for tomographic bins {params['tomo_bin1']} and {params['tomo_bin2']}" + ) + tomo_bin1_idx = data[params["key_tomo_bin_col"]] == params["tomo_bin1"] + tomo_bin2_idx = data[params["key_tomo_bin_col"]] == params["tomo_bin2"] + + g1 = data[params["key_e1"]] * params["sign_e1"] + g2 = data[params["key_e2"]] * params["sign_e2"] + + cat1 = treecorr.Catalog( + ra=data[params["key_ra"]][tomo_bin1_idx], + dec=data[params["key_dec"]][tomo_bin1_idx], + g1=g1[tomo_bin1_idx], + g2=g2[tomo_bin1_idx], + w=data["w"][tomo_bin1_idx], + ra_units=coord_units, + dec_units=coord_units, + ) + if params["tomo_bin1"] != params["tomo_bin2"]: + cat2 = treecorr.Catalog( + ra=data[params["key_ra"]][tomo_bin2_idx], + dec=data[params["key_dec"]][tomo_bin2_idx], + g1=g1[tomo_bin2_idx], + g2=g2[tomo_bin2_idx], + w=data["w"][tomo_bin2_idx], + ra_units=coord_units, + dec_units=coord_units, + ) + else: + if params["verbose"]: + print("Calculating non-tomographic 2PCF") + g1 = data[params["key_e1"]] * params["sign_e1"] + g2 = data[params["key_e2"]] * params["sign_e2"] + cat1 = treecorr.Catalog( + ra=data[params["key_ra"]], + dec=data[params["key_dec"]], + g1=g1, + g2=g2, + w=data["w"], + ra_units=coord_units, + dec_units=coord_units, + ) # Set treecorr config info for correlation sep_units = "arcmin" TreeCorrConfig = { @@ -159,7 +203,7 @@ def main(argv=None): # Compute correlation if params["verbose"]: print("Correlating...") - gg.process(cat, cat) + gg.process(cat1, cat2=cat2) # Write to file if params["verbose"]: diff --git a/src/sp_validation/cosmo_val/cosebis.py b/src/sp_validation/cosmo_val/cosebis.py index aa4657ba..87c36ac2 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,61 @@ 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( + 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[bin_key] = 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[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[bin_key] = next(iter(results[bin_key].values())) return results @@ -152,6 +172,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 @@ -218,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, @@ -228,10 +249,16 @@ 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, ) + 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 @@ -265,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/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/pure_eb.py b/src/sp_validation/cosmo_val/pure_eb.py index 7524074e..f585b0e3 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,13 +62,15 @@ 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 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 @@ -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( + version, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config, + ) + ggs_int = self.calculate_2pcf_version( + version, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config_int, + ) # Get redshift distribution if using analytic covariance + # LG TO-DO: check tomographic redshift distribution handling z_dist = ( np.column_stack(self.get_redshift(version)) if cov_path_int is not None @@ -120,15 +134,20 @@ def calculate_pure_eb( ) # Delegate to b_modes module - results = calculate_pure_eb_correlation( - gg=gg, - gg_int=gg_int, - var_method=var_method, - cov_path_int=cov_path_int, - cosmo_cov=cosmo_cov, - n_samples=n_samples, - z_dist=z_dist, - ) + results = { + bin_key: None for bin_key in ggs.keys() + } # Initialize results dictionary + + for bin_key in ggs.keys(): + results[bin_key] = calculate_pure_eb_correlation( + gg=ggs[bin_key], + gg_int=ggs_int[bin_key], + var_method=var_method, + cov_path_int=cov_path_int, + cosmo_cov=cosmo_cov, + n_samples=n_samples, + z_dist=z_dist, + ) return results diff --git a/src/sp_validation/cosmo_val/real_space.py b/src/sp_validation/cosmo_val/real_space.py index 76d05d0a..96b7a077 100644 --- a/src/sp_validation/cosmo_val/real_space.py +++ b/src/sp_validation/cosmo_val/real_space.py @@ -9,394 +9,174 @@ 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(self, ver, npatch=None, save_fits=False, **treecorr_config): + def calculate_2pcf_version( + self, + ver, + npatch=None, + compute_tomography=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. + 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"``. """ - self.print_magenta(f"Computing {ver} ξ±") - npatch = npatch or self.npatch treecorr_config = { **self._binning(**treecorr_config), "var_method": "jackknife" if int(npatch) > 1 else "shot", } - gg = treecorr.GGCorrelation(treecorr_config) + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + 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." + ) + else: + self.print_magenta(f"Computing non-tomographic ξ± for {ver}.") + tomo_bin_pairs = [("all", "all")] - # If the output file already exists, skip the calculation - 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" - ) + ggs = {f"tomo_bin_{b1}_tomo_bin_{b2}": None for b1, b2 in tomo_bin_pairs} - if os.path.exists(out_fname): - self.print_done(f"Skipping 2PCF calculation, {out_fname} exists") - gg.read(out_fname) + # LG TO-DO: Change to sacc_io method - 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") + patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") - # Use patch file if it exists - patch_file = self._output_path(f"{ver}_patches_npatch={npatch}.dat") + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) + with self.results[ver].temporarily_read_data(): + g1, g2 = self._calibrated_g(ver) + w = self._read_shear_cols(ver, "w_col") - cat_gal = treecorr.Catalog( - ra=self.results[ver].dat_shear["RA"], - dec=self.results[ver].dat_shear["Dec"], - g1=g1, - g2=g2, - w=w, + 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_gal.write_patch_centers(patch_file) + # 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_gal) - gg.write(out_fname, write_patch_results=True, write_cov=True) - - # 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, - ) + gg.process(cat_gal1, cat2=cat_gal2) + ggs[f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = gg - # 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, - ) + self.print_done(f"Done 2PCF for {ver}.") - # Add correlation object to class - if not hasattr(self, "cat_ggs"): - self.cat_ggs = {} - self.cat_ggs[ver] = gg + return ggs - self.print_done("Done 2PCF") + def calculate_2pcf( + self, + npatch=None, + compute_tomography=False, + **treecorr_config, + ): + """ + Calculate the 2PCF ξ± for every catalog version in ``self.versions``. - return gg + 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. - def plot_2pcf(self): - # Plot of n_pairs - plt.subplots(ncols=1, nrows=1) - for ver in self.versions: - self.calculate_2pcf(ver) - plt.plot( - self.cat_ggs[ver].meanr, - self.cat_ggs[ver].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].meanr * cs_plots.dx(idx, fx=1.05, nx=len(ver)), - self.cat_ggs[ver].xip, - yerr=np.sqrt(self.cat_ggs[ver].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].meanr * cs_plots.dx(idx, fx=1.05, nx=len(ver)), - self.cat_ggs[ver].xim, - yerr=np.sqrt(self.cat_ggs[ver].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].meanr, - self.cat_ggs[ver].xip * self.cat_ggs[ver].meanr, - yerr=np.sqrt(self.cat_ggs[ver].varxip) * self.cat_ggs[ver].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].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, - 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].meanr * cs_plots.dx(idx, len(ver)), - self.cat_ggs[ver].xip, - yerr=np.sqrt(self.cat_ggs[ver].varxim), - label=r"$\xi_+$", - ls="solid", - color="green", - ) - plt.errorbar( - self.cat_ggs[ver].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].meanr * cs_plots.dx(idx, len(ver)), - self.cat_ggs[ver].xip + self.xi_psf_sys[ver]["mean"], - yerr=np.sqrt( - self.cat_ggs[ver].varxip + self.xi_psf_sys[ver]["var"] - ), - label=r"$\xi_+ + \xi^{\rm psf}_{+, {\rm sys}}$", - ls="dashdot", - color="magenta", - ) + Parameters: + npatch (int, optional): Number of patches to use; defaults to the + instance's `npatch` attribute. - 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] - - 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 - ) + compute_tomography (bool, optional): Whether to compute tomographic + correlations. Defaults to False. - 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, + **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: + self.cat_ggs[ver] = self.calculate_2pcf_version( + ver, + npatch=npatch, + compute_tomography=compute_tomography, + **treecorr_config, ) - 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}") + # LG TO-DO: No longer writing out text file, change to sacc_io method + + return self.cat_ggs def calculate_aperture_mass_dispersion( - self, theta_min=0.3, theta_max=200, nbins=500, nbins_map=15, npatch=25 + self, + theta_min=0.3, + theta_max=200, + nbins=500, + nbins_map=15, + npatch=25, + compute_tomography=False, ): - self.print_start("Computing aperture-mass dispersion") self._map2 = {} theta_map = np.geomspace(theta_min * 5, theta_max / 2, nbins_map) @@ -405,53 +185,76 @@ def calculate_aperture_mass_dispersion( treecorr_config = self._binning(theta_min, theta_max, nbins) for ver in self.versions: - self.print_magenta(ver) + if compute_tomography: + tomo_bin_ids, tomo_bin_pairs = self._get_tomo_bins(ver) + 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}.") - gg = treecorr.GGCorrelation(treecorr_config) + tomo_bin_pairs = [("all", "all")] - 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"), + self._map2.setdefault(ver, {}) + cat_gal = fits.getdata(self.cc[ver]["shear"]["path"]) + + # LG TO-DO: Change to sacc_io method + 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) + + # 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, ) - gg.process(cat_gal) - gg.write(out_fname) - del cat_gal - del g1 - del g2 + gg.process(cat_gal1, cat2=cat_gal2) - 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") + mapsq, mapsq_im, mxsq, mxsq_im, varmapsq = gg.calculateMapSq( + R=theta_map, + m2_uform="Schneider", + ) + self._map2[ver][f"tomo_bin_{bin1}_tomo_bin_{bin2}"] = { + "mapsq": mapsq, + "mapsq_im": mapsq_im, + "mxsq": mxsq, + "mxsq_im": mxsq_im, + "varmapsq": varmapsq, + } + self.print_done(f"Done aperture-mass dispersion for {ver}.") @property def map2(self): @@ -459,64 +262,383 @@ def map2(self): self.calculate_aperture_mass_dispersion() return self._map2 - 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] - 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, + # 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 + ) + + # 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", ) - 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] - y = [np.abs(self.map2[ver][mode]) for ver in self.versions] - yerr = [np.sqrt(self.map2[ver]["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_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, + ), ) - cs_plots.savefig(out_path, close_fig=False) - cs_plots.show() - self.print_done(f"log-scale {mode} plot saved to {out_path}") + 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 _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, + ) diff --git a/src/sp_validation/tests/test_cosmo_val.py b/src/sp_validation/tests/test_cosmo_val.py index f50992d4..9536cbee 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(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