updating notebook with examples and adding time_variable_name parameter

Author: JeanVanDyk

causalpy/experiments/interrupted_time_series.py

@@ -48,40 +48,50 @@ def data_preprocessing(self, data, treatment_time, model):
         model.set_time_range(treatment_time, data)
         return data
 
-    def data_postprocessing(self, data, idata, treatment_time, pre_y, pre_X):
+    def data_postprocessing(self, model, data, idata, treatment_time, pre_y, pre_X):
         """
         Postprocess data based on the inferred treatment time for further analysis and plotting.
         """
+        # --- Getting the time_variable_name ---
+        time_variable_name = model.get_time_variable_name()
+
         # --- Inferred treatment time ---
         treatment_time_mean = idata.posterior["treatment_time"].mean().item()
         inferred_treatment_time = int(treatment_time_mean)
-        idx_treatment_time = data.index[data["t"] == inferred_treatment_time][0]
+        idx_treatment_time = data.index[
+            data[time_variable_name] == inferred_treatment_time
+        ][0]
 
         # --- HDI bounds (credible interval) ---
         hdi_bounds = az.hdi(idata, var_names=["treatment_time"])[
             "treatment_time"
         ].values
         hdi_start_time = int(hdi_bounds[0])
-        indice = data.index.get_loc(data.index[data["t"] == hdi_start_time][0])
+        indice = data.index.get_loc(
+            data.index[data[time_variable_name] == hdi_start_time][0]
+        )
 
         # --- Slicing ---
-        datapre = data[data["t"] < hdi_start_time]
-        datapost = data[data["t"] >= hdi_start_time]
+        datapre = data[data[time_variable_name] < hdi_start_time]
+        datapost = data[data[time_variable_name] >= hdi_start_time]
 
         truncated_y = pre_y.isel(obs_ind=slice(0, indice))
         truncated_X = pre_X.isel(obs_ind=slice(0, indice))
 
         return datapre, datapost, truncated_y, truncated_X, idx_treatment_time
 
-    def plot_intervention_line(self, ax, idata, datapost, treatment_time):
+    def plot_intervention_line(self, ax, model, idata, datapost, treatment_time):
         """
         Plot a vertical line at the inferred treatment time, along with a shaded area
         representing the Highest Density Interval (HDI) of the inferred time.
         """
+        # --- Getting the time_variable_name ---
+        time_variable_name = model.get_time_variable_name()
+
         # Extract the HDI (uncertainty interval) of the treatment time
         hdi = az.hdi(idata, var_names=["treatment_time"])["treatment_time"].values
-        x1 = datapost.index[datapost["t"] == int(hdi[0])][0]
-        x2 = datapost.index[datapost["t"] == int(hdi[1])][0]
+        x1 = datapost.index[datapost[time_variable_name] == int(hdi[0])][0]
+        x2 = datapost.index[datapost[time_variable_name] == int(hdi[1])][0]
 
         for i in [0, 1, 2]:
             ymin, ymax = ax[i].get_ylim()
@@ -119,7 +129,7 @@ def plot_treated_counterfactual(
             plot_hdi_kwargs={"color": "yellowgreen"},
         )
         handles.append((h_line, h_patch))
-        labels.append("treated counterfactual")
+        labels.append("Treated counterfactual")
 
 
 class HandlerKTT:
@@ -135,7 +145,7 @@ def data_preprocessing(self, data, treatment_time, model):
         # Use only data before treatment for training the model
         return data[data.index < treatment_time]
 
-    def data_postprocessing(self, data, idata, treatment_time, pre_y, pre_X):
+    def data_postprocessing(self, model, data, idata, treatment_time, pre_y, pre_X):
         """
         Split data into pre- and post-treatment periods using the known treatment time.
         """
@@ -147,7 +157,7 @@ def data_postprocessing(self, data, idata, treatment_time, pre_y, pre_X):
             treatment_time,
         )
 
-    def plot_intervention_line(self, ax, idata, datapost, treatment_time):
+    def plot_intervention_line(self, model, ax, idata, datapost, treatment_time):
         """
         Plot a vertical line at the known treatment time on provided axes.
         """
@@ -276,7 +286,7 @@ def __init__(
         # Postprocessing with handler
         self.datapre, self.datapost, self.pre_y, self.pre_X, self.treatment_time = (
             self.handler.data_postprocessing(
-                data, idata, treatment_time, self.pre_y, self.pre_X
+                self.model, data, idata, treatment_time, self.pre_y, self.pre_X
             )
         )
 
@@ -443,7 +453,7 @@ def _bayesian_plot(
 
         # Plot vertical line marking treatment time (with HDI if it's inferred)
         self.handler.plot_intervention_line(
-            ax, self.idata, self.datapost, self.treatment_time
+            ax, self.model, self.idata, self.datapost, self.treatment_time
         )
 
         ax[0].legend(

causalpy/pymc_models.py

@@ -566,10 +566,13 @@ class InterventionTimeEstimator(PyMCModel):
         Inference ...
     """
 
-    def __init__(self, treatment_type_effect=None, sample_kwargs=None):
+    def __init__(
+        self, time_variable_name: str, treatment_type_effect=None, sample_kwargs=None
+    ):
         """
         Initializes the InterventionTimeEstimator model.
 
+        :param time_variable_name: name of the column representing time among the covariates
         :param treatment_type_effect: Optional dictionary that specifies prior parameters for the
             intervention effects. Expected keys are:
                 - "level": [mu, sigma]
@@ -579,9 +582,10 @@ def __init__(self, treatment_type_effect=None, sample_kwargs=None):
             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
+
         if treatment_type_effect is None:
             treatment_type_effect = {}
-
         if not isinstance(treatment_type_effect, dict):
             raise TypeError("treatment_type_effect must be a dictionary.")
 
@@ -609,7 +613,7 @@ def build_model(self, X, y, coords):
         with self:
             self.add_coords(coords)
 
-            t = pm.Data("t", X.sel(coeffs="t"), dims="obs_ind")
+            t = pm.Data("t", X.sel(coeffs=self.time_variable_name), dims="obs_ind")
             X = pm.Data("X", X, dims=["obs_ind", "coeffs"])
             y = pm.Data("y", y, dims="obs_ind")
 
@@ -751,7 +755,11 @@ def _data_setter(self, X) -> None:
         new_no_of_observations = X.shape[0]
         with self:
             pm.set_data(
-                {"X": X, "t": X.sel(coeffs="t"), "y": np.zeros(new_no_of_observations)},
+                {
+                    "X": X,
+                    "t": X.sel(coeffs=self.time_variable_name),
+                    "y": np.zeros(new_no_of_observations),
+                },
                 coords={"obs_ind": np.arange(new_no_of_observations)},
             )
 
@@ -771,12 +779,18 @@ def set_time_range(self, time_range, data):
         :param time_range: tuple or None
             If not None, a tuple of two values (start_label, end_label) that correspond
             to index labels in the 't' column of the `data` DataFrame
-        :param data: pandas.DataFrame which contains a column "t".
+        :param data: pandas.DataFrame.
         """
         if time_range is None:
-            self.time_range = data["t"].min(), data["t"].max()
+            self.time_range = (
+                data[self.time_variable_name].min(),
+                data[self.time_variable_name].max(),
+            )
         else:
             self.time_range = (
-                data["t"].loc[time_range[0]],
-                data["t"].loc[time_range[1]],
+                data[self.time_variable_name].loc[time_range[0]],
+                data[self.time_variable_name].loc[time_range[1]],
             )
+
+    def get_time_variable_name(self):
+        return self.time_variable_name