{
"cells": [
{
"cell_type": "markdown",
"metadata": {},
"source": [
"# Sharp regression discontinuity with scikit-learn models"
]
},
{
"cell_type": "code",
"execution_count": 1,
"metadata": {},
"outputs": [],
"source": [
"from sklearn.gaussian_process import GaussianProcessRegressor\n",
"from sklearn.gaussian_process.kernels import ExpSineSquared, WhiteKernel\n",
"from sklearn.linear_model import LinearRegression\n",
"\n",
"import causalpy as cp"
]
},
{
"cell_type": "code",
"execution_count": 2,
"metadata": {},
"outputs": [],
"source": [
"%config InlineBackend.figure_format = 'retina'"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Load data"
]
},
{
"cell_type": "code",
"execution_count": 3,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" x | \n",
" y | \n",
" treated | \n",
"
\n",
" \n",
" \n",
" \n",
" | 0 | \n",
" -0.932739 | \n",
" -0.091919 | \n",
" False | \n",
"
\n",
" \n",
" | 1 | \n",
" -0.930778 | \n",
" -0.382663 | \n",
" False | \n",
"
\n",
" \n",
" | 2 | \n",
" -0.929110 | \n",
" -0.181786 | \n",
" False | \n",
"
\n",
" \n",
" | 3 | \n",
" -0.907419 | \n",
" -0.288245 | \n",
" False | \n",
"
\n",
" \n",
" | 4 | \n",
" -0.882469 | \n",
" -0.420811 | \n",
" False | \n",
"
\n",
" \n",
"
\n",
"
"
]
},
"metadata": {
"image/png": {
"height": 491,
"width": 731
}
},
"output_type": "display_data"
}
],
"source": [
"fig, ax = result.plot()"
]
},
{
"cell_type": "code",
"execution_count": 6,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Difference in Differences experiment\n",
"Formula: y ~ 1 + x + treated\n",
"Running variable: x\n",
"Threshold on running variable: 0.5\n",
"\n",
"Results:\n",
"Discontinuity at threshold = 0.19\n",
"\n",
"\n",
"Model coefficients:\n",
" Intercept \t 0\n",
" treated[T.True]\t 0.19\n",
" x \t 1.23\n"
]
}
],
"source": [
"result.summary(round_to=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Linear, main-effects, and interaction model"
]
},
{
"cell_type": "code",
"execution_count": 7,
"metadata": {},
"outputs": [],
"source": [
"result = cp.RegressionDiscontinuity(\n",
" data,\n",
" formula=\"y ~ 1 + x + treated + x:treated\",\n",
" model=LinearRegression(),\n",
" treatment_threshold=0.5,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 8,
"metadata": {},
"outputs": [],
"source": [
"result.plot();"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"Though we can see that this does not give a good fit of the data almost certainly overestimates the discontinuity at threshold."
]
},
{
"cell_type": "code",
"execution_count": 9,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Difference in Differences experiment\n",
"Formula: y ~ 1 + x + treated + x:treated\n",
"Running variable: x\n",
"Threshold on running variable: 0.5\n",
"\n",
"Results:\n",
"Discontinuity at threshold = 0.92\n",
"\n",
"\n",
"Model coefficients:\n",
" Intercept \t 0\n",
" treated[T.True] \t 2.47\n",
" x \t 1.32\n",
" x:treated[T.True]\t -3.11\n"
]
}
],
"source": [
"result.summary(round_to=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using a bandwidth\n",
"\n",
"One way how we could deal with this is to use the `bandwidth` kwarg. This will only fit the model to data within a certain bandwidth of the threshold. If $x$ is the running variable, then the model will only be fitted to data where $threshold - bandwidth \\le x \\le threshold + bandwidth$."
]
},
{
"cell_type": "code",
"execution_count": 10,
"metadata": {},
"outputs": [],
"source": [
"result = cp.RegressionDiscontinuity(\n",
" data,\n",
" formula=\"y ~ 1 + x + treated + x:treated\",\n",
" model=LinearRegression(),\n",
" treatment_threshold=0.5,\n",
" bandwidth=0.3,\n",
")\n",
"\n",
"result.plot();"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"We could even go crazy and just fit intercepts for the data close to the threshold. But clearly this will involve more estimation error as we are using less data."
]
},
{
"cell_type": "code",
"execution_count": 11,
"metadata": {},
"outputs": [],
"source": [
"result = cp.RegressionDiscontinuity(\n",
" data,\n",
" formula=\"y ~ 1 + treated\",\n",
" model=LinearRegression(),\n",
" treatment_threshold=0.5,\n",
" bandwidth=0.3,\n",
")\n",
"\n",
"result.plot();"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using Gaussian Processes\n",
"\n",
"Now we will demonstrate how to use a scikit-learn model. "
]
},
{
"cell_type": "code",
"execution_count": 12,
"metadata": {},
"outputs": [],
"source": [
"kernel = 1.0 * ExpSineSquared(1.0, 5.0) + WhiteKernel(1e-1)\n",
"result = cp.RegressionDiscontinuity(\n",
" data,\n",
" formula=\"y ~ 1 + x + treated\",\n",
" model=GaussianProcessRegressor(kernel=kernel),\n",
" treatment_threshold=0.5,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 13,
"metadata": {},
"outputs": [],
"source": [
"fig, ax = result.plot()"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Using basis splines"
]
},
{
"cell_type": "code",
"execution_count": 14,
"metadata": {},
"outputs": [],
"source": [
"result = cp.RegressionDiscontinuity(\n",
" data,\n",
" formula=\"y ~ 1 + bs(x, df=6) + treated\",\n",
" model=LinearRegression(),\n",
" treatment_threshold=0.5,\n",
")"
]
},
{
"cell_type": "code",
"execution_count": 15,
"metadata": {},
"outputs": [],
"source": [
"fig, ax = result.plot()"
]
},
{
"cell_type": "code",
"execution_count": 16,
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Difference in Differences experiment\n",
"Formula: y ~ 1 + bs(x, df=6) + treated\n",
"Running variable: x\n",
"Threshold on running variable: 0.5\n",
"\n",
"Results:\n",
"Discontinuity at threshold = 0.41\n",
"\n",
"\n",
"Model coefficients:\n",
" Intercept \t 0\n",
" treated[T.True]\t 0.407\n",
" bs(x, df=6)[0] \t -0.594\n",
" bs(x, df=6)[1] \t -1.07\n",
" bs(x, df=6)[2] \t 0.278\n",
" bs(x, df=6)[3] \t 1.65\n",
" bs(x, df=6)[4] \t 1.03\n",
" bs(x, df=6)[5] \t 0.567\n"
]
}
],
"source": [
"result.summary(round_to=3)"
]
},
{
"cell_type": "markdown",
"metadata": {},
"source": [
"## Effect Summary Reporting\n",
"\n",
"For decision-making, you often need a concise summary of the causal effect. The `effect_summary()` method provides a decision-ready report with key statistics. Note that for Regression Discontinuity, the effect is a single scalar (the discontinuity at the threshold), similar to Difference-in-Differences.\n",
"\n",
":::{note}\n",
"**OLS vs PyMC Models**: When using OLS models (scikit-learn), the `effect_summary()` provides **confidence intervals** and **p-values** (frequentist inference), rather than the posterior distributions, HDI intervals, and tail probabilities provided by PyMC models (Bayesian inference). OLS tables include: mean, CI_lower, CI_upper, and p_value, but do not include median, tail probabilities (P(effect>0)), or ROPE probabilities.\n",
":::\n"
]
},
{
"cell_type": "code",
"execution_count": 17,
"metadata": {},
"outputs": [
{
"data": {
"text/html": [
"\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" mean | \n",
" ci_lower | \n",
" ci_upper | \n",
" p_value | \n",
"
\n",
" \n",
" \n",
" \n",
" | discontinuity | \n",
" 0.407683 | \n",
" 0.388323 | \n",
" 0.427043 | \n",
" 0.0 | \n",
"
\n",
" \n",
"
\n",
"
\n",
"\n",
"
\n",
" \n",
" \n",
" | \n",
" mean | \n",
" ci_lower | \n",
" ci_upper | \n",
" p_value | \n",
"
\n",
" \n",
" \n",
" \n",
" | discontinuity | \n",
" 0.407683 | \n",
" 0.388323 | \n",
" 0.427043 | \n",
" 0.0 | \n",
"
\n",
" \n",
"
\n",
"