changing column index restriction to label restriction

Author: JeanVanDyk

causalpy/experiments/interrupted_time_series.py

@@ -88,11 +88,12 @@ def __init__(
         self.treatment_time = treatment_time
         # set experiment type - usually done in subclasses
         self.expt_type = "Pre-Post Fit"
+        # set if the model is supposed to infer the treatment_time
+        self.infer_treatment_time = isinstance(self.treatment_time, (type(None), tuple))
 
-        # Set the data according to if the model is
-        if treatment_time is None or isinstance(treatment_time, tuple):
+        # Set the data according to if the model is fitted on the whole bunch or not
+        if self.infer_treatment_time:
             self.datapre = data
-            self.model.set_time_range(self.treatment_time, self.datapre)
         else:
             # split data in to pre and post intervention
             self.datapre = data[data.index < self.treatment_time]
@@ -107,6 +108,12 @@ def __init__(
         self.labels = X.design_info.column_names
         self.pre_y, self.pre_X = np.asarray(y), np.asarray(X)
 
+        # Setting the time range in which the model infers treatment_time
+        # Setting the timeline index so that the model can keep of time track between predicts
+        if self.infer_treatment_time:
+            self.model.set_time_range(self.treatment_time, self.datapre)
+            self.model.set_timeline(self.labels.index("t"))
+
         # fit the model to the observed (pre-intervention) data
         if isinstance(self.model, PyMCModel):
             COORDS = {"coeffs": self.labels, "obs_ind": np.arange(self.pre_X.shape[0])}
@@ -119,17 +126,15 @@ def __init__(
         # score the goodness of fit to the pre-intervention data
         self.score = self.model.score(X=self.pre_X, y=self.pre_y)
 
-        if treatment_time is None or isinstance(treatment_time, tuple):
+        if self.infer_treatment_time:
             # We're getting the inferred switchpoint as one of the values of the timeline, from the last column
             switchpoint = int(
                 az.extract(idata, group="posterior", var_names="switchpoint")
                 .mean("sample")
                 .values
             )
-
             # we're getting the associated index of that switchpoint
-            last_column = data.columns[-1]
-            self.treatment_time = data[data[last_column] == switchpoint].index[0]
+            self.treatment_time = data[data["t"] == switchpoint].index[0]
 
             # We're getting datapre as intended for prediction
             self.datapre = data[data.index < self.treatment_time]
@@ -162,11 +167,13 @@ def __init__(
 
     def input_validation(self, data, treatment_time, model):
         """Validate the input data and model formula for correctness"""
-        if isinstance(treatment_time, (type(None), tuple)) and not hasattr(
-            model, "set_time_range"
-        ):
+        if treatment_time is None and not hasattr(model, "set_time_range"):
+            raise ModelException(
+                "If treatment_time is None, provided model must have a 'set_time_range' method"
+            )
+        if isinstance(treatment_time, tuple) and not hasattr(model, "set_time_range"):
             raise ModelException(
-                "If treatment_time is None or a tuple, provided model must have a 'set_time_range' method"
+                "If treatment_time is a tuple, provided model must have a 'set_time_range' method"
             )
         if isinstance(data.index, pd.DatetimeIndex) and not isinstance(
             treatment_time, (pd.Timestamp, tuple, type(None))

causalpy/pymc_models.py

@@ -539,8 +539,9 @@ class InterventionTimeEstimator(PyMCModel):
         >>> labels = X.design_info.column_names
         >>> _y, _X = np.asarray(y), np.asarray(X)
         >>> COORDS = {"coeffs":labels, "obs_ind": np.arange(_X.shape[0])}
-        >>> model = ITE(sample_kwargs={"draws" : 10, "tune":10, "progressbar":False}) # For a quick overview. Remove sample_kwargs parameter for better performance
-        >>> model.set_time_range(None)
+        >>> model = ITE(sample_kwargs={"draws" : 10, "tune":10, "progressbar":False})
+        >>> model.set_time_range(None, data)
+        >>> model.set_timeline(-1)
         >>> model.fit(X=_X, y=_y, coords=COORDS)
         Inference ...
     """
@@ -549,7 +550,7 @@ def __init__(self, priors={}, sample_kwargs=None):
         super().__init__(sample_kwargs)
         self.priors = priors
 
-    def build_model(self, X, t, y, coords):
+    def build_model(self, X, y, coords):
         """
         Defines the PyMC model
 
@@ -568,7 +569,7 @@ def build_model(self, X, t, y, coords):
         with self:
             self.add_coords(coords)
 
-            t = pm.Data("t", t, dims="obs_ind")
+            t = pm.Data("t", X[:, self.timeline], dims="obs_ind")
             X = pm.Data("X", X, dims=["obs_ind", "coeffs"])
             y = pm.Data("y", y[:, 0], dims="obs_ind")
             lower_bound = pm.Data("lower_bound", self.time_range[0])
@@ -644,10 +645,9 @@ 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.
         random_seed = self.sample_kwargs.get("random_seed", None)
-        t = X[:, -1]
         if self.time_range is None:
-            self.time_range = (t.min(), t.max())
-        self.build_model(X, t, y, coords)
+            self.time_range = (X[:, self.timeline].min(), X[:, self.timeline].max())
+        self.build_model(X, y, coords)
         with self:
             self.idata = pm.sample(max_treedepth=15, **self.sample_kwargs)
             self.idata.extend(pm.sample_prior_predictive(random_seed=random_seed))
@@ -669,7 +669,7 @@ 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)
-        t = X[:, -1]
+        t = X[:, self.timeline]
         self._data_setter(X, t)
         with self:  # sample with new input data
             post_pred = pm.sample_posterior_predictive(
@@ -693,3 +693,21 @@ def _data_setter(self, X, t) -> None:
                 {"X": X, "t": t, "y": np.zeros(new_no_of_observations)},
                 coords={"obs_ind": np.arange(new_no_of_observations)},
             )
+
+    def set_time_range(self, time_range, data):
+        """
+        Set time_range.
+        """
+        if time_range is None:
+            self.time_range = time_range
+        else:
+            self.time_range = (
+                data["t"].loc[time_range[0]],
+                data["t"].loc[time_range[1]],
+            )
+
+    def set_timeline(self, index):
+        """
+        Set the index of the timeline in the given covariates
+        """
+        self.timeline = index