Two point data (#432)

* Reorganizing two-point types in a different module.
* Adding optional data to two_point dataclasses.
* Updated usage of two_point dataclasses.
* Removed required data_vector when computing theory_vector.
* Testing new constructors.
* Comparing new and old constructors.
* Testing new measurement code.
* New constructor for gauss family for creating ready likelihoods.
* New tutorial for two point factories.
* Constructors to create ready.
* Added SHEAR_MINUS and SHEAR_PLUS.
* Created factories to create both ready and not-ready objects.
* Updated tests to new types.
* Adding likelihood creation to tutorial.
* Adding support for multiple measurements in the same bin.
* Added likelihood comparison on two_point_factories.
* Splitting metadata test file.
* Updating test to deal with enum str representations in python 3.10.
* Reorganizing fixtures in conftest.py.
* Testing ell and theta wrong shape check.
* Testing multiple data_types.
* Testing inverting tracer order.
* Testing name type matching.
* Testing naming convention enforcing.
* Test the consistence checks for a set of two-point dataclasses.
* Testing indices and covariance consistency.
* Testing sacc covariance matrix presence.
* Testing ambiguous types.
* Reorganized imports and added missing tests to TwoPointCWindow.
* Updated tutorials to match new module names.
* Test unsupported type exception.
* Testing metadata only constructors.
* Testing ConstGaussian create ready.
* Fixing testing covariance size.
* Testing the use of non ready statistics.

firecrown/generators/inferred_galaxy_zdist.py

@@ -12,16 +12,19 @@
 
 from numcosmo_py import Ncm
 
-from firecrown.metadata.two_point import (
+from firecrown.metadata.two_point_types import (
     InferredGalaxyZDist,
+    ALL_MEASUREMENT_TYPES,
+)
+from firecrown.metadata.two_point_types import (
+    make_measurements_dict,
     Measurement,
     Galaxies,
     CMB,
     Clusters,
-    ALL_MEASUREMENT_TYPES,
-    make_measurement_dict,
 )
 
+
 BinsType = TypedDict("BinsType", {"edges": npt.NDArray, "sigma_z": float})
 
 Y1_ALPHA = 0.94
@@ -159,7 +162,7 @@ def binned_distribution(
         sigma_z: float,
         z: npt.NDArray,
         name: str,
-        measurement: Measurement,
+        measurements: set[Measurement],
         use_autoknot: bool = False,
         autoknots_reltol: float = 1.0e-4,
         autoknots_abstol: float = 1.0e-15,
@@ -171,7 +174,7 @@ def binned_distribution(
         :param sigma_z: The resolution parameter
         :param z: The redshifts at which to evaluate the distribution
         :param name: The name of the distribution
-        :param measured_type: The measured type of the distribution
+        :param measurements: The set of measurements of the distribution
         :param use_autoknot: Whether to use the NotAKnot algorithm of NumCosmo
         :param autoknots_reltol: The relative tolerance for the NotAKnot algorithm
         :param autoknots_abstol: The absolute tolerance for the NotAKnot algorithm
@@ -217,7 +220,7 @@ def _P(z, _):
             )
 
         return InferredGalaxyZDist(
-            bin_name=name, z=z_knots, dndz=dndz, measurement=measurement
+            bin_name=name, z=z_knots, dndz=dndz, measurements=measurements
         )
 
 
@@ -250,31 +253,40 @@ def generate(self) -> npt.NDArray:
 Grid1D = LinearGrid1D | RawGrid1D
 
 
-def make_measurement(value: Measurement | dict[str, Any]) -> Measurement:
+def make_measurements(
+    value: set[Measurement] | list[dict[str, Any]]
+) -> set[Measurement]:
     """Create a Measurement object from a dictionary."""
-    if isinstance(value, ALL_MEASUREMENT_TYPES):
+    if isinstance(value, set) and all(
+        isinstance(v, ALL_MEASUREMENT_TYPES) for v in value
+    ):
         return value
 
-    if not isinstance(value, dict):
-        raise ValueError(f"Invalid Measurement: {value} is not a dictionary")
+    measurements: set[Measurement] = set()
+    for measurement_dict in value:
+        if not isinstance(measurement_dict, dict):
+            raise ValueError(f"Invalid Measurement: {value} is not a dictionary")
 
-    if "subject" not in value:
-        raise ValueError("Invalid Measurement: dictionary does not contain 'subject'")
-
-    subject = value["subject"]
-
-    match subject:
-        case "Galaxies":
-            return Galaxies[value["property"]]
-        case "CMB":
-            return CMB[value["property"]]
-        case "Clusters":
-            return Clusters[value["property"]]
-        case _:
+        if "subject" not in measurement_dict:
             raise ValueError(
-                f"Invalid Measurement: subject: '{subject}' is not recognized"
+                "Invalid Measurement: dictionary does not contain 'subject'"
             )
 
+        subject = measurement_dict["subject"]
+
+        match subject:
+            case "Galaxies":
+                measurements.update({Galaxies[measurement_dict["property"]]})
+            case "CMB":
+                measurements.update({CMB[measurement_dict["property"]]})
+            case "Clusters":
+                measurements.update({Clusters[measurement_dict["property"]]})
+            case _:
+                raise ValueError(
+                    f"Invalid Measurement: subject: '{subject}' is not recognized"
+                )
+    return measurements
+
 
 class ZDistLSSTSRDBin(BaseModel):
     """LSST Inferred galaxy redshift distributions in bins."""
@@ -286,16 +298,16 @@ class ZDistLSSTSRDBin(BaseModel):
     sigma_z: float
     z: Annotated[Grid1D, Field(union_mode="left_to_right")]
     bin_name: str
-    measurement: Annotated[Measurement, BeforeValidator(make_measurement)]
+    measurements: Annotated[set[Measurement], BeforeValidator(make_measurements)]
     use_autoknot: bool = False
     autoknots_reltol: float = 1.0e-4
     autoknots_abstol: float = 1.0e-15
 
-    @field_serializer("measurement")
+    @field_serializer("measurements")
     @classmethod
-    def serialize_measurement(cls, value: Measurement) -> dict:
+    def serialize_measurements(cls, value: set[Measurement]) -> list[dict]:
         """Serialize the Measurement."""
-        return make_measurement_dict(value)
+        return make_measurements_dict(value)
 
     def generate(self, zdist: ZDistLSSTSRD) -> InferredGalaxyZDist:
         """Generate the inferred galaxy redshift distribution in bins."""
@@ -305,7 +317,7 @@ def generate(self, zdist: ZDistLSSTSRD) -> InferredGalaxyZDist:
             sigma_z=self.sigma_z,
             z=self.z.generate(),
             name=self.bin_name,
-            measurement=self.measurement,
+            measurements=self.measurements,
             use_autoknot=self.use_autoknot,
             autoknots_reltol=self.autoknots_reltol,
             autoknots_abstol=self.autoknots_abstol,
@@ -339,7 +351,7 @@ def generate(self) -> list[InferredGalaxyZDist]:
             sigma_z=Y1_LENS_BINS["sigma_z"],
             z=RawGrid1D(values=[0.0, 3.0]),
             bin_name=f"lens_{zpl:.1f}_{zpu:.1f}_y1",
-            measurement=Galaxies.COUNTS,
+            measurements={Galaxies.COUNTS},
             use_autoknot=True,
             autoknots_reltol=1.0e-5,
         )
@@ -358,7 +370,7 @@ def generate(self) -> list[InferredGalaxyZDist]:
             sigma_z=Y1_SOURCE_BINS["sigma_z"],
             z=RawGrid1D(values=[0.0, 3.0]),
             bin_name=f"source_{zpl:.1f}_{zpu:.1f}_y1",
-            measurement=Galaxies.SHEAR_E,
+            measurements={Galaxies.SHEAR_E},
             use_autoknot=True,
             autoknots_reltol=1.0e-5,
         )
@@ -377,7 +389,7 @@ def generate(self) -> list[InferredGalaxyZDist]:
             sigma_z=Y10_LENS_BINS["sigma_z"],
             z=RawGrid1D(values=[0.0, 3.0]),
             bin_name=f"lens_{zpl:.1f}_{zpu:.1f}_y10",
-            measurement=Galaxies.COUNTS,
+            measurements={Galaxies.COUNTS},
             use_autoknot=True,
             autoknots_reltol=1.0e-5,
         )
@@ -396,7 +408,7 @@ def generate(self) -> list[InferredGalaxyZDist]:
             sigma_z=Y10_SOURCE_BINS["sigma_z"],
             z=RawGrid1D(values=[0.0, 3.0]),
             bin_name=f"source_{zpl:.1f}_{zpu:.1f}_y10",
-            measurement=Galaxies.SHEAR_E,
+            measurements={Galaxies.SHEAR_E},
             use_autoknot=True,
             autoknots_reltol=1.0e-5,
         )

firecrown/likelihood/gaussfamily.py

@@ -156,6 +156,16 @@ def __init__(
         self.theory_vector: None | npt.NDArray[np.double] = None
         self.data_vector: None | npt.NDArray[np.double] = None
 
+    @classmethod
+    def create_ready(
+        cls, statistics: Sequence[Statistic], covariance: npt.NDArray[np.float64]
+    ) -> GaussFamily:
+        """Create a GaussFamily object in the READY state."""
+        obj = cls(statistics)
+        obj._set_covariance(covariance)
+        obj.state = State.READY
+        return obj
+
     @enforce_states(
         initial=State.READY,
         terminal=State.UPDATED,
@@ -199,12 +209,26 @@ def read(self, sacc_data: sacc.Sacc) -> None:
             )
             raise RuntimeError(msg)
 
+        for stat in self.statistics:
+            stat.read(sacc_data)
+
         covariance = sacc_data.covariance.dense
 
+        self._set_covariance(covariance)
+
+    def _set_covariance(self, covariance):
+        """Set the covariance matrix.
+
+        This method is used to set the covariance matrix and perform the
+        necessary calculations to prepare the likelihood for computation.
+        """
         indices_list = []
         data_vector_list = []
         for stat in self.statistics:
-            stat.read(sacc_data)
+            if not stat.statistic.ready:
+                raise RuntimeError(
+                    f"The statistic {stat.statistic} is not ready to be used."
+                )
             if stat.statistic.sacc_indices is None:
                 raise RuntimeError(
                     f"The statistic {stat.statistic} has no sacc_indices."
@@ -216,6 +240,19 @@ def read(self, sacc_data: sacc.Sacc) -> None:
         data_vector = np.concatenate(data_vector_list)
         cov = np.zeros((len(indices), len(indices)))
 
+        largest_index = np.max(indices)
+
+        if not (
+            covariance.ndim == 2
+            and covariance.shape[0] == covariance.shape[1]
+            and largest_index < covariance.shape[0]
+        ):
+            raise ValueError(
+                f"The covariance matrix has shape {covariance.shape}, "
+                f"but the expected shape is at least "
+                f"{(largest_index + 1, largest_index + 1)}."
+            )
+
         for new_i, old_i in enumerate(indices):
             for new_j, old_j in enumerate(indices):
                 cov[new_i, new_j] = covariance[old_i, old_j]

firecrown/likelihood/number_counts.py

@@ -487,9 +487,9 @@ class PhotoZShiftFactory(BaseModel):
         Field(description="The type of the systematic."),
     ] = "PhotoZShiftFactory"
 
-    def create(self, inferred_zdist: InferredGalaxyZDist) -> PhotoZShift:
+    def create(self, bin_name: str) -> PhotoZShift:
         """Create a PhotoZShift object with the given tracer name."""
-        return PhotoZShift(inferred_zdist.bin_name)
+        return PhotoZShift(bin_name)
 
     def create_global(self) -> PhotoZShift:
         """Create a PhotoZShift object with the given tracer name."""
@@ -506,9 +506,9 @@ class LinearBiasSystematicFactory(BaseModel):
         Field(description="The type of the systematic."),
     ] = "LinearBiasSystematicFactory"
 
-    def create(self, inferred_zdist: InferredGalaxyZDist) -> LinearBiasSystematic:
+    def create(self, bin_name: str) -> LinearBiasSystematic:
         """Create a LinearBiasSystematic object with the given tracer name."""
-        return LinearBiasSystematic(inferred_zdist.bin_name)
+        return LinearBiasSystematic(bin_name)
 
     def create_global(self) -> LinearBiasSystematic:
         """Create a LinearBiasSystematic object with the given tracer name."""
@@ -525,9 +525,9 @@ class PTNonLinearBiasSystematicFactory(BaseModel):
         Field(description="The type of the systematic."),
     ] = "PTNonLinearBiasSystematicFactory"
 
-    def create(self, inferred_zdist: InferredGalaxyZDist) -> PTNonLinearBiasSystematic:
+    def create(self, bin_name: str) -> PTNonLinearBiasSystematic:
         """Create a PTNonLinearBiasSystematic object with the given tracer name."""
-        return PTNonLinearBiasSystematic(inferred_zdist.bin_name)
+        return PTNonLinearBiasSystematic(bin_name)
 
     def create_global(self) -> PTNonLinearBiasSystematic:
         """Create a PTNonLinearBiasSystematic object with the given tracer name."""
@@ -544,11 +544,9 @@ class MagnificationBiasSystematicFactory(BaseModel):
         Field(description="The type of the systematic."),
     ] = "MagnificationBiasSystematicFactory"
 
-    def create(
-        self, inferred_zdist: InferredGalaxyZDist
-    ) -> MagnificationBiasSystematic:
+    def create(self, bin_name: str) -> MagnificationBiasSystematic:
         """Create a MagnificationBiasSystematic object with the given tracer name."""
-        return MagnificationBiasSystematic(inferred_zdist.bin_name)
+        return MagnificationBiasSystematic(bin_name)
 
     def create_global(self) -> MagnificationBiasSystematic:
         """Create a MagnificationBiasSystematic object with the given tracer name."""
@@ -565,14 +563,12 @@ class ConstantMagnificationBiasSystematicFactory(BaseModel):
         Field(description="The type of the systematic."),
     ] = "ConstantMagnificationBiasSystematicFactory"
 
-    def create(
-        self, inferred_zdist: InferredGalaxyZDist
-    ) -> ConstantMagnificationBiasSystematic:
+    def create(self, bin_name: str) -> ConstantMagnificationBiasSystematic:
         """Create a ConstantMagnificationBiasSystematic object.
 
         Use the inferred_zdist to create the systematic.
         """
-        return ConstantMagnificationBiasSystematic(inferred_zdist.bin_name)
+        return ConstantMagnificationBiasSystematic(bin_name)
 
     def create_global(self) -> ConstantMagnificationBiasSystematic:
         """Create a ConstantMagnificationBiasSystematic object.
@@ -618,7 +614,7 @@ def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCounts:
             return self._cache[inferred_zdist_id]
 
         systematics: list[SourceGalaxySystematic[NumberCountsArgs]] = [
-            systematic_factory.create(inferred_zdist)
+            systematic_factory.create(inferred_zdist.bin_name)
             for systematic_factory in self.per_bin_systematics
         ]
         systematics.extend(self._global_systematics_instances)
@@ -627,3 +623,22 @@ def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCounts:
         self._cache[inferred_zdist_id] = nc
 
         return nc
+
+    def create_from_metadata_only(
+        self,
+        sacc_tracer: str,
+    ) -> NumberCounts:
+        """Create an WeakLensing object with the given tracer name and scale."""
+        sacc_tracer_id = hash(sacc_tracer)  # Improve this
+        if sacc_tracer_id in self._cache:
+            return self._cache[sacc_tracer_id]
+        systematics: list[SourceGalaxySystematic[NumberCountsArgs]] = [
+            systematic_factory.create(sacc_tracer)
+            for systematic_factory in self.per_bin_systematics
+        ]
+        systematics.extend(self._global_systematics_instances)
+
+        nc = NumberCounts(sacc_tracer=sacc_tracer, systematics=systematics)
+        self._cache[sacc_tracer_id] = nc
+
+        return nc