emutrust 0.1.0__py3-none-any.whl

emutrust/__init__.py

@@ -0,0 +1,23 @@
+"""emutrust: neural network emulators for slow scientific models, with
+calibrated uncertainty you can trust in Bayesian inference."""
+
+from .emulator import Emulator
+from .likelihood import GaussianLikelihood
+from .runner import TrainingSet, load_runs, run_model
+from .space import ParameterSpace, latin_hypercube
+from .validation import ValidationReport, validate
+
+__version__ = "0.1.0"
+
+__all__ = [
+    "Emulator",
+    "GaussianLikelihood",
+    "ParameterSpace",
+    "TrainingSet",
+    "ValidationReport",
+    "latin_hypercube",
+    "load_runs",
+    "run_model",
+    "validate",
+    "__version__",
+]

emutrust/emulator.py

@@ -0,0 +1,311 @@
+"""Neural network emulators with ensemble uncertainty."""
+
+from __future__ import annotations
+
+import copy
+from collections.abc import Mapping, Sequence
+from pathlib import Path
+from typing import Union
+
+import numpy as np
+import torch
+from torch import nn
+
+from .runner import TrainingSet
+
+PathLike = Union[str, Path]
+
+
+class Standardizer:
+    """Shift and scale each column to mean 0 and standard deviation 1.
+
+    Neural networks train best when every input and output has a similar
+    size; without this, a parameter like H0 (~70) would swamp one like
+    omch2 (~0.12).
+    """
+
+    def __init__(self, mean: np.ndarray, std: np.ndarray):
+        self.mean = np.asarray(mean, dtype=float)
+        self.std = np.asarray(std, dtype=float)
+
+    @classmethod
+    def fit(cls, data: np.ndarray) -> Standardizer:
+        return cls(data.mean(axis=0), data.std(axis=0))
+
+    def transform(self, data: np.ndarray) -> np.ndarray:
+        # A column that never changes has std 0; it maps to 0 instead of
+        # dividing by zero.
+        return (data - self.mean) / np.where(self.std > 0, self.std, 1.0)
+
+    def inverse(self, data: np.ndarray) -> np.ndarray:
+        # Multiplying by std 0 returns a constant column exactly, whatever
+        # the network outputs for it.
+        return data * self.std + self.mean
+
+
+def _should_use_log(Y: np.ndarray) -> bool:
+    """Use log outputs when every output is positive and spans >2 decades."""
+    if np.any(Y <= 0):
+        return False
+    return bool(np.max(np.log10(Y.max(axis=0) / Y.min(axis=0))) > 2.0)
+
+
+def _make_network(n_in: int, n_out: int, hidden: Sequence[int]) -> nn.Sequential:
+    layers: list[nn.Module] = []
+    width = n_in
+    for h in hidden:
+        layers += [nn.Linear(width, h), nn.SiLU()]
+        width = h
+    layers.append(nn.Linear(width, n_out))
+    return nn.Sequential(*layers)
+
+
+class Emulator:
+    """A fast stand-in for a slow model, built from an ensemble of networks.
+
+    Several networks are trained on the same training set, each starting
+    from different random weights and seeing a different random split of
+    the data. Where the training data pin the answer down, they agree;
+    where data are thin or the function is hard to learn, they disagree.
+    The spread of their predictions is the emulator's estimate of its own
+    error.
+
+    Example
+    -------
+    >>> emu = Emulator.fit(training, n_members=5, seed=0)
+    >>> mean, std = emu.predict({"H0": 67.0, "omch2": 0.12, "ns": 0.96})
+    """
+
+    def __init__(
+        self,
+        members: list[nn.Sequential],
+        x_scaler: Standardizer,
+        y_scaler: Standardizer,
+        log_outputs: bool,
+        names: tuple[str, ...],
+        config: dict,
+        std_scale: float = 1.0,
+    ):
+        self.members = members
+        self.x_scaler = x_scaler
+        self.y_scaler = y_scaler
+        self.log_outputs = log_outputs
+        self.names = tuple(names)
+        self.config = config
+        #: Multiplies every predicted std. 1.0 means the raw ensemble spread;
+        #: set it from a validation report if the spread is too small or large.
+        self.std_scale = float(std_scale)
+        for net in self.members:
+            net.eval()
+
+    @property
+    def bounds(self) -> tuple[np.ndarray, np.ndarray]:
+        """Lowest and highest value of each parameter in the training set.
+
+        The emulator has only seen the model inside this box; predictions
+        outside it are extrapolations and should not be trusted.
+        """
+        return np.asarray(self.config["train_low"]), np.asarray(self.config["train_high"])
+
+    @property
+    def n_members(self) -> int:
+        return len(self.members)
+
+    # ------------------------------------------------------------------ fit
+    @classmethod
+    def fit(
+        cls,
+        training: TrainingSet,
+        n_members: int = 5,
+        hidden: Sequence[int] = (128, 128, 128),
+        epochs: int = 2000,
+        batch_size: int = 64,
+        learning_rate: float = 1e-3,
+        val_fraction: float = 0.1,
+        patience: int = 200,
+        log_outputs: bool | None = None,
+        seed: int | None = None,
+        verbose: bool = False,
+    ) -> Emulator:
+        """Train an ensemble of ``n_members`` networks on a training set.
+
+        ``log_outputs`` controls whether the networks learn ``log10`` of the
+        outputs. ``None`` (the default) decides automatically: log is used
+        when every output is positive and spans more than two orders of
+        magnitude, as power spectra do.
+
+        Each member trains until its error on its own held-out
+        ``val_fraction`` of the data has not improved for ``patience``
+        epochs, and keeps the weights from its best epoch ("early stopping",
+        which prevents memorising the training points).
+        """
+        X = np.asarray(training.X, dtype=float)
+        Y = np.asarray(training.Y, dtype=float)
+        if n_members < 1:
+            raise ValueError("n_members must be at least 1")
+        if len(X) < 10:
+            raise ValueError(f"Need at least 10 training points, got {len(X)}")
+        if not 0 < val_fraction < 0.5:
+            raise ValueError("val_fraction must be between 0 and 0.5")
+
+        if log_outputs is None:
+            log_outputs = _should_use_log(Y)
+        elif log_outputs and np.any(Y <= 0):
+            raise ValueError("log_outputs=True needs strictly positive outputs")
+        Yt = np.log10(Y) if log_outputs else Y
+
+        x_scaler = Standardizer.fit(X)
+        y_scaler = Standardizer.fit(Yt)
+        Xs = torch.tensor(x_scaler.transform(X), dtype=torch.float32)
+        Ys = torch.tensor(y_scaler.transform(Yt), dtype=torch.float32)
+
+        rng = np.random.default_rng(seed)
+        members = []
+        for m in range(n_members):
+            member_seed = int(rng.integers(2**31))
+            net, best = _train_member(
+                Xs,
+                Ys,
+                hidden,
+                epochs,
+                batch_size,
+                learning_rate,
+                val_fraction,
+                patience,
+                member_seed,
+            )
+            if verbose:
+                print(f"  member {m + 1}/{n_members}: best validation loss {best:.2e}")
+            members.append(net)
+
+        config = {
+            "hidden": list(hidden),
+            "n_inputs": X.shape[1],
+            "n_outputs": Y.shape[1],
+            "n_train": len(X),
+            "train_low": X.min(axis=0).tolist(),
+            "train_high": X.max(axis=0).tolist(),
+        }
+        return cls(members, x_scaler, y_scaler, bool(log_outputs), training.names, config)
+
+    # -------------------------------------------------------------- predict
+    def _as_array(self, x) -> np.ndarray:
+        if isinstance(x, Mapping):
+            missing = set(self.names) - set(x)
+            if missing:
+                raise ValueError(f"Missing parameters: {sorted(missing)}")
+            x = [x[n] for n in self.names]
+        x = np.atleast_2d(np.asarray(x, dtype=float))
+        if x.shape[1] != len(self.names):
+            raise ValueError(f"Expected {len(self.names)} parameters, got {x.shape[1]}")
+        return x
+
+    def predict_members(self, x) -> np.ndarray:
+        """Each member's prediction, shape ``(n_members, n_points, n_outputs)``."""
+        xs = torch.tensor(self.x_scaler.transform(self._as_array(x)), dtype=torch.float32)
+        with torch.no_grad():
+            out = np.stack([net(xs).numpy() for net in self.members]).astype(float)
+        out = self.y_scaler.inverse(out)
+        return 10.0**out if self.log_outputs else out
+
+    def predict(self, x, return_std: bool = True):
+        """Predict outputs at one point (a dict or a 1-D array) or many points.
+
+        Returns ``(mean, std)``: the ensemble average and the spread between
+        members (times :attr:`std_scale`), both in the model's original output
+        units. A single input point gives 1-D arrays; many points give 2-D
+        arrays.
+        """
+        single = isinstance(x, Mapping) or np.ndim(x) == 1
+        preds = self.predict_members(x)
+        mean = preds.mean(axis=0)
+        std = preds.std(axis=0, ddof=1) if self.n_members > 1 else np.zeros_like(mean)
+        std = std * self.std_scale
+        if single:
+            mean, std = mean[0], std[0]
+        return (mean, std) if return_std else mean
+
+    def __call__(self, x):
+        return self.predict(x, return_std=False)
+
+    # ----------------------------------------------------------- save/load
+    def save(self, path: PathLike) -> None:
+        """Save the emulator to a single file."""
+        torch.save(
+            {
+                "members": [net.state_dict() for net in self.members],
+                "x_scaler": (self.x_scaler.mean, self.x_scaler.std),
+                "y_scaler": (self.y_scaler.mean, self.y_scaler.std),
+                "log_outputs": self.log_outputs,
+                "names": self.names,
+                "config": self.config,
+                "std_scale": self.std_scale,
+            },
+            path,
+        )
+
+    @classmethod
+    def load(cls, path: PathLike) -> Emulator:
+        """Load an emulator saved with :meth:`save`."""
+        state = torch.load(path, weights_only=False)
+        cfg = state["config"]
+        members = []
+        for sd in state["members"]:
+            net = _make_network(cfg["n_inputs"], cfg["n_outputs"], cfg["hidden"])
+            net.load_state_dict(sd)
+            members.append(net)
+        return cls(
+            members,
+            Standardizer(*state["x_scaler"]),
+            Standardizer(*state["y_scaler"]),
+            state["log_outputs"],
+            tuple(state["names"]),
+            cfg,
+            state.get("std_scale", 1.0),
+        )
+
+
+def _train_member(Xs, Ys, hidden, epochs, batch_size, lr, val_fraction, patience, seed):
+    torch.manual_seed(seed)
+    gen = torch.Generator().manual_seed(seed)
+    n = len(Xs)
+    perm = torch.randperm(n, generator=gen)
+    n_val = max(1, int(round(val_fraction * n)))
+    val_idx, train_idx = perm[:n_val], perm[n_val:]
+    Xtr, Ytr, Xval, Yval = Xs[train_idx], Ys[train_idx], Xs[val_idx], Ys[val_idx]
+
+    net = _make_network(Xs.shape[1], Ys.shape[1], hidden)
+    opt = torch.optim.Adam(net.parameters(), lr=lr)
+    sched = torch.optim.lr_scheduler.ReduceLROnPlateau(
+        opt, factor=0.5, patience=max(patience // 4, 1)
+    )
+    loss_fn = nn.MSELoss()
+
+    best_loss, best_state, since_best = float("inf"), None, 0
+    for _ in range(epochs):
+        net.train()
+        order = torch.randperm(len(Xtr), generator=gen)
+        for start in range(0, len(Xtr), batch_size):
+            idx = order[start : start + batch_size]
+            opt.zero_grad()
+            loss = loss_fn(net(Xtr[idx]), Ytr[idx])
+            loss.backward()
+            opt.step()
+
+        net.eval()
+        with torch.no_grad():
+            val_loss = loss_fn(net(Xval), Yval).item()
+        sched.step(val_loss)
+        if val_loss < best_loss:
+            best_loss, best_state, since_best = val_loss, copy.deepcopy(net.state_dict()), 0
+        else:
+            since_best += 1
+            if since_best >= patience:
+                break
+
+    net.load_state_dict(best_state)
+    net.eval()
+    return net, best_loss
+
+
+__all__ = ["Emulator", "Standardizer"]

emutrust/likelihood.py

@@ -0,0 +1,183 @@
+"""Likelihoods that include the emulator's own uncertainty."""
+
+from __future__ import annotations
+
+from typing import Union
+
+import numpy as np
+
+from .emulator import Emulator
+from .validation import ValidationReport
+
+_LOG_2PI = np.log(2 * np.pi)
+
+EmulatorError = Union[str, ValidationReport, None]
+
+
+class GaussianLikelihood:
+    """Gaussian likelihood of observed data, with an emulator as the model.
+
+    For parameters ``theta`` the emulator predicts the model output
+    ``mu(theta)``. The likelihood compares the data with ``mu`` using a
+    covariance that adds the data's noise and the emulator's own error::
+
+        C = C_noise + C_emulator
+
+    so an imperfect emulator widens the result instead of silently
+    shrinking the error bars. ``emulator_error`` chooses ``C_emulator``:
+
+    ``ValidationReport`` (recommended)
+        The error pattern measured on validation runs the emulator never
+        saw, including correlations between outputs. Emulator errors are
+        usually systematic (too high at one output means too high at the
+        next), and with many data points that correlated part dominates;
+        this option captures it.
+    ``"ensemble"``
+        The spread of the ensemble members at ``theta``, treating each
+        output's error as independent. Needs no extra model runs, but can
+        underestimate correlated errors when there are many outputs.
+    ``None``
+        Treat the emulator as exact. Only for comparison: this is the
+        overconfident likelihood that ``emutrust`` exists to avoid.
+
+    For validation-measured error, ``error_variance_kept`` sets how much of
+    the measured error variance is described by its main correlated
+    patterns (default 99.9%); the rest is kept as independent per-output
+    error. This makes each evaluation much faster with a negligible change
+    in the result. Use ``1.0`` to keep every pattern.
+
+    Parameters outside the range the emulator was trained on get a
+    log-likelihood of ``-inf``, so a sampler never relies on extrapolation.
+
+    Example
+    -------
+    >>> report = validate(emu, validation_runs)
+    >>> like = GaussianLikelihood(emu, data, noise_std, emulator_error=report)
+    >>> sampler = emcee.EnsembleSampler(32, like.ndim, like, vectorize=True)
+    """
+
+    def __init__(
+        self,
+        emulator: Emulator,
+        data: np.ndarray,
+        noise_std: float | np.ndarray,
+        emulator_error: EmulatorError = "ensemble",
+        error_variance_kept: float = 0.999,
+    ):
+        self.emulator = emulator
+        self.data = np.asarray(data, dtype=float).ravel()
+        n_out = emulator.config["n_outputs"]
+        if self.data.size != n_out:
+            raise ValueError(f"data has {self.data.size} values but the emulator predicts {n_out}")
+        noise = np.broadcast_to(np.asarray(noise_std, dtype=float), self.data.shape)
+        if np.any(noise <= 0):
+            raise ValueError("noise_std must be positive")
+        self.noise_var = noise**2
+
+        if isinstance(emulator_error, ValidationReport):
+            if emulator_error.names != emulator.names:
+                raise ValueError("The validation report belongs to a different emulator")
+            self._kind, matrix = emulator_error.error_covariance()
+            self._set_error_modes(matrix, error_variance_kept)
+            self.emulator_error = "validation"
+        elif emulator_error in ("ensemble", None):
+            self.emulator_error = emulator_error
+        else:
+            raise ValueError("emulator_error must be a ValidationReport, 'ensemble' or None")
+        self.low, self.high = emulator.bounds
+
+    def _set_error_modes(self, matrix: np.ndarray, kept: float) -> None:
+        """Split the measured error covariance into its main patterns.
+
+        The emulator's errors are dominated by a few patterns across the
+        outputs. Keeping the largest eigenvectors that together explain a
+        fraction ``kept`` of the error variance, and treating the small
+        remainder as independent per-output error, gives almost exactly the
+        same likelihood far faster: with ``K`` patterns and ``n`` outputs,
+        each evaluation costs about ``n*K**2`` instead of ``n**3``. The
+        remainder is kept on the diagonal, so every output's total error
+        variance is preserved exactly.
+        """
+        if not 0 < kept <= 1:
+            raise ValueError("error_variance_kept must be in (0, 1]")
+        values, vectors = np.linalg.eigh(matrix)
+        values, vectors = values[::-1], vectors[:, ::-1]
+        values = np.clip(values, 0.0, None)
+        total = values.sum()
+        if total == 0:
+            k = 0
+        else:
+            k = int(np.searchsorted(np.cumsum(values) / total, kept - 1e-12) + 1)
+            k = min(k, int(np.sum(values > 1e-14 * values[0])))
+        self._modes = vectors[:, :k]  # (n, K)
+        self._mode_var = values[:k]  # (K,)
+        kept_diag = np.sum(self._modes**2 * self._mode_var, axis=1)
+        self._remainder_var = np.clip(np.diag(matrix) - kept_diag, 0.0, None)  # (n,)
+
+    @property
+    def n_error_modes(self) -> int:
+        """Number of error patterns used (validation-measured error only)."""
+        return self._modes.shape[1] if self.emulator_error == "validation" else 0
+
+    @property
+    def ndim(self) -> int:
+        return len(self.emulator.names)
+
+    def in_bounds(self, theta: np.ndarray) -> np.ndarray:
+        theta = np.atleast_2d(theta)
+        return np.all((theta >= self.low) & (theta <= self.high), axis=1)
+
+    def log_likelihood(self, theta: np.ndarray) -> np.ndarray | float:
+        """Log-likelihood at one parameter vector, or at each row of many."""
+        theta = np.asarray(theta, dtype=float)
+        single = theta.ndim == 1
+        theta = np.atleast_2d(theta)
+        out = np.full(len(theta), -np.inf)
+        ok = self.in_bounds(theta) & np.all(np.isfinite(theta), axis=1)
+        if ok.any():
+            out[ok] = self._log_gaussian(theta[ok])
+        return float(out[0]) if single else out
+
+    def log_prob(self, theta: np.ndarray) -> np.ndarray | float:
+        """Log-posterior with a flat prior over the emulator's training range.
+
+        Accepts one parameter vector or a batch of them, so it can be used
+        with ``emcee.EnsembleSampler(..., vectorize=True)`` for speed.
+        """
+        return self.log_likelihood(theta)
+
+    __call__ = log_prob
+
+    def _log_gaussian(self, theta: np.ndarray) -> np.ndarray:
+        mean, std = self.emulator.predict(theta)
+        resid = self.data - mean
+
+        if self.emulator_error != "validation":
+            var = self.noise_var + (std**2 if self.emulator_error == "ensemble" else 0.0)
+            return -0.5 * np.sum(resid**2 / var + np.log(var) + _LOG_2PI, axis=1)
+
+        # C = D + A diag(lam) A^T, with D diagonal (noise plus the small
+        # remainder of the emulator error) and A the K main error patterns.
+        # For relative errors the patterns scale with the prediction.
+        n = resid.shape[1]
+        scale = mean if self._kind == "relative" else np.ones_like(mean)
+        d = self.noise_var + scale**2 * self._remainder_var  # (P, n)
+        quad = np.sum(resid**2 / d, axis=1)
+        logdet = np.sum(np.log(d), axis=1)
+        if self._modes.shape[1] > 0:
+            # Woodbury identity and matrix determinant lemma: only a K x K
+            # matrix per point needs factorising.
+            a = scale[:, :, None] * self._modes  # (P, n, K)
+            a_d_t = np.swapaxes(a / d[:, :, None], 1, 2)  # (P, K, n)
+            small = a_d_t @ a  # (P, K, K), batched matrix product (BLAS)
+            k = small.shape[1]
+            small[:, np.arange(k), np.arange(k)] += 1.0 / self._mode_var
+            chol = np.linalg.cholesky(small)
+            u = np.linalg.solve(chol, a_d_t @ resid[..., None])[..., 0]
+            quad -= np.sum(u**2, axis=1)
+            logdet += 2 * np.sum(np.log(np.diagonal(chol, axis1=1, axis2=2)), axis=1)
+            logdet += np.sum(np.log(self._mode_var))
+        return -0.5 * (quad + logdet + n * _LOG_2PI)
+
+
+__all__ = ["GaussianLikelihood"]

emutrust/runner.py

@@ -0,0 +1,151 @@
+"""Run a slow model over many parameter points, with checkpointing."""
+
+from __future__ import annotations
+
+import os
+import time
+from collections.abc import Mapping
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Union
+
+import numpy as np
+
+from .space import ParameterSpace, as_space
+
+Model = Callable[[dict], Any]
+PathLike = Union[str, "os.PathLike[str]"]
+
+
+@dataclass
+class TrainingSet:
+    """Inputs ``X`` (n, n_params) and model outputs ``Y`` (n, n_outputs)."""
+
+    X: np.ndarray
+    Y: np.ndarray
+    names: tuple[str, ...]
+    seconds_per_run: float
+
+    def __len__(self) -> int:
+        return len(self.X)
+
+
+def run_model(
+    model: Model,
+    points: np.ndarray,
+    space: ParameterSpace | Mapping[str, tuple[float, float]],
+    checkpoint: PathLike | None = None,
+    progress: bool = True,
+) -> TrainingSet:
+    """Evaluate ``model`` at every row of ``points``.
+
+    ``model`` is any Python function that takes a ``{name: value}`` dict and
+    returns an array (or list, or number) of outputs. Every call must return
+    the same number of outputs.
+
+    If ``checkpoint`` is a file path, results are saved after every run, and
+    calling ``run_model`` again with the same points resumes where it stopped
+    instead of starting over. This matters when each run takes seconds or
+    minutes and the session can crash or time out.
+    """
+    space = as_space(space)
+    points = np.atleast_2d(np.asarray(points, dtype=float))
+    if points.shape[1] != space.dim:
+        raise ValueError(
+            f"points have {points.shape[1]} columns but the space has {space.dim} parameters"
+        )
+    if not np.all(space.contains(points)):
+        raise ValueError("Some points lie outside the parameter ranges.")
+
+    path = Path(checkpoint) if checkpoint is not None else None
+    outputs: list[np.ndarray] = []
+    elapsed = 0.0
+    if path is not None and path.exists():
+        outputs, elapsed = _load_checkpoint(path, points, space)
+        if progress and outputs:
+            print(f"Resuming from checkpoint: {len(outputs)}/{len(points)} runs already done.")
+
+    n_out = outputs[0].size if outputs else None
+    for i in range(len(outputs), len(points)):
+        t0 = time.perf_counter()
+        y = np.asarray(model(space.to_dict(points[i])), dtype=float).ravel()
+        elapsed += time.perf_counter() - t0
+
+        if n_out is None:
+            n_out = y.size
+        if y.size != n_out:
+            raise ValueError(
+                f"Run {i} returned {y.size} outputs, but earlier runs returned {n_out}."
+            )
+        if not np.all(np.isfinite(y)):
+            raise ValueError(
+                f"Run {i} returned non-finite values (NaN or inf) at {space.to_dict(points[i])}."
+            )
+        outputs.append(y)
+
+        if path is not None:
+            _save_checkpoint(path, points, outputs, elapsed, space)
+        if progress:
+            print(f"  run {i + 1}/{len(points)} done ({elapsed / len(outputs):.2f} s per run)")
+
+    return TrainingSet(
+        X=points.copy(),
+        Y=np.vstack(outputs),
+        names=space.names,
+        seconds_per_run=elapsed / max(len(outputs), 1),
+    )
+
+
+def _save_checkpoint(path: Path, points, outputs, elapsed, space) -> None:
+    # Write to a temporary file and then rename it, so a crash in the middle
+    # of saving can never leave a half-written, corrupted checkpoint behind.
+    path.parent.mkdir(parents=True, exist_ok=True)
+    tmp = path.with_name(path.name + ".tmp")
+    with open(tmp, "wb") as f:
+        np.savez(
+            f,
+            points=points,
+            outputs=np.vstack(outputs),
+            elapsed=elapsed,
+            names=np.array(space.names),
+        )
+    os.replace(tmp, path)
+
+
+def _load_checkpoint(path: Path, points, space):
+    with np.load(path) as data:
+        saved_points = data["points"]
+        saved_outputs = data["outputs"]
+        elapsed = float(data["elapsed"])
+        saved_names = tuple(str(n) for n in data["names"])
+    if (
+        saved_names != space.names
+        or saved_points.shape != points.shape
+        or not np.allclose(saved_points, points)
+    ):
+        raise ValueError(
+            f"The checkpoint at {path} belongs to a different set of points or parameters. "
+            "Use a new checkpoint path, or delete the old file to start over."
+        )
+    return [row for row in saved_outputs], elapsed
+
+
+def load_runs(path: PathLike) -> TrainingSet:
+    """Load completed runs from a ``run_model`` checkpoint file.
+
+    Useful for analysing saved runs without importing or re-running the
+    slow model. Only finished runs are returned, so a partly complete
+    checkpoint gives a shorter training set.
+    """
+    with np.load(path) as data:
+        outputs = data["outputs"]
+        n = len(outputs)
+        return TrainingSet(
+            X=data["points"][:n].copy(),
+            Y=outputs.copy(),
+            names=tuple(str(name) for name in data["names"]),
+            seconds_per_run=float(data["elapsed"]) / max(n, 1),
+        )
+
+
+__all__ = ["TrainingSet", "load_runs", "run_model"]