Bayesian Model Criticism

This tutorial compares Bayesian models using predictive scores, empirical coverage, and score-based model weights.

The model is a conjugate Bayesian linear regression implemented in NumPy. It uses a closed-form posterior. The study uses the examples extra for scoring, design, optimization and plots.

Run it yourself

The full runnable script is at examples/bayesian_study.py.

uv run --extra examples python examples/bayesian_study.py

The problem

We want to fit a Bayesian regression \(y = a + bx + \varepsilon\) where \(\varepsilon \sim \mathcal{N}(0,\sigma_\text{true}^2)\) and the true coefficients are known (\(a=2\), \(b=3\), \(\sigma_\text{true}=0.5\)).

The study examines how prior variance, assumed noise, and sample size affect predictive scores and empirical coverage. Each configuration sets three factors:

Factor	Meaning
prior_var	Variance of the Normal prior on \((a, b)\)
noise_scale	Assumed observation noise \(\sigma\) (may disagree with truth)
n_obs	Number of training observations

Changing noise_scale affects posterior uncertainty and predictive intervals. A tight zero-centered prior (prior_var small) pulls coefficient estimates toward zero. More observations can reduce parameter uncertainty, but predictive quality also depends on the prior and noise assumptions.

Ground-truth model

# True data-generating process:  y = a + b*x + eps,  eps ~ N(0, sigma_true^2)
TRUE_A = 2.0  # intercept
TRUE_B = 3.0  # slope
SIGMA_TRUE = 0.5  # observation noise std

N_TEST = 50  # test locations for scoring
RNG_SEED = 42

rng = np.random.default_rng(RNG_SEED)
X_TEST = np.linspace(0.0, 1.0, N_TEST)
Y_TEST = TRUE_A + TRUE_B * X_TEST + rng.normal(0.0, SIGMA_TRUE, N_TEST)

Conjugate posterior

Because we use a Normal prior with known noise variance, the posterior over regression coefficients is available in closed form. Given design matrix \(\mathbf{X}\) and prior precision \(\boldsymbol{\Lambda}_0 = \sigma_\text{prior}^{-2}\mathbf{I}\):

\[ \boldsymbol{\Sigma}_n = \left(\boldsymbol{\Lambda}_0 + \frac{\mathbf{X}^\top\mathbf{X}}{\sigma^2}\right)^{-1} \qquad \boldsymbol{\mu}_n = \boldsymbol{\Sigma}_n \frac{\mathbf{X}^\top\mathbf{y}}{\sigma^2} \]

We draw posterior predictive samples by projecting coefficient draws through the test design matrix and adding observation noise.

def bayesian_regression(
    x_train: np.ndarray,
    y_train: np.ndarray,
    prior_var: float,
    noise_scale: float,
    n_samples: int = 500,
    seed: int = RNG_SEED,
) -> tuple[np.ndarray, np.ndarray]:
    """Fit a conjugate Bayesian linear regression and draw predictions.

    Prior:  beta = (a, b) ~ N(0, prior_var * I)
    Likelihood:  y | x, beta ~ N(X @ beta, noise_scale^2 * I)

    Args:
        x_train: Training inputs, shape (n,).
        y_train: Training targets, shape (n,).
        prior_var: Prior variance for each coefficient.
        noise_scale: Assumed observation noise standard deviation.
        n_samples: Number of posterior predictive draws.
        seed: Random seed.

    Returns:
        Tuple of (posterior_mean_at_test, predictive_samples_at_test).
        posterior_mean_at_test has shape (N_TEST,).
        predictive_samples_at_test has shape (N_TEST, n_samples).
    """
    gen = np.random.default_rng(seed)

    # Design matrices
    x_design = np.column_stack([np.ones_like(x_train), x_train])  # (n, 2)
    x_star = np.column_stack([np.ones(N_TEST), X_TEST])  # (N_TEST, 2)

    sigma2 = noise_scale**2
    prior_precision = np.eye(2) / prior_var  # (2, 2)

    # Conjugate posterior:  beta | y ~ N(mu_n, Sigma_n)
    gram = x_design.T @ x_design  # X^T X
    posterior_cov = np.linalg.inv(prior_precision + gram / sigma2)
    posterior_mean = posterior_cov @ (x_design.T @ y_train / sigma2)

    # Posterior predictive at test locations
    pred_mean = x_star @ posterior_mean  # (N_TEST,)

    # Draw beta samples, then project to predictions and add noise
    beta_samples = gen.multivariate_normal(
        posterior_mean,
        posterior_cov,
        size=n_samples,
    )  # (n_samples, 2)
    pred_samples = (x_star @ beta_samples.T) + gen.normal(
        0.0, noise_scale, (N_TEST, n_samples)
    )

    return pred_mean, pred_samples

Simulator and scorer

The simulator generates training data from the true model, fits the conjugate posterior, and returns posterior predictive samples at held-out test locations. The scorer evaluates five complementary metrics:

• CRPS (Continuous Ranked Probability Score), a univariate proper scoring rule.
• Energy score, a multivariate proper scoring rule complementary to CRPS.
• Coverage (95 %), the fraction of test points whose outcome falls inside the 95 % posterior predictive interval.
• RMSE, root mean squared error of the posterior mean.
• MAE, mean absolute error of the posterior mean.

Marginal CRPS/coverage use (test_location, member) samples. Joint energy uses (member, variable), so the scorer transposes the same array.

class BayesianRegressionSimulator:
    """Generate training data and compute the Bayesian posterior.

    Each config specifies prior hyperparameters and sample size.
    The "truth" is the held-out test set; "observations" are the
    posterior predictive samples at those test points.

    Args:
        n_samples: Number of posterior predictive draws.  Fewer draws
            give faster but noisier score estimates — useful as a cheap
            surrogate in a multi-fidelity workflow.
    """

    def __init__(self, n_samples: int = 500) -> None:
        """Initialise the simulator.

        Args:
            n_samples: Number of posterior predictive draws.
        """
        self.n_samples = n_samples

    def generate(self, config: dict[str, Any]) -> tuple[Any, Any]:
        """Draw training data, fit posterior, and return test predictions.

        Args:
            config: Must contain prior_var, noise_scale, n_obs.

        Returns:
            Tuple of (y_test, observation dict with predictive samples
            and posterior mean predictions).
        """
        gen = np.random.default_rng(RNG_SEED)
        n_obs = round(config["n_obs"])

        # Draw training data from the true model
        x_train = gen.uniform(0.0, 1.0, n_obs)
        y_train = TRUE_A + TRUE_B * x_train + gen.normal(0.0, SIGMA_TRUE, n_obs)

        # Compute conjugate posterior
        pred_mean, pred_samples = bayesian_regression(
            x_train,
            y_train,
            prior_var=config["prior_var"],
            noise_scale=config["noise_scale"],
            n_samples=self.n_samples,
        )

        observations = {
            "pred_mean": pred_mean,
            "pred_samples": pred_samples,
            "n_obs": n_obs,
        }
        return Y_TEST, observations


class BayesianRegressionScorer:
    """Score posterior predictions with proper scoring rules."""

    def score(
        self,
        truth: Any,
        observations: Any,
        config: dict[str, Any],
    ) -> dict[str, float]:
        """Compute CRPS, energy, 95 percent coverage, RMSE, and MAE.

        Args:
            truth: True test-set values (y_test).
            observations: Dict with pred_samples and pred_mean arrays.
            config: Hyperparameter dict (unused).

        Returns:
            Scores for crps, energy, coverage_95, rmse, and mae.
        """
        crps_val = score("crps", observations["pred_samples"], truth)
        # Energy scores the joint test vector: (ensemble member, variable),
        # whereas CRPS and coverage use (observation, ensemble member).
        energy_val = score("energy", observations["pred_samples"].T, truth)
        cov95 = score("coverage", observations["pred_samples"], truth, level=0.95)
        rmse_val = score("rmse", observations["pred_mean"], truth)
        mae_val = score("mae", observations["pred_mean"], truth)
        return {
            "crps": crps_val,
            "energy": energy_val,
            "coverage_95": cov95,
            "rmse": rmse_val,
            "mae": mae_val,
        }

Observables and factors

The example scales energy and coverage by 0.5 and MAE by 0.3. Positive scaling preserves Pareto membership; it does not express a preference during non-dominated sorting. A wide interval can achieve high empirical coverage while CRPS penalizes its spread. Use explicit weighted-sum or preference policies to choose among the five-objective alternatives.

observables = [
    Observable("crps", Direction.MINIMIZE),
    Observable("energy", Direction.MINIMIZE, weight=0.5),
    Observable("coverage_95", Direction.MAXIMIZE, weight=0.5),
    Observable("rmse", Direction.MINIMIZE),
    Observable("mae", Direction.MINIMIZE, weight=0.3),
]

factors = [
    Factor("prior_var", FactorType.CONTINUOUS, bounds=(0.1, 10.0)),
    Factor("noise_scale", FactorType.CONTINUOUS, bounds=(0.1, 2.0)),
    Factor("n_obs", FactorType.CONTINUOUS, bounds=(10.0, 200.0)),
]

Annotations

An Annotation attaches external metadata. This example uses a linear cost model for the number of observations:

compute_cost = Annotation(
    name="compute_cost",
    lookup=lambda n: float(n) * 0.01,
    key="n_obs",
)

Constraints

A Constraint defines a hard feasibility bound on an observable. feasibility_filter builds a phase filter that drops any design violating the constraints:

# Require at least 90% empirical coverage at the 95% nominal level
min_coverage = Constraint(
    name="min_coverage",
    observable="coverage_95",
    op=">=",
    threshold=0.90,
)

Factor screening

Before running the full grid, Morris screening identifies which factors influence the scores most. reduce_factors drops any factor whose \(\mu^*\) falls below a threshold:

    def run_fn(config: dict[str, Any]) -> dict[str, float]:
        """Compose simulator + scorer for screening.

        Returns:
            Score dictionary from the scorer.
        """
        config = {**config, "n_obs": round(config["n_obs"])}
        truth, obs = world.generate(config)
        return scorer.score(truth, obs, config)

    importance = screen(run_fn, factors, method="morris", n_trajectories=8, seed=42)
    print("Morris screening (mu_star):")
    for name, vals in importance.items():
        print(f"  {name}: {vals}")

    reduced = reduce_factors(factors, importance, threshold=0.01)
    print(f"\nImportant factors (threshold=0.01): {[f.name for f in reduced]}")

Grid evaluation

A 60-point Latin hypercube covers the three-factor space. run_grid evaluates every configuration and collects scores and annotations into a ResultsTable:

    grid = build_grid(factors, method="lhs", n_samples=60, seed=42)
    for cfg in grid:
        cfg["n_obs"] = round(cfg["n_obs"])

    print(f"\nLHS grid: {len(grid)} configurations")

    results = run_grid(
        world=world,
        scorer=scorer,
        grid=grid,
        observables=observables,
        annotations=[compute_cost],
    )

Pareto analysis

    print(f"\nPareto front: {len(front_idx)} / {results.scores.shape[0]} designs")

    header = (
        f"{'prior_var':>10s}  {'noise':>6s}  {'n_obs':>5s}  "
        f"{'CRPS':>6s}  {'Energy':>7s}  {'Cov95':>6s}  "
        f"{'RMSE':>6s}  {'MAE':>6s}  {'Cost':>5s}"
    )
    print(f"\n{header}")
    print("-" * len(header))
    for i in front_idx:
        cfg = results.configs[i]
        crps_val = results.scores[i, 0]
        energy_val = results.scores[i, 1]
        cov = results.scores[i, 2]
        rmse_val = results.scores[i, 3]
        mae_val = results.scores[i, 4]
        cost = results.annotations[i, 0] if results.annotations is not None else 0.0
        print(
            f"{cfg['prior_var']:10.2f}  {cfg['noise_scale']:6.2f}  "
            f"{cfg['n_obs']:5d}  "
            f"{crps_val:6.3f}  {energy_val:7.3f}  {cov:6.3f}  "
            f"{rmse_val:6.3f}  {mae_val:6.3f}  {cost:5.1f}"
        )

Pareto front scatter matrix

There are five objectives. plot_front displays pairwise projections of the first three (CRPS, energy and coverage), while marking membership of the full five-objective front. Projection alone does not determine that membership. Compare the five objectives jointly when interpreting a model choice.

Parallel coordinates

Each line shows a configuration. Compare its prior variance and assumed noise with its predictive scores and empirical coverage.

CRPS strip plot

Quality indicators

hypervolume measures the dominated volume of the Pareto front; igd_plus (inverted generational distance plus) measures how close the front is to a reference set. Both depend on the chosen objectives, scaling, and references:

    directions = [o.direction for o in observables]
    weights = [o.weight for o in observables]

    front_scores = results.scores[front_idx]

    # Build a synthetic reference front from per-objective best values
    n_obj = front_scores.shape[1]
    ideal = np.tile(np.median(front_scores, axis=0), (n_obj, 1))
    for j, d in enumerate(directions):
        ideal[j, j] = (
            front_scores[:, j].min()
            if d == Direction.MINIMIZE
            else front_scores[:, j].max()
        )

    ref_point = np.array([
        front_scores[:, j].max()
        if d == Direction.MINIMIZE
        else front_scores[:, j].min()
        for j, d in enumerate(directions)
    ])
    hv = hypervolume(front_scores, ref_point, directions, weights)
    igd = igd_plus(front_scores, ideal, directions, weights)
    print(f"\nHypervolume: {hv:.4f}")
    print(f"IGD+:        {igd:.4f}")

Feasibility filtering

Apply Constraint + feasibility_filter to keep only designs meeting hard bounds (here, coverage_95 >= 0.90):

    all_idx = np.arange(results.scores.shape[0])
    feas_filter = feasibility_filter(constraints=[min_coverage])
    kept = feas_filter(results, observables)
    print("\nFeasibility filter (coverage_95 >= 0.90):")
    print(f"  {len(all_idx)} total -> {len(kept)} feasible designs")

Weighted-sum filtering

weighted_sum_filter scalarises objectives into a single score via min-max normalisation and weighted combination, then keeps the top-K designs. Useful when you want a simple ranking without full Pareto analysis:

    ws_filter = weighted_sum_filter(
        weights={"crps": 1.0, "coverage_95": 0.5, "rmse": 0.8},
        k=8,
    )
    kept = ws_filter(results, observables)
    print("\nWeighted-sum filter (top 8 by scalarised score):")
    for i in kept[:5]:
        cfg = results.configs[i]
        print(
            f"  prior_var={cfg['prior_var']:.2f}, "
            f"noise={cfg['noise_scale']:.2f}, "
            f"n_obs={cfg['n_obs']}"
        )
    if len(kept) > 5:
        print(f"  ... and {len(kept) - 5} more")

Sobol grid

build_grid(method="sobol") generates a quasi-random Sobol design:

    sobol_grid = build_grid(factors, method="sobol", n_samples=40, seed=0)
    for cfg in sobol_grid:
        cfg["n_obs"] = round(cfg["n_obs"])

    print(f"\nSobol grid: {len(sobol_grid)} configurations")
    sobol_results = run_grid(
        world=world,
        scorer=scorer,
        grid=sobol_grid,
        observables=observables,
    )
    directions = [o.direction for o in observables]
    weights = [o.weight for o in observables]
    sobol_front = extract_front(sobol_results.scores, directions, weights)
    print(f"Sobol Pareto front: {len(sobol_front)} designs")

Study workflow

The Study class orchestrates multi-phase workflows. Study.summary() returns per-phase statistics and Study.stack() computes stacking weights directly:

    world = BayesianRegressionSimulator()
    scorer = BayesianRegressionScorer()

    grid = build_grid(factors, method="lhs", n_samples=60, seed=42)
    for cfg in grid:
        cfg["n_obs"] = round(cfg["n_obs"])

    study = Study(
        world=world,
        scorer=scorer,
        observables=observables,
        phases=[
            Phase(
                name="explore",
                grid=grid,
                filter_fn=feasibility_filter(
                    constraints=[min_coverage],
                ),
            ),
            Phase(name="rank", grid="carry"),
        ],
        annotations=[compute_cost],
    )
    study.run()

    # Study.summary() — per-phase statistics
    summary = study.summary()
    for phase_name, stats in summary.items():
        print(f"\n  Phase '{phase_name}':")
        print(f"    trials={stats['n_trials']}, front={stats['n_front']}")
        for obs_name, rng in stats["observable_ranges"].items():
            print(f"    {obs_name}: [{rng['min']:.4f}, {rng['max']:.4f}]")

    # Study.stack() — convenience method for score-based stacking
    stacking_w = study.stack("rank", maximize=False)
    print("Study.stack() weights (top entries):")
    for i, w in enumerate(stacking_w):
        if w > 0.01:
            print(f"  design {i}: {w:.3f}")

Score-based stacking

stack_scores minimizes a weighted average of the supplied per-model scores. Here those inputs are squared errors. This is a linear objective in the weights, so the solution selects models with the best mean input score, splitting weight uniformly across exactly tied best means. Selection is independent of positive score rescaling, subject to floating precision. See choosing ensemble weights for the supported inputs and tie policy. It does not optimize the squared error of the combined prediction or promise improvement over every model. The example computes that prediction and its RMSE separately. These weights use the same test values shown in the comparison; an independent evaluation set is needed to assess performance after selection:

    # Build a score matrix from the front: per-test-point squared errors
    front_predictions = []
    score_matrix_rows = []
    for i in front_idx:
        cfg = results.configs[i]
        _, obs = world.generate(cfg)
        front_predictions.append(obs["pred_mean"])
        sq_errors = (obs["pred_mean"] - Y_TEST) ** 2
        score_matrix_rows.append(sq_errors)

    score_mat = np.array(score_matrix_rows)  # (n_front, n_test)
    stacking_weights = stack_scores(score_mat, maximize=False)
    print("\nWeights from mean per-model squared errors:")
    for idx, w in zip(front_idx, stacking_weights, strict=True):
        if w > 0.01:
            cfg = results.configs[idx]
            print(
                f"  prior_var={cfg['prior_var']:.2f}, "
                f"noise={cfg['noise_scale']:.2f}, "
                f"n_obs={cfg['n_obs']}: w={w:.3f}"
            )

    ens_mean = ensemble_predict(front_predictions, stacking_weights)
    ens_rmse = float(np.sqrt(np.mean((ens_mean - Y_TEST) ** 2)))
    best_single = float(results.scores[front_idx, 3].min())  # rmse column
    print(f"\nBest single-model RMSE: {best_single:.4f}")
    print(f"Stacked ensemble RMSE:  {ens_rmse:.4f}")

Calibration assessment

coverage_curve evaluates empirical coverage at many nominal levels. A calibrated predictive distribution should track the diagonal across suitable repeated data. This illustration uses one test set after model selection; it is an empirical coverage plot, not a powered SBC or calibration certification:

    best_crps_idx = front_idx[int(np.argmin(results.scores[front_idx, 0]))]
    best_cfg = results.configs[best_crps_idx]
    _, best_obs = world.generate(best_cfg)
    nominal, empirical = coverage_curve(best_obs["pred_samples"], Y_TEST)
    pv = best_cfg["prior_var"]
    print(f"\nCalibration (best CRPS, prior_var={pv:.2f}):")
    for level in (0.5, 0.9, 0.95):
        closest = empirical[np.argmin(np.abs(nominal - level))]
        print(f"  Nominal {level:.0%}  -> Empirical {closest:.1%}")

Persistence

save_results writes the ResultsTable to disk; load_results reconstructs it. Useful for expensive studies where you want to separate evaluation from analysis:

    with tempfile.TemporaryDirectory() as tmp:
        save_results(results, tmp)
        loaded = load_results(tmp)
        n_saved = loaded.scores.shape[0]
        print(f"\nSaved and reloaded {n_saved} results successfully.")

Multi-fidelity workflow

MCMC fitting, fine-mesh numerical solvers, and agent-based models can be expensive to evaluate. A multi-phase study can screen designs with a cheaper calculation and re-evaluate a selected subset with a more detailed one.

Phase supports this via optional world and scorer overrides. When set, a phase uses its own simulator instead of the Study-level default. Here the first simulator draws 50 posterior samples, while the second phase draws 2 000:

    # Cheap surrogate: only 50 posterior draws (fast, noisier scores)
    cheap_world = BayesianRegressionSimulator(n_samples=50)
    # Expensive model: 2000 posterior draws (slow, precise scores)
    expensive_world = BayesianRegressionSimulator(n_samples=2000)

    # Phase-level scorer override: use a separate scorer for screening
    screen_scorer = BayesianRegressionScorer()
    validate_scorer = BayesianRegressionScorer()

    grid = build_grid(factors, method="lhs", n_samples=60, seed=42)
    for cfg in grid:
        cfg["n_obs"] = round(cfg["n_obs"])

    study = Study(
        world=expensive_world,
        scorer=validate_scorer,
        observables=observables,
        phases=[
            # Phase 1: screen with cheap world AND its own scorer
            Phase(
                name="screen",
                grid=grid,
                world=cheap_world,
                scorer=screen_scorer,
                filter_fn=top_k_pareto_filter(k=10),
            ),
            # Phase 2: validate top 10 with the expensive model + default scorer
            Phase(name="validate", grid="carry"),
        ],
        annotations=[compute_cost],
    )
    study.run()

    screen_r = study.results("screen")
    validate_r = study.results("validate")
    print("\nMulti-fidelity study:")
    print(f"  Screen phase:   {screen_r.scores.shape[0]} designs (50 draws)")
    print(f"  Validate phase: {validate_r.scores.shape[0]} designs (2000 draws)")

    directions = [o.direction for o in observables]
    weights = [o.weight for o in observables]
    front_idx = extract_front(validate_r.scores, directions, weights)
    print(f"  Final Pareto front: {len(front_idx)} designs")

The Study orchestrates both phases: the first screens 60 designs with the cheap surrogate and keeps the top 10 by Pareto rank, then the second re-evaluates those 10 designs with the expensive model.

This pattern applies whenever fidelity is a computational strategy rather than a design factor:

• Epidemiology. Screen surveillance designs with a deterministic model and re-evaluate a subset with a stochastic agent-based model.
• Engineering. Explore with a coarse mesh and re-evaluate a subset on a finer mesh.
• Forecast model assessment. Use approximate inference for screening and HMC for a later phase.

Comparisons across fidelities need compatible targets and explicit model assumptions. A later phase alone does not validate those assumptions.

What to try next

• Try stack_bayesian() on models that expose log-likelihood (requires arviz).
• Add more phases to the Study, such as a third phase that re-evaluates with even higher fidelity.
• Experiment with different weighted_sum_filter weight vectors to explore different trade-off preferences.