Skip to content

Refactor PyMC time series models to use xarray API #577

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 7 commits into
base: main
Choose a base branch
from
+1,335 −1,373
Open

Refactor PyMC time series models to use xarray API #577

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
0961b30
Refactor PyMC time series models to use xarray API
cetagostini Nov 26, 2025
5ee6651
Merge branch 'main' into cetagostini/bsts_tech_debt_v0
drbenvincent Nov 29, 2025
d06b0c6
Refactor time series model API to use xarray coordinates
drbenvincent Nov 29, 2025
387f748
get tests passing
drbenvincent Nov 29, 2025
b0b39ac
minor tweaks to dev notebook
drbenvincent Nov 29, 2025
d5700e8
increase code coverage
drbenvincent Nov 29, 2025
a15e17f
delete refactoring concerns doc
drbenvincent Nov 29, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 2 additions & 0 deletions .gitignore
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -14,3 +14,5 @@ dist/
docs/build/
docs/jupyter_execute/
docs/source/api/generated/

.cursor/
103 changes: 15 additions & 88 deletions causalpy/experiments/interrupted_time_series.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -27,11 +27,7 @@

from causalpy.custom_exceptions import BadIndexException
from causalpy.plot_utils import get_hdi_to_df, plot_xY
from causalpy.pymc_models import (
BayesianBasisExpansionTimeSeries,
PyMCModel,
StateSpaceTimeSeries,
)
from causalpy.pymc_models import PyMCModel
from causalpy.utils import round_num

from .base import BaseExperiment
Expand Down Expand Up @@ -202,27 +198,15 @@ def __init__(
)

# fit the model to the observed (pre-intervention) data
# All PyMC models now accept xr.DataArray with consistent API
if isinstance(self.model, PyMCModel):
is_bsts_like = isinstance(
self.model, (BayesianBasisExpansionTimeSeries, StateSpaceTimeSeries)
)

if is_bsts_like:
# BSTS/StateSpace models expect numpy arrays and datetime coords
X_fit = self.pre_X.values if self.pre_X.shape[1] > 0 else None # type: ignore[attr-defined]
y_fit = self.pre_y.isel(treated_units=0).values # type: ignore[attr-defined]
pre_coords: dict[str, Any] = {"datetime_index": self.datapre.index}
if X_fit is not None:
pre_coords["coeffs"] = list(self.labels)
self.model.fit(X=X_fit, y=y_fit, coords=pre_coords)
else:
# General PyMC models expect xarray with treated_units
COORDS = {
"coeffs": self.labels,
"obs_ind": np.arange(self.pre_X.shape[0]),
"treated_units": ["unit_0"],
}
self.model.fit(X=self.pre_X, y=self.pre_y, coords=COORDS)
COORDS: dict[str, Any] = {
"coeffs": self.labels,
"obs_ind": np.arange(self.pre_X.shape[0]),
"treated_units": ["unit_0"],
"datetime_index": self.datapre.index, # For time series models
}
self.model.fit(X=self.pre_X, y=self.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))
Expand All @@ -231,85 +215,28 @@ def __init__(

# score the goodness of fit to the pre-intervention data
if isinstance(self.model, PyMCModel):
is_bsts_like = isinstance(
self.model, (BayesianBasisExpansionTimeSeries, StateSpaceTimeSeries)
)
if is_bsts_like:
X_score = self.pre_X.values if self.pre_X.shape[1] > 0 else None # type: ignore[attr-defined]
y_score = self.pre_y.isel(treated_units=0).values # type: ignore[attr-defined]
score_coords: dict[str, Any] = {"datetime_index": self.datapre.index}
if X_score is not None:
score_coords["coeffs"] = list(self.labels)
self.score = self.model.score(X=X_score, y=y_score, coords=score_coords)
else:
self.score = self.model.score(X=self.pre_X, y=self.pre_y)
self.score = self.model.score(X=self.pre_X, y=self.pre_y)
elif isinstance(self.model, RegressorMixin):
self.score = self.model.score(
X=self.pre_X, y=self.pre_y.isel(treated_units=0)
)

# get the model predictions of the observed (pre-intervention) data
if isinstance(self.model, PyMCModel):
is_bsts_like = isinstance(
self.model, (BayesianBasisExpansionTimeSeries, StateSpaceTimeSeries)
)
if is_bsts_like:
X_pre_predict = self.pre_X.values if self.pre_X.shape[1] > 0 else None # type: ignore[attr-defined]
pre_pred_coords: dict[str, Any] = {"datetime_index": self.datapre.index}
self.pre_pred = self.model.predict(
X=X_pre_predict, coords=pre_pred_coords
)
if not isinstance(self.pre_pred, az.InferenceData):
self.pre_pred = az.InferenceData(posterior_predictive=self.pre_pred)
else:
self.pre_pred = self.model.predict(X=self.pre_X)
self.pre_pred = self.model.predict(X=self.pre_X)
elif isinstance(self.model, RegressorMixin):
self.pre_pred = self.model.predict(X=self.pre_X)

# calculate the counterfactual (post period)
if isinstance(self.model, PyMCModel):
is_bsts_like = isinstance(
self.model, (BayesianBasisExpansionTimeSeries, StateSpaceTimeSeries)
)
if is_bsts_like:
X_post_predict = (
self.post_X.values if self.post_X.shape[1] > 0 else None # type: ignore[attr-defined]
)
post_pred_coords: dict[str, Any] = {
"datetime_index": self.datapost.index
}
self.post_pred = self.model.predict(
X=X_post_predict, coords=post_pred_coords, out_of_sample=True
)
if not isinstance(self.post_pred, az.InferenceData):
self.post_pred = az.InferenceData(
posterior_predictive=self.post_pred
)
else:
self.post_pred = self.model.predict(X=self.post_X)
self.post_pred = self.model.predict(X=self.post_X, out_of_sample=True)
elif isinstance(self.model, RegressorMixin):
self.post_pred = self.model.predict(X=self.post_X)

# calculate impact - use appropriate y data format for each model type
# calculate impact - all PyMC models now use 2D data with treated_units
if isinstance(self.model, PyMCModel):
is_bsts_like = isinstance(
self.model, (BayesianBasisExpansionTimeSeries, StateSpaceTimeSeries)
)
if is_bsts_like:
pre_y_for_impact = self.pre_y.isel(treated_units=0)
post_y_for_impact = self.post_y.isel(treated_units=0)
self.pre_impact = self.model.calculate_impact(
pre_y_for_impact, self.pre_pred
)
self.post_impact = self.model.calculate_impact(
post_y_for_impact, self.post_pred
)
else:
# PyMC models with treated_units use 2D data
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(self.pre_y, self.pre_pred)
self.post_impact = self.model.calculate_impact(self.post_y, self.post_pred)
elif isinstance(self.model, RegressorMixin):
# SKL models work with 1D data
self.pre_impact = self.model.calculate_impact(
Expand Down
Loading