Introducing BinPairSelector class system for bin pair selection. (#552)

* Introduce BinPairSelector class system for bin combinations.
* Reorganize all tutorials.
* New tutorial for Selectors.
---------
Co-authored-by: Marc Paterno <[email protected]>

Author: vitenti

firecrown/generators/_inferred_galaxy_zdist.py

@@ -546,10 +546,10 @@ def get_lsst_y1_lens_harmonic_bin_collection() -> ZDistLSSTSRDBinCollection:
                 zpu=zpu,
                 sigma_z=y1_lens_bins["sigma_z"],
                 z=RawGrid1D(values=[0.0, 3.5]),
-                bin_name=f"lens_{zpl:.1f}_{zpu:.1f}_y1",
+                bin_name=f"lsst_y1_lens{i}",
                 measurements={Galaxies.COUNTS},
             )
-            for zpl, zpu in pairwise(y1_lens_bins["edges"])
+            for i, (zpl, zpu) in enumerate(pairwise(y1_lens_bins["edges"]))
         ],
     )
 
@@ -568,10 +568,10 @@ def get_lsst_y1_source_harmonic_bin_collection() -> ZDistLSSTSRDBinCollection:
                 zpu=zpu,
                 sigma_z=y1_source_bins["sigma_z"],
                 z=RawGrid1D(values=[0.0, 3.5]),
-                bin_name=f"source_{zpl:.1f}_{zpu:.1f}_y1",
+                bin_name=f"lsst_y1_source{i}",
                 measurements={Galaxies.SHEAR_E},
             )
-            for zpl, zpu in pairwise(y1_source_bins["edges"])
+            for i, (zpl, zpu) in enumerate(pairwise(y1_source_bins["edges"]))
         ],
     )
 
@@ -590,10 +590,10 @@ def get_lsst_y10_lens_harmonic_bin_collection() -> ZDistLSSTSRDBinCollection:
                 zpu=zpu,
                 sigma_z=y10_lens_bins["sigma_z"],
                 z=RawGrid1D(values=[0.0, 3.5]),
-                bin_name=f"lens_{zpl:.1f}_{zpu:.1f}_y10",
+                bin_name=f"lsst_y10_lens{i}",
                 measurements={Galaxies.COUNTS},
             )
-            for zpl, zpu in pairwise(y10_lens_bins["edges"])
+            for i, (zpl, zpu) in enumerate(pairwise(y10_lens_bins["edges"]))
         ],
     )
 
@@ -612,10 +612,10 @@ def get_lsst_y10_source_harmonic_bin_collection() -> ZDistLSSTSRDBinCollection:
                 zpu=zpu,
                 sigma_z=y10_source_bins["sigma_z"],
                 z=RawGrid1D(values=[0.0, 3.5]),
-                bin_name=f"source_{zpl:.1f}_{zpu:.1f}_y10",
+                bin_name=f"lsst_y10_source{i}",
                 measurements={Galaxies.SHEAR_E},
             )
-            for zpl, zpu in pairwise(y10_source_bins["edges"])
+            for i, (zpl, zpu) in enumerate(pairwise(y10_source_bins["edges"]))
         ],
     )
 

firecrown/metadata_functions/__init__.py

@@ -27,6 +27,7 @@
     maybe_enforce_window,
 )
 from firecrown.metadata_functions._combination_utils import (
+    make_binned_two_point_filtered,
     make_all_photoz_bin_combinations,
     make_all_photoz_bin_combinations_with_cmb,
     make_cmb_galaxy_combinations_only,
@@ -53,6 +54,7 @@
     "extract_all_photoz_bin_combinations",
     "extract_window_function",
     "maybe_enforce_window",
+    "make_binned_two_point_filtered",
     "make_all_photoz_bin_combinations",
     "make_all_photoz_bin_combinations_with_cmb",
     "make_cmb_galaxy_combinations_only",

firecrown/metadata_functions/_combination_utils.py

@@ -1,4 +1,11 @@
-"""Utilities for creating combinations of photo-z bins and measurements."""
+"""Utilities for creating combinations of tomographic bins and measurements.
+
+This module provides functions to generate two-point correlation combinations from
+tomographic redshift bins. It supports:
+- All possible galaxy-galaxy correlations
+- CMB-galaxy cross-correlations
+- Filtered combinations based on bin pair selectors
+"""
 
 from itertools import product, chain
 
@@ -7,37 +14,95 @@
 import firecrown.metadata_types as mdt
 
 
+def _validate_list_of_inferred_galaxy_zdists(
+    inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
+) -> None:
+    """Validate that tomographic bin names are unique.
+
+    :param inferred_galaxy_zdists: List of tomographic bins to validate.
+
+    :raises ValueError: If any bin names appear more than once in the list.
+    """
+    bin_names_set = set()
+    # Produce a list of duplicates
+    bin_names = []
+    for igz in inferred_galaxy_zdists:
+        if igz.bin_name in bin_names_set:
+            bin_names.append(igz.bin_name)
+        else:
+            bin_names_set.add(igz.bin_name)
+
+    if bin_names:
+        raise ValueError(
+            f"Duplicate inferred galaxy z distribution bin names found: {bin_names}"
+        )
+
+
 def make_all_photoz_bin_combinations(
     inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
 ) -> list[mdt.TwoPointXY]:
-    """Extract the two-point function metadata from a sacc file."""
-    bin_combinations = [
-        mdt.TwoPointXY(
-            x=igz1, y=igz2, x_measurement=x_measurement, y_measurement=y_measurement
-        )
-        for igz1, igz2 in product(inferred_galaxy_zdists, repeat=2)
-        for x_measurement, y_measurement in product(
-            igz1.measurements, igz2.measurements
-        )
-        if mdt.measurement_is_compatible(x_measurement, y_measurement)
+    """Create all possible two-point correlation combinations for galaxy bins.
+
+    This function generates all possible pairs of (bin, measurement) combinations,
+    keeping only those where the measurements are compatible. For auto-correlations
+    (same measurement type), only unique pairs are kept to avoid duplicates
+    (e.g., only bin0-bin1, not both bin0-bin1 and bin1-bin0).
+
+    :param inferred_galaxy_zdists: List of tomographic redshift bins with their
+        associated measurement types.
+
+    :return: List of all valid TwoPointXY combinations.
+
+    :raises ValueError: If duplicate bin names are found in inferred_galaxy_zdists.
+    """
+    _validate_list_of_inferred_galaxy_zdists(inferred_galaxy_zdists)
+    expanded = [
+        (igz, m) for igz in inferred_galaxy_zdists for m in igz.measurement_list
+    ]
+
+    # Create all combinations of the expanded list, keeping only compatible ones
+    # and avoiding duplicates in the case of correlations of the same type
+    all_xy = [
+        mdt.TwoPointXY(x=igz1, y=igz2, x_measurement=m1, y_measurement=m2)
+        for (igz1, m1), (igz2, m2) in product(expanded, repeat=2)
+        if mdt.measurement_is_compatible(m1, m2)
+        and ((m1 != m2) or (igz2.bin_name >= igz1.bin_name))
     ]
 
-    return bin_combinations
+    # Reorder expanded to have alphabetical order considering first measurements, then
+    # bin names.
+    return sorted(
+        all_xy,
+        key=lambda xy: (
+            xy.x_measurement,
+            xy.y_measurement,
+            xy.x.bin_name,
+            xy.y.bin_name,
+        ),
+    )
 
 
 def make_all_photoz_bin_combinations_with_cmb(
     inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
     cmb_tracer_name: str = "cmb_convergence",
     include_cmb_auto: bool = False,
 ) -> list[mdt.TwoPointXY]:
-    """Create all galaxy combinations plus mdt.CMB-galaxy cross-correlations.
+    """Create all galaxy-galaxy and CMB-galaxy correlation combinations.
 
-    :param inferred_galaxy_zdists: List of galaxy redshift bins
-    :param cmb_tracer_name: Name of the mdt.CMB tracer
-    :param include_cmb_auto: Whether to include mdt.CMB auto-correlation
-        (default: False)
-    :return: List of all XY combinations including mdt.CMB-galaxy crosses
+    This function generates all possible two-point correlations including both
+    galaxy-galaxy auto/cross-correlations and CMB-galaxy cross-correlations.
+
+    :param inferred_galaxy_zdists: List of galaxy redshift bins with their associated
+        measurement types.
+    :param cmb_tracer_name: Name to assign to the CMB tracer (default:
+        "cmb_convergence").
+    :param include_cmb_auto: Whether to include CMB auto-correlation (default: False).
+
+    :return: Combined list of galaxy-galaxy and CMB-galaxy correlation combinations.
+
+    :raises ValueError: If duplicate bin names are found in inferred_galaxy_zdists.
     """
+    _validate_list_of_inferred_galaxy_zdists(inferred_galaxy_zdists)
     # Get all galaxy-galaxy combinations first
     galaxy_combinations = make_all_photoz_bin_combinations(inferred_galaxy_zdists)
     all_combinations = galaxy_combinations + make_cmb_galaxy_combinations_only(
@@ -52,12 +117,24 @@ def make_cmb_galaxy_combinations_only(
     cmb_tracer_name: str = "cmb_convergence",
     include_cmb_auto: bool = False,
 ) -> list[mdt.TwoPointXY]:
-    """Create only mdt.CMB-galaxy cross-correlations.
+    """Create only CMB-galaxy cross-correlations.
+
+    This function generates cross-correlations between CMB convergence and galaxy
+    measurements, optionally including the CMB auto-correlation. It does NOT include
+    any galaxy-galaxy correlations.
+
+    :param inferred_galaxy_zdists: List of galaxy redshift bins with their
+        associated measurement types.
+    :param cmb_tracer_name: Name to assign to the CMB tracer (default:
+        "cmb_convergence").
+    :param include_cmb_auto: Whether to include CMB auto-correlation (default: False).
 
-    :param inferred_galaxy_zdists: List of galaxy redshift bins
-    :param cmb_tracer_name: Name of the mdt.CMB tracer
-    :return: List of mdt.CMB-galaxy cross-correlation XY combinations only
+    :return: List of CMB-galaxy cross-correlation combinations (and optionally CMB
+        auto).
+
+    :raises ValueError: If duplicate bin names are found in inferred_galaxy_zdists.
     """
+    _validate_list_of_inferred_galaxy_zdists(inferred_galaxy_zdists)
     # Create a mock mdt.CMB "bin"
     cmb_bin = mdt.InferredGalaxyZDist(
         bin_name=cmb_tracer_name,
@@ -94,3 +171,35 @@ def make_cmb_galaxy_combinations_only(
         )
 
     return cmb_galaxy_combinations
+
+
+def make_binned_two_point_filtered(
+    inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
+    bin_pair_selector: mdt.BinPairSelector,
+) -> list[mdt.TwoPointXY]:
+    """Create two-point correlations filtered by a bin pair selector.
+
+    This function generates all possible bin combinations and then filters them using
+    the provided selector, keeping only pairs that satisfy the selection criteria
+    (e.g., auto-correlations only, specific measurements, neighboring bins).
+
+    :param inferred_galaxy_zdists: List of tomographic redshift bins with their
+        associated measurement types.
+    :param bin_pair_selector: Selector defining which bin pairs to include.
+
+    :return: List of TwoPointXY combinations that pass the selector's criteria.
+
+    :raises ValueError: If duplicate bin names are found in inferred_galaxy_zdists.
+
+    Example:
+        # Get only auto-correlations of source measurements
+        selector = AutoNameBinPairSelector() & SourceBinPairSelector()
+        combinations = make_binned_two_point_filtered(bins, selector)
+    """
+    _validate_list_of_inferred_galaxy_zdists(inferred_galaxy_zdists)
+    all_bin_combinations = make_all_photoz_bin_combinations(inferred_galaxy_zdists)
+    return [
+        xy
+        for xy in all_bin_combinations
+        if bin_pair_selector.keep((xy.x, xy.y), (xy.x_measurement, xy.y_measurement))
+    ]

firecrown/metadata_functions/_extraction.py

@@ -458,10 +458,19 @@ def extract_all_harmonic_metadata_indices(
 
 def extract_all_harmonic_metadata(
     sacc_data: sacc.Sacc,
-    allow_mixed_types: bool = False,
     allowed_data_type: None | list[str] = None,
+    allow_mixed_types: bool = False,
+    bin_pair_selector: None | mdt.BinPairSelector = None,
 ) -> list[mdt.TwoPointHarmonic]:
-    """Extract the two-point function metadata and data from a sacc file."""
+    """Extract two-point harmonic-space metadata and data from a SACC file.
+
+    :param sacc_data: The SACC object containing tracers and data points.
+    :param allowed_data_type: Optional list of SACC data type strings to include.
+        If None, all harmonic-space data types are extracted.
+    :param bin_pair_selector: Optional selector to filter which bin pairs to include.
+        If None, all valid bin pairs are returned.
+    :return: List of TwoPointHarmonic objects with metadata and ell values.
+    """
     inferred_galaxy_zdists_dict = {
         igz.bin_name: igz
         for igz in extract_all_tracers_inferred_galaxy_zdists(
@@ -478,6 +487,13 @@ def extract_all_harmonic_metadata(
 
         XY = make_two_point_xy(inferred_galaxy_zdists_dict, tracer_names, dt)
 
+        # Apply bin pair selector if provided
+        if bin_pair_selector is not None:
+            if not bin_pair_selector.keep(
+                (XY.x, XY.y), (XY.x_measurement, XY.y_measurement)
+            ):
+                continue
+
         ells, _, indices = sacc_data.get_ell_cl(
             data_type=dt,
             tracer1=tracer_names[0],
@@ -501,10 +517,19 @@ def extract_all_harmonic_metadata(
 
 def extract_all_real_metadata(
     sacc_data: sacc.Sacc,
-    allow_mixed_types: bool = False,
     allowed_data_type: None | list[str] = None,
+    allow_mixed_types: bool = False,
+    bin_pair_selector: None | mdt.BinPairSelector = None,
 ) -> list[mdt.TwoPointReal]:
-    """Extract the two-point function metadata and data from a sacc file."""
+    """Extract two-point real-space metadata and data from a SACC file.
+
+    :param sacc_data: The SACC object containing tracers and data points.
+    :param allowed_data_type: Optional list of SACC data type strings to include.
+        If None, all real-space data types are extracted.
+    :param bin_pair_selector: Optional selector to filter which bin pairs to include.
+        If None, all valid bin pairs are returned.
+    :return: List of TwoPointReal objects with metadata and theta values.
+    """
     inferred_galaxy_zdists_dict = {
         igz.bin_name: igz
         for igz in extract_all_tracers_inferred_galaxy_zdists(
@@ -521,6 +546,13 @@ def extract_all_real_metadata(
 
         XY = make_two_point_xy(inferred_galaxy_zdists_dict, tracer_names, dt)
 
+        # Apply bin pair selector if provided
+        if bin_pair_selector is not None:
+            if not bin_pair_selector.keep(
+                (XY.x, XY.y), (XY.x_measurement, XY.y_measurement)
+            ):
+                continue
+
         t1, t2 = tracer_names
         thetas, _, _ = sacc_data.get_theta_xi(
             data_type=dt, tracer1=t1, tracer2=t2, return_cov=False, return_ind=True
@@ -532,14 +564,31 @@ def extract_all_real_metadata(
 
 
 def extract_all_photoz_bin_combinations(
-    sacc_data: sacc.Sacc, allow_mixed_types: bool = False
+    sacc_data: sacc.Sacc,
+    allow_mixed_types: bool = False,
+    bin_pair_selector: None | mdt.BinPairSelector = None,
 ) -> list[mdt.TwoPointXY]:
-    """Extracts the two-point function metadata from a sacc file."""
+    """Extract all two-point bin pair combinations from a SACC file.
+
+    :param sacc_data: The SACC object containing tracers and data points.
+    :param bin_pair_selector: Optional selector to filter which bin pairs to include.
+        If None, all valid bin pairs are returned.
+    :return: List of TwoPointXY objects representing valid bin pair combinations.
+    """
     inferred_galaxy_zdists = extract_all_tracers_inferred_galaxy_zdists(
         sacc_data, allow_mixed_types
     )
     bin_combinations = make_all_photoz_bin_combinations(inferred_galaxy_zdists)
 
+    if bin_pair_selector is not None:
+        bin_combinations = [
+            xy
+            for xy in bin_combinations
+            if bin_pair_selector.keep(
+                (xy.x, xy.y), (xy.x_measurement, xy.y_measurement)
+            )
+        ]
+
     return bin_combinations