gemini

Author: juanitorduz

causalpy/data/datasets.py

@@ -43,12 +43,12 @@
 }
 
 
-def _get_data_home() -> pathlib.PosixPath:
+def _get_data_home() -> pathlib.Path:
     """Return the path of the data directory"""
     return pathlib.Path(cp.__file__).parents[1] / "causalpy" / "data"
 
 
-def load_data(dataset: str = None) -> pd.DataFrame:
+def load_data(dataset: str | None = None) -> pd.DataFrame:
     """Loads the requested dataset and returns a pandas DataFrame.
 
     :param dataset: The desired dataset to load

causalpy/experiments/diff_in_diff.py

@@ -96,6 +96,7 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(model=model)
+        self.causal_impact: xr.DataArray | float | None
         # rename the index to "obs_ind"
         data.index.name = "obs_ind"
         self.data = data
@@ -213,6 +214,7 @@ def __init__(
 
         # calculate causal impact
         if isinstance(self.model, PyMCModel):
+            assert self.model.idata is not None
             # This is the coefficient on the interaction term
             coeff_names = self.model.idata.posterior.coords["coeffs"].data
             for i, label in enumerate(coeff_names):
@@ -395,7 +397,7 @@ def _plot_causal_impact_arrow(results, ax):
         labels = ["Control group"]
 
         # Plot model fit to treatment group
-        time_points = self.x_pred_control[self.time_variable_name].values
+        time_points = self.x_pred_treatment[self.time_variable_name].values
         h_line, h_patch = plot_xY(
             time_points,
             self.y_pred_treatment["posterior_predictive"].mu.isel(treated_units=0),

causalpy/experiments/interrupted_time_series.py

@@ -93,6 +93,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(model=model)
+        self.pre_y: xr.DataArray
+        self.post_y: xr.DataArray
         # rename the index to "obs_ind"
         data.index.name = "obs_ind"
         self.input_validation(data, treatment_time)

causalpy/experiments/prepostnegd.py

@@ -98,6 +98,10 @@ def __init__(
         **kwargs,
     ):
         super().__init__(model=model)
+        self.causal_impact: xr.DataArray
+        self.pred_xi: np.ndarray
+        self.pred_untreated: az.InferenceData
+        self.pred_treated: az.InferenceData
         self.data = data
         self.expt_type = "Pretest/posttest Nonequivalent Group Design"
         self.formula = formula
@@ -140,6 +144,7 @@ def __init__(
         else:
             raise ValueError("Model type not recognized")
 
+        assert self.model.idata is not None
         # Calculate the posterior predictive for the treatment and control for an
         # interpolated set of pretest values
         # get the model predictions of the observed data

causalpy/plot_utils.py

@@ -24,10 +24,11 @@
 import xarray as xr
 from matplotlib.collections import PolyCollection
 from matplotlib.lines import Line2D
+from pandas.api.extensions import ExtensionArray
 
 
 def plot_xY(
-    x: Union[pd.DatetimeIndex, np.array],
+    x: Union[pd.DatetimeIndex, np.ndarray, pd.Index, pd.Series, ExtensionArray],
     Y: xr.DataArray,
     ax: plt.Axes,
     plot_hdi_kwargs: Optional[Dict[str, Any]] = None,

causalpy/pymc_models.py

@@ -91,7 +91,7 @@ class PyMCModel(pm.Model):
     Inference data...
     """
 
-    default_priors = {}
+    default_priors: Dict[str, Prior] = {}
 
     def priors_from_data(self, X, y) -> Dict[str, Any]:
         """
@@ -236,6 +236,7 @@ def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None:
         self.build_model(X, y, coords)
         with self:
             self.idata = pm.sample(**self.sample_kwargs)
+            assert self.idata is not None
             self.idata.extend(pm.sample_prior_predictive(random_seed=random_seed))
             self.idata.extend(
                 pm.sample_posterior_predictive(
@@ -349,6 +350,8 @@ def calculate_cumulative_impact(self, impact):
         return impact.cumsum(dim="obs_ind")
 
     def print_coefficients(self, labels, round_to=None) -> None:
+        assert self.idata is not None
+
         def print_row(
             max_label_length: int, name: str, coeff_samples: xr.DataArray, round_to: int
         ) -> None: