#53 adding more docstring info

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -26,7 +26,14 @@ def __init__(self, prediction_model=None, **kwargs):
 class TimeSeriesExperiment(ExperimentalDesign):
     """A class to analyse time series quasi-experiments"""
 
-    def __init__(self, data, treatment_time, formula, prediction_model=None, **kwargs):
+    def __init__(
+        self,
+        data: pd.DataFrame,
+        treatment_time: int,
+        formula: str,
+        prediction_model=None,
+        **kwargs,
+    ) -> None:
         super().__init__(prediction_model=prediction_model, **kwargs)
         self.treatment_time = treatment_time
         # split data in to pre and post intervention
@@ -111,6 +118,8 @@ def plot(self):
 
 
 class SyntheticControl(TimeSeriesExperiment):
+    """A wrapper around the TimeSeriesExperiment class"""
+
     def plot(self):
         """Plot the results"""
         fig, ax = super().plot()
@@ -121,6 +130,8 @@ def plot(self):
 
 
 class InterruptedTimeSeries(TimeSeriesExperiment):
+    """A wrapper around the TimeSeriesExperiment class"""
+
     pass
 
 

causalpy/pymc_models.py

@@ -39,7 +39,7 @@ def predict(self, X):
         return post_pred
 
     def score(self, X, y):
-        """Score the predictions $R^2$ given inputs X and outputs y."""
+        """Score the predictions :math:`R^2` given inputs ``X`` and outputs ``y``."""
         yhat = self.predict(X)
         yhat = az.extract(yhat, group="posterior_predictive", var_names="y_hat").mean(
             dim="sample"

causalpy/skl_experiments.py

@@ -359,7 +359,7 @@ def __init__(
         )
 
     def _is_treated(self, x):
-        """Returns ``True`` if `x` is greater than or equal to the treatment threshold.
+        """Returns ``True`` if ``x`` is greater than or equal to the treatment threshold.
 
         .. warning::
 
@@ -401,6 +401,7 @@ def plot(self):
         return (fig, ax)
 
     def summary(self):
+        """Print text output summarising the results"""
         print("Difference in Differences experiment")
         print(f"Formula: {self.formula}")
         print(f"Running variable: {self.running_variable_name}")