Removing time variable

Author: JeanVanDyk

causalpy/experiments/interrupted_time_series.py

@@ -93,15 +93,11 @@ def data_postprocessing(
         # --- Return ---
         res = {}
 
-        # --- Retrieve timeline and inferred treatment time ---
-        time_var = model.time_variable_name
-        timeline = data[time_var]
-
         tt_samples = idata.posterior["treatment_time"].values
         tt_mean = int(tt_samples.mean().item())
 
         # Actual timestamp (index) corresponding to inferred treatment
-        tt = data[timeline == tt_mean].index[0]
+        tt = data.index[tt_mean]
         # Index of the inferred treatment time in the data
         tt_idx = data.index.get_loc(tt)
         res["treatment_time"] = tt
@@ -128,9 +124,13 @@ def data_postprocessing(
 
         # --- Create a mask to isolate post-treatment period ---
         # Timeline reshaped to match broadcasting with treatment time
-        timeline_reshape = timeline.values.reshape(1, 1, -1)
+        timeline = [
+            [[i for i in range(len(data))] for _ in range(len(tt_samples[0]))]
+            for _ in range(len(tt_samples))
+        ]
+        timeline_broadcast = np.array(timeline)
         tt_broadcast = tt_samples[:, :, None].astype(int)
-        mask = (timeline_reshape >= tt_broadcast).astype(int)
+        mask = (timeline_broadcast >= tt_broadcast).astype(int)
 
         # --- Compute cumulative post-treatment impact ---
         post_impact = impact * mask
@@ -184,12 +184,10 @@ def plot_intervention_line(
         Draw a vertical line at the inferred treatment time and shade the HDI interval around it.
         """
         data = pd.concat([datapre, datapost])
-        timeline = data[model.time_variable_name]
-
         # Extract the HDI (uncertainty interval) of the treatment time
         hdi = az.hdi(idata, var_names=["treatment_time"])["treatment_time"].values
-        x1 = data[timeline == int(hdi[0])].index[0]
-        x2 = data[timeline == int(hdi[1])].index[0]
+        x1 = data.index[int(hdi[0])]
+        x2 = data.index[int(hdi[1])]
 
         for i in [0, 1, 2]:
             ymin, ymax = ax[i].get_ylim()

causalpy/pymc_models.py

@@ -560,23 +560,21 @@ class InterventionTimeEstimator(PyMCModel):
         ...     coords={"obs_ind": data.index},
         ...     )
         >>> COORDS = {"coeffs":labels, "obs_ind": np.arange(_X.shape[0])}
-        >>> model = ITE(time_variable_name="t", treatment_effect_type="level", sample_kwargs={"draws" : 10, "tune":10, "progressbar":False})
+        >>> model = ITE(treatment_effect_type="level", sample_kwargs={"draws" : 10, "tune":10, "progressbar":False})
         >>> model.set_time_range(None, data)
         >>> model.fit(X=_X, y=_y, coords=COORDS)
         Inference ...
     """
 
     def __init__(
         self,
-        time_variable_name: str,
         treatment_effect_type: str | list[str],
         treatment_effect_param=None,
         sample_kwargs=None,
     ):
         """
         Initializes the InterventionTimeEstimator model.
 
-        :param time_variable_name: name of the column representing time among the covariates
         :param treatment_effect_type: Optional dictionary that specifies prior parameters for the
             intervention effects. Expected keys are:
                 - "level": [mu, sigma]
@@ -586,7 +584,6 @@ def __init__(
             If the associated list is incomplete, default values will be used.
         :param sample_kwargs: Optional dictionary of arguments passed to pm.sample().
         """
-        self.time_variable_name = time_variable_name
 
         super().__init__(sample_kwargs)
 
@@ -657,7 +654,7 @@ def build_model(self, X, y, coords):
         with self:
             self.add_coords(coords)
 
-            t = pm.Data("t", X.sel(coeffs=self.time_variable_name), dims="obs_ind")
+            t = pm.Data("t", np.arange(len(X)), dims="obs_ind")
             X = pm.Data("X", X, dims=["obs_ind", "coeffs"])
             y = pm.Data("y", y, dims="obs_ind")
 
@@ -768,7 +765,7 @@ def _data_setter(self, X) -> None:
             pm.set_data(
                 {
                     "X": X,
-                    "t": X.sel(coeffs=self.time_variable_name),
+                    "t": np.arange(len(X)),
                     "y": np.zeros(new_no_of_observations),
                 },
                 coords={"obs_ind": np.arange(new_no_of_observations)},
@@ -794,11 +791,11 @@ def set_time_range(self, time_range, data):
         """
         if time_range is None:
             self.time_range = (
-                data[self.time_variable_name].min(),
-                data[self.time_variable_name].max(),
+                0,
+                len(data),
             )
         else:
             self.time_range = (
-                data[self.time_variable_name].loc[time_range[0]],
-                data[self.time_variable_name].loc[time_range[1]],
+                data.index.get_loc(time_range[0]),
+                data.index.get_loc(time_range[1]),
             )