Merge pull request #330 from pymc-labs/failing-doctests

Fix failing doctests

Author: drbenvincent

causalpy/pymc_experiments.py

@@ -99,6 +99,11 @@ def print_coefficients(self, round_to=None) -> None:
         ...                 "progressbar": False
         ...             }),
         ...  )
+        <BLANKLINE>
+        <BLANKLINE>
+        <BLANKLINE>
+        <BLANKLINE>
+        <BLANKLINE>
         >>> result.print_coefficients(round_to=1) # doctest: +NUMBER
         Model coefficients:
         Intercept                     1, 94% HDI [1, 1]
@@ -157,6 +162,11 @@ class PrePostFit(ExperimentalDesign, PrePostFitDataValidator):
     ...         }
     ...     ),
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     >>> result.summary(round_to=1) # doctest: +NUMBER
     ==================================Pre-Post Fit==================================
     Formula: actual ~ 0 + a + g
@@ -383,6 +393,11 @@ class InterruptedTimeSeries(PrePostFit):
     ...         }
     ...     )
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     """
 
     expt_type = "Interrupted Time Series"
@@ -418,6 +433,11 @@ class SyntheticControl(PrePostFit):
     ...         }
     ...     ),
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     """
 
     expt_type = "Synthetic Control"
@@ -474,6 +494,11 @@ class DifferenceInDifferences(ExperimentalDesign, DiDDataValidator):
     ...         }
     ...     )
     ...  )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     """
 
     def __init__(
@@ -765,6 +790,11 @@ class RegressionDiscontinuity(ExperimentalDesign, RDDataValidator):
     ...     ),
     ...     treatment_threshold=0.5,
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     """
 
     def __init__(
@@ -1174,6 +1204,10 @@ class PrePostNEGD(ExperimentalDesign, PrePostNEGDDataValidator):
     ...         }
     ...     )
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
+    <BLANKLINE>
     >>> result.summary(round_to=1) # doctest: +NUMBER
     ==================Pretest/posttest Nonequivalent Group Design===================
     Formula: post ~ 1 + C(group) + pre
@@ -1404,6 +1438,8 @@ class InstrumentalVariable(ExperimentalDesign, IVDataValidator):
     ...         formula=formula,
     ...         model=InstrumentalVariableRegression(sample_kwargs=sample_kwargs),
     ... )
+    <BLANKLINE>
+    <BLANKLINE>
     """
 
     def __init__(

causalpy/pymc_models.py

@@ -72,13 +72,17 @@ class ModelBuilder(pm.Model):
     ...             }
     ... )
     >>> model.fit(X, y)
-    Inference...
+    <BLANKLINE>
+    <BLANKLINE>
+    Inference data...
     >>> X_new = rng.normal(loc=0, scale=1, size=(20,2))
     >>> model.predict(X_new)
-    Inference...
-    >>> model.score(X, y) # doctest: +NUMBER
-    r2        0.3
-    r2_std    0.0
+    <BLANKLINE>
+    Inference data...
+    >>> model.score(X, y)
+    <BLANKLINE>
+    r2        0.390344
+    r2_std    0.081135
     dtype: float64
     """
 
@@ -112,10 +116,7 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None:
 
         # Ensure random_seed is used in sample_prior_predictive() and
         # sample_posterior_predictive() if provided in sample_kwargs.
-        if "random_seed" in self.sample_kwargs:
-            random_seed = self.sample_kwargs["random_seed"]
-        else:
-            random_seed = None
+        random_seed = self.sample_kwargs.get("random_seed", None)
 
         self.build_model(X, y, coords)
         with self:
@@ -137,10 +138,17 @@ def predict(self, X):
 
         """
 
+        # Ensure random_seed is used in sample_prior_predictive() and
+        # sample_posterior_predictive() if provided in sample_kwargs.
+        random_seed = self.sample_kwargs.get("random_seed", None)
+
         self._data_setter(X)
         with self:  # sample with new input data
             post_pred = pm.sample_posterior_predictive(
-                self.idata, var_names=["y_hat", "mu"], progressbar=False
+                self.idata,
+                var_names=["y_hat", "mu"],
+                progressbar=False,
+                random_seed=random_seed,
             )
         return post_pred
 
@@ -193,7 +201,9 @@ class WeightedSumFitter(ModelBuilder):
     >>> y = np.asarray(sc['actual']).reshape((sc.shape[0], 1))
     >>> wsf = WeightedSumFitter(sample_kwargs={"progressbar": False})
     >>> wsf.fit(X,y)
-    Inference ...
+    <BLANKLINE>
+    <BLANKLINE>
+    Inference data...
     """  # noqa: W605
 
     def build_model(self, X, y, coords):
@@ -249,7 +259,9 @@ class LinearRegression(ModelBuilder):
     ...                 'obs_indx': np.arange(rd.shape[0])
     ...                },
     ... )
-    Inference...
+    <BLANKLINE>
+    <BLANKLINE>
+    Inference data...
     """  # noqa: W605
 
     def build_model(self, X, y, coords):
@@ -301,6 +313,8 @@ class InstrumentalVariableRegression(ModelBuilder):
     ...                  "eta": 2,
     ...                  "lkj_sd": 2,
     ...              })
+    <BLANKLINE>
+    <BLANKLINE>
     Inference data...
     """