# Copyright 2022 - 2026 The PyMC Labs Developers
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Backend adapters for model fitting, prediction, and coefficients."""
from __future__ import annotations
import copy
import warnings
from abc import ABC, abstractmethod
from typing import Any, Literal
import numpy as np
import pandas as pd
import xarray as xr
from sklearn.base import RegressorMixin, clone
from sklearn.metrics import r2_score
from causalpy._arviz_compat import hdi_bounds
from causalpy.constants import HDI_PROB
from causalpy.pymc_forecast_models import PyMCForecastModel
from causalpy.pymc_models import PyMCModel
from causalpy.skl_models import create_causalpy_compatible_class
from causalpy.utils import round_num
BackendKind = Literal["pymc", "sklearn", "pymc-forecast"]
[docs]
def build_coords(
coeffs: list[str] | tuple[str, ...],
n_obs: int,
*,
treated_units: tuple[str, ...] | list[str] = ("unit_0",),
**extra: Any,
) -> dict[str, Any]:
"""Build the standard PyMC coordinate dict for regression experiments.
Parameters
----------
coeffs : list of str or tuple of str
Coefficient / predictor names for the ``coeffs`` coord.
n_obs : int
Number of observations; used to build ``obs_ind`` as ``np.arange(n_obs)``.
treated_units : list of str or tuple of str, default ``("unit_0",)``
Names for the treated-unit dimension of ``y``.
**extra
Additional coordinate entries merged into the result (e.g.
``datetime_index`` for ITS).
"""
return {
"coeffs": list(coeffs),
"obs_ind": np.arange(n_obs),
"treated_units": list(treated_units),
**extra,
}
def _extract_mu(prediction: xr.DataTree) -> xr.DataArray:
"""Extract response-scale ``mu`` from a DataTree prediction container."""
mu = prediction["posterior_predictive"]["mu"].transpose(
"chain", "draw", "obs_ind", "treated_units"
)
# Enforce the canonical container: stray non-dim coords (e.g. the
# state-space backend's `observed_state`) would otherwise leak into
# downstream impact containers and break coordinate equality checks.
return mu.drop_vars([name for name in mu.coords if name not in mu.dims])
def _sklearn_array(value: Any) -> np.ndarray:
"""Coerce xarray or array-like inputs to a numpy array for sklearn."""
if isinstance(value, xr.DataArray):
return np.asarray(value.data)
return np.asarray(value)
def _sklearn_y(y: Any) -> np.ndarray:
"""Coerce outcome arrays to sklearn's preferred 1D shape when possible.
Collapses a single trailing treated-units column to 1D. Genuine multi-output
``y`` (>1 column) is passed through unchanged; experiments whose sklearn
backend cannot fit multiple outcomes (e.g. synthetic control's
``WeightedProportion``) must reject that case upstream at construction.
"""
arr = _sklearn_array(y)
if arr.ndim == 2 and arr.shape[1] == 1:
return np.squeeze(arr, axis=1)
return arr
def _canonical_pymc_coefficients(posterior: xr.Dataset) -> xr.DataArray:
"""Normalize supported PyMC coefficient variables to the canonical contract."""
coefficient_names = ("beta", "b", "beta_z")
coefficient_name = next(
(name for name in coefficient_names if name in posterior), None
)
if coefficient_name is None:
raise ValueError(
"PyMC posterior must expose one of 'beta', 'b', or 'beta_z' "
"as design-matrix coefficients."
)
coefficients = posterior[coefficient_name]
label_dims = ("coeffs", "covariates", "instruments", "outcome_coeffs")
label_dim = next((dim for dim in label_dims if dim in coefficients.dims), None)
if label_dim is None:
raise ValueError(
"PyMC coefficient draws must include one of "
f"{label_dims!r}, got dims={coefficients.dims!r}."
)
if label_dim != "coeffs":
coefficients = coefficients.rename({label_dim: "coeffs"})
required_dims = {"chain", "draw", "coeffs"}
if not required_dims.issubset(coefficients.dims):
raise ValueError(
"PyMC coefficient draws must include dimensions "
f"{required_dims!r}, got dims={coefficients.dims!r}."
)
unexpected_dims = set(coefficients.dims) - required_dims - {"treated_units"}
if unexpected_dims:
raise ValueError(
"PyMC coefficient draws include unsupported dimensions "
f"{unexpected_dims!r}."
)
dims = ["chain", "draw", "coeffs"]
if "treated_units" in coefficients.dims:
dims.append("treated_units")
coefficients = coefficients.transpose(*dims).rename("coefficients")
return coefficients.drop_vars(
[name for name in coefficients.coords if name not in coefficients.dims]
)
def _print_coefficients(
coefficients: xr.DataArray,
labels: list[str],
round_to: int | None,
) -> None:
"""Print a coefficient container without dispatching on backend type."""
coefficients = coefficients.sel(coeffs=labels)
with_uncertainty = coefficients.sizes["chain"] * coefficients.sizes["draw"] > 1
treated_units: list[Any] = (
list(coefficients.coords["treated_units"].values)
if "treated_units" in coefficients.dims
else [None]
)
max_label_length = max(len(name) for name in labels)
print("Model coefficients:")
for unit in treated_units:
if len(treated_units) > 1:
print(f"\nTreated unit: {unit}")
unit_coefficients = (
coefficients.sel(treated_units=unit) if unit is not None else coefficients
)
for name in labels:
samples = unit_coefficients.sel(coeffs=name)
formatted_name = f"{name:<{max_label_length}}"
mean = round_num(float(samples.mean()), round_to)
if with_uncertainty:
lower, upper = hdi_bounds(samples, prob=HDI_PROB)
value = (
f"{mean}, {HDI_PROB * 100:.0f}% HDI "
f"[{round_num(float(lower), round_to)}, "
f"{round_num(float(upper), round_to)}]"
)
else:
value = mean
print(f" {formatted_name} {value}")
[docs]
class ModelAdapter(ABC):
"""Experiment-agnostic wrapper around a CausalPy statistical backend."""
@property
@abstractmethod
def model(self) -> PyMCModel | RegressorMixin | PyMCForecastModel:
"""The underlying model instance."""
@property
@abstractmethod
def kind(self) -> BackendKind:
"""Backend identifier."""
@property
def is_bayesian(self) -> bool:
"""Whether the backend is Bayesian (PyMC or pymc-forecast)."""
return self.kind in ("pymc", "pymc-forecast")
@property
def is_ols(self) -> bool:
"""Whether the backend is OLS/sklearn."""
return self.kind == "sklearn"
@property
def supports_idata(self) -> bool:
"""Whether the backend exposes an inference-result DataTree."""
return self.kind in ("pymc", "pymc-forecast")
@property
@abstractmethod
def idata(self) -> xr.DataTree | None:
"""Return a fitted inference-result DataTree when supported, otherwise ``None``."""
[docs]
def require_idata(self) -> xr.DataTree:
"""Return fitted inference-result DataTree or raise an explicit capability error."""
if not self.supports_idata:
raise TypeError(
f"{type(self).__name__} does not support InferenceData/DataTree results."
)
idata = self.idata
if idata is None:
raise RuntimeError("Model has not been fit yet.")
return idata
[docs]
@abstractmethod
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> Any:
"""Fit the model with backend-appropriate conventions.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Outcome vector or matrix.
coords : dict, optional
Coordinate metadata for PyMC models. Ignored by sklearn backends.
"""
[docs]
@abstractmethod
def predict(
self,
X: Any,
*,
coords: dict[str, Any] | None = None,
out_of_sample: bool = False,
) -> xr.DataArray:
"""Return expected outcomes with canonical prediction dimensions.
Every backend returns the same container: response-scale expected
outcomes as an :class:`xarray.DataArray` with dimensions
``("chain", "draw", "obs_ind", "treated_units")``. Point-estimate
backends (sklearn) return singleton ``chain``/``draw`` dimensions —
a point estimate is a posterior with one atom.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix for which to generate predictions.
coords : dict, optional
Coordinate metadata for Bayesian backends.
out_of_sample : bool, default False
Whether predictions are out-of-sample. Used by PyMC backends only.
Returns
-------
xr.DataArray
Expected outcomes with dimensions ``("chain", "draw", "obs_ind",
"treated_units")``.
"""
[docs]
@abstractmethod
def score(
self, X: Any, y: Any, *, coords: dict[str, Any] | None = None
) -> pd.Series:
"""Return per-unit :math:`R^2` scores in the canonical container.
Every backend returns a :class:`pandas.Series` with one
``unit_{i}_r2`` entry per treated unit. Backends with posterior
dispersion also include ``unit_{i}_r2_std`` entries.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Observed outcomes.
coords : dict, optional
Coordinate metadata for Bayesian backends.
Returns
-------
pd.Series
Per-treated-unit :math:`R^2` values and optional posterior
standard deviations.
"""
[docs]
@abstractmethod
def coefficients(self) -> xr.DataArray:
"""Return model coefficients with canonical coefficient dimensions.
Every supported backend returns an :class:`xarray.DataArray` with
dimensions ``("chain", "draw", "coeffs")`` and an optional trailing
``"treated_units"`` dimension. Point-estimate backends return singleton
``chain`` and ``draw`` dimensions.
"""
[docs]
def print_coefficients(
self, labels: list[str], round_to: int | None = None
) -> None:
"""Print model coefficients with labels.
Parameters
----------
labels : list of str
Coefficient names aligned with the fitted model.
round_to : int, optional
Number of significant figures to round to.
"""
_print_coefficients(self.coefficients(), labels, round_to)
[docs]
class PyMCModelAdapter(ModelAdapter):
"""Adapter for :class:`~causalpy.pymc_models.PyMCModel` backends.
Parameters
----------
model : PyMCModel
Fitted or unfitted PyMC backend model.
"""
[docs]
def __init__(self, model: PyMCModel) -> None:
self._model = model
@property
def model(self) -> PyMCModel:
"""The underlying PyMC model."""
return self._model
@property
def kind(self) -> BackendKind:
"""Backend identifier."""
return "pymc"
@property
def idata(self) -> xr.DataTree | None:
"""Return the model's DataTree when fitted."""
return self._model.idata
[docs]
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> xr.DataTree:
"""Fit the PyMC model.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Outcome vector or matrix.
coords : dict, optional
Coordinate metadata for the PyMC model.
"""
if isinstance(X, dict) and isinstance(y, dict):
return self._model.fit_mapping(X=X, y=y, coords=coords)
if isinstance(X, dict) or isinstance(y, dict):
raise TypeError("X and y must either both be mappings or both be arrays")
return self._model.fit(X=X, y=y, coords=coords)
[docs]
def predict(
self,
X: Any,
*,
coords: dict[str, Any] | None = None,
out_of_sample: bool = False,
) -> xr.DataArray:
"""Predict expected outcomes using the PyMC model.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix for which to generate predictions.
coords : dict, optional
Coordinate metadata for the PyMC model.
out_of_sample : bool, default False
Whether predictions are out-of-sample.
Returns
-------
xr.DataArray
Posterior draws of ``mu`` with canonical prediction dimensions.
"""
return _extract_mu(
self._model.predict(X=X, coords=coords, out_of_sample=out_of_sample)
)
[docs]
def score(
self, X: Any, y: Any, *, coords: dict[str, Any] | None = None
) -> pd.Series:
"""Score predictions from the PyMC model.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Observed outcomes.
coords : dict, optional
Coordinate metadata for the PyMC model.
"""
return self._model.score(X=X, y=y, coords=coords)
[docs]
def coefficients(self) -> xr.DataArray:
"""Return posterior coefficient draws in the canonical container."""
if self._model.idata is None:
raise RuntimeError("Model has not been fit yet.")
return _canonical_pymc_coefficients(self._model.idata.posterior)
[docs]
class SklearnModelAdapter(ModelAdapter):
"""Adapter for sklearn :class:`~sklearn.base.RegressorMixin` backends.
Parameters
----------
model : RegressorMixin
CausalPy-compatible sklearn backend model.
"""
[docs]
def __init__(self, model: RegressorMixin) -> None:
self._model = model
self._coeffs: np.ndarray | None = None
self._treated_units: np.ndarray | None = None
@property
def model(self) -> RegressorMixin:
"""The underlying sklearn model."""
return self._model
@property
def kind(self) -> BackendKind:
"""Backend identifier."""
return "sklearn"
@property
def idata(self) -> None:
"""Return ``None`` because sklearn models have no inference-result DataTree."""
return None
[docs]
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> Any:
"""Fit the sklearn model.
Parameters
----------
X : array-like
Predictor matrix.
y : array-like
Outcome vector or matrix.
coords : dict, optional
Ignored for sklearn backends.
"""
X_array = _sklearn_array(X)
if isinstance(X, xr.DataArray) and "coeffs" in X.coords:
self._coeffs = np.asarray(X.coords["coeffs"])
else:
self._coeffs = np.asarray([f"coeff_{i}" for i in range(X_array.shape[1])])
if isinstance(y, xr.DataArray) and "treated_units" in y.coords:
self._treated_units = np.asarray(y.coords["treated_units"])
else:
self._treated_units = None
return self._model.fit(X=X_array, y=_sklearn_y(y))
[docs]
def predict(
self,
X: Any,
*,
coords: dict[str, Any] | None = None,
out_of_sample: bool = False,
) -> xr.DataArray:
"""Return point predictions as singleton posterior draws.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix for which to generate predictions.
coords : dict, optional
Ignored for sklearn backends.
out_of_sample : bool, default False
Ignored for sklearn backends.
Returns
-------
xr.DataArray
Point predictions with canonical prediction dimensions and
singleton ``chain``/``draw`` dimensions.
"""
values = np.asarray(self._model.predict(X=_sklearn_array(X)))
if values.ndim == 1:
values = values[:, None]
if values.ndim != 2:
raise ValueError(
"Expected sklearn predictions with shape (obs,) or "
f"(obs, treated_units), got {values.shape}."
)
obs_ind = (
X.get_index("obs_ind")
if isinstance(X, xr.DataArray) and "obs_ind" in X.coords
else np.arange(values.shape[0])
)
treated_units = (
self._treated_units
if self._treated_units is not None
else np.asarray([f"unit_{i}" for i in range(values.shape[1])])
)
if len(treated_units) != values.shape[1]:
raise ValueError(
"Prediction output columns do not match the treated units used for fit."
)
return xr.DataArray(
values[None, None, :, :],
dims=("chain", "draw", "obs_ind", "treated_units"),
coords={
"chain": [0],
"draw": [0],
"obs_ind": obs_ind,
"treated_units": treated_units,
},
)
[docs]
def score(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
sample_weight: Any | None = None,
multioutput: Literal["raw_values"] = "raw_values",
force_finite: bool = True,
) -> pd.Series:
"""Return per-output :math:`R^2` scores from the sklearn model.
Parameters
----------
X : array-like
Predictor matrix.
y : array-like
Observed outcomes.
coords : dict, optional
Ignored for sklearn backends.
sample_weight : array-like, optional
Sample weights passed to :func:`sklearn.metrics.r2_score`.
multioutput : {"raw_values"}, default "raw_values"
The required aggregation mode. Per-unit scores require the raw
value for each output.
force_finite : bool, default True
Whether to replace non-finite scores for constant targets, passed to
:func:`sklearn.metrics.r2_score`.
Returns
-------
pd.Series
One ``unit_{i}_r2`` entry per output. Point estimates carry no
dispersion entries.
"""
if multioutput != "raw_values":
raise ValueError(
"SklearnModelAdapter.score() requires "
'multioutput="raw_values" for the canonical per-unit score contract.'
)
scores = np.atleast_1d(
r2_score(
_sklearn_y(y),
self._model.predict(X=_sklearn_array(X)),
sample_weight=sample_weight,
multioutput=multioutput,
force_finite=force_finite,
)
)
return pd.Series(
{f"unit_{i}_r2": float(score) for i, score in enumerate(scores)}
)
[docs]
def coefficients(self) -> xr.DataArray:
"""Return fitted sklearn coefficients as singleton posterior draws."""
values = np.asarray(self._model.coef_)
n_coeffs = values.shape[-1]
coeffs = (
self._coeffs
if self._coeffs is not None
else np.asarray([f"coeff_{i}" for i in range(n_coeffs)])
)
if len(coeffs) != n_coeffs:
raise ValueError(
"Coefficient output does not match the predictors used for fit."
)
if values.ndim == 1:
if self._treated_units is None:
return xr.DataArray(
values[None, None, :],
dims=("chain", "draw", "coeffs"),
coords={"chain": [0], "draw": [0], "coeffs": coeffs},
name="coefficients",
)
if len(self._treated_units) != 1:
raise ValueError(
"Coefficient output columns do not match the treated units "
"used for fit."
)
values = values[None, :]
elif values.ndim != 2:
raise ValueError(
"Expected sklearn coefficients with shape (coeffs,) or "
f"(treated_units, coeffs), got {values.shape}."
)
treated_units = (
self._treated_units
if self._treated_units is not None
else np.asarray([f"unit_{i}" for i in range(values.shape[0])])
)
if len(treated_units) != values.shape[0]:
raise ValueError(
"Coefficient output rows do not match the treated units used for fit."
)
return xr.DataArray(
values.T[None, None, :, :],
dims=("chain", "draw", "coeffs", "treated_units"),
coords={
"chain": [0],
"draw": [0],
"coeffs": coeffs,
"treated_units": treated_units,
},
name="coefficients",
)
[docs]
class PyMCForecastAdapter(ModelAdapter):
"""Adapter for :class:`~causalpy.pymc_forecast_models.PyMCForecastModel`
backends.
The wrapped model already speaks CausalPy's Bayesian conventions
(``mu``/``y_hat`` posterior-predictive output on ``obs_ind`` /
``treated_units`` coords), so this adapter is pure delegation.
Parameters
----------
model : PyMCForecastModel
Wrapped ``pymc_forecast`` backend model.
"""
[docs]
def __init__(self, model: PyMCForecastModel) -> None:
self._model = model
@property
def model(self) -> PyMCForecastModel:
"""The underlying pymc-forecast wrapper."""
return self._model
@property
def kind(self) -> BackendKind:
"""Backend identifier."""
return "pymc-forecast"
@property
def idata(self) -> xr.DataTree | None:
"""Return the model's DataTree when fitted."""
return self._model.idata
[docs]
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> xr.DataTree:
"""Fit the forecasting model on the pre-period.
Parameters
----------
X : xarray.DataArray
Design matrix with dims ``["obs_ind", "coeffs"]``.
y : xarray.DataArray
Outcome with dims ``["obs_ind", "treated_units"]``.
coords : dict, optional
Coordinate metadata; ignored (real coordinates are read from
``X`` and ``y``).
"""
return self._model.fit(X=X, y=y, coords=coords)
[docs]
def predict(
self,
X: Any,
*,
coords: dict[str, Any] | None = None,
out_of_sample: bool = False,
) -> xr.DataArray:
"""Predict in-sample or forecast the counterfactual.
Parameters
----------
X : xarray.DataArray
Design matrix for which to generate predictions.
coords : dict, optional
Coordinate metadata accepted by the forecasting backend but not
used by its forecasting implementation.
out_of_sample : bool, default False
``True`` draws the post-period counterfactual via the model's
forecasting path.
Returns
-------
xr.DataArray
Posterior draws of ``mu`` with canonical prediction dimensions.
"""
return _extract_mu(
self._model.predict(X=X, coords=coords, out_of_sample=out_of_sample)
)
[docs]
def score(
self, X: Any, y: Any, *, coords: dict[str, Any] | None = None
) -> pd.Series:
"""Score in-sample predictions with the Bayesian :math:`R^2`.
Parameters
----------
X : xarray.DataArray
Design matrix.
y : xarray.DataArray
Observed outcomes.
coords : dict, optional
Coordinate metadata accepted by the forecasting backend but not
used by its scoring implementation.
"""
return self._model.score(X=X, y=y, coords=coords)
[docs]
def coefficients(self) -> xr.DataArray:
"""Forecasting models have no design-matrix coefficients."""
raise NotImplementedError(
"pymc-forecast models do not expose design-matrix coefficients; "
"inspect the fitted posterior via `.idata` instead."
)
[docs]
def print_coefficients(
self, labels: list[str], round_to: int | None = None
) -> None:
"""Print posterior summaries of the model's scalar parameters.
Parameters
----------
labels : list of str
Design-matrix labels; ignored by forecasting models.
round_to : int, optional
Number of significant figures to round to.
"""
self._model.print_coefficients(labels, round_to)
def _prepare_sklearn_model(model: RegressorMixin) -> RegressorMixin:
"""Clone, augment, and validate a sklearn estimator for CausalPy."""
try:
model = clone(model)
except TypeError:
model = copy.deepcopy(model)
model = create_causalpy_compatible_class(model)
if getattr(model, "fit_intercept", False):
warnings.warn(
f"{type(model).__name__} had fit_intercept=True, but CausalPy "
"requires fit_intercept=False because the intercept is already "
"included in the design matrix by patsy. A cloned copy of the "
"model with fit_intercept=False will be used; the original "
"instance is unchanged.",
UserWarning,
stacklevel=3,
)
model.fit_intercept = False
return model
[docs]
def make_model_adapter(
model: PyMCModel | RegressorMixin | PyMCForecastModel | None,
*,
default_model_class: type[PyMCModel] | None,
supports_bayes: bool,
supports_ols: bool,
supports_pymc_forecast: bool = False,
) -> ModelAdapter:
"""Resolve, validate, and wrap a model in a backend adapter.
Parameters
----------
model : PyMCModel, RegressorMixin, PyMCForecastModel, or None
User-supplied model instance, or ``None`` to use the default.
default_model_class : type[PyMCModel] or None
PyMC model class used when ``model`` is ``None``.
supports_bayes : bool
Whether the experiment supports Bayesian backends.
supports_ols : bool
Whether the experiment supports OLS/sklearn backends.
supports_pymc_forecast : bool, default False
Whether the experiment supports pymc-forecast backends.
Returns
-------
ModelAdapter
Backend-specific adapter wrapping the resolved model.
"""
if isinstance(model, RegressorMixin):
model = _prepare_sklearn_model(model)
if model is None and default_model_class is not None:
model = default_model_class()
if model is None:
raise ValueError("model not set or passed.")
if isinstance(model, PyMCModel):
if not supports_bayes:
raise ValueError("Bayesian models not supported.")
return PyMCModelAdapter(model)
if isinstance(model, RegressorMixin):
if not supports_ols:
raise ValueError("OLS models not supported.")
return SklearnModelAdapter(model)
if isinstance(model, PyMCForecastModel):
if not supports_pymc_forecast:
raise ValueError("pymc-forecast models not supported.")
return PyMCForecastAdapter(model)
raise ValueError("Unsupported model type")
