# 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 experiment model fitting and prediction."""
from __future__ import annotations
import copy
import warnings
from abc import ABC, abstractmethod
from typing import Any, Literal
import arviz as az
import numpy as np
import xarray as xr
from sklearn.base import RegressorMixin, clone
from causalpy.pymc_models import PyMCModel
from causalpy.skl_models import create_causalpy_compatible_class
BackendKind = Literal["pymc", "sklearn"]
[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 _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
[docs]
class ModelAdapter(ABC):
"""Experiment-agnostic wrapper around a CausalPy statistical backend."""
@property
@abstractmethod
def model(self) -> PyMCModel | RegressorMixin:
"""The underlying model instance."""
@property
@abstractmethod
def kind(self) -> BackendKind:
"""Backend identifier."""
@property
def is_bayesian(self) -> bool:
"""Whether the backend is Bayesian (PyMC)."""
return self.kind == "pymc"
@property
def is_ols(self) -> bool:
"""Whether the backend is OLS/sklearn."""
return self.kind == "sklearn"
@property
@abstractmethod
def idata(self) -> az.InferenceData:
"""Return InferenceData for Bayesian models."""
[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,
*,
out_of_sample: bool = False,
**kwargs: Any,
) -> Any:
"""Predict with backend-appropriate conventions.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix for which to generate predictions.
out_of_sample : bool, default False
Whether predictions are out-of-sample. Used by PyMC backends only.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
[docs]
@abstractmethod
def score(self, X: Any, y: Any, **kwargs: Any) -> Any:
"""Score predictions against observed outcomes.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Observed outcomes.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
[docs]
@abstractmethod
def coefficients(self) -> np.ndarray:
"""Return point estimates of model coefficients."""
[docs]
@abstractmethod
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.
"""
[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) -> az.InferenceData:
"""Return the model's InferenceData object."""
return self._model.idata
[docs]
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> az.InferenceData:
"""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.
"""
return self._model.fit(X=X, y=y, coords=coords)
[docs]
def predict(
self,
X: Any,
*,
out_of_sample: bool = False,
**kwargs: Any,
) -> Any:
"""Predict using the PyMC model.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix for which to generate predictions.
out_of_sample : bool, default False
Whether predictions are out-of-sample.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
return self._model.predict(X=X, out_of_sample=out_of_sample, **kwargs)
[docs]
def score(self, X: Any, y: Any, **kwargs: Any) -> Any:
"""Score predictions from the PyMC model.
Parameters
----------
X : array-like or xarray.DataArray
Predictor matrix.
y : array-like or xarray.DataArray
Observed outcomes.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
return self._model.score(X=X, y=y, **kwargs)
[docs]
def coefficients(self) -> np.ndarray:
"""Return posterior mean coefficients."""
if self._model.idata is None:
raise RuntimeError("Model has not been fit yet.")
beta = self._model.idata.posterior["beta"]
return beta.mean(dim=["chain", "draw"]).values
[docs]
def print_coefficients(
self, labels: list[str], round_to: int | None = None
) -> None:
"""Print PyMC model coefficients.
Parameters
----------
labels : list of str
Coefficient names aligned with the fitted model.
round_to : int, optional
Number of significant figures to round to.
"""
self._model.print_coefficients(labels, round_to)
[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
@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) -> az.InferenceData:
"""OLS models do not expose InferenceData."""
raise AttributeError("OLS models do not have idata.")
[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.
"""
return self._model.fit(X=_sklearn_array(X), y=_sklearn_y(y))
[docs]
def predict(
self,
X: Any,
*,
out_of_sample: bool = False,
**kwargs: Any,
) -> Any:
"""Predict using the sklearn model.
Parameters
----------
X : array-like
Predictor matrix for which to generate predictions.
out_of_sample : bool, default False
Ignored for sklearn backends.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
return self._model.predict(X=_sklearn_array(X), **kwargs)
[docs]
def score(self, X: Any, y: Any, **kwargs: Any) -> Any:
"""Score predictions from the sklearn model.
Parameters
----------
X : array-like
Predictor matrix.
y : array-like
Observed outcomes.
**kwargs
Additional keyword arguments forwarded to the underlying model.
"""
return self._model.score(X=_sklearn_array(X), y=_sklearn_y(y), **kwargs)
[docs]
def coefficients(self) -> np.ndarray:
"""Return fitted sklearn coefficients."""
return self._model.get_coeffs()
[docs]
def print_coefficients(
self, labels: list[str], round_to: int | None = None
) -> None:
"""Print sklearn model coefficients.
Parameters
----------
labels : list of str
Coefficient names aligned with the fitted model.
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 | None,
*,
default_model_class: type[PyMCModel] | None,
supports_bayes: bool,
supports_ols: bool,
) -> ModelAdapter:
"""Resolve, validate, and wrap a model in a backend adapter.
Parameters
----------
model : PyMCModel, RegressorMixin, 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.
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)
raise ValueError("Unsupported model type")
