Removing the overriding of fit and calculate_impact, adding a test and fixing a bug

Author: JeanVanDyk

causalpy/pymc_models.py

@@ -699,25 +699,6 @@ def build_model(self, X, y, coords):
             # Likelihodd of the base time series and the intervention's effect
             pm.Normal("y_ts", mu=mu_ts, sigma=sigma, observed=y, dims="obs_ind")
 
-    def fit(self, X, y, coords: Optional[Dict[str, Any]] = None) -> None:
-        """Draw samples from posterior, prior predictive, and posterior predictive
-        distributions, placing them in the model's idata attribute.
-        """
-
-        # 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)
-        self.build_model(X, y, coords)
-        with self:
-            self.idata = pm.sample(**self.sample_kwargs)
-            self.idata.extend(pm.sample_prior_predictive(random_seed=random_seed))
-            self.idata.extend(
-                pm.sample_posterior_predictive(
-                    self.idata, progressbar=False, random_seed=random_seed
-                )
-            )
-        return self.idata
-
     def predict(self, X):
         """
         Predict data given input data `X`
@@ -731,19 +712,20 @@ def predict(self, X):
         random_seed = self.sample_kwargs.get("random_seed", None)
         self._data_setter(X)
         with self:  # sample with new input data
-            post_pred = pm.sample_posterior_predictive(
+            pp = pm.sample_posterior_predictive(
                 self.idata,
                 var_names=["y_hat", "y_ts", "mu", "mu_ts", "mu_in"],
                 progressbar=False,
                 random_seed=random_seed,
             )
-        return post_pred
 
-    def calculate_impact(
-        self, y_true: xr.DataArray, y_pred: az.InferenceData
-    ) -> xr.DataArray:
-        impact = y_true.data - y_pred["posterior_predictive"]["y_hat"]
-        return impact.transpose(..., "obs_ind")
+        # TODO: This is a bit of a hack. Maybe it could be done properly in _data_setter?
+        if isinstance(X, xr.DataArray):
+            pp["posterior_predictive"] = pp["posterior_predictive"].assign_coords(
+                obs_ind=X.obs_ind
+            )
+
+        return pp
 
     def _data_setter(self, X) -> None:
         """
@@ -770,7 +752,7 @@ def score(self, X, y) -> pd.Series:
         mu_ts = self.predict(X)
         mu_ts = az.extract(mu_ts, group="posterior_predictive", var_names="mu_ts").T
         # Note: First argument must be a 1D array
-        return r2_score(y.data, mu_ts)
+        return r2_score(y.data, mu_ts.data)
 
     def set_time_range(self, time_range, data):
         """

causalpy/tests/test_integration_pymc_examples.py

@@ -402,6 +402,63 @@ def test_its():
     )
 
 
+@pytest.mark.integration
+def test_its_no_treatment_time():
+    """
+    Test Interrupted Time-Series experiment on COVID data with an unknown treatment time.
+
+    Loads data and checks:
+    1. data is a dataframe
+    2. causalpy.InterruptedtimeSeries returns correct type
+    3. the correct number of MCMC chains exists in the posterior inference data
+    4. the correct number of MCMC draws exists in the posterior inference data
+    5. the method get_plot_data returns a DataFrame with expected columns
+    """
+
+    df = (
+        cp.load_data("covid")
+        .assign(date=lambda x: pd.to_datetime(x["date"]))
+        .set_index("date")
+    )
+    result = cp.InterruptedTimeSeries(
+        df,
+        None,
+        formula="standardize(deaths) ~ 0 + t + C(month) + standardize(temp)",  # noqa E501
+        model=cp.pymc_models.InterventionTimeEstimator(
+            time_variable_name="t",
+            treatment_type_effect={"impulse": []},
+            sample_kwargs=sample_kwargs,
+        ),
+    )
+    assert isinstance(df, pd.DataFrame)
+    assert isinstance(result, cp.InterruptedTimeSeries)
+    assert len(result.idata.posterior.coords["chain"]) == sample_kwargs["chains"]
+    assert len(result.idata.posterior.coords["draw"]) == sample_kwargs["draws"]
+    result.summary()
+    fig, ax = result.plot()
+    assert isinstance(fig, plt.Figure)
+    # For multi-panel plots, ax should be an array of axes
+    assert isinstance(ax, np.ndarray) and all(
+        isinstance(item, plt.Axes) for item in ax
+    ), "ax must be a numpy.ndarray of plt.Axes"
+    # Test get_plot_data with default parameters
+    plot_data = result.get_plot_data()
+    assert isinstance(plot_data, pd.DataFrame), (
+        "The returned object is not a pandas DataFrame"
+    )
+    expected_columns = [
+        "prediction",
+        "pred_hdi_lower_94",
+        "pred_hdi_upper_94",
+        "impact",
+        "impact_hdi_lower_94",
+        "impact_hdi_upper_94",
+    ]
+    assert set(expected_columns).issubset(set(plot_data.columns)), (
+        f"DataFrame is missing expected columns {expected_columns}"
+    )
+
+
 @pytest.mark.integration
 def test_its_covid():
     """