start embracing xarray to handle broadcasting

Author: drbenvincent

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,23 @@ 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 model
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_ind": np.arange(self.X.shape[0])}
@@ -190,6 +209,8 @@ def __init__(
         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,37 @@ 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):

causalpy/experiments/synthetic_control.py

@@ -95,7 +95,7 @@ def __init__(
         self.datapost = data[data.index >= self.treatment_time]
 
         # split data into the 4 quadrants (pre/post, control/treated) and store as
-        # xarray DataArray objects.
+        # xarray.DataArray objects.
         # NOTE: if we have renamed/ensured the index is named "obs_ind", then it will
         # make constructing the xarray DataArray objects easier.
         self.datapre_control = xr.DataArray(

causalpy/pymc_models.py

@@ -163,16 +163,12 @@ def score(self, X, y) -> pd.Series:
 
         """
         mu = self.predict(X)
-        mu = az.extract(mu, group="posterior_predictive", var_names="mu").T.values
-        # Note: First argument must be a 1D array
-        return r2_score(y.flatten(), mu)
+        mu = az.extract(mu, group="posterior_predictive", var_names="mu").T
+        return r2_score(y.data, mu.data)
 
     def calculate_impact(
-        self, y_true: xr.DataArray | np.ndarray, y_pred: az.InferenceData
+        self, y_true: xr.DataArray, y_pred: az.InferenceData
     ) -> xr.DataArray:
-        if isinstance(y_true, np.ndarray):
-            y_true = xr.DataArray(y_true, dims=["obs_ind"])
-
         impact = y_true - y_pred["posterior_predictive"]["y_hat"]
         return impact.transpose(..., "obs_ind")
 
@@ -240,7 +236,7 @@ def build_model(self, X, y, coords):
         with self:
             self.add_coords(coords)
             X = pm.Data("X", X, dims=["obs_ind", "coeffs"])
-            y = pm.Data("y", y[:, 0], dims="obs_ind")
+            y = pm.Data("y", y, dims="obs_ind")
             beta = pm.Normal("beta", 0, 50, dims="coeffs")
             sigma = pm.HalfNormal("sigma", 1)
             mu = pm.Deterministic("mu", pm.math.dot(X, beta), dims="obs_ind")