QRFs for Conformalized Quantile Regression#
This example demonstrates the use of a quantile regression forest (QRF) to construct reliable prediction intervals using conformalized quantile regression (CQR). While QRFs can estimate quantiles, they may require additional calibration to provide reliable interval estimates. CQR calibrates the intervals on data held out from training, which guarantees that their coverage (i.e., the percentage of samples that actually fall within their prediction interval) is at least the target level, on average. We also calibrate them with the out-of-bag (OOB) predictions of the training samples, which lets the QRF fit on all of the training data and gives narrower intervals, but does not guarantee the coverage. Since the coverage on one split of the data varies by chance, the plots show each sample’s interval from one repetition of 4-fold cross-validation, and the mean coverage and interval width over five repetitions (20 splits). This example is adapted from “Prediction intervals: Quantile Regression Forests” by Carl McBride Ellis.
import altair as alt
import numpy as np
import pandas as pd
from sklearn import datasets
from sklearn.model_selection import RepeatedKFold, train_test_split
from quantile_forest import RandomForestQuantileRegressor
random_seed = 0
random_state = np.random.RandomState(random_seed)
n_samples = 900
n_folds = 4 # the number of folds of the cross-validation
n_repeats = 5 # the number of times that the cross-validation is repeated
coverages = np.linspace(0.1, 0.9, num=9).round(1).tolist() # the "coverage level"
# Load the California Housing dataset.
X, y = datasets.fetch_california_housing(return_X_y=True, as_frame=True)
perm = random_state.permutation(min(len(X), n_samples))
X = X.iloc[perm]
y = y.iloc[perm]
def predict_intervals(qrf, X, **kwargs):
"""Predict the interval at each coverage level, as an array of shape (n, n_coverages, 2)."""
quantiles = [q for coverage in coverages for q in [(1 - coverage) / 2, (1 + coverage) / 2]]
return qrf.predict(X, quantiles=quantiles, **kwargs).reshape(len(X), len(coverages), 2)
def conformalize(y_pis, y_pis_calib, y_calib):
"""Calibrate prediction intervals with CQR, given those of calibration samples."""
y_pis = y_pis.copy()
for i, coverage in enumerate(coverages):
# Calculate the conformity scores on the calibration data.
conf_scores = np.maximum(y_pis_calib[:, i, 0] - y_calib, y_calib - y_pis_calib[:, i, 1])
# Get the ceil((n + 1) * coverage)-th smallest of the n conformity scores,
# `s`, which guarantees a coverage of at least `coverage`.
n_calib = len(y_calib)
level = min(1, np.ceil((n_calib + 1) * coverage) / n_calib)
s = np.quantile(conf_scores, level, method="higher")
# Subtract `s` from the lower quantile and add it to the upper quantile.
y_pis[:, i, 0] -= s
y_pis[:, i, 1] += s
return y_pis
def predict_strategies(X_train, X_test, y_train, random_state=None):
"""Predict the point predictions and prediction intervals of each strategy."""
qrf = RandomForestQuantileRegressor(max_samples_leaf=None, random_state=random_state)
qrf.fit(X_train, y_train)
# Calculate the point predictions and prediction intervals on the test data.
y_pred = qrf.predict(X_test, quantiles="mean", aggregate_leaves_first=False)
y_pis = predict_intervals(qrf, X_test)
# For CQR, fit on half of the training data and calibrate on the other half.
X_fit, X_calib, y_fit, y_calib = train_test_split(
X_train, y_train, test_size=0.5, random_state=random_state
)
qrf_cqr = RandomForestQuantileRegressor(max_samples_leaf=None, random_state=random_state)
qrf_cqr.fit(X_fit, y_fit)
y_pred_cqr = qrf_cqr.predict(X_test, quantiles="mean", aggregate_leaves_first=False)
y_pis_cqr = conformalize(
predict_intervals(qrf_cqr, X_test), predict_intervals(qrf_cqr, X_calib), y_calib
)
# For CQR with OOB calibration, calibrate on the training data, with the
# intervals of each sample predicted by the trees that did not draw it.
y_pis_oob = predict_intervals(qrf, X_train, oob_score=True)
y_pis_cqr_oob = conformalize(y_pis, y_pis_oob, y_train)
return {
"qrf": (y_pred, y_pis),
"cqr": (y_pred_cqr, y_pis_cqr),
"cqr_oob": (y_pred, y_pis_cqr_oob),
}
# Predict with each strategy on each split of the cross-validation.
cv = RepeatedKFold(n_splits=n_folds, n_repeats=n_repeats, random_state=random_seed)
dfs, metrics = [], []
for split, (train_index, test_index) in enumerate(cv.split(X)):
X_train, X_test = X.iloc[train_index], X.iloc[test_index]
y_train, y_test = y.iloc[train_index], y.iloc[test_index].to_numpy()
for strategy, (y_pred, y_pis) in predict_strategies(
X_train, X_test, y_train, random_state=random_seed + split
).items():
y_pred_low, y_pred_upp = y_pis[:, :, 0], y_pis[:, :, 1]
# Calculate the coverage and mean width of the intervals at each level.
within = (y_pred_low <= y_test[:, np.newaxis]) & (y_test[:, np.newaxis] <= y_pred_upp)
metrics.append(
pd.DataFrame(
{
"strategy": strategy,
"target": coverages,
"coverage": within.mean(axis=0),
"width": (y_pred_upp - y_pred_low).mean(axis=0) * 100_000,
}
)
)
# Keep the intervals of the first repetition, in which each sample is tested once.
if split < n_folds:
y_values = {"y_test": y_test, "y_pred": y_pred}
for i, coverage in enumerate(coverages):
y_values[f"y_pred_low_{round(coverage * 100)}"] = y_pred_low[:, i]
y_values[f"y_pred_upp_{round(coverage * 100)}"] = y_pred_upp[:, i]
dfs.append(
pd.DataFrame(y_values)
.pipe(lambda x: (x * 100_000).round().astype(int)) # convert to dollars
.assign(strategy=strategy)
)
df = pd.concat(dfs, ignore_index=True)
df_metrics = pd.concat(metrics).groupby(["strategy", "target"], as_index=False).mean()
def plot_prediction_intervals_by_strategy(df, df_metrics):
"""Plot prediction intervals by interval estimate strategy."""
# Single shared slider for varying the target coverage level.
slider = alt.binding_range(name="Coverage Target: ", min=0.1, max=0.9, step=0.1)
coverage_val = alt.param(name="coverage", value=0.9, bind=slider)
# Drop-down menu for choosing the calibration data of the CQR intervals.
select = alt.binding_select(
name="CQR Calibration: ",
options=["cqr", "cqr_oob"],
labels=["Held-Out Data", "Out-of-Bag Predictions"],
)
calibration_val = alt.param(name="calibration", value="cqr", bind=select)
domain = [
int(np.min((df[["y_test", "y_pred"]].min(axis=0)))), # min of all axes
int(np.max((df[["y_test", "y_pred"]].max(axis=0)))), # max of all axes
]
def _plot_prediction_intervals(df, df_metrics, filter_strategy="true", suffix=""):
click = alt.selection_point(name=f"click{suffix}", bind="legend", fields=["y_label"])
tooltip = [
alt.Tooltip("y_test:Q", format="$,d", title="True Price"),
alt.Tooltip("y_pred:Q", format="$,d", title="Predicted Price"),
alt.Tooltip("y_pred_low:Q", format="$,d", title="Predicted Lower Price"),
alt.Tooltip("y_pred_upp:Q", format="$,d", title="Predicted Upper Price"),
alt.Tooltip("y_label:N", title="Within Interval"),
]
base = (
alt.Chart(df)
.transform_filter(filter_strategy)
.transform_calculate(
y_pred_low="datum['y_pred_low_' + round(coverage * 100)]",
y_pred_upp="datum['y_pred_upp_' + round(coverage * 100)]",
)
.transform_calculate(
y_label=(
"((datum.y_test >= datum.y_pred_low) & (datum.y_test <= datum.y_pred_upp))"
" ? 'Yes' : 'No'"
)
)
)
bar = base.mark_bar(width=1).encode(
x=alt.X("y_pred:Q", scale=alt.Scale(domain=domain, padding=0), title=""),
y=alt.Y("y_pred_low:Q", scale=alt.Scale(domain=domain, padding=0), title=""),
y2=alt.Y2("y_pred_upp:Q", title=None),
color=alt.condition(click, alt.value("#e0f2ff"), alt.value("lightgray")),
opacity=alt.condition(click, alt.value(0.8), alt.value(0)),
tooltip=tooltip,
)
tick = base.mark_tick(orient="horizontal", thickness=1, width=5).encode(
x=alt.X("y_pred:Q", title=""),
color=alt.value("#006aff"),
opacity=alt.condition(click, alt.value(0.4), alt.value(0)),
)
tick_low = tick.encode(y=alt.Y("y_pred_low:Q", scale=alt.Scale(clamp=True), title=""))
tick_upp = tick.encode(y=alt.Y("y_pred_upp:Q", scale=alt.Scale(clamp=True), title=""))
circle = (
base.add_params(click)
.mark_circle(size=15)
.encode(
x=alt.X(
"y_pred:Q",
axis=alt.Axis(format="$~s"),
scale=alt.Scale(domain=domain, nice=False),
title="Predicted Prices",
),
y=alt.Y(
"y_test:Q",
axis=alt.Axis(format="$~s"),
scale=alt.Scale(domain=domain, nice=False),
title="True Prices",
),
color=alt.condition(
click,
alt.Color(
"y_label:N",
scale=alt.Scale(domain=["Yes", "No"], range=["#f2a619", "red"]),
title="Within Interval",
),
alt.value("lightgray"),
),
opacity=alt.condition(click, alt.value(1), alt.value(0)),
tooltip=tooltip,
)
)
diagonal = (
alt.Chart(pd.DataFrame({"var1": domain, "var2": domain}))
.mark_line(color="gray", opacity=0.4, strokeDash=[2, 2])
.encode(
x=alt.X("var1:Q"),
y=alt.Y("var2:Q"),
)
)
# The mean coverage and interval width over all of the splits.
text = (
alt.Chart(df_metrics)
.transform_filter(filter_strategy)
.transform_filter("round(datum.target * 100) == round(coverage * 100)")
.transform_calculate(
coverage_text=(
"'Mean Coverage: ' + format(datum.coverage * 100, '.1f') + '%'"
" + ' (target = ' + format(datum.target * 100, '.1f') + '%)'"
),
width_text="'Mean Interval Width: ' + format(datum.width, '$,d')",
)
.mark_text(align="left", baseline="top", color="gray")
.encode(x=alt.value(5))
)
text_coverage = text.encode(y=alt.value(5), text=alt.Text("coverage_text:N"))
text_width = text.encode(y=alt.value(18), text=alt.Text("width_text:N"))
chart = bar + tick_low + tick_upp + circle + diagonal + text_coverage + text_width
return chart.properties(height=250, width=300)
# The CQR plot shows the intervals of the calibration chosen in the drop-down menu.
panel_qrf = _plot_prediction_intervals(
df.query("strategy == 'qrf'"), df_metrics.query("strategy == 'qrf'"), suffix="_qrf"
).properties(title="Quantile Regression Forest (QRF)")
panel_cqr = _plot_prediction_intervals(
df.query("strategy != 'qrf'"),
df_metrics.query("strategy != 'qrf'"),
"datum.strategy == calibration",
suffix="_cqr",
).properties(title="Conformalized Quantile Regression (CQR)")
chart = alt.hconcat(panel_qrf, panel_cqr).add_params(coverage_val, calibration_val)
return chart
chart = plot_prediction_intervals_by_strategy(df, df_metrics)
chart