#53 tidy up and add some docstrings

Author: drbenvincent

CONTRIBUTING.md

@@ -29,6 +29,8 @@ pip install -r requirements-dev.txt
 pip install -r requirements-docs.txt
 ```
 
+It may also be necessary to [install](https://pandoc.org/installing.html) `pandoc`. On a mac, I run `brew install pandoc`.
+
 5. You may also need to run this to get pre-commit checks working
 
 ```
@@ -60,6 +62,14 @@ make clean && make html
 
 Docs are built in `docs/_build`, but these docs are _not_ committed to the GitHub repository due to `.gitignore`.
 
+## Remote documentation
+
+Documentation is hosted on https://causalpy.readthedocs.io/. New remote builds are triggered automatically whenever there is an update to the `main` branch.
+
+The `.readthedocs.yaml` file contains the configurations for the remote build.
+
+If there are autodoc issues/errors in remote builds of the docs, we need to add all package dependencies (in `requirements.txt`) into the list `autodoc_mock_imports` in `docs/config.py`.
+
 ## New releases [work in progress]
 
 1. Bump the release version in `causalpy/version.py`. This is automatically read by `setup.py` and `docs/config.py`.

causalpy/plot_utils.py

@@ -2,7 +2,7 @@
 
 
 def plot_xY(x, Y, ax, plot_hdi_kwargs=dict()):
-    """Utility function to plot HDI intervals"""
+    """Utility function to plot HDI intervals."""
     quantiles = Y.quantile(
         (0.025, 0.25, 0.5, 0.75, 0.975), dim=("chain", "draw")
     ).transpose()

causalpy/pymc_experiments.py

@@ -12,6 +12,8 @@
 
 
 class ExperimentalDesign:
+    """Base class"""
+
     prediction_model = None
 
     def __init__(self, prediction_model=None, **kwargs):
@@ -22,6 +24,8 @@ 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):
         super().__init__(prediction_model=prediction_model, **kwargs)
         self.treatment_time = treatment_time
@@ -71,6 +75,7 @@ def __init__(self, data, treatment_time, formula, prediction_model=None, **kwarg
         self.post_impact_cumulative = self.post_impact.cumsum(dim="obs_ind")
 
     def plot(self):
+        """Plot the results"""
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
         # pre-intervention period
@@ -107,6 +112,7 @@ def plot(self):
 
 class SyntheticControl(TimeSeriesExperiment):
     def plot(self):
+        """Plot the results"""
         fig, ax = super().plot()
         # plot control units as well
         ax[0].plot(self.datapre.index, self.pre_X, "-", c=[0.8, 0.8, 0.8], zorder=1)
@@ -119,7 +125,12 @@ class InterruptedTimeSeries(TimeSeriesExperiment):
 
 
 class DifferenceInDifferences(ExperimentalDesign):
-    """Note: there is no pre/post intervention data distinction for DiD, we fit all the data available."""
+    """A class to analyse data from Difference in Difference settings.
+
+    .. Note::
+
+        There is no pre/post intervention data distinction for DiD, we fit all the data available.
+    """
 
     def __init__(
         self,
@@ -183,6 +194,7 @@ def __init__(
         )
 
     def plot(self):
+        """Plot the results"""
         fig, ax = plt.subplots()
 
         # Plot raw data
@@ -272,7 +284,13 @@ def plot(self):
 
 
 class RegressionDiscontinuity(ExperimentalDesign):
-    """Note: there is no pre/post intervention data distinction, we fit all the data available."""
+    """
+    A class to analyse regression discontinuity experiments.
+
+    .. Note::
+
+    There is no pre/post intervention data distinction for DiD, we fit all the data available.
+    """
 
     def __init__(
         self,
@@ -337,9 +355,16 @@ def __init__(
         )
 
     def _is_treated(self, x):
+        """Returns true is `x` is greater than or equal to the treatment threshold
+
+        .. admonition::
+
+           Assumes treatment is given to those ABOVE the treatment threshold.
+        """
         return np.greater_equal(x, self.treatment_threshold)
 
     def plot(self):
+        """Plot the results"""
         fig, ax = plt.subplots()
         # Plot raw data
         sns.scatterplot(
@@ -373,6 +398,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}")

causalpy/pymc_models.py

@@ -21,6 +21,7 @@ def _data_setter(self, X):
             pm.set_data({"X": X})
 
     def fit(self, X, y, coords):
+        """Draw samples from posterior, prior predictive, and posterior predictive distributions."""
         self.build_model(X, y, coords)
         with self.model:
             self.idata = pm.sample()
@@ -29,6 +30,7 @@ def fit(self, X, y, coords):
         return self.idata
 
     def predict(self, X):
+        """Predict data given input data `X`"""
         self._data_setter(X)
         with self.model:  # sample with new input data
             post_pred = pm.sample_posterior_predictive(
@@ -37,6 +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."""
         yhat = self.predict(X)
         yhat = az.extract(yhat, group="posterior_predictive", var_names="y_hat").mean(
             dim="sample"
@@ -48,6 +51,7 @@ class WeightedSumFitter(ModelBuilder):
     """Used for synthetic control experiments"""
 
     def build_model(self, X, y, coords):
+        """Defines the PyMC model"""
         with self:
             self.add_coords(coords)
             n_predictors = X.shape[1]
@@ -60,7 +64,10 @@ def build_model(self, X, y, coords):
 
 
 class LinearRegression(ModelBuilder):
+    """Custom PyMC model for linear regression"""
+
     def build_model(self, X, y, coords):
+        """Defines the PyMC model"""
         with self:
             self.add_coords(coords)
             X = pm.MutableData("X", X, dims=["obs_ind", "coeffs"])

docs/conf.py

@@ -14,31 +14,6 @@
 
 sys.path.insert(0, os.path.abspath("../"))
 
-
-# Need to mock the package dependencies in order for autodoc to work when docs are remotely built on readthedocs
-# from mock import Mock as MagicMock
-
-
-# class Mock(MagicMock):
-#     @classmethod
-#     def __getattr__(cls, name):
-#         return MagicMock()
-
-
-# MOCK_MODULES = [
-#     "arviz",
-#     "matplotlib",
-#     "numpy",
-#     "pandas",
-#     "patsy",
-#     "pymc",
-#     "scipy",
-#     "seaborn",
-#     "sklearn",
-#     "xarray",
-# ]
-# sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES)
-
 # autodoc_mock_imports
 # This avoids autodoc breaking when it can't find packages imported in the code.
 # https://www.sphinx-doc.org/en/master/usage/extensions/autodoc.html#confval-autodoc_mock_imports

requirements-docs.txt

@@ -1,6 +1,5 @@
 ipykernel
 linkify-it-py
-# mock
 myst_parser
 nbsphinx
 pathlib