# 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.custom_exceptions import (
GroupNotSampledException,
PriorPredictiveNotSupportedException,
)
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 _design_fingerprint, 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.DataTree) -> xr.DataArray:
"""Normalize supported PyMC coefficient variables to the canonical contract."""
variables = posterior if isinstance(posterior, xr.Dataset) else posterior.dataset
coefficient_names = ("beta", "b", "beta_z")
coefficient_name = next(
(name for name in coefficient_names if name in variables), None
)
if coefficient_name is None:
raise ValueError(
"PyMC posterior must expose one of 'beta', 'b', or 'beta_z' "
"as design-matrix coefficients."
)
coefficients = variables[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,
group: Literal["prior", "posterior"] = "posterior",
) -> 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.
group : {"prior", "posterior"}, default "posterior"
Draw group to condition forward sampling on. Bayesian backends
reproduce the prediction machinery using the requested group's
draws; point-estimate backends only ever have the (implicit)
posterior atom.
Returns
-------
xr.DataArray
Expected outcomes with dimensions ``("chain", "draw", "obs_ind",
"treated_units")``.
"""
@property
def is_built(self) -> bool:
"""Whether the backend's model graph / design state is constructed."""
return False
@property
def has_posterior(self) -> bool:
"""Whether posterior draws are available on this backend."""
return False
@property
def has_prior(self) -> bool:
"""Whether prior draws are available on this backend."""
return False
@property
def supports_prior_predictive(self) -> bool:
"""Whether this backend can sample a prior predictive phase."""
return False
[docs]
def build(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> None:
"""Construct the backend's graph/design state without sampling.
Idempotent by skipping when already built. Bayesian backends merge
data-driven priors and construct the PyMC graph; sklearn backends
record the design matrices for :meth:`sample_posterior`.
Parameters
----------
X : array-like or xarray.DataArray or mapping
Predictor matrix in the backend's expected form.
y : array-like or xarray.DataArray or mapping
Outcome vector or matrix in the backend's expected form.
coords : dict, optional
Coordinate metadata for PyMC models. Ignored by sklearn backends.
"""
raise NotImplementedError(
f"{type(self).__name__} does not support deferred construction."
)
[docs]
def sample_prior_predictive(self, **kwargs: Any) -> None:
"""Sample the prior predictive phase.
Other Parameters
----------------
**kwargs
Forwarded to the backend's prior-predictive sampler, overriding
the model's stored ``prior_sample_kwargs`` for this call only.
Raises
------
PriorPredictiveNotSupportedException
For backends without a prior predictive phase.
"""
raise PriorPredictiveNotSupportedException(
f"The {type(self.model).__name__} backend does not support prior "
"predictive sampling."
)
[docs]
@abstractmethod
def sample_posterior(self, **kwargs: Any) -> Any:
"""Sample the posterior phase with backend-appropriate conventions.
Other Parameters
----------------
**kwargs
Forwarded to the backend's posterior sampler, overriding the
model's stored ``sample_kwargs`` for this call only.
"""
[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, *, group: Literal["prior", "posterior"] = "posterior"
) -> 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.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Draw group to read coefficient variables from. ``"prior"`` is how
coefficient-estimand experiments compute prior-phase effects.
"""
[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
@property
def is_built(self) -> bool:
"""Whether the PyMC graph has been constructed."""
return bool(getattr(self._model, "_built", False))
@property
def has_posterior(self) -> bool:
"""Whether posterior draws are available."""
idata = self._model.idata
return idata is not None and "posterior" in idata.children
@property
def has_prior(self) -> bool:
"""Whether prior draws are available."""
idata = self._model.idata
return idata is not None and "prior" in idata.children
@property
def supports_prior_predictive(self) -> bool:
"""Whether the wrapped PyMC model exposes a prior predictive phase."""
return bool(getattr(self._model, "supports_prior_predictive", True))
[docs]
def build(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> None:
"""Merge priors and construct the PyMC graph without sampling.
Idempotent: skips when the graph already exists. Mapping inputs are
routed to the model's mapping-aware build path.
Parameters
----------
X : xarray.DataArray or dict of str to xarray.DataArray
Predictor matrix or mapping inputs.
y : xarray.DataArray or dict of str to xarray.DataArray
Outcome matrix or mapping inputs.
coords : dict, optional
Coordinate metadata for the PyMC model.
"""
if isinstance(X, dict) and isinstance(y, dict):
self._model.build_mapping(X=X, y=y, coords=coords)
return
if isinstance(X, dict) or isinstance(y, dict):
raise TypeError("X and y must either both be mappings or both be arrays")
self._model.build(X=X, y=y, coords=coords)
[docs]
def sample_prior_predictive(self, **kwargs: Any) -> None:
"""Sample the prior predictive phase on the built graph.
Other Parameters
----------------
**kwargs
Forwarded to the wrapped model, overriding its stored
``prior_sample_kwargs`` for this call only.
"""
if not self.supports_prior_predictive:
raise PriorPredictiveNotSupportedException(
f"The {type(self._model).__name__} backend does not support "
"prior predictive sampling."
)
self._model.sample_prior_predictive(**kwargs)
[docs]
def sample_posterior(self, **kwargs: Any) -> xr.DataTree:
"""Sample the posterior phase on the built graph.
Other Parameters
----------------
**kwargs
Forwarded to the wrapped model, overriding its stored
``sample_kwargs`` for this call only.
"""
return self._model.sample_posterior(**kwargs)
[docs]
def fit(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> xr.DataTree:
"""Fit the PyMC model (build + prior phase + posterior phase).
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,
group: Literal["prior", "posterior"] = "posterior",
) -> 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.
group : {"prior", "posterior"}, default "posterior"
Draw group to condition forward sampling on. The returned draws
always land in the ``posterior_predictive`` group (PyMC's output
location regardless of conditioning group) but carry the
conditioning group's ``chain``/``draw`` sizes.
Returns
-------
xr.DataArray
Forward draws of ``mu`` with canonical prediction dimensions.
Notes
-----
Reading ``posterior_predictive`` here is correct for prior-conditioned
output too; do not "fix" it to read ``idata[group]``.
"""
return _extract_mu(
self._model.predict(
X=X, coords=coords, out_of_sample=out_of_sample, group=group
)
)
[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, *, group: Literal["prior", "posterior"] = "posterior"
) -> xr.DataArray:
"""Return coefficient draws from the requested group, canonically.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Draw group to read coefficient variables from. The prior group is
how coefficient-estimand experiments compute prior-phase effects.
"""
idata = self._model.idata
if idata is None or group not in idata.children:
call = "fit()" if group == "posterior" else "sample_prior_predictive()"
raise GroupNotSampledException(
f"No {group!r} draws are available on this model. Call {call} first.",
group=group,
)
return _canonical_pymc_coefficients(idata[group])
[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
self._fit_inputs: tuple[Any, Any] | None = None
self._fit_fingerprint: tuple | None = None
self._is_fitted: bool = False
@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
@property
def is_built(self) -> bool:
"""Whether design matrices have been recorded via :meth:`build`."""
return self._fit_inputs is not None
@property
def has_posterior(self) -> bool:
"""Whether the sklearn model has been fitted."""
return self._is_fitted
[docs]
def build(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> None:
"""Record the design matrices for :meth:`sample_posterior`.
sklearn has no graph to construct; building only captures the fit
inputs. Idempotent by skipping when already recorded.
Parameters
----------
X : array-like
Predictor matrix.
y : array-like
Outcome vector or matrix.
coords : dict, optional
Ignored for sklearn backends.
"""
if self._fit_inputs is not None:
if _design_fingerprint(X, y) != self._fit_fingerprint:
raise RuntimeError(
"This backend is already built with different inputs. "
"Design matrices are recorded exactly once per instance; "
"assign a fresh model instead of rebuilding."
)
return
self._fit_inputs = (X, y)
self._fit_fingerprint = _design_fingerprint(X, y)
[docs]
def sample_prior_predictive(self, **kwargs: Any) -> None:
"""Raise: point-estimate backends have no prior predictive phase.
Other Parameters
----------------
**kwargs
Ignored; the capability error is raised unconditionally.
"""
raise PriorPredictiveNotSupportedException(
f"The {type(self._model).__name__} backend does not support "
"prior predictive sampling."
)
[docs]
def sample_posterior(self, **kwargs: Any) -> Any:
"""Fit the sklearn model on the recorded design matrices.
Other Parameters
----------------
**kwargs
Not supported by point-estimate backends; any override attempt
raises ``TypeError``.
"""
if self._fit_inputs is None:
raise RuntimeError(
"Design matrices have not been recorded. Call build(X, y) — "
"or an experiment's build() / fit(), which auto-call it — "
"before sampling."
)
if kwargs:
raise TypeError(
"Point-estimate backends accept no sampling overrides; got "
f"{sorted(kwargs)!r}"
)
X, y = self._fit_inputs
result = self.fit(X=X, y=y)
self._is_fitted = True
return result
[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
result = self._model.fit(X=X_array, y=_sklearn_y(y))
self._is_fitted = True
return result
[docs]
def predict(
self,
X: Any,
*,
coords: dict[str, Any] | None = None,
out_of_sample: bool = False,
group: Literal["prior", "posterior"] = "posterior",
) -> 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.
group : {"prior", "posterior"}, default "posterior"
Only the implicit posterior atom exists on point-estimate
backends; requesting ``"prior"`` raises the group-not-sampled
error naming :meth:`sample_prior_predictive`.
Returns
-------
xr.DataArray
Point predictions with canonical prediction dimensions and
singleton ``chain``/``draw`` dimensions.
"""
if group != "posterior":
raise GroupNotSampledException(
"No 'prior' draws are available on a point-estimate backend. "
"Call sample_prior_predictive() first — which itself raises "
"PriorPredictiveNotSupportedException for this backend — or "
"use a Bayesian model for prior checks.",
group=group,
)
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, *, group: Literal["prior", "posterior"] = "posterior"
) -> xr.DataArray:
"""Return fitted sklearn coefficients as singleton posterior draws.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Only the implicit posterior atom exists on point-estimate
backends; requesting ``"prior"`` raises the group-not-sampled
error naming :meth:`sample_prior_predictive`.
"""
if group != "posterior":
raise GroupNotSampledException(
"No 'prior' draws are available on a point-estimate backend. "
"Call sample_prior_predictive() first — which itself raises "
"PriorPredictiveNotSupportedException for this backend.",
group=group,
)
if not self._is_fitted:
raise GroupNotSampledException(
"No posterior draws are available on this model. Call fit() first.",
group="posterior",
)
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
self._fit_inputs: tuple[Any, Any, dict[str, Any] | None] | None = None
self._fit_fingerprint: tuple | None = None
@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
@property
def is_built(self) -> bool:
"""Whether fit inputs have been recorded via :meth:`build`."""
return self._fit_inputs is not None
@property
def has_posterior(self) -> bool:
"""Whether posterior draws are available."""
idata = self._model.idata
return idata is not None and "posterior" in idata.children
[docs]
def build(
self,
X: Any,
y: Any,
*,
coords: dict[str, Any] | None = None,
) -> None:
"""Record the fit inputs for :meth:`sample_posterior`.
The wrapped forecaster exposes no graph-construction step of its own,
so building only captures the inputs. Idempotent by skipping when
already recorded.
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; forwarded to :meth:`fit`.
"""
if self._fit_inputs is not None:
if _design_fingerprint(X, y) != self._fit_fingerprint:
raise RuntimeError(
"This backend is already built with different inputs. "
"Fit inputs are recorded exactly once per instance; "
"assign a fresh model instead of rebuilding."
)
return
self._fit_inputs = (X, y, coords)
self._fit_fingerprint = _design_fingerprint(X, y)
[docs]
def sample_prior_predictive(self, **kwargs: Any) -> None:
"""Raise: the forecaster exposes no prior-drawing path upstream.
Other Parameters
----------------
**kwargs
Ignored; the capability error is raised unconditionally.
"""
raise PriorPredictiveNotSupportedException(
"The PyMCForecastModel backend does not support prior predictive "
"sampling; the installed pymc-forecast API exposes no prior draw "
"path."
)
[docs]
def sample_posterior(self, **kwargs: Any) -> xr.DataTree:
"""Fit the forecaster on the recorded inputs.
Other Parameters
----------------
**kwargs
Not supported by the forecasting backend; any override attempt
raises ``TypeError``.
"""
if self._fit_inputs is None:
raise RuntimeError(
"Fit inputs have not been recorded. Call build(X, y) — or an "
"experiment's build() / fit(), which auto-call it — before "
"sampling."
)
if kwargs:
raise TypeError(
"pymc-forecast backends accept no sampling overrides; got "
f"{sorted(kwargs)!r}"
)
X, y, coords = self._fit_inputs
return self._model.fit(X=X, y=y, coords=coords)
[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,
group: Literal["prior", "posterior"] = "posterior",
) -> 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.
group : {"prior", "posterior"}, default "posterior"
Only posterior draws exist on this backend; requesting
``"prior"`` raises the group-not-sampled error naming
:meth:`sample_prior_predictive`.
Returns
-------
xr.DataArray
Posterior draws of ``mu`` with canonical prediction dimensions.
"""
if group != "posterior":
raise GroupNotSampledException(
"No 'prior' draws are available on a pymc-forecast backend. "
"Call sample_prior_predictive() first — which itself raises "
"PriorPredictiveNotSupportedException for this backend.",
group=group,
)
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, *, group: Literal["prior", "posterior"] = "posterior"
) -> xr.DataArray:
"""Forecasting models have no design-matrix coefficients.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Ignored; no coefficient container exists on this backend.
"""
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")
