TwoPoint metadata, constructiors for ready objects, metadata serializ…

…ation. (#399)

* Creating constructiors for ready objects.
* Added a set of dataclasses to describe TwoPoint sources metadata.
* Defining MeasuredTypes.
* Adding translation between metadata types and SACC strings
* Refactoring to reduce complexity of read method.
* Using (str, Enum) since StrEnum is not available before python 3.11. 
* kw_only option for dataclasses is only available on python >= 3.10.
* Refactor _compute_theory_vector into _compute_theory_vector_harmonic and _compute_theory_vector_real_space.
* Testing measured types compatibilities.
* Testing type related exceptions.
* Using MagicMock in a way compatible with py39 -- py312.
* TwoPoint metadata testing.
* Testing SACC metadata extraction.
* Better type detection, more tests.
* Testing cosmosis output in the datablock.
* Checking two point harmonic and real.
* Testing create ready and source factories.
* Testing factory cache.
* Testing two_point windows and cuts.
* Testing more error cases.
* Testing data and configuration warnings.
* Add serialization to and from YAML
* Make TwoPoint serializable, add tests and fixtures
* Add serialization for TwoPointCells
* Add serialization support for Window
* Add method to test optional numpy arrays for equality
* Add serialization for TwoPointCWindow
* Add serialization to TwoPointXiTheta
* Add restriction to glib version
* Add testing for Window equality method

---------

Co-authored-by: Sandro Dias Pinto Vitenti <vitenti@uel.br>

environment.yml

@@ -16,6 +16,7 @@ dependencies:
   - flake8-docstrings
   - fuzzywuzzy
   - getdist
+  - glib <= 2.78.4
   - idna
   - matplotlib-base
   - more-itertools

firecrown/connector/cosmosis/likelihood.py

@@ -98,9 +98,9 @@ def execute(self, sample: cosmosis.datablock) -> int:
         )
         self.map.set_params_from_cosmosis(cosmological_params)
 
-        ccl_args = self.map.calculate_ccl_args(sample)
-
-        ccl_cosmo = ccl.CosmologyCalculator(**self.map.asdict(), **ccl_args)
+        ccl_cosmo = ccl.CosmologyCalculator(
+            **self.map.asdict(), **self.map.calculate_ccl_args(sample)
+        )
 
         # TODO: Future development will need to capture elements that get put into the
         # datablock. This probably will be in a different "physics module" and not in
@@ -156,22 +156,31 @@ def execute(self, sample: cosmosis.datablock) -> int:
         # some method in the Statistic base class should be called here. For
         # statistics other than TwoPoint, the base class implementation should
         # do nothing.
-        for stat in self.likelihood.statistics:
+        for gstat in self.likelihood.statistics:
+            stat = gstat.statistic
+
             if isinstance(stat, TwoPoint):
                 assert stat.sacc_tracers is not None
                 tracer = f"{stat.sacc_tracers[0]}_{stat.sacc_tracers[1]}"
 
-                sample.put(
-                    "data_vector",
-                    f"ell_or_theta_{stat.sacc_data_type}_{tracer}",
-                    stat.ell_or_theta_,
-                )
-                sample.put(
+                if stat.ells is not None:
+                    sample.put_int_array_1d(
+                        "data_vector",
+                        f"ell_{stat.sacc_data_type}_{tracer}",
+                        stat.ells,
+                    )
+                elif stat.thetas is not None:
+                    sample.put_double_array_1d(
+                        "data_vector",
+                        f"theta_{stat.sacc_data_type}_{tracer}",
+                        stat.thetas,
+                    )
+                sample.put_double_array_1d(
                     "data_vector",
                     f"theory_{stat.sacc_data_type}_{tracer}",
                     stat.get_theory_vector(),
                 )
-                sample.put(
+                sample.put_double_array_1d(
                     "data_vector",
                     f"data_{stat.sacc_data_type}_{tracer}",
                     stat.get_data_vector(),

firecrown/likelihood/gauss_family/statistic/source/number_counts.py

@@ -1,7 +1,7 @@
 """Number counts source and systematics."""
 
 from __future__ import annotations
-from typing import Optional, final
+from typing import Optional, final, Sequence
 from dataclasses import dataclass, replace
 from abc import abstractmethod
 
@@ -27,6 +27,7 @@
     DerivedParameterCollection,
 )
 from .....updatable import UpdatableCollection
+from .....metadata.two_point import InferredGalaxyZDist
 
 __all__ = ["NumberCounts"]
 
@@ -267,7 +268,9 @@ def __init__(
         has_rsd: bool = False,
         derived_scale: bool = False,
         scale: float = 1.0,
-        systematics: Optional[list[SourceGalaxySystematic[NumberCountsArgs]]] = None,
+        systematics: Optional[
+            Sequence[SourceGalaxySystematic[NumberCountsArgs]]
+        ] = None,
     ):
         """Initialize the NumberCounts object.
 
@@ -293,6 +296,33 @@ def __init__(
         self.current_tracer_args: Optional[NumberCountsArgs] = None
         self.tracer_args: NumberCountsArgs
 
+    @classmethod
+    def create_ready(
+        cls,
+        inferred_zdist: InferredGalaxyZDist,
+        has_rsd: bool = False,
+        derived_scale: bool = False,
+        scale: float = 1.0,
+        systematics: Optional[list[SourceGalaxySystematic[NumberCountsArgs]]] = None,
+    ) -> NumberCounts:
+        """Create a NumberCounts object with the given tracer name and scale."""
+        obj = cls(
+            sacc_tracer=inferred_zdist.bin_name,
+            systematics=systematics,
+            has_rsd=has_rsd,
+            derived_scale=derived_scale,
+            scale=scale,
+        )
+        obj.tracer_args = NumberCountsArgs(
+            scale=obj.scale,
+            z=inferred_zdist.z,
+            dndz=inferred_zdist.dndz,
+            bias=None,
+            mag_bias=None,
+        )
+
+        return obj
+
     @final
     def _update_source(self, params: ParamsMap):
         """Perform any updates necessary after the parameters have being updated.
@@ -400,3 +430,42 @@ def get_scale(self):
         """Return the scale for this source."""
         assert self.current_tracer_args
         return self.current_tracer_args.scale
+
+
+class NumberCountsSystematicFactory:
+    """Factory class for NumberCountsSystematic objects."""
+
+    @abstractmethod
+    def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCountsSystematic:
+        """Create a NumberCountsSystematic object with the given tracer name."""
+
+
+class NumberCountsFactory:
+    """Factory class for NumberCounts objects."""
+
+    def __init__(
+        self,
+        per_bin_systematics: list[NumberCountsSystematicFactory],
+        global_systematics: Sequence[NumberCountsSystematic],
+    ) -> None:
+        self.per_bin_systematics: list[NumberCountsSystematicFactory] = (
+            per_bin_systematics
+        )
+        self.global_systematics: Sequence[NumberCountsSystematic] = global_systematics
+        self.cache: dict[str, NumberCounts] = {}
+
+    def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCounts:
+        """Create a NumberCounts object with the given tracer name and scale."""
+        if inferred_zdist.bin_name in self.cache:
+            return self.cache[inferred_zdist.bin_name]
+
+        systematics: list[SourceGalaxySystematic[NumberCountsArgs]] = [
+            systematic_factory.create(inferred_zdist)
+            for systematic_factory in self.per_bin_systematics
+        ]
+        systematics.extend(self.global_systematics)
+
+        nc = NumberCounts.create_ready(inferred_zdist, systematics=systematics)
+        self.cache[inferred_zdist.bin_name] = nc
+
+        return nc

firecrown/likelihood/gauss_family/statistic/source/source.py

@@ -274,7 +274,7 @@ def __init__(
         self,
         *,
         sacc_tracer: str,
-        systematics: Optional[list[SourceGalaxySystematic]] = None,
+        systematics: Optional[Sequence[SourceGalaxySystematic]] = None,
     ):
         """Initialize the SourceGalaxy object.
 

firecrown/likelihood/gauss_family/statistic/source/weak_lensing.py

@@ -1,7 +1,7 @@
 """Weak lensing source and systematics."""
 
 from __future__ import annotations
-from typing import Optional, final
+from typing import Optional, final, Sequence
 from dataclasses import dataclass, replace
 from abc import abstractmethod
 
@@ -24,6 +24,7 @@
     ParamsMap,
 )
 from .....modeling_tools import ModelingTools
+from .....metadata.two_point import InferredGalaxyZDist
 
 __all__ = ["WeakLensing"]
 
@@ -214,7 +215,7 @@ def __init__(
         *,
         sacc_tracer: str,
         scale: float = 1.0,
-        systematics: Optional[list[SourceGalaxySystematic[WeakLensingArgs]]] = None,
+        systematics: Optional[Sequence[SourceGalaxySystematic[WeakLensingArgs]]] = None,
     ):
         """Initialize the WeakLensing object.
 
@@ -233,6 +234,20 @@ def __init__(
         self.current_tracer_args: Optional[WeakLensingArgs] = None
         self.tracer_args: WeakLensingArgs
 
+    @classmethod
+    def create_ready(
+        cls,
+        inferred_zdist: InferredGalaxyZDist,
+        systematics: Optional[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.tracer_args = WeakLensingArgs(
+            scale=obj.scale, z=inferred_zdist.z, dndz=inferred_zdist.dndz, ia_bias=None
+        )
+
+        return obj
+
     @final
     def _update_source(self, params: ParamsMap):
         """Implementation of Source interface `_update_source`.
@@ -296,3 +311,42 @@ def get_scale(self):
         """Returns the scales for this Source."""
         assert self.current_tracer_args
         return self.current_tracer_args.scale
+
+
+class WeakLensingSystematicFactory:
+    """Factory class for WeakLensingSystematic objects."""
+
+    @abstractmethod
+    def create(self, inferred_zdist: InferredGalaxyZDist) -> WeakLensingSystematic:
+        """Create a WeakLensingSystematic object with the given tracer name."""
+
+
+class WeakLensingFactory:
+    """Factory class for WeakLensing objects."""
+
+    def __init__(
+        self,
+        per_bin_systematics: list[WeakLensingSystematicFactory],
+        global_systematics: Sequence[WeakLensingSystematic],
+    ) -> None:
+        self.per_bin_systematics: list[WeakLensingSystematicFactory] = (
+            per_bin_systematics
+        )
+        self.global_systematics: Sequence[WeakLensingSystematic] = global_systematics
+        self.cache: dict[str, WeakLensing] = {}
+
+    def create(self, inferred_zdist: InferredGalaxyZDist) -> WeakLensing:
+        """Create a WeakLensing object with the given tracer name and scale."""
+        if inferred_zdist.bin_name in self.cache:
+            return self.cache[inferred_zdist.bin_name]
+
+        systematics: list[SourceGalaxySystematic[WeakLensingArgs]] = [
+            systematic_factory.create(inferred_zdist)
+            for systematic_factory in self.per_bin_systematics
+        ]
+        systematics.extend(self.global_systematics)
+
+        wl = WeakLensing.create_ready(inferred_zdist, systematics)
+        self.cache[inferred_zdist.bin_name] = wl
+
+        return wl