Uncertainty quantification¶
iBART is a selection method, so its biggest uncertainty is which descriptors
get chosen, not just the scatter around a fixed model. jax_ibart.uncertainty
has tools for both. We use a noisy synthetic problem in which only a*b drives
the response (c and d are irrelevant).
import numpy as np
import matplotlib.pyplot as plt
from jax_ibart import IBart, uncertainty
rng = np.random.default_rng(0)
n = 240
a, b, c, d = (rng.uniform(0.5, 2.0, n) for _ in range(4))
X = np.column_stack([a, b, c, d]); names = ["a", "b", "c", "d"]
y = 5.0 * (a * b) + rng.normal(0, 0.4, n) # noisy
# a factory makes a fresh, identically configured estimator on demand
factory = lambda: IBart(backend="diff", schedule=("b",), max_dim=1, seed=0, steps=900)
1. Conformal prediction intervals¶
Split conformal prediction wraps the entire fit, selection included, and calibrates on held-out data. That gives a finite-sample coverage guarantee even after selection. We check the realized coverage on a fresh test set.
idx = rng.permutation(n); te, tr = idx[:80], idx[80:]
conf = uncertainty.conformal_intervals(factory, X[tr], y[tr], names,
alpha=0.10, cal_frac=0.4, seed=0)
cov = uncertainty.empirical_coverage(conf, X[te], y[te])
se = np.sqrt(0.9 * 0.1 / len(te))
print(f"target coverage: 0.90 realized: {cov:.2f} half-width: {conf.half_width:.3f}")
print(f"sampling noise on a coverage estimate from {len(te)} test points: +/- {se:.2f} (1 s.e.)")
target coverage: 0.90 realized: 0.85 half-width: 0.599 sampling noise on a coverage estimate from 80 test points: +/- 0.03 (1 s.e.)
lo, point, hi = conf.interval(X[te])
order = np.argsort(point)
inside = (y[te] >= lo) & (y[te] <= hi)
fig, ax = plt.subplots(figsize=(7, 3.6))
xs = np.arange(len(order))
ax.fill_between(xs, lo[order], hi[order], alpha=0.25, label="90% conformal band")
ax.plot(xs, point[order], lw=1, label="prediction")
ax.scatter(xs[inside[order]], y[te][order][inside[order]], s=12, color="k",
label="true (inside)", zorder=3)
ax.scatter(xs[~inside[order]], y[te][order][~inside[order]], s=18, color="C3",
label="true (outside)", zorder=3)
ax.set_xlabel("test point (sorted by prediction)"); ax.set_ylabel("y")
ax.legend(fontsize=8); fig.tight_layout(); plt.show()
The realized coverage is within sampling noise of the 90% target: with 80 test points, one standard error on a coverage estimate is about 0.03.
2. Selection stability¶
The most informative check for a selection method is how often each feature and
each descriptor is chosen across bootstrap resamples. The true drivers a and
b should dominate, and c and d should rarely or never appear.
stab = uncertainty.selection_stability(factory, X, y, names, n_boot=20, seed=0)
ff = {nm: stab["feature_freq"].get(nm, 0.0) for nm in names}
fig, ax = plt.subplots(figsize=(5, 3))
ax.bar(list(ff), list(ff.values()))
ax.set_ylabel("inclusion frequency"); ax.set_ylim(0, 1.05)
ax.set_title(f"Feature selection stability ({stab['n_boot']} bootstrap fits)")
fig.tight_layout(); plt.show()
print("features :", ff)
print("descriptors:", stab["descriptor_freq"])
features : {'a': 1.0, 'b': 1.0, 'c': 0.0, 'd': 0.0}
descriptors: {'a*b': 1.0}
Here the answer is clear-cut: every bootstrap fit chose a*b. On a real
problem, a descriptor chosen in only a fraction of resamples is a warning that
the data do not pin it down.
3. Out-of-sample RMSE distribution¶
Following the papers' protocol, we report the RMSE over many random train/test splits as a distribution rather than a single number. The noise level is 0.4, so that is the best RMSE any model can achieve here.
cv = uncertainty.repeated_split_rmse(factory, X, y, names, n_splits=20, seed=0)
fig, ax = plt.subplots(figsize=(5, 3))
ax.hist(cv["rmse"], bins=12)
ax.axvline(cv["mean"], color="k", ls="--", label=f"mean {cv['mean']:.2f}")
ax.axvline(0.4, color="C3", ls=":", label="noise level 0.4")
ax.set_xlabel("out-of-sample RMSE"); ax.set_ylabel("count"); ax.legend()
fig.tight_layout(); plt.show()
print(f"RMSE = {cv['mean']:.3f} +/- {cv['std']:.3f} (90% CI [{cv['q05']:.2f}, {cv['q95']:.2f}])")
RMSE = 0.409 +/- 0.046 (90% CI [0.33, 0.47])
4. Coefficient intervals (classical, optimistic)¶
Classical OLS confidence intervals on the final model are fast but optimistic: they condition on the selected descriptors and ignore the selection step. Use the conformal intervals (section 1) for honest predictive coverage.
model = factory().fit(X, y, names)
ols = uncertainty.ols_intervals(model, X, y)
for term in ols["coefficients"]:
print(f"{term['term']:>10}: {term['coef']:+.3f} "
f"[{term['ci_low']:+.3f}, {term['ci_high']:+.3f}]")
intercept: +0.004 [-0.093, +0.101]
a*b: +4.986 [+4.934, +5.038]
Takeaway: report the conformal band for honest prediction intervals, and the
selection-stability frequencies to show which descriptors are trustworthy.
Together they capture the uncertainty a single point estimate hides. See also
uncertainty.bart_posterior_intervals for Bayesian posterior predictive bands.