Update docstrings again (#429)

* Add type annotations
* Update docstrings

firecrown/connector/cobaya/ccl.py

@@ -11,15 +11,15 @@
 
 from cobaya.theory import Theory
 
-from firecrown.connector.mapping import mapping_builder
+from firecrown.connector.mapping import mapping_builder, MappingCAMB
 
 
 class CCLConnector(Theory):
     """A class implementing cobaya.theory.Theory."""
 
     input_style: str | None = None
 
-    def initialize(self):
+    def initialize(self) -> None:
         """Initialize a CCLConnector object.
 
         Required by Cobaya.
@@ -28,14 +28,18 @@ def initialize(self):
         Cobaya does not allow us to override __init__.
         """
         assert self.input_style
-        self.map = mapping_builder(input_style=self.input_style)
+        # We have to do some extra type-fiddling here because mapping_builder
+        # has a declared return type of the base class.
+        new_mapping = mapping_builder(input_style=self.input_style)
+        assert isinstance(new_mapping, MappingCAMB)
+        self.map = new_mapping
 
         self.a_bg = np.linspace(0.1, 1.0, 50)
         self.z_bg = 1.0 / self.a_bg - 1.0
         self.z_Pk = np.arange(0.0, 6.0, 1)
         self.Pk_kmax = 1.0
 
-    def initialize_with_params(self):
+    def initialize_with_params(self) -> None:
         """Complete the initialization of a CCLConnector object.
 
         Required by Cobaya.
@@ -44,7 +48,7 @@ def initialize_with_params(self):
         that point. This version has nothing to do.
         """
 
-    def initialize_with_provider(self, provider):
+    def initialize_with_provider(self, provider) -> None:
         """Set the object's provider.
 
         Required by Cobaya.
@@ -53,7 +57,7 @@ def initialize_with_provider(self, provider):
         """
         self.provider = provider
 
-    def get_can_provide_params(self):
+    def get_can_provide_params(self) -> list[str]:
         """Return the list of params provided.
 
         Required by Cobaya.
@@ -70,7 +74,7 @@ def get_can_support_params(self) -> list[str]:
         """
         return self.map.get_params_names()
 
-    def get_allow_agnostic(self):
+    def get_allow_agnostic(self) -> bool:
         """Is it allowed to pass all unassigned input parameters to this component.
 
         Required by Cobaya.
@@ -100,7 +104,7 @@ def get_requirements(
 
         return pyccl_calculator_requires
 
-    def must_provide(self, **requirements):
+    def must_provide(self, **requirements) -> None:
         """Required by Cobaya.
 
         This version does nothing.

firecrown/connector/cobaya/likelihood.py

@@ -41,27 +41,35 @@ def initialize(self):
         )
 
     def initialize_with_params(self) -> None:
-        """Required by Cobaya.
+        """Complete the initialization of a LikelihoodConnector object.
+
+        Required by Cobaya.
 
         This version has nothing to do.
         """
 
     def initialize_with_provider(self, provider) -> None:
-        """Required by Cobaya.
+        """Set the obejct's provider.
 
-        Sets instance's provided to the given provider.
+        Required by Cobaya.
+
+        :param provider: A Cobaya provider.
         """
         self.provider = provider
 
     def get_can_provide_params(self) -> list[str]:
-        """Required by Cobaya.
+        """Return the list of params provided.
+
+        Required by Cobaya.
 
         Returns an empty list.
         """
         return self.derived_parameters
 
-    def get_allow_agnostic(self):
-        """Required by Cobaya.
+    def get_allow_agnostic(self) -> bool:
+        """Is it allowed to pass all unassigned input parameters to this component.
+
+        Required by Cobaya.
 
         Return False.
         """
@@ -70,10 +78,13 @@ def get_allow_agnostic(self):
     def get_requirements(
         self,
     ) -> dict[str, None | dict[str, npt.NDArray[np.float64]] | dict[str, object]]:
-        """Required by Cobaya.
+        """Returns a dictionary.
 
         Returns a dictionary with keys corresponding the contained likelihood's
         required parameter, plus "pyccl". All values are None.
+
+        Required by Cobaya.
+        :return: a dictionary
         """
         likelihood_requires: dict[
             str, None | dict[str, npt.NDArray[np.float64]] | dict[str, object]
@@ -87,16 +98,17 @@ def get_requirements(
 
         return likelihood_requires
 
-    def must_provide(self, **requirements):
+    def must_provide(self, **requirements) -> None:
         """Required by Cobaya.
 
         This version does nothing.
         """
 
     def logp(self, **params_values) -> float:
-        """Required by Cobaya.
+        """Return the log of the calculated likelihood.
 
-        Return the (log) calculated likelihood.
+        Required by Cobaya.
+        :params values: The values of the parameters to use.
         """
         pyccl = self.provider.get_pyccl()
 

firecrown/connector/cosmosis/likelihood.py

@@ -22,26 +22,32 @@
 
 
 def extract_section(sample: cosmosis.datablock, section: str) -> NamedParameters:
-    """Extract all the parameters from the name datablock section into a dictionary."""
+    """Extract all the parameters from the name datablock section into a dictionary.
+
+    :param sample: the CosmoSiS datablock to query
+    :param section: the name of the section desired
+    :return: a dictionary of the parameters in the section
+    """
     if not sample.has_section(section):
         raise RuntimeError(f"Datablock section `{section}' does not exist.")
     sec_dict = {name: sample[section, name] for _, name in sample.keys(section=section)}
     return NamedParameters(sec_dict)
 
 
 class FirecrownLikelihood:
-    """CosmoSIS likelihood module for calculating Firecrown likelihood.
+    """CosmoSIS likelihood module for calculating a Firecrown likelihood.
 
     In this simplest implementation, we have only a single module. This module
     is responsible for calling CCL to perform theory calculations, based on the
     output of CAMB, and also for calculating the data likelihood based on this
     theory.
-
-    :param config: current CosmoSIS datablock
     """
 
-    def __init__(self, config: cosmosis.datablock):
-        """Create the FirecrownLikelihood object from the given configuration."""
+    def __init__(self, config: cosmosis.datablock) -> None:
+        """Create the FirecrownLikelihood object from the given configuration.
+
+        :param config: the datablock the configuration
+        """
         likelihood_source = config.get_string(option_section, "likelihood_source", "")
         if likelihood_source == "":
             likelihood_source = config[option_section, "firecrown_config"]
@@ -67,9 +73,13 @@ def __init__(self, config: cosmosis.datablock):
             print(f"The Firecrown likelihood needs a required parameter: {err}")
             print("*" * 30)
             raise
-        self.map: MappingCosmoSIS = mapping_builder(
+        # We have to do some extra type-fiddling here because mapping_builder
+        # has a declared return type of the base class.
+        new_mapping = mapping_builder(
             input_style="CosmoSIS", require_nonlinear_pk=require_nonlinear_pk
         )
+        assert isinstance(new_mapping, MappingCosmoSIS)
+        self.map = new_mapping
 
         # If sampling_sections  is empty, but we have required parameters, then
         # we have a configuration problem, and ParamsMap can never be built
@@ -92,7 +102,11 @@ def __init__(self, config: cosmosis.datablock):
                 raise RuntimeError(msg)
 
     def execute(self, sample: cosmosis.datablock) -> int:
-        """This is the method called for each sample generated by the sampler."""
+        """This is the method called for each sample generated by the sampler.
+
+        :param sample: the sample generated by the sampler
+        :return: 0
+        """
         cosmological_params: NamedParameters = extract_section(
             sample, "cosmological_parameters"
         )
@@ -130,13 +144,16 @@ def execute(self, sample: cosmosis.datablock) -> int:
         self.tools.reset()
         return 0
 
-    def special_gauss_family_handling(self, sample):
+    def special_gauss_family_handling(self, sample: cosmosis.datablock) -> None:
         """Special handling for the GaussFamily likelihood.
 
         We need to save concatenated data vector and inverse covariance to
         enable support for the CosmoSIS Fisher sampler. This can only work
         for likelihoods that have these quantities. Currently, this is only
         GaussFamily.
+
+        :param sample: the sample generated by the sampler
+        :return: None
         """
         assert isinstance(self.likelihood, GaussFamily)
         sample.put(
@@ -166,8 +183,16 @@ def special_gauss_family_handling(self, sample):
             if isinstance(stat, TwoPoint):
                 self.handle_twopoint_statistic(sample, stat)
 
-    def handle_twopoint_statistic(self, sample, stat):
-        """Handle the TwoPoint statistic for the GaussFamily likelihood."""
+    def handle_twopoint_statistic(
+        self, sample: cosmosis.datablock, stat: TwoPoint
+    ) -> None:
+        """Handle the TwoPoint statistic for the GaussFamily likelihood.
+
+        This puts the theory and data vectors in the data block.
+
+        :param sample: the sample generated by the sampler
+        :param stat: a TwoPoint statistic
+        """
         assert stat.sacc_tracers is not None
         tracer = f"{stat.sacc_tracers[0]}_{stat.sacc_tracers[1]}"
         if stat.ells is not None:
@@ -194,7 +219,11 @@ def handle_twopoint_statistic(self, sample, stat):
         )
 
     def update_likelihood_and_tools(self, firecrown_params: ParamsMap) -> None:
-        """Update the likelihood and tools with the new parameters."""
+        """Update the likelihood and tools with the new parameters.
+
+        :param firecrown_params: the new parameters
+        :return: None
+        """
         try:
             self.likelihood.update(firecrown_params)
             self.tools.update(firecrown_params)
@@ -207,6 +236,9 @@ def form_error_message(self, exc: MissingSamplerParameterError) -> str:
 
         This error message will also include when that parameter should have been
         supplied by the sampler.
+
+        :param exc: the missing parameter error
+        :return: the error message
         """
         msg = (
             "A required parameter was not found in any of the "
@@ -224,7 +256,11 @@ def form_error_message(self, exc: MissingSamplerParameterError) -> str:
         return msg
 
     def calculate_firecrown_params(self, sample: cosmosis.datablock) -> ParamsMap:
-        """Calculate the ParamsMap for this sample."""
+        """Calculate the ParamsMap for this sample.
+
+        :param sample: the sample generated by the sampler
+        :return: a ParamsMap with the firecrown parameters
+        """
         firecrown_params = ParamsMap()
         for section in self.sampling_sections:
             section_params = extract_section(sample, section)
@@ -245,9 +281,10 @@ def calculate_firecrown_params(self, sample: cosmosis.datablock) -> ParamsMap:
 def setup(config: cosmosis.datablock) -> FirecrownLikelihood:
     """Setup hook for a CosmoSIS module.
 
-    Returns an instance of
-    class FirecrownLikelihood. The same object will be passed to the CosmoSIS
-    execute hook.
+    The returned object will be passed to the CosmoSIS execute hook.
+
+    :param config: the datablock the configuration
+    :return: an instance of class FirecrownLikelihood
     """
     return FirecrownLikelihood(config)
 
@@ -257,10 +294,17 @@ def execute(sample: cosmosis.datablock, instance: FirecrownLikelihood) -> int:
 
     Return 0 on success. The parameter `sample` represents the current MCMC sample;
     `instance` is the FirecrownLikelihood object created by `setup`.
+
+    :param sample: the sample generated by the sampler
+    :param instance: the FirecrownLikelihood object
+    :return: the status of the call to the module's execute function
     """
     return instance.execute(sample)
 
 
 def cleanup(_) -> int:
-    """Cleanup hook for a CosmoSIS module. This one has nothing to do."""
+    """Cleanup hook for a CosmoSIS module. This one has nothing to do.
+
+    :return: 0
+    """
     return 0