Derived parameters support (#189)

* Initial version of DerivedParameters framework.

* Finished first working version of DerivedParameters framework. Included support for derived parameters in the CosmoSIS connector and included its use in the des_y1_3x2pt example.

* Improving iterable in DerivedParameterCollection. Fixed minor flake8/mypy/... complaints. Added black --check to the CI.

* Fixed pylint issues in tests.

* Update documentation of _get_derived_parameters

Co-authored-by: Marc Paterno <paterno@fnal.gov>

.github/workflows/ci.yml

@@ -67,6 +67,11 @@ jobs:
         shell: bash -l {0}
         run: python -m pip install cobaya
         if: steps.cache.outputs.cache-hit != 'true'      
+      - name: Running black check
+        shell: bash -l {0}
+        run: |
+          black --check firecrown
+          black --check tests
       - name: Running flake8
         shell: bash -l {0}
         run: flake8 firecrown

environment.yml

@@ -4,6 +4,7 @@ channels:
 dependencies:
   - cosmosis
   - cosmosis-build-standard-library
+  - black
   - flake8
   - mypy
   - pylint

examples/des_y1_3x2pt/des_y1_3x2pt.ini

@@ -17,6 +17,7 @@ likelihoods = firecrown
 quiet = T
 debug = T
 timing = T
+extra_output = TwoPoint/NumberCountsScale_lens0 TwoPoint/NumberCountsScale_lens1 TwoPoint/NumberCountsScale_lens2 TwoPoint/NumberCountsScale_lens3 TwoPoint/NumberCountsScale_lens4
 
 [consistency]
 file = ${CSL_DIR}/utility/consistency/consistency_interface.py
@@ -51,6 +52,7 @@ save_dir = des_y1_3x2pt_output
 
 [metropolis]
 samples = 1000
+nsteps = 1
 
 [emcee]
 walkers = 64

examples/des_y1_3x2pt/des_y1_3x2pt.py

@@ -60,7 +60,7 @@
     """
         The source is created and saved (temporarely in the sources dict).
     """
-    sources[f"lens{i}"] = nc.NumberCounts(sacc_tracer=f"lens{i}", systematics=[pzshift])
+    sources[f"lens{i}"] = nc.NumberCounts(sacc_tracer=f"lens{i}", systematics=[pzshift], derived_scale=True)
 
 """
     Now that we have all sources we can instantiate all the two-point

firecrown/connector/cosmosis/likelihood.py

@@ -75,10 +75,15 @@ def execute(self, sample: cosmosis.datablock):
 
         self.likelihood.update(firecrown_params)
         loglike = self.likelihood.compute_loglike(cosmo)
+        derived_params_collection = self.likelihood.get_derived_parameters()
+        assert derived_params_collection is not None
         self.likelihood.reset()
 
         sample.put_double(section_names.likelihoods, "firecrown_like", loglike)
 
+        for section, name, val in derived_params_collection:
+            sample.put(section, name, val)
+
         # Save concatenated data vector and inverse covariance to enable support
         # for the CosmoSIS fisher sampler.
         sample.put(

firecrown/likelihood/gauss_family/gauss_family.py

@@ -21,7 +21,7 @@
 from ..likelihood import Likelihood
 from ...updatable import UpdatableCollection
 from .statistic.statistic import Statistic
-from ...parameters import ParamsMap, RequiredParameters
+from ...parameters import ParamsMap, RequiredParameters, DerivedParameterCollection
 
 
 class GaussFamily(Likelihood):
@@ -39,7 +39,7 @@ def __init__(self, statistics: List[Statistic]):
         self.inv_cov: Optional[np.ndarray] = None
 
     def read(self, sacc_data: sacc.Sacc) -> None:
-        """Read the covariance matrirx for this likelihood from the SACC file."""
+        """Read the covariance matrix for this likelihood from the SACC file."""
 
         _sd = sacc_data.copy()
         inds_list = []
@@ -59,21 +59,24 @@ def read(self, sacc_data: sacc.Sacc) -> None:
     @final
     def compute_chisq(self, cosmo: pyccl.Cosmology) -> float:
         """Calculate and return the chi-squared for the given cosmology."""
-        residuals = []
-        theory_vector = []
-        data_vector = []
+        residuals_list: List[np.ndarray] = []
+        theory_vector_list: List[np.ndarray] = []
+        data_vector_list: List[np.ndarray] = []
         for stat in self.statistics:
             data, theory = stat.compute(cosmo)
-            residuals.append(np.atleast_1d(data - theory))
-            theory_vector.append(np.atleast_1d(theory))
-            data_vector.append(np.atleast_1d(data))
+            residuals_list.append(np.atleast_1d(data - theory))
+            theory_vector_list.append(np.atleast_1d(theory))
+            data_vector_list.append(np.atleast_1d(data))
+
+        residuals = np.concatenate(residuals_list, axis=0)
+        self.predicted_data_vector: np.ndarray = np.concatenate(theory_vector_list)
+        self.measured_data_vector: np.ndarray = np.concatenate(data_vector_list)
 
-        residuals = np.concatenate(residuals, axis=0)
-        self.predicted_data_vector = np.concatenate(theory_vector)
-        self.measured_data_vector = np.concatenate(data_vector)
         # pylint: disable-next=C0103
         x = scipy.linalg.solve_triangular(self.cholesky, residuals, lower=True)
-        return np.dot(x, x)
+        chisq = np.dot(x, x)
+
+        return chisq
 
     @final
     def _update(self, params: ParamsMap) -> None:
@@ -93,6 +96,15 @@ def _reset(self) -> None:
         self._reset_gaussian_family()
         self.statistics.reset()
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        derived_parameters = (
+            self._get_derived_parameters_gaussian_family()
+            + self.statistics.get_derived_parameters()
+        )
+
+        return derived_parameters
+
     @abstractmethod
     def _update_gaussian_family(self, params: ParamsMap) -> None:
         """Abstract method to update GaussianFamily state. Must be implemented by all
@@ -120,3 +132,7 @@ def required_parameters(self) -> RequiredParameters:
     @abstractmethod
     def required_parameters_gaussian_family(self):
         """Required parameters for GaussFamily subclasses."""
+
+    @abstractmethod
+    def _get_derived_parameters_gaussian_family(self) -> DerivedParameterCollection:
+        """Get derived parameters for GaussFamily subclasses."""

firecrown/likelihood/gauss_family/gaussian.py

@@ -13,7 +13,7 @@
 import pyccl
 
 from .gauss_family import GaussFamily
-from ...parameters import ParamsMap, RequiredParameters
+from ...parameters import ParamsMap, RequiredParameters, DerivedParameterCollection
 
 
 class ConstGaussian(GaussFamily):
@@ -35,3 +35,7 @@ def _reset_gaussian_family(self):
     @final
     def required_parameters_gaussian_family(self):
         return RequiredParameters([])
+
+    @final
+    def _get_derived_parameters_gaussian_family(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])

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

@@ -13,7 +13,13 @@
 
 from .source import Source
 from .source import Systematic
-from .....parameters import ParamsMap, RequiredParameters, parameter_get_full_name
+from .....parameters import (
+    ParamsMap,
+    RequiredParameters,
+    parameter_get_full_name,
+    DerivedParameterScalar,
+    DerivedParameterCollection,
+)
 from .....updatable import UpdatableCollection
 
 __all__ = ["NumberCounts"]
@@ -87,6 +93,10 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def apply(
         self, cosmo: pyccl.Cosmology, tracer_arg: NumberCountsArgs
     ) -> NumberCountsArgs:
@@ -165,6 +175,10 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def apply(
         self, cosmo: pyccl.Cosmology, tracer_arg: NumberCountsArgs
     ) -> NumberCountsArgs:
@@ -227,6 +241,10 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def apply(self, cosmo: pyccl.Cosmology, tracer_arg: NumberCountsArgs):
         """Apply a shift to the photo-z distribution of a source."""
 
@@ -260,6 +278,7 @@ def __init__(
         sacc_tracer: str,
         has_rsd: bool = False,
         has_mag_bias: bool = False,
+        derived_scale: bool = False,
         scale: float = 1.0,
         systematics: Optional[List[NumberCountsSystematic]] = None,
     ):
@@ -268,6 +287,7 @@ def __init__(
         self.sacc_tracer = sacc_tracer
         self.has_rsd = has_rsd
         self.has_mag_bias = has_mag_bias
+        self.derived_scale = derived_scale
 
         self.systematics = UpdatableCollection([])
         if systematics:
@@ -313,6 +333,24 @@ def required_parameters(self) -> RequiredParameters:
             )
         return rp + self.systematics.required_parameters()
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        if self.derived_scale:
+            assert self.current_tracer_args is not None
+            derived_scale = DerivedParameterScalar(
+                "TwoPoint",
+                f"NumberCountsScale_{self.sacc_tracer}",
+                self.current_tracer_args.scale,
+            )
+            derived_parameters = DerivedParameterCollection([derived_scale])
+        else:
+            derived_parameters = DerivedParameterCollection([])
+        derived_parameters = (
+            derived_parameters + self.systematics.get_derived_parameters()
+        )
+
+        return derived_parameters
+
     def _read(self, sacc_data):
         """Read the data for this source from the SACC file.
 

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

@@ -13,7 +13,12 @@
 
 from .source import Source
 from .source import Systematic
-from .....parameters import ParamsMap, RequiredParameters, parameter_get_full_name
+from .....parameters import (
+    ParamsMap,
+    RequiredParameters,
+    parameter_get_full_name,
+    DerivedParameterCollection,
+)
 from .....updatable import UpdatableCollection
 
 __all__ = ["WeakLensing"]
@@ -82,6 +87,10 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def apply(self, cosmo: pyccl.Cosmology, tracer_arg: WeakLensingArgs):
         """Apply multiplicative shear bias to a source. The `scale_` of the
         source is multiplied by `(1 + m)`.
@@ -159,6 +168,10 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def apply(
         self, cosmo: pyccl.Cosmology, tracer_arg: WeakLensingArgs
     ) -> WeakLensingArgs:
@@ -208,6 +221,12 @@ def required_parameters(self) -> RequiredParameters:
             [parameter_get_full_name(self.sacc_tracer, pn) for pn in self.params_names]
         )
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        derived_parameters = DerivedParameterCollection([])
+
+        return derived_parameters
+
     def apply(self, cosmo: pyccl.Cosmology, tracer_arg: WeakLensingArgs):
         """Apply a shift to the photo-z distribution of a source."""
 
@@ -263,6 +282,14 @@ def _reset_source(self) -> None:
     def required_parameters(self) -> RequiredParameters:
         return self.systematics.required_parameters()
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        derived_parameters = DerivedParameterCollection([])
+        derived_parameters = (
+            derived_parameters + self.systematics.get_derived_parameters()
+        )
+        return derived_parameters
+
     def _read(self, sacc_data):
         """Read the data for this source from the SACC file.
 

firecrown/likelihood/gauss_family/statistic/supernova.py

@@ -2,15 +2,15 @@
 """
 
 from __future__ import annotations
-from typing import Tuple, final
+from typing import Optional, Tuple, final
 
 import numpy as np
 
 import pyccl
 import sacc
 
 from .statistic import Statistic
-from ....parameters import ParamsMap, RequiredParameters
+from ....parameters import ParamsMap, RequiredParameters, DerivedParameterCollection
 
 
 class Supernova(Statistic):
@@ -23,7 +23,7 @@ def __init__(self, sacc_tracer):
 
         self.sacc_tracer = sacc_tracer
         self.data_vector = None
-        self.a = None  # pylint: disable-msg=invalid-name
+        self.a: Optional[np.ndarray] = None  # pylint: disable-msg=invalid-name
         self.M = None  # pylint: disable-msg=invalid-name
 
     def read(self, sacc_data: sacc.Sacc):
@@ -55,6 +55,10 @@ def required_parameters(self) -> RequiredParameters:
         """
         return RequiredParameters(["m"])
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])
+
     def compute(self, cosmo: pyccl.Cosmology) -> Tuple[np.ndarray, np.ndarray]:
         """Compute a two-point statistic from sources."""
 

firecrown/likelihood/gauss_family/statistic/two_point.py

@@ -14,7 +14,7 @@
 
 from .statistic import Statistic
 from .source.source import Source, Systematic
-from ....parameters import ParamsMap, RequiredParameters
+from ....parameters import ParamsMap, RequiredParameters, DerivedParameterCollection
 
 # only supported types are here, any thing else will throw
 # a value error
@@ -198,6 +198,13 @@ def _reset(self) -> None:
     def required_parameters(self) -> RequiredParameters:
         return self.source0.required_parameters() + self.source1.required_parameters()
 
+    @final
+    def _get_derived_parameters(self) -> DerivedParameterCollection:
+        derived_parameters = DerivedParameterCollection([])
+        derived_parameters = derived_parameters + self.source0.get_derived_parameters()
+        derived_parameters = derived_parameters + self.source1.get_derived_parameters()
+        return derived_parameters
+
     def read(self, sacc_data):
         """Read the data for this statistic from the SACC file.
 

firecrown/likelihood/gauss_family/student_t.py

@@ -11,7 +11,7 @@
 
 from .gauss_family import GaussFamily
 from .statistic.statistic import Statistic
-from ...parameters import ParamsMap, RequiredParameters
+from ...parameters import ParamsMap, RequiredParameters, DerivedParameterCollection
 
 
 class StudentT(GaussFamily):
@@ -50,3 +50,7 @@ def _reset_gaussian_family(self):
     @final
     def required_parameters_gaussian_family(self):
         return RequiredParameters([])
+
+    @final
+    def _get_derived_parameters_gaussian_family(self) -> DerivedParameterCollection:
+        return DerivedParameterCollection([])

firecrown/likelihood/likelihood.py

@@ -62,7 +62,7 @@ def get_params_names(self) -> Optional[List[str]]:
 
     @abstractmethod
     def read(self, sacc_data: sacc.Sacc):
-        """Read the covariance matrirx for this likelihood from the SACC file."""
+        """Read the covariance matrix for this likelihood from the SACC file."""
 
     @abstractmethod
     def compute_loglike(self, cosmo: pyccl.Cosmology) -> float: