Generalize Tracer Framework and Add CMB Lensing Support

firecrown/app/sacc/_utils.py

@@ -18,7 +18,7 @@
 )
 
 
-def mean_std_tracer(tracer: mdt.InferredGalaxyZDist):
+def mean_std_tracer(tracer: mdt.TomographicBin):
     """Compute the mean and standard deviation of a tracer.
 
     :param tracer: The galaxy redshift distribution tracer to analyze.

firecrown/generators/_inferred_galaxy_zdist.py

@@ -12,7 +12,7 @@
 from scipy.special import erf, erfc, gamma  # pylint: disable=no-name-in-module
 
 from firecrown.metadata_functions import make_measurements, make_measurements_dict
-from firecrown.metadata_types import Galaxies, InferredGalaxyZDist, Measurement
+from firecrown.metadata_types import Galaxies, TomographicBin, Measurement
 
 
 class BinsType(TypedDict):
@@ -360,7 +360,7 @@ def binned_distribution(
         z: npt.NDArray,
         name: str,
         measurements: set[Measurement],
-    ) -> InferredGalaxyZDist:
+    ) -> TomographicBin:
         """Generate the inferred galaxy redshift distribution in bins.
 
         :param zpl: The lower bound of the integration
@@ -411,7 +411,7 @@ def _P(z, _):
                 / norma
             )
 
-        return InferredGalaxyZDist(
+        return TomographicBin(
             bin_name=name, z=z_knots, dndz=dndz, measurements=measurements
         )
 
@@ -463,7 +463,7 @@ def serialize_measurements(cls, value: set[Measurement]) -> list[dict]:
         """Serialize the Measurement."""
         return make_measurements_dict(value)
 
-    def generate(self, zdist: ZDistLSSTSRD) -> InferredGalaxyZDist:
+    def generate(self, zdist: ZDistLSSTSRD) -> TomographicBin:
         """Generate the inferred galaxy redshift distribution in bins."""
         return zdist.binned_distribution(
             zpl=self.zpl,
@@ -489,7 +489,7 @@ class ZDistLSSTSRDBinCollection(BaseModel):
     autoknots_reltol: float = 1.0e-4
     autoknots_abstol: float = 1.0e-15
 
-    def generate(self) -> list[InferredGalaxyZDist]:
+    def generate(self) -> list[TomographicBin]:
         """Generate the inferred galaxy redshift distributions in bins."""
         zdist = ZDistLSSTSRD(
             alpha=self.alpha,

firecrown/likelihood/_cmb.py

@@ -7,7 +7,7 @@
 from pydantic import BaseModel, ConfigDict, PrivateAttr
 
 from firecrown.likelihood._base import Source, Tracer
-from firecrown.metadata_types import InferredGalaxyZDist, TypeSource
+from firecrown.metadata_types import ProjectedField, CMBLensing, TypeSource
 from firecrown.modeling_tools import ModelingTools
 
 
@@ -113,14 +113,15 @@ def model_post_init(self, _, /) -> None:
         """Initialize the CMBConvergenceFactory."""
         self._cache: dict[int, CMBConvergence] = {}
 
-    def create(self, inferred_galaxy_zdist: InferredGalaxyZDist) -> CMBConvergence:
+    def create(self, cmb_tracer: ProjectedField) -> CMBConvergence:
         """Create a CMBConvergence object with the given inferred galaxy z distribution.
 
-        :param inferred_galaxy_zdist: the inferred galaxy redshift distribution
+        :param cmb_tracer: the inferred galaxy redshift distribution
         :return: a fully initialized CMBConvergence object
         """
+        assert isinstance(cmb_tracer, CMBLensing)
         # Use the bin_name as the tracer identifier
-        sacc_tracer = inferred_galaxy_zdist.bin_name
+        sacc_tracer = cmb_tracer.bin_name
         tracer_id = hash(sacc_tracer)
 
         if tracer_id in self._cache:

firecrown/likelihood/_two_point.py

@@ -37,7 +37,6 @@
     CMB_TYPES,
     GALAXY_LENS_TYPES,
     GALAXY_SOURCE_TYPES,
-    InferredGalaxyZDist,
     Measurement,
     TracerNames,
     TwoPointCorrelationSpace,
@@ -280,6 +279,11 @@ def _from_metadata_single_base(
 
         :return: A TwoPoint statistic.
         """
+        # metadata.XY.x/Y are typed as Tracer (protocol). In this code path we
+        # expect concrete TomographicBin instances (with redshift arrays). Use
+        # runtime isinstance checks and raise an informative error if this is
+        # not the case — this both documents the assumption and narrows the
+        # type for the type checker.
         source0 = use_source_factory(
             metadata.XY.x, metadata.XY.x_measurement, tp_factory
         )
@@ -863,21 +867,19 @@ def from_metadata(
 
 
 def use_source_factory(
-    inferred_galaxy_zdist: InferredGalaxyZDist,
+    tomographic_bin: mdt.ProjectedField,
     measurement: Measurement,
     tp_factory: TwoPointFactory,
 ) -> WeakLensing | NumberCounts | CMBConvergence:
     """Apply the factory to the inferred galaxy redshift distribution."""
-    if measurement not in inferred_galaxy_zdist.measurements:
+    if measurement not in tomographic_bin.measurements:
         raise ValueError(
             f"Measurement {measurement} not found in inferred galaxy redshift "
-            f"distribution {inferred_galaxy_zdist.bin_name}!"
+            f"distribution {tomographic_bin.bin_name}!"
         )
 
-    source_factory = tp_factory.get_factory(
-        measurement, inferred_galaxy_zdist.type_source
-    )
-    source = source_factory.create(inferred_galaxy_zdist)
+    source_factory = tp_factory.get_factory(measurement, tomographic_bin.type_source)
+    source = source_factory.create(tomographic_bin)
     return source
 
 

firecrown/likelihood/_weak_lensing.py

@@ -26,7 +26,7 @@
     SourceGalaxySystematic,
     Tracer,
 )
-from firecrown.metadata_types import InferredGalaxyZDist, TypeSource
+from firecrown.metadata_types import ProjectedField, TomographicBin, TypeSource
 from firecrown.modeling_tools import ModelingTools
 from firecrown.updatable import ParamsMap
 
@@ -462,14 +462,17 @@ def __init__(
     @classmethod
     def create_ready(
         cls,
-        inferred_zdist: InferredGalaxyZDist,
+        tomographic_bin: TomographicBin,
         systematics: None | list[SourceGalaxySystematic[WeakLensingArgs]] = None,
     ) -> WeakLensing:
         """Create a WeakLensing object with the given tracer name and scale."""
-        obj = cls(sacc_tracer=inferred_zdist.bin_name, systematics=systematics)
+        obj = cls(sacc_tracer=tomographic_bin.bin_name, systematics=systematics)
         # pylint: disable=unexpected-keyword-arg
         obj.tracer_args = WeakLensingArgs(
-            scale=obj.scale, z=inferred_zdist.z, dndz=inferred_zdist.dndz, ia_bias=None
+            scale=obj.scale,
+            z=tomographic_bin.z,
+            dndz=tomographic_bin.dndz,
+            ia_bias=None,
         )
         # pylint: enable=unexpected-keyword-arg
 
@@ -580,7 +583,7 @@ class MultiplicativeShearBiasFactory(BaseModel):
     def create(self, bin_name: str) -> MultiplicativeShearBias:
         """Create a MultiplicativeShearBias object.
 
-        :param inferred_zdist: The inferred galaxy redshift distribution for
+        :param tomographic_bin: The inferred galaxy redshift distribution for
             the created MultiplicativeShearBias object.
         :return: The created MultiplicativeShearBias object.
         """
@@ -609,7 +612,7 @@ class LinearAlignmentSystematicFactory(BaseModel):
     def create(self, bin_name: str) -> LinearAlignmentSystematic:
         """Create a LinearAlignmentSystematic object.
 
-        :param inferred_zdist: The inferred galaxy redshift distribution for
+        :param tomographic_bin: The inferred galaxy redshift distribution for
             the created LinearAlignmentSystematic object.
         :return: The created LinearAlignmentSystematic object.
         """
@@ -637,7 +640,7 @@ class TattAlignmentSystematicFactory(BaseModel):
     def create(self, bin_name: str) -> TattAlignmentSystematic:
         """Create a TattAlignmentSystematic object.
 
-        :param inferred_zdist: The inferred galaxy redshift distribution for
+        :param tomographic_bin: The inferred galaxy redshift distribution for
             the created TattAlignmentSystematic object.
         :return: The created TattAlignmentSystematic object.
         """
@@ -687,19 +690,20 @@ def model_post_init(self, _, /) -> None:
             for wl_systematic_factory in self.global_systematics
         ]
 
-    def create(self, inferred_zdist: InferredGalaxyZDist) -> WeakLensing:
+    def create(self, tomographic_bin: ProjectedField) -> WeakLensing:
         """Create a WeakLensing object with the given tracer name and scale."""
-        inferred_zdist_id = id(inferred_zdist)
+        assert isinstance(tomographic_bin, TomographicBin)
+        inferred_zdist_id = id(tomographic_bin)
         if inferred_zdist_id in self._cache:
             return self._cache[inferred_zdist_id]
 
         systematics: list[SourceGalaxySystematic[WeakLensingArgs]] = [
-            systematic_factory.create(inferred_zdist.bin_name)
+            systematic_factory.create(tomographic_bin.bin_name)
             for systematic_factory in self.per_bin_systematics
         ]
         systematics.extend(self._global_systematics_instances)
 
-        wl = WeakLensing.create_ready(inferred_zdist, systematics)
+        wl = WeakLensing.create_ready(tomographic_bin, systematics)
         self._cache[inferred_zdist_id] = wl
 
         return wl

firecrown/likelihood/number_counts/_factories.py

@@ -20,7 +20,7 @@
 )
 from firecrown.likelihood._base import SourceGalaxySystematic
 from firecrown.likelihood.number_counts._args import NumberCountsArgs
-from firecrown.metadata_types import InferredGalaxyZDist, TypeSource
+from firecrown.metadata_types import ProjectedField, TomographicBin, TypeSource
 
 
 class LinearBiasSystematicFactory(BaseModel):
@@ -155,24 +155,25 @@ def model_post_init(self, _, /) -> None:
             for nc_systematic_factory in self.global_systematics
         ]
 
-    def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCounts:
+    def create(self, tomographic_bin: ProjectedField) -> NumberCounts:
         """Create a NumberCounts object with the given tracer name and scale.
 
-        :param inferred_zdist: the inferred redshift distribution
+        :param tomographic_bin: the inferred redshift distribution
         :return: a fully initialized NumberCounts object
         """
-        inferred_zdist_id = id(inferred_zdist)
+        assert isinstance(tomographic_bin, TomographicBin)
+        inferred_zdist_id = id(tomographic_bin)
         if inferred_zdist_id in self._cache:
             return self._cache[inferred_zdist_id]
 
         systematics: list[SourceGalaxySystematic[NumberCountsArgs]] = [
-            systematic_factory.create(inferred_zdist.bin_name)
+            systematic_factory.create(tomographic_bin.bin_name)
             for systematic_factory in self.per_bin_systematics
         ]
         systematics.extend(self._global_systematics_instances)
 
         nc = NumberCounts.create_ready(
-            inferred_zdist, systematics=systematics, has_rsd=self.include_rsd
+            tomographic_bin, systematics=systematics, has_rsd=self.include_rsd
         )
         self._cache[inferred_zdist_id] = nc
 

firecrown/likelihood/number_counts/_source.py

@@ -16,7 +16,7 @@
     SourceGalaxySystematic,
     Tracer,
 )
-from firecrown.metadata_types import InferredGalaxyZDist
+from firecrown.metadata_types import TomographicBin
 from firecrown.modeling_tools import ModelingTools
 from firecrown.updatable import (
     DerivedParameter,
@@ -69,7 +69,7 @@ def __init__(
     @classmethod
     def create_ready(
         cls,
-        inferred_zdist: InferredGalaxyZDist,
+        tomographic_bin: TomographicBin,
         has_rsd: bool = False,
         derived_scale: bool = False,
         scale: float = 1.0,
@@ -80,7 +80,7 @@ def create_ready(
         This is the recommended way to create a NumberCounts object. It creates
         a fully initialized object.
 
-        :param inferred_zdist: the inferred redshift distribution
+        :param tomographic_bin: the inferred redshift distribution
         :param has_rsd: whether to include RSD in the tracer
         :param derived_scale: whether to include a derived parameter for the scale
             of the tracer
@@ -89,7 +89,7 @@ def create_ready(
         :return: a fully initialized NumberCounts object
         """
         obj = cls(
-            sacc_tracer=inferred_zdist.bin_name,
+            sacc_tracer=tomographic_bin.bin_name,
             systematics=systematics,
             has_rsd=has_rsd,
             derived_scale=derived_scale,
@@ -98,8 +98,8 @@ def create_ready(
         # pylint: disable=unexpected-keyword-arg
         obj.tracer_args = NumberCountsArgs(
             scale=obj.scale,
-            z=inferred_zdist.z,
-            dndz=inferred_zdist.dndz,
+            z=tomographic_bin.z,
+            dndz=tomographic_bin.dndz,
             bias=None,
             mag_bias=None,
         )

firecrown/metadata_functions/_combination_utils.py

@@ -9,8 +9,6 @@
 
 from itertools import product, chain
 
-import numpy as np
-
 import firecrown.metadata_types as mdt
 
 
@@ -39,7 +37,7 @@ def _validate_list_of_inferred_galaxy_zdists(
 
 
 def make_all_photoz_bin_combinations(
-    inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
+    inferred_galaxy_zdists: list[mdt.TomographicBin],
 ) -> list[mdt.TwoPointXY]:
     """Create all possible two-point correlation combinations for galaxy bins.
 
@@ -83,7 +81,7 @@ def make_all_photoz_bin_combinations(
 
 
 def make_all_photoz_bin_combinations_with_cmb(
-    inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
+    inferred_galaxy_zdists: list[mdt.TomographicBin],
     cmb_tracer_name: str = "cmb_convergence",
     include_cmb_auto: bool = False,
 ) -> list[mdt.TwoPointXY]:
@@ -113,7 +111,7 @@ def make_all_photoz_bin_combinations_with_cmb(
 
 
 def make_cmb_galaxy_combinations_only(
-    inferred_galaxy_zdists: list[mdt.InferredGalaxyZDist],
+    inferred_galaxy_zdists: list[mdt.TomographicBin],
     cmb_tracer_name: str = "cmb_convergence",
     include_cmb_auto: bool = False,
 ) -> list[mdt.TwoPointXY]:
@@ -136,19 +134,24 @@ def make_cmb_galaxy_combinations_only(
     """
     _validate_list_of_inferred_galaxy_zdists(inferred_galaxy_zdists)
     # Create a mock mdt.CMB "bin"
-    cmb_bin = mdt.InferredGalaxyZDist(
+    cmb_bin = mdt.CMBLensing(
         bin_name=cmb_tracer_name,
-        z=np.array([1100.0]),
-        dndz=np.array([1.0]),
+        z_lss=1100.0,
         measurements={mdt.CMB.CONVERGENCE},
         type_source=mdt.TypeSource.DEFAULT,
     )
 
     cmb_galaxy_combinations = []
 
+    x: mdt.ProjectedField
+    y: mdt.ProjectedField
     for galaxy_bin in inferred_galaxy_zdists:
-        galaxy_bin_type = list(product([galaxy_bin], list(galaxy_bin.measurements)))
-        cmb_bin_type = list(product([cmb_bin], list(cmb_bin.measurements)))
+        galaxy_bin_type: list[tuple[mdt.ProjectedField, mdt.Measurement]] = list(
+            product([galaxy_bin], list(galaxy_bin.measurements))
+        )
+        cmb_bin_type: list[tuple[mdt.ProjectedField, mdt.Measurement]] = list(
+            product([cmb_bin], list(cmb_bin.measurements))
+        )
         for (x, m1), (y, m2) in chain(
             product(galaxy_bin_type, cmb_bin_type),
             product(cmb_bin_type, galaxy_bin_type),

firecrown/metadata_functions/_extraction.py

@@ -20,7 +20,7 @@
 def extract_all_tracers_inferred_galaxy_zdists(
     sacc_data: sacc.Sacc,
     allow_mixed_types: bool = False,
-) -> list[mdt.InferredGalaxyZDist]:
+) -> list[mdt.TomographicBin]:
     """Extracts the two-point function metadata from a Sacc object.
 
     The Sacc object contains a set of tracers (one-dimensional bins) and data
@@ -44,7 +44,7 @@ def extract_all_tracers_inferred_galaxy_zdists(
                 )
 
     return [
-        mdt.InferredGalaxyZDist(
+        mdt.TomographicBin(
             bin_name=tracer.name,
             z=tracer.z,
             dndz=tracer.nz,