From bd7047635cecdae3f08f34833d5e93e51b6010b2 Mon Sep 17 00:00:00 2001
From: Marc Paterno <paterno@fnal.gov>
Date: Tue, 13 Aug 2024 13:45:49 -0500
Subject: [PATCH] Yet even more docstrings (#445)

* Improve docstrings and type annotations
* Fix bug in NumberCounts property computation
---
 firecrown/likelihood/__init__.py              |   2 +-
 .../binned_cluster_number_counts.py           |  43 ++++--
 firecrown/likelihood/gauss_family/__init__.py |   2 +
 .../gauss_family/statistic/source/__init__.py |   2 +
 firecrown/likelihood/gaussfamily.py           |  85 ++++++++++--
 firecrown/likelihood/gaussian.py              |  13 +-
 firecrown/likelihood/likelihood.py            |  90 +++++++++---
 firecrown/likelihood/number_counts.py         | 129 +++++++++++++-----
 firecrown/likelihood/supernova.py             |   2 +-
 firecrown/likelihood/two_point.py             |   5 +-
 .../gauss_family/statistic/test_statistic.py  |  10 +-
 11 files changed, 294 insertions(+), 89 deletions(-)

diff --git a/firecrown/likelihood/__init__.py b/firecrown/likelihood/__init__.py
index 636d36e76..0a6eede37 100644
--- a/firecrown/likelihood/__init__.py
+++ b/firecrown/likelihood/__init__.py
@@ -1,4 +1,4 @@
-"""Classes used to represent likelihoods and functions to support them.
+"""Classes used to represent likelihoods, and functions to support them.
 
 Subpackages contain specific likelihood implementations, e.g., Gaussian and Student-t.
 The submodule :mod:`firecrown.likelihood.likelihood` contain the abstract base class for
diff --git a/firecrown/likelihood/binned_cluster_number_counts.py b/firecrown/likelihood/binned_cluster_number_counts.py
index 2732134b5..00ab0398d 100644
--- a/firecrown/likelihood/binned_cluster_number_counts.py
+++ b/firecrown/likelihood/binned_cluster_number_counts.py
@@ -1,8 +1,4 @@
-"""This module holds classes needed to predict the binned cluster number counts.
-
-The binned cluster number counts statistic predicts the number of galaxy
-clusters within a single redshift and mass bin.
-"""
+"""Binned cluster number counts statistic support."""
 
 from __future__ import annotations
 
@@ -26,11 +22,7 @@
 
 
 class BinnedClusterNumberCounts(Statistic):
-    """The Binned Cluster Number Counts statistic.
-
-    This class will make a prediction for the number of clusters in a z, mass bin
-    and compare that prediction to the data provided in the sacc file.
-    """
+    """A statistic representing the number of clusters in a z, mass bin."""
 
     def __init__(
         self,
@@ -39,6 +31,13 @@ def __init__(
         cluster_recipe: ClusterRecipe,
         systematics: None | list[SourceSystematic] = None,
     ):
+        """Initialize this statistic.
+
+        :param cluster_properties: The cluster observables to use.
+        :param survey_name: The name of the survey to use.
+        :param cluster_recipe: The cluster recipe to use.
+        :param systematics: The systematics to apply to this statistic.
+        """
         super().__init__()
         self.systematics = systematics or []
         self.theory_vector: None | TheoryVector = None
@@ -50,7 +49,10 @@ def __init__(
         self.bins: list[SaccBin] = []
 
     def read(self, sacc_data: sacc.Sacc) -> None:
-        """Read the data for this statistic and mark it as ready for use."""
+        """Read the data for this statistic and mark it as ready for use.
+
+        :param sacc_data: The data in the sacc format.
+        """
         # Build the data vector and indices needed for the likelihood
         if self.cluster_properties == ClusterProperty.NONE:
             raise ValueError("You must specify at least one cluster property.")
@@ -77,12 +79,19 @@ def read(self, sacc_data: sacc.Sacc) -> None:
         super().read(sacc_data)
 
     def get_data_vector(self) -> DataVector:
-        """Gets the statistic data vector."""
+        """Gets the statistic data vector.
+
+        :return: The statistic data vector.
+        """
         assert self.data_vector is not None
         return self.data_vector
 
     def _compute_theory_vector(self, tools: ModelingTools) -> TheoryVector:
-        """Compute a statistic from sources, concrete implementation."""
+        """Compute a statistic from sources, concrete implementation.
+
+        :param tools: The modeling tools used to compute the statistic.
+        :return: The computed statistic.
+        """
         assert tools.cluster_abundance is not None
 
         theory_vector_list: list[float] = []
@@ -116,6 +125,9 @@ def get_binned_cluster_property(
         Using the data from the sacc file, this function evaluates the likelihood for
         a single point of the parameter space, and returns the predicted mean mass of
         the clusters in each bin.
+
+        :param cluster_counts: The number of clusters in each bin.
+        :param cluster_properties: The cluster observables to use.
         """
         assert tools.cluster_abundance is not None
 
@@ -124,8 +136,6 @@ def get_binned_cluster_property(
             total_observable = self.cluster_recipe.evaluate_theory_prediction(
                 tools.cluster_abundance, this_bin, self.sky_area, cluster_properties
             )
-            cluster_counts.append(counts)
-
             mean_observable = total_observable / counts
             mean_values.append(mean_observable)
 
@@ -137,6 +147,9 @@ def get_binned_cluster_counts(self, tools: ModelingTools) -> list[float]:
         Using the data from the sacc file, this function evaluates the likelihood for
         a single point of the parameter space, and returns the predicted number of
         clusters in each bin.
+
+        :param tools: The modeling tools used to compute the statistic.
+        :return: The number of clusters in each bin.
         """
         assert tools.cluster_abundance is not None
 
diff --git a/firecrown/likelihood/gauss_family/__init__.py b/firecrown/likelihood/gauss_family/__init__.py
index 9c0fa90a1..34815d384 100644
--- a/firecrown/likelihood/gauss_family/__init__.py
+++ b/firecrown/likelihood/gauss_family/__init__.py
@@ -1 +1,3 @@
+"""Backward compatibility support for deprecated directory structure."""
+
 # flake8: noqa
diff --git a/firecrown/likelihood/gauss_family/statistic/source/__init__.py b/firecrown/likelihood/gauss_family/statistic/source/__init__.py
index 9c0fa90a1..34815d384 100644
--- a/firecrown/likelihood/gauss_family/statistic/source/__init__.py
+++ b/firecrown/likelihood/gauss_family/statistic/source/__init__.py
@@ -1 +1,3 @@
+"""Backward compatibility support for deprecated directory structure."""
+
 # flake8: noqa
diff --git a/firecrown/likelihood/gaussfamily.py b/firecrown/likelihood/gaussfamily.py
index 80ff5bbca..265163822 100644
--- a/firecrown/likelihood/gaussfamily.py
+++ b/firecrown/likelihood/gaussfamily.py
@@ -1,4 +1,4 @@
-"""Support for the family of Gaussian likelihood."""
+"""Support for the family of Gaussian likelihoods."""
 
 from __future__ import annotations
 
@@ -30,7 +30,12 @@
 
 
 class State(Enum):
-    """The states used in GaussFamily."""
+    """The states used in GaussFamily.
+
+    GaussFamily and all subclasses enforce a statemachine behavior based on
+    these states to ensure that the necessary initialization and setup is done
+    in the correct order.
+    """
 
     INITIALIZED = 1
     READY = 2
@@ -62,6 +67,12 @@ def enforce_states(
     If terminal is None the state of the object is not modified.
     If terminal is not None and the call to the wrapped method returns
     normally the state of the object is set to terminal.
+
+    :param initial: The initial states allowable for the wrapped method
+    :param terminal: The terminal state ensured for the wrapped method. None
+        indicates no state change happens.
+    :param failure_message: The failure message for the AssertionError raised
+    :return: The wrapped method
     """
     initials: list[State]
     if isinstance(initial, list):
@@ -74,6 +85,9 @@ def decorator_enforce_states(func: Callable[P, T]) -> Callable[P, T]:
 
         This closure is what actually contains the values of initials, terminal, and
         failure_message.
+
+        :param func: The method to be wrapped
+        :return: The wrapped method
         """
 
         @wraps(func)
@@ -132,8 +146,11 @@ class GaussFamily(Likelihood):
     def __init__(
         self,
         statistics: Sequence[Statistic],
-    ):
-        """Initialize the base class parts of a GaussFamily object."""
+    ) -> None:
+        """Initialize the base class parts of a GaussFamily object.
+
+        :param statistics: A list of statistics to be include in chisquared calculations
+        """
         super().__init__()
         self.state: State = State.INITIALIZED
         if len(statistics) == 0:
@@ -160,7 +177,12 @@ def __init__(
     def create_ready(
         cls, statistics: Sequence[Statistic], covariance: npt.NDArray[np.float64]
     ) -> GaussFamily:
-        """Create a GaussFamily object in the READY state."""
+        """Create a GaussFamily object in the READY state.
+
+        :param statistics: A list of statistics to be include in chisquared calculations
+        :param covariance: The covariance matrix of the statistics
+        :return: A ready GaussFamily object
+        """
         obj = cls(statistics)
         obj._set_covariance(covariance)
         obj.state = State.READY
@@ -178,6 +200,8 @@ def _update(self, _: ParamsMap) -> None:
         for its own reasons must be sure to do what this does: check the state
         at the start of the method, and change the state at the end of the
         method.
+
+        :param _: a ParamsMap object, not used
         """
 
     @enforce_states(
@@ -201,7 +225,10 @@ def _reset(self) -> None:
         failure_message="read() must only be called once",
     )
     def read(self, sacc_data: sacc.Sacc) -> None:
-        """Read the covariance matrix for this likelihood from the SACC file."""
+        """Read the covariance matrix for this likelihood from the SACC file.
+
+        :param sacc_data: The SACC data object to be read
+        """
         if sacc_data.covariance is None:
             msg = (
                 f"The {type(self).__name__} likelihood requires a covariance, "
@@ -216,11 +243,13 @@ def read(self, sacc_data: sacc.Sacc) -> None:
 
         self._set_covariance(covariance)
 
-    def _set_covariance(self, covariance):
+    def _set_covariance(self, covariance: npt.NDArray[np.float64]) -> None:
         """Set the covariance matrix.
 
         This method is used to set the covariance matrix and perform the
         necessary calculations to prepare the likelihood for computation.
+
+        :param covariance: The covariance matrix for this likelihood
         """
         indices_list = []
         data_vector_list = []
@@ -276,6 +305,7 @@ def get_cov(
         :param statistic: The statistic for which the sub-covariance matrix
             should be returned. If not specified, return the covariance of all
             statistics.
+        :return: The covariance matrix (or portion thereof)
         """
         assert self.cov is not None
         if statistic is None:
@@ -301,7 +331,10 @@ def get_cov(
         failure_message="read() must be called before get_data_vector()",
     )
     def get_data_vector(self) -> npt.NDArray[np.float64]:
-        """Get the data vector from all statistics in the right order."""
+        """Get the data vector from all statistics in the right order.
+
+        :return: The data vector
+        """
         assert self.data_vector is not None
         return self.data_vector
 
@@ -315,6 +348,7 @@ def compute_theory_vector(self, tools: ModelingTools) -> npt.NDArray[np.float64]
         """Computes the theory vector using the current instance of pyccl.Cosmology.
 
         :param tools: Current ModelingTools object
+        :return: The computed theory vector
         """
         theory_vector_list: list[npt.NDArray[np.float64]] = [
             stat.compute_theory_vector(tools) for stat in self.statistics
@@ -329,7 +363,10 @@ def compute_theory_vector(self, tools: ModelingTools) -> npt.NDArray[np.float64]
         "get_theory_vector()",
     )
     def get_theory_vector(self) -> npt.NDArray[np.float64]:
-        """Get the theory vector from all statistics in the right order."""
+        """Get the already-computed theory vector from all statistics.
+
+        :return: The theory vector, with all statistics in the right order
+        """
         assert (
             self.theory_vector is not None
         ), "theory_vector is None after compute_theory_vector() has been called"
@@ -343,7 +380,14 @@ def get_theory_vector(self) -> npt.NDArray[np.float64]:
     def compute(
         self, tools: ModelingTools
     ) -> tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
-        """Calculate and return both the data and theory vectors."""
+        """Calculate and return both the data and theory vectors.
+
+        This method is dprecated and will be removed in a future version of Firecrown.
+
+        :param tools: the ModelingTools to be used in the calculation of the
+            theory vector
+        :return: a tuple containing the data vector and the theory vector
+        """
         warnings.warn(
             "The use of the `compute` method on Statistic is deprecated."
             "The Statistic objects should implement `get_data` and "
@@ -359,7 +403,12 @@ def compute(
         failure_message="update() must be called before compute_chisq()",
     )
     def compute_chisq(self, tools: ModelingTools) -> float:
-        """Calculate and return the chi-squared for the given cosmology."""
+        """Calculate and return the chi-squared for the given cosmology.
+
+        :param tools: the ModelingTools to be used in the calculation of the
+            theory vector
+        :return: the chi-squared
+        """
         theory_vector: npt.NDArray[np.float64]
         data_vector: npt.NDArray[np.float64]
         residuals: npt.NDArray[np.float64]
@@ -386,6 +435,10 @@ def get_sacc_indices(
         """Get the SACC indices of the statistic or list of statistics.
 
         If no statistic is given, get the indices of all statistics of the likelihood.
+
+        :param statistics: The statistic or list of statistics for which the
+            SACC indices are desired
+        :return: The SACC indices
         """
         if statistic is None:
             statistic = [stat.statistic for stat in self.statistics]
@@ -409,7 +462,15 @@ def get_sacc_indices(
     def make_realization(
         self, sacc_data: sacc.Sacc, add_noise: bool = True, strict: bool = True
     ) -> sacc.Sacc:
-        """Create a new realization of the model."""
+        """Create a new realization of the model.
+
+        :param sacc_data: The SACC data object containing the covariance matrix
+            to be read
+        :param add_noise: If True, add noise to the realization.
+        :param strict: If True, check that the indices of the realization cover
+            all the indices of the SACC data object.
+        :return: The SACC data object containing the new realization
+        """
         sacc_indices = self.get_sacc_indices()
 
         if add_noise:
diff --git a/firecrown/likelihood/gaussian.py b/firecrown/likelihood/gaussian.py
index f503e6b6d..dddd044af 100644
--- a/firecrown/likelihood/gaussian.py
+++ b/firecrown/likelihood/gaussian.py
@@ -13,12 +13,19 @@
 class ConstGaussian(GaussFamily):
     """A Gaussian log-likelihood with a constant covariance matrix."""
 
-    def compute_loglike(self, tools: ModelingTools):
-        """Compute the log-likelihood."""
+    def compute_loglike(self, tools: ModelingTools) -> float:
+        """Compute the log-likelihood.
+
+        :params tools: The modeling tools used to compute the likelihood.
+        :return: The log-likelihood.
+        """
         return -0.5 * self.compute_chisq(tools)
 
     def make_realization_vector(self) -> np.ndarray:
-        """Create a new realization of the model."""
+        """Create a new (randomized) realization of the model.
+
+        :return: A new realization of the model
+        """
         theory_vector = self.get_theory_vector()
         assert self.cholesky is not None
         new_data_vector = theory_vector + np.dot(
diff --git a/firecrown/likelihood/likelihood.py b/firecrown/likelihood/likelihood.py
index 9ace48cdf..2aa751552 100644
--- a/firecrown/likelihood/likelihood.py
+++ b/firecrown/likelihood/likelihood.py
@@ -20,7 +20,7 @@
 a *data vector* and a *covariance matrix*, which must be present in the
 :class:`Sacc` object given to the :meth:`read` method.
 
-The theory predictions that are used in the calcluation of a likelihood are
+The theory predictions that are used in the calculation of a likelihood are
 expected to change for different calls to the :meth:`compute_loglike` method.
 In order to prepare a :class:`Likelihood` object for each call to
 :meth:`compute_loglike`, the following sequence of calls must be made (note
@@ -75,24 +75,32 @@ class Likelihood(Updatable):
     Concrete subclasses represent specific likelihood forms (e.g. gaussian with
     constant covariance matrix, or Student's t, etc.).
 
-    Concrete subclasses must have an implementation of both *read* and
-    *compute_loglike*. Note that abstract subclasses of Likelihood might implement
+    Concrete subclasses must have an implementation of both :meth:`read` and
+    :meth:`compute_loglike`. Note that abstract subclasses of Likelihood might implement
     these methods, and provide other abstract methods for their subclasses to implement.
     """
 
     def __init__(self, parameter_prefix: None | str = None) -> None:
-        """Default initialization for a base Likelihood object."""
+        """Default initialization for a base Likelihood object.
+
+        :params parameter_prefix: The prefix to prepend to all parameter names
+        """
         super().__init__(parameter_prefix=parameter_prefix)
 
     @abstractmethod
     def read(self, sacc_data: sacc.Sacc) -> None:
-        """Read the covariance matrix for this likelihood from the SACC file."""
+        """Read the covariance matrix for this likelihood from the SACC file.
+
+        :param sacc_data: The SACC data object to be read
+        """
 
     def make_realization_vector(self) -> npt.NDArray[np.float64]:
         """Create a new realization of the model.
 
         This new realization uses the previously computed theory vector and covariance
         matrix.
+
+        :return: the new realization of the theory vector
         """
         raise NotImplementedError(
             "This class does not implement make_realization_vector."
@@ -111,11 +119,17 @@ def make_realization(
             only the theory vector.
         :param strict: If True, check that the indices of the realization cover
             all the indices of the SACC data object.
+
+        :return: the new SACC object containing the new realization
         """
 
     @abstractmethod
     def compute_loglike(self, tools: ModelingTools) -> float:
-        """Compute the log-likelihood of generic CCL data."""
+        """Compute the log-likelihood of generic CCL data.
+
+        :param tools: the ModelingTools to be used in calculating the likelihood
+        :return: the log-likelihood
+        """
 
 
 class NamedParameters:
@@ -142,14 +156,22 @@ def __init__(
             ]
         ) = None,
     ):
-        """Initialize the object from the supplied mapping of values."""
+        """Initialize the object from the supplied mapping of values.
+
+        :param mapping: the mapping from strings to values used for initialization
+        """
         if mapping is None:
             self.data = {}
         else:
             self.data = dict(mapping)
 
     def get_bool(self, name: str, default_value: None | bool = None) -> bool:
-        """Return the named parameter as a bool."""
+        """Return the named parameter as a bool.
+
+        :param name: the name of the parameter to be returned
+        :param default_value: the default value if the parameter is not found
+        :return: the value of the parameter (or the default value)
+        """
         if default_value is None:
             val = self.data[name]
         else:
@@ -159,7 +181,12 @@ def get_bool(self, name: str, default_value: None | bool = None) -> bool:
         return val
 
     def get_string(self, name: str, default_value: None | str = None) -> str:
-        """Return the named parameter as a string."""
+        """Return the named parameter as a string.
+
+        :param name: the name of the parameter to be returned
+        :param default_value: the default value if the parameter is not found
+        :return: the value of the parameter (or the default value)
+        """
         if default_value is None:
             val = self.data[name]
         else:
@@ -169,7 +196,12 @@ def get_string(self, name: str, default_value: None | str = None) -> str:
         return val
 
     def get_int(self, name: str, default_value: None | int = None) -> int:
-        """Return the named parameter as an int."""
+        """Return the named parameter as an int.
+
+        :param name: the name of the parameter to be returned
+        :param default_value: the default value if the parameter is not found
+        :return: the value of the parameter (or the default value)
+        """
         if default_value is None:
             val = self.data[name]
         else:
@@ -179,7 +211,12 @@ def get_int(self, name: str, default_value: None | int = None) -> int:
         return val
 
     def get_float(self, name: str, default_value: None | float = None) -> float:
-        """Return the named parameter as a float."""
+        """Return the named parameter as a float.
+
+        :param name: the name of the parameter to be returned
+        :param default_value: the default value if the parameter is not found
+        :return: the value of the parameter (or the default value)
+        """
         if default_value is None:
             val = self.data[name]
         else:
@@ -189,7 +226,11 @@ def get_float(self, name: str, default_value: None | float = None) -> float:
         return val
 
     def get_int_array(self, name: str) -> npt.NDArray[np.int64]:
-        """Return the named parameter as a numpy array of int."""
+        """Return the named parameter as a numpy array of int.
+
+        :param name: the name of the parameter to be returned
+        :return: the value of the parameter
+        """
         tmp = self.data[name]
         assert isinstance(tmp, np.ndarray)
         val = tmp.view(dtype=np.int64)
@@ -197,7 +238,11 @@ def get_int_array(self, name: str) -> npt.NDArray[np.int64]:
         return val
 
     def get_float_array(self, name: str) -> npt.NDArray[np.float64]:
-        """Return the named parameter as a numpy array of float."""
+        """Return the named parameter as a numpy array of float.
+
+        :param name: the name of the parameter to be returned
+        :return: the value of the parameter
+        """
         tmp = self.data[name]
         assert isinstance(tmp, np.ndarray)
         val = tmp.view(dtype=np.float64)
@@ -209,7 +254,10 @@ def to_set(
     ) -> set[
         str | int | bool | float | npt.NDArray[np.int64] | npt.NDArray[np.float64]
     ]:
-        """Return the contained data as a set."""
+        """Return the contained data as a set.
+
+        :return: the value of the parameter as a set
+        """
         return set(self.data)
 
     def set_from_basic_dict(
@@ -219,7 +267,10 @@ def set_from_basic_dict(
             str | float | int | bool | Sequence[float] | Sequence[int] | Sequence[bool],
         ],
     ) -> None:
-        """Set the contained data from a dictionary of basic types."""
+        """Set the contained data from a dictionary of basic types.
+
+        :param basic_dict: the mapping from strings to values used for initialization
+        """
         for key, value in basic_dict.items():
             if isinstance(value, (str, float, int, bool)):
                 self.data = dict(self.data, **{key: value})
@@ -245,7 +296,10 @@ def convert_to_basic_dict(
         str,
         str | float | int | bool | Sequence[float] | Sequence[int] | Sequence[bool],
     ]:
-        """Convert a NamedParameters object to a dictionary of basic types."""
+        """Convert a NamedParameters object to a dictionary of built-in types.
+
+        :return: a dictionary containing the parameters as built-in Python types
+        """
         basic_dict: dict[
             str,
             str | float | int | bool | Sequence[float] | Sequence[int] | Sequence[bool],
@@ -281,6 +335,7 @@ def load_likelihood_from_module_type(
     :param module: a loaded module
     :param build_parameters: a NamedParameters object containing the factory
         function parameters
+    :return : a tuple of the likelihood and the modeling tools
     """
     if not hasattr(module, "build_likelihood"):
         if not hasattr(module, "likelihood"):
@@ -336,6 +391,7 @@ def load_likelihood_from_script(
     :param filename: script filename
     :param build_parameters: a NamedParameters object containing the factory
         function parameters
+    :return : a tuple of the likelihood and the modeling tools
     """
     _, file_extension = os.path.splitext(filename)
 
@@ -384,6 +440,7 @@ def load_likelihood_from_module(
     :param module: module name
     :param build_parameters: a NamedParameters object containing the factory
         function parameters
+    :return : a tuple of the likelihood and the modeling tools
     """
     try:
         mod = importlib.import_module(module)
@@ -406,6 +463,7 @@ def load_likelihood(
     :param likelihood_name: script filename or module name
     :param build_parameters: a NamedParameters object containing the factory
         function parameters
+    :return : a tuple of the likelihood and the modeling tools
     """
     try:
         return load_likelihood_from_script(likelihood_name, build_parameters)
diff --git a/firecrown/likelihood/number_counts.py b/firecrown/likelihood/number_counts.py
index 48b3a3ae9..a3cfa1c56 100644
--- a/firecrown/likelihood/number_counts.py
+++ b/firecrown/likelihood/number_counts.py
@@ -45,14 +45,23 @@ class NumberCountsArgs(SourceGalaxyArgs):
 class NumberCountsSystematic(SourceGalaxySystematic[NumberCountsArgs]):
     """Abstract base class for systematics for Number Counts sources.
 
-    Derived classes must implement :python`apply` with the correct signature.
+    Derived classes must implement :meth:`apply` with the correct signature.
     """
 
     @abstractmethod
     def apply(
         self, tools: ModelingTools, tracer_arg: NumberCountsArgs
     ) -> NumberCountsArgs:
-        """Apply method to include systematics in the tracer_arg."""
+        """Apply method to include systematics in the tracer_arg.
+
+        This does not modify the supplied tracer_arg; it returns a new
+        one that has been updated.
+
+        :param tools: the Modeling tools used to update the tracer_arg
+        :param tracer_arg: the original NumberCountsArgs to which to apply the
+             systematic
+        :return: the updated NumberCountsArgs
+        """
 
 
 class PhotoZShift(SourceGalaxyPhotoZShift[NumberCountsArgs]):
@@ -107,12 +116,10 @@ def apply(
     ) -> NumberCountsArgs:
         """Apply a linear bias systematic.
 
-        Parameters
-        ----------
-        cosmo : Cosmology
-            A Cosmology object.
-        tracer_arg : NumberCountsArgs
-            The source to which apply the shear bias.
+        :param tools: the ModelingTools used to update the tracer_arg
+        :param tracer_arg: a NumberCountsArgs object with values to be updated
+
+        :return: the updated NumberCountsArgs object
         """
         ccl_cosmo = tools.get_ccl_cosmology()
         pref = ((1.0 + tracer_arg.z) / (1.0 + self.z_piv)) ** self.alphaz
@@ -224,10 +231,10 @@ def apply(
     ) -> NumberCountsArgs:
         """Apply a magnification bias systematic.
 
-        :param tools: a ModelingTools object
-        :param tracer_arg: a NumberCountsArgs object
+        :param tools: currently unused, by required by the interface
+        :param tracer_arg: a NumberCountsArgs object with values to be updated
 
-        :return: a NumberCountsArgs object
+        :return: an updated NumberCountsArgs object
         """
         z_bar = self.z_c + self.z_m * (self.r_lim - 24.0)
         # The slope of log(n_tot(z,r_lim)) with respect to r_lim
@@ -286,7 +293,7 @@ def apply(
         :param tools: currently unused, but required by interface
         :param tracer_arg: a NumberCountsArgs object with values to be updated
 
-        :return: the updated NumberCountsArgs object
+        :return: an updated NumberCountsArgs object
         """
         return replace(
             tracer_arg,
@@ -344,7 +351,19 @@ def create_ready(
         scale: float = 1.0,
         systematics: None | list[SourceGalaxySystematic[NumberCountsArgs]] = None,
     ) -> NumberCounts:
-        """Create a NumberCounts object with the given tracer name and scale."""
+        """Create a NumberCounts object with the given tracer name and scale.
+
+        This is the recommended way to create a NumberCounts object. It creates
+        a fully initialized object.
+
+        :param inferred_zdist: the inferred redshift distribution
+        :param has_rsd: whether to include RSD in the tracer
+        :param derived_scale: whether to include a derived parameter for the scale
+            of the tracer
+        :param scale: the initial scale of the tracer
+        :param systematics: a list of systematics to apply to the tracer
+        :return: a fully initialized NumberCounts object
+        """
         obj = cls(
             sacc_tracer=inferred_zdist.bin_name,
             systematics=systematics,
@@ -365,15 +384,21 @@ def create_ready(
         return obj
 
     @final
-    def _update_source(self, params: ParamsMap):
+    def _update_source(self, params: ParamsMap) -> None:
         """Perform any updates necessary after the parameters have being updated.
 
         This implementation must update all contained Updatable instances.
+
+        :param params: the parameters to be used for the update
         """
         self.systematics.update(params)
 
     @final
     def _get_derived_parameters(self) -> DerivedParameterCollection:
+        """Return the derived parameters for this source.
+
+        :return: the derived parameters
+        """
         if self.derived_scale:
             assert self.current_tracer_args is not None
             derived_scale = DerivedParameter(
@@ -387,13 +412,10 @@ def _get_derived_parameters(self) -> DerivedParameterCollection:
 
         return derived_parameters
 
-    def _read(self, sacc_data):
+    def _read(self, sacc_data) -> None:
         """Read the data for this source from the SACC file.
 
-        Parameters
-        ----------
-        sacc_data : sacc.Sacc
-            The data in the sacc format.
+        :param sacc_data: The data in the sacc format.
         """
         # pylint: disable=unexpected-keyword-arg
         self.tracer_args = NumberCountsArgs(
@@ -409,7 +431,11 @@ def _read(self, sacc_data):
     def create_tracers(
         self, tools: ModelingTools
     ) -> tuple[list[Tracer], NumberCountsArgs]:
-        """Create the tracers for this source."""
+        """Create the tracers for this source.
+
+        :param tools: the ModelingTools used to create the tracers
+        :return: a tuple of tracers and the updated tracer_args
+        """
         tracer_args = self.tracer_args
         tracer_args = replace(tracer_args, bias=self.bias * np.ones_like(tracer_args.z))
 
@@ -471,8 +497,11 @@ def create_tracers(
 
         return tracers, tracer_args
 
-    def get_scale(self):
-        """Return the scale for this source."""
+    def get_scale(self) -> float:
+        """Return the scale for this source.
+
+        :return: the scale for this source.
+        """
         assert self.current_tracer_args
         return self.current_tracer_args.scale
 
@@ -488,11 +517,18 @@ class PhotoZShiftFactory(BaseModel):
     ] = "PhotoZShiftFactory"
 
     def create(self, bin_name: str) -> PhotoZShift:
-        """Create a PhotoZShift object with the given tracer name."""
+        """Create a PhotoZShift object with the given tracer name for a bin.
+
+        :param bin_name: the name of the bin
+        :return: the created PhotoZShift object
+        """
         return PhotoZShift(bin_name)
 
     def create_global(self) -> PhotoZShift:
-        """Create a PhotoZShift object with the given tracer name."""
+        """Required by the interface, but raises an error.
+
+        PhotoZShift systematics cannot be global.
+        """
         raise ValueError("PhotoZShift cannot be global.")
 
 
@@ -526,11 +562,15 @@ class PTNonLinearBiasSystematicFactory(BaseModel):
     ] = "PTNonLinearBiasSystematicFactory"
 
     def create(self, bin_name: str) -> PTNonLinearBiasSystematic:
-        """Create a PTNonLinearBiasSystematic object with the given tracer name."""
+        """Create a PTNonLinearBiasSystematic object with the given tracer name.
+
+        :param bin_name: the name of the bin
+        :return: the created PTNonLinearBiasSystematic object
+        """
         return PTNonLinearBiasSystematic(bin_name)
 
     def create_global(self) -> PTNonLinearBiasSystematic:
-        """Create a PTNonLinearBiasSystematic object with the given tracer name."""
+        """Create a global PTNonLinearBiasSystematic object."""
         return PTNonLinearBiasSystematic()
 
 
@@ -545,11 +585,18 @@ class MagnificationBiasSystematicFactory(BaseModel):
     ] = "MagnificationBiasSystematicFactory"
 
     def create(self, bin_name: str) -> MagnificationBiasSystematic:
-        """Create a MagnificationBiasSystematic object with the given tracer name."""
+        """Create a MagnificationBiasSystematic object with the given tracer name.
+
+        :param bin_name: the name of the bin
+        :return: the created MagnificationBiasSystematic object
+        """
         return MagnificationBiasSystematic(bin_name)
 
     def create_global(self) -> MagnificationBiasSystematic:
-        """Create a MagnificationBiasSystematic object with the given tracer name."""
+        """Required by the interface, but raises an error.
+
+        MagnificationBiasSystematic systematics cannot be global.
+        """
         raise ValueError("MagnificationBiasSystematic cannot be global.")
 
 
@@ -566,14 +613,15 @@ class ConstantMagnificationBiasSystematicFactory(BaseModel):
     def create(self, bin_name: str) -> ConstantMagnificationBiasSystematic:
         """Create a ConstantMagnificationBiasSystematic object.
 
-        Use the inferred_zdist to create the systematic.
+        :param bin_name: the name of the bin
+        :return: the created ConstantMagnificationBiasSystematic object
         """
         return ConstantMagnificationBiasSystematic(bin_name)
 
     def create_global(self) -> ConstantMagnificationBiasSystematic:
-        """Create a ConstantMagnificationBiasSystematic object.
+        """Required by the interface, but raises an error.
 
-        Use the inferred_zdist to create the systematic.
+        ConstantMagnificationBiasSystematic systematics cannot be global.
         """
         raise ValueError("ConstantMagnificationBiasSystematic cannot be global.")
 
@@ -599,8 +647,11 @@ class NumberCountsFactory(BaseModel):
     per_bin_systematics: Sequence[NumberCountsSystematicFactory]
     global_systematics: Sequence[NumberCountsSystematicFactory]
 
-    def model_post_init(self, __context) -> None:
-        """Initialize the NumberCountsFactory."""
+    def model_post_init(self, _) -> None:
+        """Initialize the NumberCountsFactory.
+
+        :param _: required by the interface but not used
+        """
         self._cache: dict[int, NumberCounts] = {}
         self._global_systematics_instances = [
             nc_systematic_factory.create_global()
@@ -608,7 +659,11 @@ def model_post_init(self, __context) -> None:
         ]
 
     def create(self, inferred_zdist: InferredGalaxyZDist) -> NumberCounts:
-        """Create a NumberCounts object with the given tracer name and scale."""
+        """Create a NumberCounts object with the given tracer name and scale.
+
+        :param inferred_zdist: the inferred redshift distribution
+        :return: a fully initialized NumberCounts object
+        """
         inferred_zdist_id = id(inferred_zdist)
         if inferred_zdist_id in self._cache:
             return self._cache[inferred_zdist_id]
@@ -628,7 +683,11 @@ def create_from_metadata_only(
         self,
         sacc_tracer: str,
     ) -> NumberCounts:
-        """Create an WeakLensing object with the given tracer name and scale."""
+        """Create an WeakLensing object with the given tracer name and scale.
+
+        :param sacc_tracer: the name of the tracer
+        :return: a fully initialized NumberCounts object
+        """
         sacc_tracer_id = hash(sacc_tracer)  # Improve this
         if sacc_tracer_id in self._cache:
             return self._cache[sacc_tracer_id]
diff --git a/firecrown/likelihood/supernova.py b/firecrown/likelihood/supernova.py
index 40ee63e34..278d05f1c 100644
--- a/firecrown/likelihood/supernova.py
+++ b/firecrown/likelihood/supernova.py
@@ -23,7 +23,7 @@
 
 
 class Supernova(Statistic):
-    """A supernova statistic.
+    """A statistic representing the distance modulus for a single supernova.
 
     This statistic that applies an additive shift M to a supernova's distance
     modulus.
diff --git a/firecrown/likelihood/two_point.py b/firecrown/likelihood/two_point.py
index 0f1becbd9..bfc77a706 100644
--- a/firecrown/likelihood/two_point.py
+++ b/firecrown/likelihood/two_point.py
@@ -289,7 +289,10 @@ def use_source_factory_metadata_only(
 
 
 class TwoPoint(Statistic):
-    """A two-point statistic.
+    """A statistic that represents the correlation between two measurements.
+
+    If the same source is used twice in the same TwoPoint object, this produces
+    an autocorrelation.
 
     For example, shear correlation function, galaxy-shear correlation function, etc.
 
diff --git a/tests/likelihood/gauss_family/statistic/test_statistic.py b/tests/likelihood/gauss_family/statistic/test_statistic.py
index 0738ab303..217410034 100644
--- a/tests/likelihood/gauss_family/statistic/test_statistic.py
+++ b/tests/likelihood/gauss_family/statistic/test_statistic.py
@@ -21,7 +21,7 @@ def test_vector_create():
     assert vals.dtype == np.float64
     assert vals.shape == (10,)
     for cls in VECTOR_CLASSES:
-        result = cls.create(vals)  # type: ignore
+        result = cls.create(vals)
         assert isinstance(result, cls)
         assert result.shape == (10,)
         assert np.array_equal(vals, result)
@@ -32,7 +32,7 @@ def test_vector_from_list():
     assert isinstance(vals, list)
     assert len(vals) == 4
     for cls in VECTOR_CLASSES:
-        result = cls.from_list(vals)  # type: ignore
+        result = cls.from_list(vals)
         assert isinstance(result, cls)
         assert result.shape == (4,)
         for i, val in enumerate(vals):
@@ -41,7 +41,7 @@ def test_vector_from_list():
 
 def test_vector_slicing():
     for cls in VECTOR_CLASSES:
-        vec = cls.create(np.random.random_sample((12,)))  # type: ignore
+        vec = cls.create(np.random.random_sample((12,)))
         assert isinstance(vec, cls)
         middle_part = vec[3:6]
         assert middle_part.shape == (3,)
@@ -50,7 +50,7 @@ def test_vector_slicing():
 
 def test_vector_copying():
     for cls in VECTOR_CLASSES:
-        vec = cls.create(np.random.random_sample((12,)))  # type: ignore
+        vec = cls.create(np.random.random_sample((12,)))
         assert isinstance(vec, cls)
         vec_copy = vec.copy()
         assert vec_copy is not vec
@@ -70,7 +70,7 @@ def test_ufunc_on_vector():
     data = np.array([0.0, 0.25, 0.50])
     expected = np.sin(data)
     for cls in VECTOR_CLASSES:
-        vec = cls.create(data)  # type: ignore
+        vec = cls.create(data)
         result = np.sin(vec)
         assert isinstance(result, cls)
         assert np.array_equal(result, expected)