Merge branch 'main' into did-banks

Author: drbenvincent

causalpy/__init__.py

@@ -13,9 +13,7 @@
 #   limitations under the License.
 import arviz as az
 
-import causalpy.pymc_experiments as pymc_experiments  # to be deprecated
 import causalpy.pymc_models as pymc_models
-import causalpy.skl_experiments as skl_experiments  # to be deprecated
 import causalpy.skl_models as skl_models
 from causalpy.skl_models import create_causalpy_compatible_class
 from causalpy.version import __version__
@@ -41,11 +39,9 @@
     "InversePropensityWeighting",
     "load_data",
     "PrePostNEGD",
-    "pymc_experiments",  # to be deprecated
     "pymc_models",
     "RegressionDiscontinuity",
     "RegressionKink",
-    "skl_experiments",  # to be deprecated
     "skl_models",
     "SyntheticControl",
 ]

causalpy/experiments/diff_in_diff.py

@@ -19,6 +19,7 @@
 import numpy as np
 import pandas as pd
 import seaborn as sns
+import xarray as xr
 from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 from sklearn.base import RegressorMixin
@@ -87,7 +88,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(model=model)
-
+        # rename the index to "obs_ind"
+        data.index.name = "obs_ind"
         self.data = data
         self.expt_type = "Difference in Differences"
         self.formula = formula
@@ -102,6 +104,21 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
+        # turn into xarray.DataArray's
+        self.X = xr.DataArray(
+            self.X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": np.arange(self.X.shape[0]),
+                "coeffs": self.labels,
+            },
+        )
+        self.y = xr.DataArray(
+            self.y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": np.arange(self.y.shape[0])},
+        )
+
         # fit model
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_ind": np.arange(self.X.shape[0])}
@@ -183,13 +200,15 @@ def __init__(
                     )
         elif isinstance(self.model, RegressorMixin):
             # This is the coefficient on the interaction term
-            # TODO: THIS IS NOT YET CORRECT ?????
+            # TODO: CHECK FOR CORRECTNESS
             self.causal_impact = (
                 self.y_pred_treatment[1] - self.y_pred_counterfactual[0]
-            )[0]
+            )
         else:
             raise ValueError("Model type not recognized")
 
+        return
+
     def input_validation(self):
         """Validate the input data and model formula for correctness"""
         if "post_treatment" not in self.formula:

causalpy/experiments/interrupted_time_series.py

@@ -20,6 +20,7 @@
 import arviz as az
 import numpy as np
 import pandas as pd
+import xarray as xr
 from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 from sklearn.base import RegressorMixin
@@ -84,6 +85,8 @@ def __init__(
         **kwargs,
     ) -> None:
         super().__init__(model=model)
+        # rename the index to "obs_ind"
+        data.index.name = "obs_ind"
         self.input_validation(data, treatment_time)
         self.treatment_time = treatment_time
         # set experiment type - usually done in subclasses
@@ -107,6 +110,33 @@ def __init__(
         )
         self.post_X = np.asarray(new_x)
         self.post_y = np.asarray(new_y)
+        # turn into xarray.DataArray's
+        self.pre_X = xr.DataArray(
+            self.pre_X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": self.datapre.index,
+                "coeffs": self.labels,
+            },
+        )
+        self.pre_y = xr.DataArray(
+            self.pre_y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": self.datapre.index},
+        )
+        self.post_X = xr.DataArray(
+            self.post_X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": self.datapost.index,
+                "coeffs": self.labels,
+            },
+        )
+        self.post_y = xr.DataArray(
+            self.post_y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": self.datapost.index},
+        )
 
         # fit the model to the observed (pre-intervention) data
         if isinstance(self.model, PyMCModel):
@@ -125,10 +155,8 @@ def __init__(
 
         # calculate the counterfactual
         self.post_pred = self.model.predict(X=self.post_X)
-        self.pre_impact = self.model.calculate_impact(self.pre_y[:, 0], self.pre_pred)
-        self.post_impact = self.model.calculate_impact(
-            self.post_y[:, 0], self.post_pred
-        )
+        self.pre_impact = self.model.calculate_impact(self.pre_y, self.pre_pred)
+        self.post_impact = self.model.calculate_impact(self.post_y, self.post_pred)
         self.post_impact_cumulative = self.model.calculate_cumulative_impact(
             self.post_impact
         )

causalpy/experiments/prepostnegd.py

@@ -21,6 +21,7 @@
 import numpy as np
 import pandas as pd
 import seaborn as sns
+import xarray as xr
 from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 from sklearn.base import RegressorMixin
@@ -111,6 +112,21 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
+        # turn into xarray.DataArray's
+        self.X = xr.DataArray(
+            self.X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": np.arange(self.X.shape[0]),
+                "coeffs": self.labels,
+            },
+        )
+        self.y = xr.DataArray(
+            self.y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": self.data.index},
+        )
+
         # fit the model to the observed (pre-intervention) data
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_ind": np.arange(self.X.shape[0])}

causalpy/experiments/regression_discontinuity.py

@@ -23,7 +23,7 @@
 from matplotlib import pyplot as plt
 from patsy import build_design_matrices, dmatrices
 from sklearn.base import RegressorMixin
-
+import xarray as xr
 from causalpy.custom_exceptions import (
     DataException,
     FormulaException,
@@ -121,6 +121,21 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
+        # turn into xarray.DataArray's
+        self.X = xr.DataArray(
+            self.X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": np.arange(self.X.shape[0]),
+                "coeffs": self.labels,
+            },
+        )
+        self.y = xr.DataArray(
+            self.y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": np.arange(self.y.shape[0])},
+        )
+
         # fit model
         if isinstance(self.model, PyMCModel):
             # fit the model to the observed (pre-intervention) data

causalpy/experiments/regression_kink.py

@@ -23,7 +23,7 @@
 import pandas as pd
 import seaborn as sns
 from patsy import build_design_matrices, dmatrices
-
+import xarray as xr
 from causalpy.plot_utils import plot_xY
 
 from .base import BaseExperiment
@@ -84,6 +84,21 @@ def __init__(
         self.y, self.X = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
+        # turn into xarray.DataArray's
+        self.X = xr.DataArray(
+            self.X,
+            dims=["obs_ind", "coeffs"],
+            coords={
+                "obs_ind": np.arange(self.X.shape[0]),
+                "coeffs": self.labels,
+            },
+        )
+        self.y = xr.DataArray(
+            self.y[:, 0],
+            dims=["obs_ind"],
+            coords={"obs_ind": np.arange(self.y.shape[0])},
+        )
+
         COORDS = {"coeffs": self.labels, "obs_ind": np.arange(self.X.shape[0])}
         self.model.fit(X=self.X, y=self.y, coords=COORDS)