Consolidate experiment design-matrix attributes into xr.Dataset

Bundle loose xr.DataArray attributes on experiment classes into
xr.Dataset objects to reduce attribute sprawl. Pre/post classes
(ITS, SC) use pre_design/post_design; formula-based classes use
a single design Dataset. Deprecated @Property accessors preserve
backward compatibility. Closes #199.

Made-with: Cursor

Author: drbenvincent

causalpy/experiments/diff_in_diff.py

@@ -15,6 +15,7 @@
 Difference in differences
 """
 
+import warnings
 from typing import Any, Literal
 
 import arviz as az
@@ -129,42 +130,60 @@ def _build_design_matrices(self) -> None:
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
         self.labels = X.design_info.column_names
-        self.y, self.X = np.asarray(y), np.asarray(X)
+        self._y_raw, self._X_raw = np.asarray(y), np.asarray(X)
         self.outcome_variable_name = y.design_info.column_names[0]
 
     def _prepare_data(self) -> None:
-        """Convert design matrices to xarray DataArrays."""
-        self.X = xr.DataArray(
-            self.X,
-            dims=["obs_ind", "coeffs"],
-            coords={
-                "obs_ind": np.arange(self.X.shape[0]),
-                "coeffs": self.labels,
-            },
+        """Bundle design matrices into an ``xr.Dataset``."""
+        n = self._X_raw.shape[0]
+        self.design = xr.Dataset(
+            {
+                "X": xr.DataArray(
+                    self._X_raw,
+                    dims=["obs_ind", "coeffs"],
+                    coords={"obs_ind": np.arange(n), "coeffs": self.labels},
+                ),
+                "y": xr.DataArray(
+                    self._y_raw,
+                    dims=["obs_ind", "treated_units"],
+                    coords={"obs_ind": np.arange(n), "treated_units": ["unit_0"]},
+                ),
+            }
         )
-        self.y = xr.DataArray(
-            self.y,
-            dims=["obs_ind", "treated_units"],
-            coords={"obs_ind": np.arange(self.y.shape[0]), "treated_units": ["unit_0"]},
+        del self._X_raw, self._y_raw
+
+    @property
+    def X(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.design['X']`` instead."""
+        warnings.warn(
+            "X is deprecated, use design['X']", DeprecationWarning, stacklevel=2
+        )
+        return self.design["X"]
+
+    @property
+    def y(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.design['y']`` instead."""
+        warnings.warn(
+            "y is deprecated, use design['y']", DeprecationWarning, stacklevel=2
         )
+        return self.design["y"]
 
     def algorithm(self) -> None:
         """Run the experiment algorithm: fit model, predict, and calculate causal impact."""
-        # fit model
+        X = self.design["X"]
+        y = self.design["y"]
+
         if isinstance(self.model, PyMCModel):
             COORDS = {
                 "coeffs": self.labels,
-                "obs_ind": np.arange(self.X.shape[0]),
+                "obs_ind": np.arange(X.shape[0]),
                 "treated_units": ["unit_0"],
             }
-            self.model.fit(X=self.X, y=self.y, coords=COORDS)
+            self.model.fit(X=X, y=y, coords=COORDS)
         elif isinstance(self.model, RegressorMixin):
-            # Ensure the intercept is part of the coefficients array rather than
-            # a separate intercept_ attribute.  See #664 / PR #693 for
-            # centralising this in BaseExperiment.
             if hasattr(self.model, "fit_intercept"):
                 self.model.fit_intercept = False
-            self.model.fit(X=self.X, y=self.y)
+            self.model.fit(X=X, y=y)
         else:
             raise ValueError("Model type not recognized")
 

causalpy/experiments/interrupted_time_series.py

@@ -15,6 +15,7 @@
 Interrupted Time Series Analysis
 """
 
+import warnings
 from typing import Any, Literal
 
 import arviz as az
@@ -139,9 +140,8 @@ def __init__(
         **kwargs: Any,
     ) -> None:
         super().__init__(model=model)
-        self.pre_y: xr.DataArray
-        self.post_y: xr.DataArray
-        # rename the index to "obs_ind"
+        self.pre_design: xr.Dataset
+        self.post_design: xr.Dataset
         data.index.name = "obs_ind"
         self.data = data
         self.input_validation(data, treatment_time, treatment_end_time)
@@ -155,96 +155,98 @@ def __init__(
 
     def _build_design_matrices(self) -> None:
         """Build design matrices for pre and post intervention periods using patsy."""
-        # set things up with pre-intervention data
         y, X = dmatrices(self.formula, self.datapre)
         self.outcome_variable_name = y.design_info.column_names[0]
         self._y_design_info = y.design_info
         self._x_design_info = X.design_info
         self.labels = X.design_info.column_names
-        self.pre_y, self.pre_X = np.asarray(y), np.asarray(X)
-        # process post-intervention data
+        self._pre_y_raw, self._pre_X_raw = np.asarray(y), np.asarray(X)
         (new_y, new_x) = build_design_matrices(
             [self._y_design_info, self._x_design_info], self.datapost
         )
-        self.post_X = np.asarray(new_x)
-        self.post_y = np.asarray(new_y)
+        self._post_X_raw = np.asarray(new_x)
+        self._post_y_raw = np.asarray(new_y)
 
     def _prepare_data(self) -> None:
-        """Convert design matrices to xarray DataArrays for pre and post periods."""
-        self.pre_X = xr.DataArray(
-            self.pre_X,
-            dims=["obs_ind", "coeffs"],
-            coords={
-                "obs_ind": self.datapre.index,
-                "coeffs": self.labels,
-            },
-        )
-        self.pre_y = xr.DataArray(
-            self.pre_y,  # Keep 2D shape
-            dims=["obs_ind", "treated_units"],
-            coords={"obs_ind": self.datapre.index, "treated_units": ["unit_0"]},
-        )
-        self.post_X = xr.DataArray(
-            self.post_X,
-            dims=["obs_ind", "coeffs"],
-            coords={
-                "obs_ind": self.datapost.index,
-                "coeffs": self.labels,
-            },
+        """Bundle design matrices into ``xr.Dataset`` objects for pre and post periods."""
+        self.pre_design = xr.Dataset(
+            {
+                "X": xr.DataArray(
+                    self._pre_X_raw,
+                    dims=["obs_ind", "coeffs"],
+                    coords={"obs_ind": self.datapre.index, "coeffs": self.labels},
+                ),
+                "y": xr.DataArray(
+                    self._pre_y_raw,
+                    dims=["obs_ind", "treated_units"],
+                    coords={
+                        "obs_ind": self.datapre.index,
+                        "treated_units": ["unit_0"],
+                    },
+                ),
+            }
         )
-        self.post_y = xr.DataArray(
-            self.post_y,  # Keep 2D shape
-            dims=["obs_ind", "treated_units"],
-            coords={"obs_ind": self.datapost.index, "treated_units": ["unit_0"]},
+        self.post_design = xr.Dataset(
+            {
+                "X": xr.DataArray(
+                    self._post_X_raw,
+                    dims=["obs_ind", "coeffs"],
+                    coords={"obs_ind": self.datapost.index, "coeffs": self.labels},
+                ),
+                "y": xr.DataArray(
+                    self._post_y_raw,
+                    dims=["obs_ind", "treated_units"],
+                    coords={
+                        "obs_ind": self.datapost.index,
+                        "treated_units": ["unit_0"],
+                    },
+                ),
+            }
         )
+        del self._pre_X_raw, self._pre_y_raw, self._post_X_raw, self._post_y_raw
 
     def algorithm(self) -> None:
         """Run the experiment algorithm: fit model, predict, and calculate causal impact."""
-        # fit the model to the observed (pre-intervention) data
-        # All PyMC models now accept xr.DataArray with consistent API
+        pre_X = self.pre_design["X"]
+        pre_y = self.pre_design["y"]
+        post_X = self.post_design["X"]
+        post_y = self.post_design["y"]
+
         if isinstance(self.model, PyMCModel):
             COORDS: dict[str, Any] = {
                 "coeffs": self.labels,
-                "obs_ind": np.arange(self.pre_X.shape[0]),
+                "obs_ind": np.arange(pre_X.shape[0]),
                 "treated_units": ["unit_0"],
-                "datetime_index": self.datapre.index,  # For time series models
+                "datetime_index": self.datapre.index,
             }
-            self.model.fit(X=self.pre_X, y=self.pre_y, coords=COORDS)
+            self.model.fit(X=pre_X, y=pre_y, coords=COORDS)
         elif isinstance(self.model, RegressorMixin):
-            # For OLS models, use 1D y data
-            self.model.fit(X=self.pre_X, y=self.pre_y.isel(treated_units=0))
+            self.model.fit(X=pre_X, y=pre_y.isel(treated_units=0))
         else:
             raise ValueError("Model type not recognized")
 
-        # score the goodness of fit to the pre-intervention data
         if isinstance(self.model, PyMCModel):
-            self.score = self.model.score(X=self.pre_X, y=self.pre_y)
+            self.score = self.model.score(X=pre_X, y=pre_y)
         elif isinstance(self.model, RegressorMixin):
-            self.score = self.model.score(
-                X=self.pre_X, y=self.pre_y.isel(treated_units=0)
-            )
+            self.score = self.model.score(X=pre_X, y=pre_y.isel(treated_units=0))
 
-        # get the model predictions of the observed (pre-intervention) data
         if isinstance(self.model, PyMCModel | RegressorMixin):
-            self.pre_pred = self.model.predict(X=self.pre_X)
+            self.pre_pred = self.model.predict(X=pre_X)
 
-        # calculate the counterfactual (post period)
         if isinstance(self.model, PyMCModel):
-            self.post_pred = self.model.predict(X=self.post_X, out_of_sample=True)
+            self.post_pred = self.model.predict(X=post_X, out_of_sample=True)
         elif isinstance(self.model, RegressorMixin):
-            self.post_pred = self.model.predict(X=self.post_X)
+            self.post_pred = self.model.predict(X=post_X)
 
-        # calculate impact - all PyMC models now use 2D data with treated_units
         if isinstance(self.model, PyMCModel):
-            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.pre_impact = self.model.calculate_impact(pre_y, self.pre_pred)
+            self.post_impact = self.model.calculate_impact(post_y, self.post_pred)
         elif isinstance(self.model, RegressorMixin):
-            # SKL models work with 1D data
             self.pre_impact = self.model.calculate_impact(
-                self.pre_y.isel(treated_units=0), self.pre_pred
+                pre_y.isel(treated_units=0), self.pre_pred
             )
             self.post_impact = self.model.calculate_impact(
-                self.post_y.isel(treated_units=0), self.post_pred
+                post_y.isel(treated_units=0), self.post_pred
             )
 
         self.post_impact_cumulative = self.model.calculate_cumulative_impact(
@@ -318,6 +320,46 @@ def datapost(self) -> pd.DataFrame:
         """
         return self.data[self.data.index >= self.treatment_time]
 
+    @property
+    def pre_X(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.pre_design['X']`` instead."""
+        warnings.warn(
+            "pre_X is deprecated, use pre_design['X']",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.pre_design["X"]
+
+    @property
+    def pre_y(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.pre_design['y']`` instead."""
+        warnings.warn(
+            "pre_y is deprecated, use pre_design['y']",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.pre_design["y"]
+
+    @property
+    def post_X(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.post_design['X']`` instead."""
+        warnings.warn(
+            "post_X is deprecated, use post_design['X']",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.post_design["X"]
+
+    @property
+    def post_y(self) -> xr.DataArray:
+        """.. deprecated:: Use ``self.post_design['y']`` instead."""
+        warnings.warn(
+            "post_y is deprecated, use post_design['y']",
+            DeprecationWarning,
+            stacklevel=2,
+        )
+        return self.post_design["y"]
+
     def _split_post_period(self) -> None:
         """Split post period into intervention and post-intervention periods.
 
@@ -627,9 +669,7 @@ def _bayesian_plot(
 
         (h,) = ax[0].plot(
             self.datapre.index,
-            self.pre_y.isel(treated_units=0)
-            if hasattr(self.pre_y, "isel")
-            else self.pre_y[:, 0],
+            self.pre_design["y"].isel(treated_units=0),
             "k.",
             label="Observations",
         )
@@ -654,9 +694,7 @@ def _bayesian_plot(
 
         ax[0].plot(
             self.datapost.index,
-            self.post_y.isel(treated_units=0)
-            if hasattr(self.post_y, "isel")
-            else self.post_y[:, 0],
+            self.post_design["y"].isel(treated_units=0),
             "k.",
         )
         # Shaded causal effect
@@ -669,9 +707,7 @@ def _bayesian_plot(
         h = ax[0].fill_between(
             self.datapost.index,
             y1=post_pred_mu,
-            y2=self.post_y.isel(treated_units=0)
-            if hasattr(self.post_y, "isel")
-            else self.post_y[:, 0],
+            y2=self.post_design["y"].isel(treated_units=0),
             color="C0",
             alpha=0.25,
         )
@@ -807,10 +843,10 @@ def _ols_plot(
 
         fig, ax = plt.subplots(3, 1, sharex=True, figsize=(7, 8))
 
-        ax[0].plot(self.datapre.index, self.pre_y, "k.")
+        ax[0].plot(self.datapre.index, self.pre_design["y"], "k.")
         ax[0].plot(self.datapre.index, self.pre_pred, c="k", label="model fit")
 
-        ax[0].plot(self.datapost.index, self.post_y, "k.")
+        ax[0].plot(self.datapost.index, self.post_design["y"], "k.")
         ax[0].plot(
             self.datapost.index,
             self.post_pred,
@@ -822,7 +858,7 @@ def _ols_plot(
         ax[0].fill_between(
             self.datapost.index,
             y1=np.squeeze(self.post_pred),
-            y2=np.squeeze(self.post_y),
+            y2=np.squeeze(self.post_design["y"]),
             color="C0",
             alpha=0.25,
             label="Causal impact",