Added state control for updatable objects. (#184)

* Added state control for updatable objects. Added testing and fixed pylint issues (old and new) on tests (including a pylintrc for tests and its inclusion on GHA).

* Removed extra parameter in _reset_source.

* Added pylint to environment.yml.

* Fixed minor typos.

.github/workflows/ci.yml

@@ -73,6 +73,9 @@ jobs:
       - name: Running mypy
         shell: bash -l {0}
         run: mypy -p firecrown --ignore-missing-imports
+      - name: Running pylint-tests
+        shell: bash -l {0}
+        run: pylint --rcfile pylintrc_for_tests --recursive=y tests
       - name: Running pytest
         shell: bash -l {0}
         run: python -m pytest -vv 

environment.yml

@@ -6,6 +6,7 @@ dependencies:
   - cosmosis-build-standard-library
   - flake8
   - mypy
+  - pylint
   - numpy>=1.16
   - pip
   - pyccl>=2.0

firecrown/connector/cobaya/likelihood.py

@@ -62,7 +62,6 @@ def get_requirements(self) -> Dict:
         required parameter, plus "pyccl". All values are None.
         """
         likelihood_requires = {"pyccl": None}
-
         required_params = self.likelihood.required_parameters()
 
         for param_name in required_params.get_params_names():
@@ -77,12 +76,14 @@ def must_provide(self, **requirements):
         """
 
     def logp(self, **params_values) -> float:
-        """Requried by Cobaya.
+        """Required by Cobaya.
 
         Return the (log) calculated likelihood.
         """
         pyccl = self.provider.get_pyccl()
 
         self.likelihood.update(ParamsMap(params_values))
+        loglike = self.likelihood.compute_loglike(pyccl)
+        self.likelihood.reset()
 
-        return self.likelihood.compute_loglike(pyccl)
+        return loglike

firecrown/connector/cosmosis/likelihood.py

@@ -74,9 +74,10 @@ def execute(self, sample: cosmosis.datablock):
         firecrown_params = self.calculate_firecrown_params(sample)
 
         self.likelihood.update(firecrown_params)
-        lnlike = self.likelihood.compute_loglike(cosmo)
+        loglike = self.likelihood.compute_loglike(cosmo)
+        self.likelihood.reset()
 
-        sample.put_double(section_names.likelihoods, "firecrown_like", lnlike)
+        sample.put_double(section_names.likelihoods, "firecrown_like", loglike)
 
         # Save concatenated data vector and inverse covariance to enable support
         # for the CosmoSIS fisher sampler.

firecrown/likelihood/gauss_family/gauss_family.py

@@ -79,10 +79,19 @@ def _update(self, params: ParamsMap):
         self.statistics.update(params)
         self._update_gaussian_family(params)
 
+    @final
+    def _reset(self):
+        self._reset_gaussian_family()
+        self.statistics.reset()
+
     @abstractmethod
     def _update_gaussian_family(self, params: ParamsMap):
         pass
 
+    @abstractmethod
+    def _reset_gaussian_family(self, params: ParamsMap):
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         stats_rp = self.statistics.required_parameters()

firecrown/likelihood/gauss_family/gaussian.py

@@ -28,6 +28,10 @@ def compute_loglike(self, cosmo: pyccl.Cosmology):
     def _update_gaussian_family(self, params: ParamsMap):
         pass
 
+    @final
+    def _reset_gaussian_family(self):
+        pass
+
     @final
     def required_parameters_gaussian_family(self):
         return RequiredParameters([])

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

@@ -71,6 +71,8 @@ class LinearBiasSystematic(NumberCountsSystematic):
     z_piv: float
 
     def __init__(self, sacc_tracer: str):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
@@ -81,6 +83,10 @@ def _update(self, params: ParamsMap):
         self.alphag = params.get_from_prefix_param(self.sacc_tracer, "alphag")
         self.z_piv = params.get_from_prefix_param(self.sacc_tracer, "z_piv")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -139,6 +145,8 @@ class MagnificationBiasSystematic(NumberCountsSystematic):
     z_m: float
 
     def __init__(self, sacc_tracer: str):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
@@ -151,6 +159,10 @@ def _update(self, params: ParamsMap):
         self.z_c = params.get_from_prefix_param(self.sacc_tracer, "z_c")
         self.z_m = params.get_from_prefix_param(self.sacc_tracer, "z_m")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -199,12 +211,18 @@ class PhotoZShift(NumberCountsSystematic):
     delta_z: float
 
     def __init__(self, sacc_tracer: str):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
     def _update(self, params: ParamsMap):
         self.delta_z = params.get_from_prefix_param(self.sacc_tracer, "delta_z")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -245,6 +263,8 @@ def __init__(
         scale=1.0,
         systematics: Optional[List[NumberCountsSystematic]] = None,
     ):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
         self.has_rsd = has_rsd
         self.has_mag_bias = has_mag_bias
@@ -270,6 +290,10 @@ def _update_source(self, params: ParamsMap):
 
         self.systematics.update(params)
 
+    @final
+    def _reset_source(self) -> None:
+        self.systematics.reset()
+
     @final
     def required_parameters(self) -> RequiredParameters:
         if self.has_mag_bias:

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

@@ -17,7 +17,7 @@
 
 
 class Systematic(Updatable):
-    """The systematic (e.g., shear biases, photo-z shifts, etc.).
+    """An abstract systematic class (e.g., shear biases, photo-z shifts, etc.).
 
     This class currently has no methods at all, because the argument types for
     the `apply` method of different subclasses are different."""
@@ -28,7 +28,7 @@ def read(self, sacc_data: sacc.Sacc):
 
 
 class Source(Updatable):
-    """The source (e.g., a sample of lenses).
+    """An abstract source class (e.g., a sample of lenses).
 
     Parameters
     ----------
@@ -60,12 +60,20 @@ def _read(self, sacc_data: sacc.Sacc):
     def _update_source(self, params: ParamsMap):
         pass
 
+    @abstractmethod
+    def _reset_source(self):
+        pass
+
     @final
     def _update(self, params: ParamsMap):
         self.cosmo_hash = None
         self.tracer = None
         self._update_source(params)
 
+    @final
+    def _reset(self) -> None:
+        self._reset_source()
+
     @abstractmethod
     def get_scale(self) -> float:
         pass

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

@@ -62,6 +62,8 @@ def __init__(self, sacc_tracer: str):
         ----------
         sacc_tracer : The name of the multiplicative bias parameter.
         """
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
@@ -70,6 +72,10 @@ def _update(self, params: ParamsMap):
         parameters."""
         self.m = params.get_from_prefix_param(self.sacc_tracer, "mult_bias")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -130,6 +136,8 @@ def __init__(self, sacc_tracer: Optional[str] = None):
         z_piv : The pivot redshift parameter for the intrinsic alignment
         parameter.
         """
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
@@ -139,6 +147,10 @@ def _update(self, params: ParamsMap):
         self.alphag = params.get_from_prefix_param(self.sacc_tracer, "alphag")
         self.z_piv = params.get_from_prefix_param(self.sacc_tracer, "z_piv")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -176,12 +188,18 @@ class PhotoZShift(WeakLensingSystematic):
     delta_z: float
 
     def __init__(self, sacc_tracer: str):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
 
     @final
     def _update(self, params: ParamsMap):
         self.delta_z = params.get_from_prefix_param(self.sacc_tracer, "delta_z")
 
+    @final
+    def _reset(self) -> None:
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(
@@ -215,6 +233,8 @@ def __init__(
         scale=1.0,
         systematics: Optional[List[WeakLensingSystematic]] = None,
     ):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
         self.scale = scale
         self.z_orig: Optional[np.ndarray] = None
@@ -233,6 +253,10 @@ def __init__(
     def _update_source(self, params: ParamsMap):
         self.systematics.update(params)
 
+    @final
+    def _reset_source(self) -> None:
+        self.systematics.reset()
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return self.systematics.required_parameters()

firecrown/likelihood/gauss_family/statistic/statistic.py

@@ -20,7 +20,7 @@
 
 
 class Statistic(Updatable):
-    """A statistic (e.g., two-point function, mass function, etc.)."""
+    """An abstract statistic class (e.g., two-point function, mass function, etc.)."""
 
     systematics: List[Systematic]
     sacc_inds: List[int]

firecrown/likelihood/gauss_family/statistic/supernova.py

@@ -38,6 +38,8 @@ def _cached_distmod(cosmo, tracers, z):
 
 class Supernova(Statistic):
     def __init__(self, sacc_tracer):
+        super().__init__()
+
         self.sacc_tracer = sacc_tracer
         self.data_vector = None
 
@@ -57,6 +59,10 @@ def read(self, sacc_data: sacc.Sacc):
     def _update(self, params: ParamsMap):
         self.M = params["m"]  # CosmoSIS makes everything lowercase
 
+    @final
+    def _reset(self):
+        pass
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return RequiredParameters(["m"])

firecrown/likelihood/gauss_family/statistic/two_point.py

@@ -163,6 +163,8 @@ def __init__(
         ell_or_theta_min=None,
         ell_or_theta_max=None,
     ):
+        super().__init__()
+
         self.sacc_data_type = sacc_data_type
         self.source0 = source0
         self.source1 = source1
@@ -192,6 +194,11 @@ def _update(self, params: ParamsMap):
         self.source0.update(params)
         self.source1.update(params)
 
+    @final
+    def _reset(self) -> None:
+        self.source0.reset()
+        self.source1.reset()
+
     @final
     def required_parameters(self) -> RequiredParameters:
         return self.source0.required_parameters() + self.source1.required_parameters()

firecrown/likelihood/gauss_family/student_t.py

@@ -48,6 +48,10 @@ def compute_loglike(self, cosmo: pyccl.Cosmology):
     def _update_gaussian_family(self, params: ParamsMap):
         pass
 
+    @final
+    def _reset_gaussian_family(self):
+        pass
+
     @final
     def required_parameters_gaussian_family(self):
         return RequiredParameters([])

firecrown/likelihood/likelihood.py

@@ -35,7 +35,9 @@ class Likelihood(Updatable):
     """
 
     def __init__(self):
-        """Deafault initialization for a base Likelihood object."""
+        """Default initialization for a base Likelihood object."""
+        super().__init__()
+
         self.params_names: Optional[List[str]] = None
         self.predicted_data_vector: Optional[npt.NDArray[np.double]] = None
         self.measured_data_vector: Optional[npt.NDArray[np.double]] = None
@@ -106,4 +108,6 @@ def load_likelihood(filename: str) -> Likelihood:
         )
 
     likelihood = mod.likelihood
+    assert isinstance(likelihood, Likelihood)
+
     return likelihood