Tomographic 2pcf calculation with treecorr

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"]:

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,