# 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.
"""Regression kink design."""
import re # noqa: I001
import warnings
from matplotlib import pyplot as plt
import numpy as np
import pandas as pd
import seaborn as sns
from patsy import ModelDesc
import xarray as xr
from causalpy.formula_utils import build_design_matrices, build_formula_matrices
from causalpy.input_data import DataFrameLike, to_pandas
from causalpy.experiments._results import KinkResult
from causalpy.experiments.model_adapter import build_coords
from causalpy.plot_utils import (
_PosteriorPlotStyle,
format_r2_score,
plot_posterior_over_x,
)
from causalpy.pymc_models import LinearRegression, PyMCModel
from causalpy.reporting import EffectSummary, _effect_summary_rkink
from causalpy.constants import HDI_PROB, LEGEND_FONT_SIZE
from .base import BaseExperiment
from typing import Any, Literal
from causalpy.utils import _is_variable_dummy_coded, round_num
from causalpy.custom_exceptions import (
DataException,
FormulaException,
)
[docs]
class RegressionKink(BaseExperiment[KinkResult]):
"""A class to analyse regression kink designs.
Parameters
----------
data : dataframe-like
Any eager dataframe Narwhals supports, such as pandas, Polars, or
PyArrow. Converted to pandas internally.
formula : str
A statistical model formula.
kink_point : float
A scalar value at which the kink occurs.
model : PyMCModel, optional
A PyMC model. Defaults to :class:`LinearRegression`.
running_variable_name : str, default "x"
The name of the running variable column.
epsilon : float, default 0.001
A small scalar for evaluating the causal impact above/below the kink.
bandwidth : float, default np.inf
Data outside of the bandwidth (relative to the kink) is not used to
fit the model.
Notes
-----
**Lazy lifecycle**
Construction only validates input and builds design matrices — nothing is
sampled. Call :meth:`fit` to run posterior inference (it returns ``self``,
so construction and fitting chain in one expression), and optionally
:meth:`sample_prior_predictive` first for prior predictive checks
(``plot(group="prior")``, ``effect_summary(group="prior")``). Results
live on ``exp.result`` / ``exp.prior_result``.
**Estimate extraction**
The class predicts the conditional expectation at ``kink_point - epsilon``, ``kink_point``, and ``kink_point + epsilon``. It forms finite-difference slopes on the left and right and stores their difference as ``gradient_change``. This is a local prediction contrast on derivatives, not a population-standardized effect.
Examples
--------
>>> import causalpy as cp
>>> df = cp.load_data("rd")
>>> kink = 0.5
>>> result = cp.RegressionKink(
... df,
... formula=f"y ~ 1 + x + I((x - {kink}) * treated)",
... kink_point=kink,
... model=cp.pymc_models.LinearRegression(
... sample_kwargs={"random_seed": 42, "progressbar": False}
... ),
... ).fit()
"""
supports_ols = False
supports_bayes = True
_default_model_class = LinearRegression
[docs]
def __init__(
self,
data: DataFrameLike,
formula: str,
kink_point: float,
model: PyMCModel | None = None,
running_variable_name: str = "x",
epsilon: float = 0.001,
bandwidth: float = np.inf,
) -> None:
super().__init__(model=model)
self.expt_type = "Regression Kink"
self.data = to_pandas(data)
self.data.index.name = "obs_ind"
self.formula = formula
self.running_variable_name = running_variable_name
self.kink_point = kink_point
self.epsilon = epsilon
self.bandwidth = bandwidth
self.input_validation()
self._build_design_matrices()
self._prepare_data()
self._prepare_prediction_grids()
def _build_design_matrices(self) -> None:
"""Build design matrices from formula and data, applying bandwidth filtering."""
if self.bandwidth is not np.inf:
fmin = self.kink_point - self.bandwidth
fmax = self.kink_point + self.bandwidth
filtered_data = self.data.query(f"{fmin} <= x <= {fmax}")
if len(filtered_data) <= 10:
warnings.warn(
f"Choice of bandwidth parameter has lead to only {len(filtered_data)} remaining datapoints. Consider increasing the bandwidth parameter.", # noqa: E501
UserWarning,
stacklevel=2,
)
y, X = build_formula_matrices(self.formula, filtered_data)
else:
y, X = build_formula_matrices(self.formula, self.data)
self._y_design_info = y.design_info
self._x_design_info = X.design_info
self.labels = X.design_info.column_names
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:
"""Bundle design matrices into an ``xr.Dataset``."""
n = self._X_raw.shape[0]
self.design = self._build_design_dataset(
self._X_raw,
self._y_raw,
obs_ind=np.arange(n),
coeffs=self.labels,
)
del self._X_raw, self._y_raw
def _prepare_prediction_grids(self) -> None:
"""Build the deterministic running-variable grid used for plotting.
Draw-independent design-stage artifact behind
``result.predictions``; computed once at configure time and never
re-assigned per draw group.
"""
if self.bandwidth is not np.inf:
fmin = self.kink_point - self.bandwidth
fmax = self.kink_point + self.bandwidth
xi = np.linspace(fmin, fmax, 200)
else:
xi = np.linspace(
np.min(self.data[self.running_variable_name]),
np.max(self.data[self.running_variable_name]),
200,
)
self.x_pred = pd.DataFrame(
{self.running_variable_name: xi, "treated": self._is_treated(xi)}
)
def _fit_inputs(self) -> tuple[Any, Any, dict[str, Any]]:
"""Return the design matrices and coordinates for model build."""
X = self.design["X"]
return (
X,
self.design["y"],
build_coords(self.labels, X.shape[0]),
)
def _finalize(self, group: Literal["prior", "posterior"]) -> None:
"""Compute the group's result bundle from its draws and assign it.
The body is the historical ``algorithm()`` with the draw group
threaded through prediction. Posterior fits score against the
observed data; prior draws are not scored (R² against observed
data is not informative under a prior).
"""
X = self.design["X"]
y = self.design["y"]
# predictions over the running-variable grid built at configure time
(new_x,) = build_design_matrices([self._x_design_info], self.x_pred)
predictions = self._model_backend.predict(X=np.asarray(new_x), group=group)
# evaluate gradient change around kink point
mu_kink_left, mu_kink, mu_kink_right = self._probe_kink_point(group=group)
gradient_change = self._eval_gradient_change(
mu_kink_left, mu_kink, mu_kink_right, self.epsilon
)
score = None
if group == "posterior":
score = self._model_backend.score(X=X, y=y)
bundle = KinkResult(
predictions=predictions,
gradient_change=gradient_change,
score=score,
)
self._assign_bundle(group, bundle)
@staticmethod
def _eval_gradient_change(
mu_kink_left: xr.DataArray,
mu_kink: xr.DataArray,
mu_kink_right: xr.DataArray,
epsilon: float,
) -> xr.DataArray:
"""Evaluate the gradient change at the kink point.
It works by evaluating the model below the kink point, at the kink point,
and above the kink point.
This is a static method for ease of testing.
"""
gradient_left = (mu_kink - mu_kink_left) / epsilon
gradient_right = (mu_kink_right - mu_kink) / epsilon
gradient_change = gradient_right - gradient_left
return gradient_change
def _probe_kink_point(
self, *, group: Literal["prior", "posterior"]
) -> tuple[xr.DataArray, xr.DataArray, xr.DataArray]:
"""Probe the kink point to evaluate the predicted outcome at the kink point and
either side."""
# Create a dataframe to evaluate predicted outcome at the kink point and either
# side
x_predict = pd.DataFrame(
{
self.running_variable_name: np.array(
[
self.kink_point - self.epsilon,
self.kink_point,
self.kink_point + self.epsilon,
]
),
"treated": np.array([0, 1, 1]),
}
)
(new_x,) = build_design_matrices([self._x_design_info], x_predict)
predicted = self._model_backend.predict(X=np.asarray(new_x), group=group)
mu_kink_left = predicted.sel(obs_ind=0)
mu_kink = predicted.sel(obs_ind=1)
mu_kink_right = predicted.sel(obs_ind=2)
return mu_kink_left, mu_kink, mu_kink_right
def _is_treated(self, x: np.ndarray | pd.Series) -> np.ndarray:
"""Returns ``True`` if `x` is greater than or equal to the treatment threshold.""" # noqa: E501
return np.greater_equal(x, self.kink_point)
[docs]
def summary(self, round_to: int | None = 2) -> None:
"""Print summary of main results and model coefficients.
Parameters
----------
round_to : int, optional
Number of decimals used to round results. Defaults to 2. Use
``None`` to return raw numbers.
"""
bundle = self.result
print(
f"""
{self.expt_type:=^80}
Formula: {self.formula}
Running variable: {self.running_variable_name}
Kink point on running variable: {self.kink_point}
Results:
Change in slope at kink point = {round_num(bundle.gradient_change.mean(), round_to)}
"""
)
self.print_coefficients(round_to)
[docs]
def plot(
self,
*,
group: Literal["prior", "posterior"] = "posterior",
round_to: int | None = 2,
ci_prob: float = HDI_PROB,
kind: Literal["ribbon", "histogram", "spaghetti"] = "ribbon",
ci_kind: Literal["hdi", "eti"] = "hdi",
num_samples: int = 50,
figsize: tuple[float, float] | None = None,
show: bool = True,
legend_kwargs: dict[str, Any] | None = None,
) -> tuple[plt.Figure, plt.Axes]:
"""Plot the regression kink results.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Which draw group to plot. ``"prior"`` renders the reduced
prior-check figure — the prior-implied fit against the observed
data only — and requires :meth:`sample_prior_predictive`;
``"posterior"`` (default) renders the full results figure and
requires :meth:`fit`.
round_to : int, optional
Number of decimals used to round numerical results in the figure
title (e.g. the Bayesian :math:`R^2`). Defaults to 2. Use
``None`` to render raw numbers.
ci_prob : float
Probability mass of the highest density interval drawn around the
posterior predictive band, and the central credible interval
reported in the figure title for the change in gradient at the
kink point. Must be in ``(0, 1]``. Defaults to
:data:`~causalpy.constants.HDI_PROB` (currently 0.94).
kind : {"ribbon", "histogram", "spaghetti"}, optional
How posterior uncertainty is rendered via
:func:`~causalpy.plot_utils.plot_posterior_over_x`. Defaults to ``"ribbon"``.
For ``"spaghetti"``, legends use draw lines rather than a shaded
band. For ``"histogram"``, uncertainty is shown as a 2D density
heatmap with a mean line overlay (no ribbon patch for legends).
ci_kind : {"hdi", "eti"}, optional
Credible interval type when ``kind="ribbon"``. Defaults to
``"hdi"``.
num_samples : int, optional
Number of posterior draws when ``kind="spaghetti"``. Defaults
to 50. Ignored for other kinds.
figsize : tuple of (float, float), optional
Width and height of the figure in inches, passed to
:func:`matplotlib.pyplot.subplots`. Defaults to ``None`` (use
matplotlib's default).
show : bool
Whether to automatically display the plot. Defaults to ``True``.
legend_kwargs : dict, optional
Keyword arguments to adjust legend placement and styling.
Supported keys: ``loc``, ``bbox_to_anchor``, ``fontsize``,
``frameon``, ``title`` (``bbox_transform`` is accepted alongside
``bbox_to_anchor``). The existing legend is modified **in
place** so that custom handles are preserved.
Returns
-------
fig : matplotlib.figure.Figure
The figure that was created.
ax : matplotlib.axes.Axes
The axes object containing the plot.
"""
return self._render_plot(
show=show,
legend_kwargs=legend_kwargs,
group=group,
round_to=round_to,
ci_prob=ci_prob,
kind=kind,
ci_kind=ci_kind,
num_samples=num_samples,
figsize=figsize,
)
def _plot(
self,
*,
group: Literal["prior", "posterior"] = "posterior",
round_to: int | None = 2,
ci_prob: float = HDI_PROB,
kind: Literal["ribbon", "histogram", "spaghetti"] = "ribbon",
ci_kind: Literal["hdi", "eti"] = "hdi",
num_samples: int = 50,
figsize: tuple[float, float] | None = None,
**kwargs: Any,
) -> tuple[plt.Figure, plt.Axes]:
"""Generate plot for regression kink designs.
Parameters
----------
group : {"prior", "posterior"}
``"prior"`` renders the reduced single-panel prior-check figure
via :meth:`_plot_prior_checks`; ``"posterior"`` renders the full
results figure.
round_to : int, optional
Number of decimals used to round results. Defaults to 2. Use ``None``
to return raw numbers.
ci_prob : float, optional
Probability mass of the highest density interval drawn around the
posterior predictive band, and the central credible interval
reported in the figure title for the change in gradient at the
kink point. Must be in ``(0, 1]``. Defaults to
:data:`~causalpy.constants.HDI_PROB` (currently 0.94).
figsize : tuple of (float, float), optional
Width and height of the figure in inches. Defaults to ``None``
(use matplotlib's default).
"""
bundle = self._require_bundle(group)
if group == "prior":
return self._plot_prior_checks(bundle=bundle)
style: _PosteriorPlotStyle = {
"ci_prob": ci_prob,
"kind": kind,
"ci_kind": ci_kind,
"num_samples": num_samples,
}
fig, ax = plt.subplots(figsize=figsize)
# Plot raw data
sns.scatterplot(
self.data,
x=self.running_variable_name,
y=self.outcome_variable_name,
c="k", # hue="treated",
ax=ax,
)
h_line, h_patch = plot_posterior_over_x(
self.x_pred[self.running_variable_name],
bundle.predictions.isel(treated_units=0),
ax=ax,
**style,
plot_hdi_kwargs={"color": "C1"},
)
handles = [(h_line, h_patch)]
labels = ["Posterior mean"]
# create strings to compose title
r2 = format_r2_score(bundle.score, round_to=round_to, context="on all data")
percentiles = bundle.gradient_change.quantile(
[(1 - ci_prob) / 2, 1 - (1 - ci_prob) / 2]
).values
ci = (
rf"$CI_{{{ci_prob * 100:.0f}\%}}$"
+ f"[{round_num(percentiles[0], round_to if round_to is not None else 2)}, {round_num(percentiles[1], round_to if round_to is not None else 2)}]"
)
grad_change = f"""
Change in gradient = {round_num(bundle.gradient_change.mean(), round_to if round_to is not None else 2)},
"""
ax.set(title=r2 + "\n" + grad_change + ci)
# Intervention line
ax.axvline(
x=self.kink_point,
ls="-",
lw=3,
color="r",
label="treatment threshold",
)
ax.legend(
handles=handles,
labels=labels,
fontsize=LEGEND_FONT_SIZE,
)
return fig, ax
def _plot_prior_checks(self, *, bundle: KinkResult) -> tuple[plt.Figure, plt.Axes]:
"""Render the reduced prior-check figure.
A prior check answers whether the prior-implied fit is plausible
against the observed data, so a single panel suffices: the observed
scatter plus the prior-implied fit line and band over the
running-variable grid, with the kink point marked.
"""
style: _PosteriorPlotStyle = {
"ci_prob": HDI_PROB,
"kind": "ribbon",
"ci_kind": "hdi",
"num_samples": 50,
}
fig, ax = plt.subplots(figsize=(7, 4))
sns.scatterplot(
self.data,
x=self.running_variable_name,
y=self.outcome_variable_name,
c="k",
ax=ax,
)
h_line, h_patch = plot_posterior_over_x(
self.x_pred[self.running_variable_name],
bundle.predictions.isel(treated_units=0),
ax=ax,
**style,
plot_hdi_kwargs={"color": "C1"},
)
ax.axvline(
x=self.kink_point,
ls="--",
lw=1.5,
color="r",
label="treatment threshold",
)
ax.legend(
handles=[(h_line, h_patch)],
labels=["Prior fit"],
fontsize=LEGEND_FONT_SIZE,
)
ax.set(title="Prior predictive check")
return fig, ax
[docs]
def effect_summary(
self,
*,
group: Literal["prior", "posterior"] = "posterior",
direction: Literal["increase", "decrease", "two-sided"] = "increase",
alpha: float = 0.05,
min_effect: float | None = None,
) -> EffectSummary:
"""
Generate a decision-ready summary of causal effects for Regression Kink.
Parameters
----------
group : {"prior", "posterior"}, default "posterior"
Which draw group to summarize. ``"prior"`` requires
:meth:`sample_prior_predictive` and produces prior-appropriate
prose — under a neutral prior, ``P(effect > 0)`` should sit near
0.5, so a tail probability far from 0.5 flags a design-matrix or
prior-specification problem rather than a causal finding.
``"posterior"`` requires :meth:`fit`.
direction : {"increase", "decrease", "two-sided"}, default="increase"
Direction for tail probability calculation (PyMC only, ignored for OLS).
alpha : float, default=0.05
Significance level for HDI/CI intervals (1-alpha confidence level).
min_effect : float, optional
Region of Practical Equivalence (ROPE) threshold (PyMC only, ignored for OLS).
Returns
-------
EffectSummary
Object with .table (DataFrame) and .text (str) attributes
"""
# Resolve the group's bundle once; helpers consume containers.
bundle = self._require_bundle(group)
# The helper applies the prior-plausibility prose prefix itself.
return _effect_summary_rkink(
bundle,
direction=direction,
alpha=alpha,
min_effect=min_effect,
group=group,
)
