bug fixing

Author: drbenvincent

causalpy/experiments/interrupted_time_series.py

@@ -125,10 +125,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/skl_models.py

@@ -16,6 +16,7 @@
 from functools import partial
 
 import numpy as np
+import xarray as xr
 from scipy.optimize import fmin_slsqp
 from sklearn.base import RegressorMixin
 from sklearn.linear_model._base import LinearModel
@@ -28,6 +29,9 @@ class ScikitLearnAdaptor:
 
     def calculate_impact(self, y_true, y_pred):
         """Calculate the causal impact of the intervention."""
+        if isinstance(y_true, np.ndarray):
+            y_true = xr.DataArray(y_true, dims=["obs_ind"])
+
         return y_true - y_pred
 
     def calculate_cumulative_impact(self, impact):