"""Generate and render the exact Bayesian regression-model blog figures. The experiment results are checked in as ``bayesian_regression_results.pkl``. By default, this script derives the total-area rate and half-dataset tail area from those results, then generates the figures. Pass ``--regenerate`` to recompute the exact conjugate-Bayesian experiment first. Run from any directory with: python bayesian_regression.py python bayesian_regression.py --regenerate """ from __future__ import annotations import argparse import pickle from dataclasses import dataclass from pathlib import Path from typing import Any, Literal import matplotlib as mpl # These assets are generated non-interactively (including in CI and Codex). # Selecting the backend before importing pyplot avoids macOS AppKit crashes. mpl.use("Agg") import matplotlib.patheffects as path_effects import matplotlib.pyplot as plt import numpy as np from matplotlib.colors import BoundaryNorm, ListedColormap from matplotlib.lines import Line2D from matplotlib.patches import Patch ASSET_DIR = Path(__file__).resolve().parent DEFAULT_RESULTS = ASSET_DIR / "bayesian_regression_results.pkl" MODEL_COLORS = ("#0072B2", "#E69F00", "#009E73") GRID_COLOR = "#D8DEE4" TEXT_COLOR = "#24313A" MUTED_TEXT_COLOR = "#65727C" NATS_TO_BITS = 1 / np.log(2) N_SAMPLES = 3000 N_VALIDATION = 512 N_FEATURES = 64 N_TRUE_FEATURES = 8 DATA_SEED = 0 LONG_HORIZON_SAMPLES = 12000 LONG_HORIZON_SEED = 20260828 LONG_PHASE_RESOLUTION = 600 DEFAULT_TRIALS = 25 EVALUATION_SIZES = np.unique( np.concatenate( ( np.arange(1, 101), np.geomspace(101, N_SAMPLES - 1, 300).astype(int), ) ) ) EVALUATION_SIZE_SET = frozenset(int(size) for size in EVALUATION_SIZES) PREQUENTIAL_HALF_WINDOW = 25 @dataclass(frozen=True) class ModelSpec: short_name: str description: str color: str feature_map: Literal["linear", "additive_trees"] num_input_features: int prior_variance: float sigma_noise: float bins_per_feature: int | None = None @property def key(self) -> str: return ( f"BayesianRegressionModel(feature_map={self.feature_map}, " f"num_input_features={self.num_input_features}, " f"bins_per_feature={self.bins_per_feature}, " f"prior_variance={self.prior_variance:g}, " f"sigma_noise={self.sigma_noise:g})" ) @property def num_features(self) -> int: if self.feature_map == "linear": return self.num_input_features if self.bins_per_feature is None: raise ValueError("Tree features require bins_per_feature") return self.num_input_features * self.bins_per_feature @property def legend_label(self) -> str: return f"{self.short_name} · {self.description}" def transform_features(self, features: np.ndarray) -> np.ndarray: """Map raw inputs to this model's fixed Bayesian regression features.""" selected = features[:, : self.num_input_features] if self.feature_map == "linear": return selected if self.bins_per_feature is None: raise ValueError("Tree features require bins_per_feature") edges = np.linspace(0, 1, self.bins_per_feature + 1)[1:-1] leaves = [ np.eye(self.bins_per_feature)[np.digitize(selected[:, index], edges)] for index in range(self.num_input_features) ] # Normalizing the concatenated one-hot vectors makes the prior-predictive # variance equal to ``prior_variance`` rather than growing with the # number of additive trees. return np.concatenate(leaves, axis=1) / np.sqrt(self.num_input_features) MODEL_SPECS = ( ModelSpec( short_name="Model A", description="linear · d=64 · tight prior / low noise", color=MODEL_COLORS[0], feature_map="linear", num_input_features=64, prior_variance=0.005, sigma_noise=0.5, ), ModelSpec( short_name="Model B", description="tree ensemble · 8×6 leaves / high noise", color=MODEL_COLORS[1], feature_map="additive_trees", num_input_features=8, bins_per_feature=6, prior_variance=5, sigma_noise=0.8, ), ModelSpec( short_name="Model C", description="linear · d=8 · tight prior / medium noise", color=MODEL_COLORS[2], feature_map="linear", num_input_features=8, prior_variance=0.005, sigma_noise=0.625, ), ) METRIC_LAYOUT = ( ("marginal_cross_entropy_val", "Final height\n(validation loss)"), ("joint_marginal_information", "Total area\n(−log evidence)"), ( "conditional_joint_marginal_information_half", "Tail area\n(−CLML, last half)", ), ( "local_prequential_loss", "Local prequential loss\n(windowed)", ), ("joint_marginal_information_rate", "Average total area\n(per example)"), ("iterative_train_loss", "Training-speed proxy\n(current loss)"), ) def configure_style() -> None: """Use a restrained sketch style while keeping the exact data legible.""" mpl.rcParams.update( { "figure.facecolor": "white", "axes.facecolor": "white", "savefig.facecolor": "white", "font.family": "sans-serif", "font.sans-serif": [ "Comic Sans MS", "Chalkboard SE", "Marker Felt", "DejaVu Sans", ], "font.size": 9.5, "axes.titlesize": 10.8, "axes.titleweight": "bold", "axes.labelsize": 10.5, "axes.labelcolor": TEXT_COLOR, "axes.edgecolor": "#9AA5AE", "axes.linewidth": 0.8, "axes.spines.top": False, "axes.spines.right": False, "xtick.color": TEXT_COLOR, "ytick.color": TEXT_COLOR, "text.color": TEXT_COLOR, "legend.frameon": False, "lines.linewidth": 1.9, "svg.fonttype": "none", } ) def load_results(path: Path) -> dict[str, dict[str, Any]]: with path.open("rb") as handle: results = pickle.load(handle) missing_models = [spec.key for spec in MODEL_SPECS if spec.key not in results] if missing_models: available = ", ".join(results) missing = ", ".join(missing_models) raise KeyError(f"Missing models: {missing}. Available models: {available}") return results def gaussian_nll( targets: np.ndarray, means: np.ndarray, variances: np.ndarray, ) -> np.ndarray: """Return independent univariate Gaussian negative log probabilities.""" safe_variances = np.maximum(variances, np.finfo(float).tiny) return 0.5 * ( np.log(2 * np.pi * safe_variances) + np.square(targets - means) / safe_variances ) def generate_data() -> tuple[ np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, ]: """Generate one noiseless regression problem shared by every model.""" rng = np.random.RandomState(DATA_SEED) features = rng.rand(N_SAMPLES + N_VALIDATION, N_FEATURES) true_weights = np.zeros(N_FEATURES) true_weights[:N_TRUE_FEATURES] = rng.randn(N_TRUE_FEATURES) + 2 targets = features @ true_weights return ( features[:N_SAMPLES], targets[:N_SAMPLES], features[N_SAMPLES : N_SAMPLES + N_VALIDATION], targets[N_SAMPLES : N_SAMPLES + N_VALIDATION], true_weights, ) def asymptotic_loss_floor(spec: ModelSpec, true_weights: np.ndarray) -> float: """Return the population posterior-predictive loss floor in bits.""" approximation_mse = 0.0 if spec.feature_map == "additive_trees": if spec.bins_per_feature is None: raise ValueError("Tree features require bins_per_feature") represented_weights = true_weights[: spec.num_input_features] approximation_mse += float( np.sum(np.square(represented_weights)) / (12 * spec.bins_per_feature**2) ) omitted_weights = true_weights[spec.num_input_features :] if np.any(omitted_weights): approximation_mse += float(np.sum(np.square(omitted_weights)) / 12) if spec.feature_map == "linear": # There is no explicit intercept, but the retained U(0, 1) # regressors can partially represent the omitted term's mean. omitted_sum = float(np.sum(omitted_weights)) approximation_mse += omitted_sum**2 / ( 4 * (1 + 3 * spec.num_input_features) ) return float( 0.5 * ( np.log(2 * np.pi * spec.sigma_noise**2) + approximation_mse / spec.sigma_noise**2 ) * NATS_TO_BITS ) def fractional_fit_losses( spec: ModelSpec, train_features: np.ndarray, train_targets: np.ndarray, val_features: np.ndarray, val_targets: np.ndarray, ) -> tuple[dict[int, float], dict[int, float]]: """Evaluate the full-data fractional-posterior training-speed proxy.""" train_features = spec.transform_features(train_features) val_features = spec.transform_features(val_features) prior_standard_deviation = np.sqrt(spec.prior_variance) whitened_train = train_features * prior_standard_deviation whitened_val = val_features * prior_standard_deviation scaled_gram = whitened_train.T @ whitened_train / spec.sigma_noise**2 scaled_rhs = whitened_train.T @ train_targets / spec.sigma_noise**2 eigenvalues, eigenvectors = np.linalg.eigh(scaled_gram) rotated_rhs = eigenvectors.T @ scaled_rhs rotated_train = whitened_train @ eigenvectors rotated_val = whitened_val @ eigenvectors squared_rotated_train = np.square(rotated_train) squared_rotated_val = np.square(rotated_val) train_losses: dict[int, float] = {} val_losses: dict[int, float] = {} for size in EVALUATION_SIZES: size = int(size) fraction = size / N_SAMPLES posterior_precision = 1 + fraction * eigenvalues posterior_mean_rotated = fraction * rotated_rhs / posterior_precision train_means = rotated_train @ posterior_mean_rotated train_variances = ( spec.sigma_noise**2 + squared_rotated_train @ (1 / posterior_precision) ) val_means = rotated_val @ posterior_mean_rotated val_variances = ( spec.sigma_noise**2 + squared_rotated_val @ (1 / posterior_precision) ) train_losses[size] = float( np.mean(gaussian_nll(train_targets, train_means, train_variances)) ) val_losses[size] = float( np.mean(gaussian_nll(val_targets, val_means, val_variances)) ) return train_losses, val_losses def sequential_trial( spec: ModelSpec, train_features: np.ndarray, train_targets: np.ndarray, val_features: np.ndarray, val_targets: np.ndarray, permutation: np.ndarray, ) -> tuple[dict[int, float], dict[int, float]]: """Compute exact evidence and posterior-predictive losses for one ordering.""" train_features = spec.transform_features(train_features) val_features = spec.transform_features(val_features) ordered_features = train_features[permutation] ordered_targets = train_targets[permutation] posterior_mean = np.zeros(spec.num_features) posterior_covariance = np.eye(spec.num_features) * spec.prior_variance val_means = np.zeros(len(val_targets)) val_variances = ( spec.sigma_noise**2 + spec.prior_variance * np.sum(np.square(val_features), axis=1) ) cumulative_nll = 0.0 joint_information: dict[int, float] = {} val_predictive_loss: dict[int, float] = {} for index in range(N_SAMPLES - 1): feature = ordered_features[index] covariance_feature = posterior_covariance @ feature predictive_variance = max( float(spec.sigma_noise**2 + feature @ covariance_feature), np.finfo(float).tiny, ) residual = float(ordered_targets[index] - feature @ posterior_mean) cumulative_nll += 0.5 * ( np.log(2 * np.pi * predictive_variance) + residual**2 / predictive_variance ) val_projection = val_features @ covariance_feature update_scale = residual / predictive_variance posterior_mean += covariance_feature * update_scale val_means += val_projection * update_scale val_variances -= np.square(val_projection) / predictive_variance posterior_covariance -= ( np.outer(covariance_feature, covariance_feature) / predictive_variance ) size = index + 1 joint_information[size] = cumulative_nll if size in EVALUATION_SIZE_SET: val_predictive_loss[size] = float( np.mean(gaussian_nll(val_targets, val_means, val_variances)) ) return joint_information, val_predictive_loss def extend_training_data( train_features: np.ndarray, train_targets: np.ndarray, true_weights: np.ndarray, ) -> tuple[np.ndarray, np.ndarray]: """Extend the original i.i.d. problem for the long-horizon phase panel.""" extension_size = LONG_HORIZON_SAMPLES - N_SAMPLES if extension_size < 0: raise ValueError("Long horizon cannot be shorter than the main experiment") rng = np.random.RandomState(LONG_HORIZON_SEED) extension_features = rng.rand(extension_size, N_FEATURES) extension_targets = extension_features @ true_weights return ( np.concatenate((train_features, extension_features)), np.concatenate((train_targets, extension_targets)), ) def long_horizon_joint_information( spec: ModelSpec, train_features: np.ndarray, train_targets: np.ndarray, num_trials: int, ) -> dict[int, float]: """Average exact evidence along continuations of the main data orderings.""" transformed_features = spec.transform_features(train_features) total = np.zeros(LONG_HORIZON_SAMPLES) extension_size = LONG_HORIZON_SAMPLES - N_SAMPLES for trial in range(num_trials): initial_order = np.random.RandomState(trial + 31).permutation(N_SAMPLES) extension_order = N_SAMPLES + np.random.RandomState( trial + 31031 ).permutation(extension_size) permutation = np.concatenate((initial_order, extension_order)) ordered_features = transformed_features[permutation] ordered_targets = train_targets[permutation] posterior_mean = np.zeros(spec.num_features) posterior_covariance = np.eye(spec.num_features) * spec.prior_variance cumulative_nll = 0.0 for index in range(LONG_HORIZON_SAMPLES): feature = ordered_features[index] covariance_feature = posterior_covariance @ feature predictive_variance = max( float(spec.sigma_noise**2 + feature @ covariance_feature), np.finfo(float).tiny, ) residual = float(ordered_targets[index] - feature @ posterior_mean) cumulative_nll += 0.5 * ( np.log(2 * np.pi * predictive_variance) + residual**2 / predictive_variance ) posterior_mean += ( covariance_feature * residual / predictive_variance ) posterior_covariance -= ( np.outer(covariance_feature, covariance_feature) / predictive_variance ) total[index] += cumulative_nll mean = total / num_trials return { size: float(mean[size - 1]) for size in range(1, LONG_HORIZON_SAMPLES + 1) } def generate_results(num_trials: int) -> dict[str, dict[str, Any]]: """Run the shared-data model comparison and return serializable metrics.""" if num_trials < 1: raise ValueError("Trial count must be positive") train_features, train_targets, val_features, val_targets, true_weights = ( generate_data() ) long_features, long_targets = extend_training_data( train_features, train_targets, true_weights, ) results: dict[str, dict[str, Any]] = {} for spec in MODEL_SPECS: iterative_train, iterative_val = fractional_fit_losses( spec, train_features, train_targets, val_features, val_targets, ) metrics: dict[str, Any] = { "iterative_train_loss": iterative_train, "iterative_val_loss": iterative_val, "joint_marginal_information": {}, "long_joint_marginal_information": long_horizon_joint_information( spec, long_features, long_targets, num_trials, ), "conditional_joint_marginal_information_half": {}, "marginal_cross_entropy_val": {}, } for trial in range(num_trials): permutation = np.random.RandomState(trial + 31).permutation(N_SAMPLES) joint, val_loss = sequential_trial( spec, train_features, train_targets, val_features, val_targets, permutation, ) metrics["joint_marginal_information"][trial] = joint metrics["marginal_cross_entropy_val"][trial] = val_loss results[spec.key] = metrics return results def save_results(results: dict[str, dict[str, Any]], path: Path) -> None: """Write generated experiment metrics for reproducible figure rendering.""" with path.open("wb") as handle: pickle.dump(results, handle, protocol=pickle.HIGHEST_PROTOCOL) def derive_area_metrics(results: dict[str, dict[str, Any]]) -> None: """Derive tail area and total-area rate without changing the source file.""" for spec in MODEL_SPECS: metrics = results[spec.key] joint_by_trial = metrics["joint_marginal_information"] conditional_by_trial: dict[Any, dict[int, float]] = {} rate_by_trial: dict[Any, dict[int, float]] = {} local_prequential_by_trial: dict[Any, dict[int, float]] = {} for trial, series in joint_by_trial.items(): conditional_series: dict[int, float] = {} rate_series: dict[int, float] = {} for size, value in series.items(): size = int(size) rate_series[size] = float(value) / size conditioning_size = size // 2 if conditioning_size in series: conditional_series[size] = float(value) - float( series[conditioning_size] ) conditional_by_trial[trial] = conditional_series rate_by_trial[trial] = rate_series max_size = max(int(size) for size in series) local_prequential_by_trial[trial] = { int(size): ( float(series[min(max_size, int(size) + PREQUENTIAL_HALF_WINDOW)]) - ( float(series[int(size) - PREQUENTIAL_HALF_WINDOW]) if int(size) > PREQUENTIAL_HALF_WINDOW else 0.0 ) ) / ( min(max_size, int(size) + PREQUENTIAL_HALF_WINDOW) - max(0, int(size) - PREQUENTIAL_HALF_WINDOW) ) for size in EVALUATION_SIZES if int(size) <= max_size } metrics["conditional_joint_marginal_information_half"] = ( conditional_by_trial ) metrics["joint_marginal_information_rate"] = rate_by_trial metrics["local_prequential_loss"] = local_prequential_by_trial def summarize_metric( metrics: dict[str, Any], metric_name: str, ) -> tuple[np.ndarray, np.ndarray, np.ndarray]: """Return x, mean, and a normal-approximation 95% interval half-width.""" source = metrics[metric_name] if not source: raise ValueError(f"Metric {metric_name!r} is empty") first_value = next(iter(source.values())) if isinstance(first_value, dict): trials = list(source.values()) sizes = np.array( sorted({int(size) for trial in trials for size in trial}), dtype=int, ) values = np.array( [ [float(trial.get(int(size), np.nan)) for size in sizes] for trial in trials ] ) mean = np.nanmean(values, axis=0) counts = np.sum(np.isfinite(values), axis=0) if len(trials) > 1: standard_error = np.nanstd(values, axis=0, ddof=1) / np.sqrt(counts) interval = 1.96 * standard_error else: interval = np.zeros_like(mean) else: sizes = np.array(sorted(int(size) for size in source), dtype=int) mean = np.array([float(source[int(size)]) for size in sizes]) interval = np.zeros_like(mean) return sizes, mean * NATS_TO_BITS, interval * NATS_TO_BITS def save_figure(fig: mpl.figure.Figure, output_dir: Path, stem: str) -> None: output_dir.mkdir(parents=True, exist_ok=True) fig.savefig(output_dir / f"{stem}.png", dpi=180, bbox_inches="tight") svg_path = output_dir / f"{stem}.svg" fig.savefig(svg_path, bbox_inches="tight", metadata={"Date": None}) svg_path.write_text( "\n".join(line.rstrip() for line in svg_path.read_text().splitlines()) + "\n" ) plt.close(fig) def plot_metric_grid( results: dict[str, dict[str, Any]], output_dir: Path, ) -> None: _, _, _, _, true_weights = generate_data() fig, axes = plt.subplots( 2, 3, figsize=(12.2, 7.2), sharex=True, ) for ax, (metric_name, title) in zip(axes.flat, METRIC_LAYOUT): metric_curves: list[tuple[np.ndarray, np.ndarray]] = [] for spec in MODEL_SPECS: sizes, mean, interval = summarize_metric( results[spec.key], metric_name, ) valid = np.isfinite(mean) & (mean > 0) ax.plot(sizes[valid], mean[valid], color=spec.color) metric_curves.append((sizes[valid], mean[valid])) if np.any(interval[valid] > 0): lower = np.maximum( mean[valid] - interval[valid], np.finfo(float).tiny, ) upper = mean[valid] + interval[valid] ax.fill_between( sizes[valid], lower, upper, color=spec.color, alpha=0.14, linewidth=0, ) if metric_name == "marginal_cross_entropy_val": for spec in MODEL_SPECS: loss_floor = asymptotic_loss_floor(spec, true_weights) ax.axhline( loss_floor, color=spec.color, linestyle=(0, (1.5, 2.5)), linewidth=1.0, alpha=0.42, zorder=1, ) for left_index, right_index in ((0, 1), (0, 2), (1, 2)): sizes, left = metric_curves[left_index] right_sizes, right = metric_curves[right_index] if not np.array_equal(sizes, right_sizes): continue differences = left - right crossings = np.flatnonzero( np.signbit(differences[1:]) != np.signbit(differences[:-1]) ) for crossing in crossings: fraction = -differences[crossing] / ( differences[crossing + 1] - differences[crossing] ) crossing_size = sizes[crossing] + fraction * ( sizes[crossing + 1] - sizes[crossing] ) crossing_loss = left[crossing] + fraction * ( left[crossing + 1] - left[crossing] ) ax.scatter( crossing_size, crossing_loss, s=23, facecolor="white", edgecolor=TEXT_COLOR, linewidth=0.8, zorder=5, ) ax.set_title(title, pad=8) ax.set_yscale("log") ax.set_xlim(0, N_SAMPLES) ax.set_xticks(np.linspace(0, N_SAMPLES, 6, dtype=int)) ax.grid(axis="y", color=GRID_COLOR, linewidth=0.7, alpha=0.72) ax.set_axisbelow(True) handles = [ Line2D([0], [0], color=spec.color, label=spec.legend_label) for spec in MODEL_SPECS ] fig.legend( handles=handles, loc="upper center", bbox_to_anchor=(0.5, 0.9), ncol=3, fontsize=9.0, columnspacing=1.8, handlelength=2.2, ) fig.suptitle( "Exact Bayesian regression models: the criteria disagree", fontsize=15.2, fontweight="bold", y=0.98, ) fig.supxlabel("Dataset size, $N$", y=0.025) fig.supylabel("Loss or code length (bits, log scale)", x=0.012) fig.subplots_adjust( left=0.08, right=0.99, bottom=0.105, top=0.745, wspace=0.18, hspace=0.40, ) save_figure(fig, output_dir, "binary_regression_information_metrics") def mean_joint_information( results: dict[str, dict[str, Any]], spec: ModelSpec, ) -> dict[int, float]: trials = results[spec.key]["joint_marginal_information"] sizes = sorted({int(size) for trial in trials.values() for size in trial}) return { size: float( np.mean( [ float(trial[size]) for trial in trials.values() if size in trial ] ) ) for size in sizes } def compute_phase_diagram( results: dict[str, dict[str, Any]], ) -> tuple[np.ndarray, int]: """Find the model with minimum tail-area loss for every valid (N, k).""" means = [mean_joint_information(results, spec) for spec in MODEL_SPECS] max_size = min(max(series) for series in means) winners = np.full((max_size + 1, max_size + 1), np.nan) cumulative = np.array( [[0.0] + [series[size] for size in range(1, max_size + 1)] for series in means] ) for dataset_size in range(1, max_size + 1): conditioning_sizes = np.arange(dataset_size - 1, -1, -1) scores = ( cumulative[:, dataset_size, None] - cumulative[:, conditioning_sizes] ) winners[1 : dataset_size + 1, dataset_size] = np.argmin(scores, axis=0) return winners, max_size def compute_long_phase_diagram( results: dict[str, dict[str, Any]], ) -> tuple[np.ndarray, int]: """Sample the long-horizon tail-area winner on a square display grid.""" means = [ { int(size): float(value) for size, value in results[spec.key][ "long_joint_marginal_information" ].items() } for spec in MODEL_SPECS ] max_size = min(max(series) for series in means) cumulative = np.array( [ [0.0] + [series[size] for size in range(1, max_size + 1)] for series in means ] ) sampled_sizes = np.unique( np.linspace(1, max_size, LONG_PHASE_RESOLUTION).astype(int) ) winners = np.full((len(sampled_sizes), len(sampled_sizes)), np.nan) for column, dataset_size in enumerate(sampled_sizes): valid = sampled_sizes <= dataset_size tail_sizes = sampled_sizes[valid] scores = ( cumulative[:, dataset_size, None] - cumulative[:, dataset_size - tail_sizes] ) winners[valid, column] = np.argmin(scores, axis=0) return winners, max_size def eventual_all_a_size(results: dict[str, dict[str, Any]]) -> int | None: """Return the first N after which Model A wins every possible split.""" means = [ { int(size): float(value) for size, value in results[spec.key][ "long_joint_marginal_information" ].items() } for spec in MODEL_SPECS ] max_size = min(max(series) for series in means) cumulative = np.array( [ [0.0] + [series[size] for size in range(1, max_size + 1)] for series in means ] ) last_non_a = 0 for dataset_size in range(1, max_size + 1): conditioning_sizes = np.arange(dataset_size - 1, -1, -1) scores = ( cumulative[:, dataset_size, None] - cumulative[:, conditioning_sizes] ) if np.any(np.argmin(scores, axis=0) != 0): last_non_a = dataset_size if last_non_a == max_size: return None return last_non_a + 1 def draw_phase_panel( ax: mpl.axes.Axes, winners: np.ndarray, max_size: int, *, title: str, ) -> None: """Draw one triangular (N, k) model-selection phase panel.""" cmap = ListedColormap([spec.color for spec in MODEL_SPECS]) cmap.set_bad("#F4F6F7") norm = BoundaryNorm([-0.5, 0.5, 1.5, 2.5], cmap.N) ax.imshow( winners, origin="lower", extent=(0.5, max_size + 0.5, 0.5, max_size + 0.5), cmap=cmap, norm=norm, interpolation="nearest", aspect="equal", rasterized=True, ) dataset_sizes = np.array([1, max_size]) half_tail = dataset_sizes / 2 ax.plot(dataset_sizes, half_tail, color="white", linewidth=3.0, alpha=0.9) ax.plot( dataset_sizes, half_tail, color=TEXT_COLOR, linewidth=0.9, alpha=0.65, ) ax.set_xlim(0, max_size) ax.set_ylim(0, max_size) ax.set_xlabel("Dataset size, $N$") ax.set_title(title, pad=10) ax.grid(False) def plot_phase_diagram( results: dict[str, dict[str, Any]], output_dir: Path, ) -> None: winners, max_size = compute_phase_diagram(results) long_winners, long_max_size = compute_long_phase_diagram(results) all_a_size = eventual_all_a_size(results) fig, axes = plt.subplots( 1, 2, figsize=(12.4, 5.8), constrained_layout=True, ) draw_phase_panel( axes[0], winners[1:, 1:], max_size, title="Finite-data detail", ) draw_phase_panel( axes[1], long_winners, long_max_size, title="Long-horizon view", ) outlined_text = [ path_effects.withStroke(linewidth=3.5, foreground="white", alpha=0.95) ] axes[0].text( max_size * 0.59, max_size * 0.295 + 9, "$k=N/2$: score last half", color=TEXT_COLOR, fontsize=9.2, rotation=26, ha="center", va="bottom", path_effects=outlined_text, ) axes[0].text( max_size * 0.52, max_size * 0.02, "$k=1$: one-step loss", color=TEXT_COLOR, fontsize=9.2, ha="center", va="bottom", path_effects=outlined_text, ) axes[0].text( max_size * 0.67, max_size * 0.70, "$k=N$: total area", color=TEXT_COLOR, fontsize=9.2, rotation=45, ha="center", va="bottom", path_effects=outlined_text, ) axes[0].text( max_size * 0.29, max_size * 0.73, "invalid: $k>N$", color=MUTED_TEXT_COLOR, fontsize=9.5, ha="center", ) # Direct labels keep the model identity readable even without relying on # the colors alone. Their positions sit in the three broad stable regions. for label, x_fraction, y_fraction in ( ("Model A", 0.90, 0.10), ("Model B", 0.14, 0.125), ("Model C", 0.78, 0.28), ): axes[0].text( max_size * x_fraction, max_size * y_fraction, label, color=TEXT_COLOR, fontsize=10, fontweight="bold", ha="center", va="center", path_effects=outlined_text, ) if all_a_size is not None: rounded_all_a_size = round(all_a_size, -2) axes[1].plot( [all_a_size, all_a_size], [0, all_a_size], color="white", linewidth=2.5, linestyle=(0, (4, 3)), alpha=0.9, ) axes[1].plot( [all_a_size, all_a_size], [0, all_a_size], color=TEXT_COLOR, linewidth=0.8, linestyle=(0, (4, 3)), alpha=0.65, ) axes[1].text( all_a_size + long_max_size * 0.025, long_max_size * 0.48, f"A wins every split\nby $N\\approx{rounded_all_a_size:,}$", color=TEXT_COLOR, fontsize=9.5, fontweight="bold", ha="left", va="center", path_effects=outlined_text, ) axes[0].legend( handles=[ Patch(facecolor=spec.color, label=spec.short_name) for spec in MODEL_SPECS ], title="Selected model (lowest score)", loc="upper left", frameon=False, ) axes[0].set_ylabel("Scored tail size, $k$") fig.suptitle( "CLML spans one-step loss, tail scores, and total evidence", fontsize=15.2, fontweight="bold", ) # This is the one SVG embedded directly by the post. Outline its lettering # so the hand-drawn font remains stable on devices that do not have it. with mpl.rc_context({"svg.fonttype": "path"}): save_figure( fig, output_dir, "binary_regression_conditional_joint_marginal_information_decision_boundary", ) def parse_args() -> argparse.Namespace: parser = argparse.ArgumentParser(description=__doc__) parser.add_argument( "--results", type=Path, default=DEFAULT_RESULTS, help="Path to the checked-in experiment results.", ) parser.add_argument( "--output-dir", type=Path, default=ASSET_DIR, help="Directory in which to write PNG and SVG figures.", ) parser.add_argument( "--regenerate", action="store_true", help="Recompute and overwrite the checked-in experiment results.", ) parser.add_argument( "--trials", type=int, default=DEFAULT_TRIALS, help="Number of data-order trials used when regenerating.", ) return parser.parse_args() def main() -> None: args = parse_args() if args.regenerate: results = generate_results(args.trials) save_results(results, args.results) else: results = load_results(args.results) derive_area_metrics(results) with plt.xkcd(scale=0.42, length=120, randomness=0.9): configure_style() plot_metric_grid(results, args.output_dir) plot_phase_diagram(results, args.output_dir) if __name__ == "__main__": main()