Issue 129: increase docstring coverage, now at 86%

Author: Jason Preszler

causalpy/custom_exceptions.py

@@ -1,21 +1,26 @@
+"""
+Custom Exceptions for CausalPy.
+"""
+
+
 class BadIndexException(Exception):
     """Custom exception used when we have a mismatch in types between the dataframe
     index and an event, typically a treatment or intervention."""
 
-    def __init__(self, message):
+    def __init__(self, message: str):
         self.message = message
 
 
 class FormulaException(Exception):
     """Exception raised given when there is some error in a user-provided model
     formula"""
 
-    def __init__(self, message):
+    def __init__(self, message: str):
         self.message = message
 
 
 class DataException(Exception):
     """Exception raised given when there is some error in user-provided dataframe"""
 
-    def __init__(self, message):
+    def __init__(self, message: str):
         self.message = message

causalpy/data/simulate_data.py

@@ -11,6 +11,9 @@
 def _smoothed_gaussian_random_walk(
     gaussian_random_walk_mu, gaussian_random_walk_sigma, N, lowess_kwargs
 ):
+    """
+    Generates Gaussian random walk data and applies LOWESS
+    """
     x = np.arange(N)
     y = norm(gaussian_random_walk_mu, gaussian_random_walk_sigma).rvs(N).cumsum()
     filtered = lowess(y, x, **lowess_kwargs)

causalpy/plot_utils.py

@@ -1,3 +1,7 @@
+"""
+Plotting utility functions.
+"""
+
 from typing import Any, Dict, Optional, Tuple, Union
 
 import arviz as az

causalpy/pymc_experiments.py

@@ -1,3 +1,13 @@
+"""
+Experiment routines for PyMC models.
+
+Includes:
+1. ExperimentalDesign base class
+2. Pre-Post Fit
+3. Synthetic Control
+4. Difference in differences
+5. Regression Discontinuity
+"""
 import warnings
 from typing import Optional, Union
 
@@ -36,7 +46,7 @@ def idata(self):
         """Access to the InferenceData object"""
         return self.model.idata
 
-    def print_coefficients(self):
+    def print_coefficients(self) -> None:
         """Prints the model coefficients"""
         print("Model coefficients:")
         coeffs = az.extract(self.idata.posterior, var_names="beta")
@@ -236,7 +246,7 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs):
 
         return (fig, ax)
 
-    def summary(self):
+    def summary(self) -> None:
         """Print text output summarising the results"""
 
         print(f"{self.expt_type:=^80}")
@@ -524,13 +534,14 @@ def _plot_causal_impact_arrow(self, ax):
             va="center",
         )
 
-    def _causal_impact_summary_stat(self):
+    def _causal_impact_summary_stat(self) -> str:
+        """Computes the mean and 94% credible interval bounds for the causal impact."""
         percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values
         ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
         causal_impact = f"{self.causal_impact.mean():.2f}, "
         return f"Causal impact = {causal_impact + ci}"
 
-    def summary(self):
+    def summary(self) -> None:
         """Print text output summarising the results"""
 
         print(f"{self.expt_type:=^80}")
@@ -716,7 +727,7 @@ def plot(self):
         )
         return (fig, ax)
 
-    def summary(self):
+    def summary(self) -> None:
         """Print text output summarising the results"""
 
         print(f"{self.expt_type:=^80}")
@@ -795,7 +806,7 @@ def __init__(
 
         # ================================================================
 
-    def _input_validation(self):
+    def _input_validation(self) -> None:
         """Validate the input data and model formula for correctness"""
         if not _series_has_2_levels(self.data[self.group_variable_name]):
             raise DataException(
@@ -856,13 +867,14 @@ def plot(self):
         ax[1].set(title="Estimated treatment effect")
         return fig, ax
 
-    def _causal_impact_summary_stat(self):
+    def _causal_impact_summary_stat(self) -> str:
+        """Computes the mean and 94% credible interval bounds for the causal impact."""
         percentiles = self.causal_impact.quantile([0.03, 1 - 0.03]).values
         ci = r"$CI_{94\%}$" + f"[{percentiles[0]:.2f}, {percentiles[1]:.2f}]"
         causal_impact = f"{self.causal_impact.mean():.2f}, "
         return f"Causal impact = {causal_impact + ci}"
 
-    def summary(self):
+    def summary(self) -> None:
         """Print text output summarising the results"""
 
         print(f"{self.expt_type:=^80}")

causalpy/pymc_models.py

@@ -40,6 +40,7 @@ def build_model(self, X, y, coords) -> None:
         raise NotImplementedError("This method must be implemented by a subclass")
 
     def _data_setter(self, X) -> None:
+        """Set data for the model."""
         with self.model:
             pm.set_data({"X": X})
 

causalpy/skl_experiments.py

@@ -1,3 +1,6 @@
+"""
+Experiments for Scikit-Learn models
+"""
 import warnings
 from typing import Optional
 
@@ -78,6 +81,7 @@ def __init__(
         self.post_impact_cumulative = np.cumsum(self.post_impact)
 
     def plot(self, counterfactual_label="Counterfactual", **kwargs):
+        """Plot experiment results"""
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
         ax[0].plot(self.datapre.index, self.pre_y, "k.")
@@ -140,9 +144,11 @@ def plot(self, counterfactual_label="Counterfactual", **kwargs):
         return (fig, ax)
 
     def get_coeffs(self):
+        """Returns model coefficients"""
         return np.squeeze(self.model.coef_)
 
     def plot_coeffs(self):
+        """Plots coefficient bar plot"""
         df = pd.DataFrame(
             {"predictor variable": self.labels, "ols_coef": self.get_coeffs()}
         )
@@ -463,6 +469,7 @@ def _is_treated(self, x):
         return np.greater_equal(x, self.treatment_threshold)
 
     def plot(self):
+        """Plot results"""
         fig, ax = plt.subplots()
         # Plot raw data
         sns.scatterplot(

causalpy/skl_models.py

@@ -1,3 +1,9 @@
+"""
+Scikit-Learn Models
+
+Includes:
+1. Weighted Proportion
+"""
 from functools import partial
 
 import numpy as np
@@ -18,9 +24,11 @@ class WeightedProportion(LinearModel, RegressorMixin):
     """
 
     def loss(self, W, X, y):
+        """Compute root mean squared loss with data X, weights W, and predictor y"""
         return np.sqrt(np.mean((y - np.dot(X, W.T)) ** 2))
 
     def fit(self, X, y):
+        """Fit model on data X with predictor y"""
         w_start = [1 / X.shape[1]] * X.shape[1]
         coef_ = fmin_slsqp(
             partial(self.loss, X=X, y=y),
@@ -34,4 +42,5 @@ def fit(self, X, y):
         return self
 
     def predict(self, X):
+        """Predict results for data X"""
         return np.dot(X, self.coef_.T)

causalpy/tests/conftest.py

@@ -4,5 +4,6 @@
 
 @pytest.fixture(scope="session")
 def rng() -> np.random.Generator:
+    """Random number generator that can persist through a pytest session"""
     seed: int = sum(map(ord, "causalpy"))
     return np.random.default_rng(seed=seed)

causalpy/tests/test_data_loading.py

@@ -19,6 +19,10 @@
 
 @pytest.mark.parametrize("dataset_name", tests)
 def test_data_loading(dataset_name):
+    """
+    Checks that test data can be loaded into data frames and that there are no
+    missing values in any column.
+    """
     df = cp.load_data(dataset_name)
     assert isinstance(df, pd.DataFrame)
     # Check that there are no missing values in any column