tscp 0.1.0.dev0__py3-none-any.whl

tscp/nn/data.py

@@ -0,0 +1,85 @@
+"""Batches of archive rows for training learned components."""
+
+import math
+from collections.abc import Iterator
+
+import torch
+from torch import Tensor
+
+from tscp.data import ResidualArchive
+from tscp.features import FeatureMap
+
+
+class RowBatches:
+    """Batches of the rows of an archive, with their features, for :func:`tscp.nn.train`.
+
+    Each pass yields dicts with:
+
+    - ``"index"``: ``[B]`` positions of the batch's rows within ``rows``;
+    - ``"residuals"``: ``[B, H, N]``;
+    - ``"observed"``: ``[B, H, N]`` bool, all True when the archive has no mask;
+    - ``"features"`` and ``"features_valid"``: ``[B, *element, F]`` and ``[B, *element]``,
+      when a feature map is given.
+
+    Features are computed once, when the batches are made, causally from the rows'
+    timeline. With ``shuffle`` the order changes on every pass, drawn from a generator
+    seeded with ``seed``, so fresh batches with the same seed give the same passes.
+
+    Args:
+        rows: The rows, such as ``train`` in a component's ``fit``.
+        features: Optional feature map.
+        batch_size: Rows per batch.
+        shuffle: Shuffle the rows on every pass.
+        seed: Seed of the shuffles.
+    """
+
+    def __init__(
+        self,
+        rows: ResidualArchive,
+        features: FeatureMap | None = None,
+        *,
+        batch_size: int = 256,
+        shuffle: bool = False,
+        seed: int = 0,
+    ):
+        if batch_size < 1:
+            raise ValueError("batch_size must be positive")
+        self.rows = rows
+        self.batch_size = int(batch_size)
+        self.shuffle = bool(shuffle)
+        self._generator = torch.Generator().manual_seed(seed)
+        self._observed = rows.observed
+        if self._observed is None:
+            self._observed = torch.ones_like(rows.residuals, dtype=torch.bool)
+        self.features: Tensor | None = None
+        self.features_valid: Tensor | None = None
+        if features is not None:
+            self.features = features(rows)
+            valid = features.valid(rows)
+            if valid is None:
+                valid = torch.ones(self.features.shape[:-1], dtype=torch.bool, device=rows.device)
+            self.features_valid = valid
+
+    def __len__(self) -> int:
+        return math.ceil(len(self.rows) / self.batch_size)
+
+    def __iter__(self) -> Iterator[dict[str, Tensor]]:
+        count = len(self.rows)
+        if self.shuffle:
+            order = torch.randperm(count, generator=self._generator).to(self.rows.device)
+        else:
+            order = torch.arange(count, device=self.rows.device)
+        for start in range(0, count, self.batch_size):
+            index = order[start : start + self.batch_size]
+            batch = {
+                "index": index,
+                "residuals": self.rows.residuals[index],
+                "observed": self._observed[index],
+            }
+            if self.features is not None:
+                batch["features"] = self.features[index]
+                batch["features_valid"] = self.features_valid[index]
+            yield batch
+
+
+__all__ = ["RowBatches"]

tscp/nn/losses.py

@@ -0,0 +1,63 @@
+"""Losses for training learned components."""
+
+import torch
+from torch import Tensor
+
+from tscp.metrics import winkler_score
+from tscp.regions import Interval
+
+
+def _mean(values: Tensor, weights: Tensor | None) -> Tensor:
+    if weights is None:
+        return values.mean()
+    weights = weights.to(values.dtype).expand_as(values)
+    return (values * weights).sum() / weights.sum().clamp(min=torch.finfo(values.dtype).tiny)
+
+
+def pinball_loss(
+    quantiles: Tensor,
+    targets: Tensor,
+    level: float | Tensor,
+    weights: Tensor | None = None,
+) -> Tensor:
+    """Mean pinball loss of predicted quantiles, minimized by the true ``level``-quantile.
+
+    ``level`` is the probability of the quantile, such as 0.9 for the 90% quantile, not a
+    miscoverage level. With ``e = target - quantile`` the loss is ``level * e`` when
+    ``e >= 0`` and ``(level - 1) * e`` otherwise.
+
+    Args:
+        quantiles: Predicted quantiles.
+        targets: Observed values, broadcastable against ``quantiles``.
+        level: Probability in (0, 1), a float or a broadcastable tensor.
+        weights: Optional nonnegative weights, such as an observed mask; the loss is then
+            their weighted mean, and 0 if every weight is 0.
+    """
+    error = targets - quantiles
+    return _mean(torch.maximum(level * error, (level - 1) * error), weights)
+
+
+def interval_score(
+    lower: Tensor,
+    upper: Tensor,
+    targets: Tensor,
+    alpha: float,
+    weights: Tensor | None = None,
+) -> Tensor:
+    """Mean interval (Winkler) score: the width plus ``2 / alpha`` times the distance missed.
+
+    It is minimized by the central ``1 - alpha`` interval of the targets' distribution,
+    and it is :func:`tscp.metrics.winkler_score` averaged, so it trains for what evaluation
+    reports.
+
+    Args:
+        lower: Lower bounds.
+        upper: Upper bounds, broadcastable against ``lower``.
+        targets: Observed values, broadcastable against the bounds.
+        alpha: Target miscoverage level.
+        weights: Optional nonnegative weights; see :func:`pinball_loss`.
+    """
+    return _mean(winkler_score(Interval(lower, upper), targets, alpha), weights)
+
+
+__all__ = ["interval_score", "pinball_loss"]

tscp/nn/training.py

@@ -0,0 +1,125 @@
+"""A small, reproducible training loop for learned components."""
+
+import contextlib
+import math
+from collections.abc import Callable, Iterable, Iterator
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor
+
+
+@contextlib.contextmanager
+def seeded(seed: int, device: torch.device | str | None = None) -> Iterator[None]:
+    """Run a block with the random state seeded, and restore the caller's afterwards.
+
+    Building a network inside ``seeded(seed)`` makes its initial weights the same on every
+    run, without changing the random numbers the rest of the program draws.
+
+    Args:
+        seed: The seed.
+        device: A CUDA device whose generator to fork and seed as well.
+    """
+    devices = [device] if device is not None and torch.device(device).type == "cuda" else []
+    with torch.random.fork_rng(devices=devices):
+        torch.manual_seed(seed)
+        yield
+
+
+@dataclass(frozen=True)
+class TrainingResult:
+    """What :func:`train` did.
+
+    Attributes:
+        losses: Mean training loss of each epoch.
+        validation_losses: Validation loss after each epoch; empty without validation.
+        best_epoch: The epoch whose state the module was left in, with validation.
+    """
+
+    losses: tuple[float, ...]
+    validation_losses: tuple[float, ...]
+    best_epoch: int | None
+
+
+def train(
+    module: torch.nn.Module,
+    loss: Callable[[torch.nn.Module, Any], Tensor],
+    batches: Iterable[Any],
+    *,
+    validation: Iterable[Any] | None = None,
+    epochs: int = 100,
+    optimizer: Callable[[list[torch.nn.Parameter]], torch.optim.Optimizer] | None = None,
+    patience: int | None = None,
+    seed: int = 0,
+) -> TrainingResult:
+    """Fit ``module`` by minimizing ``loss(module, batch)`` over the batches.
+
+    Training is reproducible without touching the caller's random state: it runs inside
+    :func:`seeded`, so dropout and shuffles draw the same numbers on every run.
+
+    With ``validation``, the mean loss over its batches is computed after every epoch, in
+    eval mode and without gradients. The module ends in the state of the epoch with the
+    lowest validation loss, and with ``patience``, training stops after that many epochs
+    without improvement.
+
+    Args:
+        module: The module to train; its parameters with ``requires_grad`` are updated.
+        loss: Maps the module and a batch to a scalar loss.
+        batches: Training batches; iterated once per epoch, so it must be re-iterable,
+            like :class:`tscp.nn.RowBatches`.
+        validation: Validation batches, also re-iterable.
+        epochs: Most passes over the training batches.
+        optimizer: Builds the optimizer from the parameters; Adam with learning rate
+            ``1e-3`` by default.
+        patience: Epochs without a better validation loss before stopping.
+        seed: Seed of the random state during training.
+    """
+    if epochs < 1:
+        raise ValueError("epochs must be positive")
+    if patience is not None and patience < 1:
+        raise ValueError("patience must be positive")
+    parameters = [parameter for parameter in module.parameters() if parameter.requires_grad]
+    if not parameters:
+        raise ValueError("the module has no parameters to train")
+    build = optimizer or (lambda params: torch.optim.Adam(params, lr=1e-3))
+    losses: list[float] = []
+    validation_losses: list[float] = []
+    best, best_epoch, best_state, waited = math.inf, None, None, 0
+    with seeded(seed, parameters[0].device):
+        step = build(parameters)
+        for epoch in range(epochs):
+            module.train()
+            total, count = 0.0, 0
+            for batch in batches:
+                step.zero_grad()
+                value = loss(module, batch)
+                value.backward()
+                step.step()
+                total, count = total + float(value.detach()), count + 1
+            if count == 0:
+                raise ValueError("the training batches are empty")
+            losses.append(total / count)
+            if validation is None:
+                continue
+            module.eval()
+            with torch.no_grad():
+                values = [float(loss(module, batch)) for batch in validation]
+            if not values:
+                raise ValueError("the validation batches are empty")
+            current = sum(values) / len(values)
+            validation_losses.append(current)
+            if current < best:
+                best, best_epoch, waited = current, epoch, 0
+                best_state = {k: v.detach().clone() for k, v in module.state_dict().items()}
+            else:
+                waited += 1
+                if patience is not None and waited >= patience:
+                    break
+    if best_state is not None:
+        module.load_state_dict(best_state)
+    module.eval()
+    return TrainingResult(tuple(losses), tuple(validation_losses), best_epoch)
+
+
+__all__ = ["TrainingResult", "seeded", "train"]

tscp/online/__init__.py

@@ -0,0 +1,6 @@
+"""Axis 5, online update: how the threshold reacts to its own misses."""
+
+from .aci import ACIUpdate
+from .base import OnlineUpdate
+
+__all__ = ["ACIUpdate", "OnlineUpdate"]

tscp/online/aci.py

@@ -0,0 +1,53 @@
+"""Adaptive conformal inference."""
+
+import torch
+from torch import Tensor
+
+from tscp.online.base import OnlineUpdate
+from tscp.quantile import ScoreDistribution
+from tscp.registry import register
+
+
+@register("update", "aci")
+class ACIUpdate(OnlineUpdate):
+    """Adaptive conformal inference (Gibbs & Candès, NeurIPS 2021).
+
+    Keeps a level ``alpha_t`` per kept element and moves it after each released outcome:
+    ``alpha_{t+1} = alpha_t + gamma * (alpha - err_t)``, where ``err_t`` is 1 when the
+    target fell outside its region. A level at or below 0 gives an unbounded region and
+    one at or above 1 an empty region, so over any sequence of ``T`` outcomes the miss
+    rate stays within ``(max(alpha_1, 1 - alpha_1) + gamma) / (gamma * T)`` of ``alpha``.
+
+    Args:
+        gamma: Step size, positive.
+    """
+
+    def __init__(self, gamma: float = 0.005):
+        super().__init__()
+        if not gamma > 0:
+            raise ValueError("gamma must be positive")
+        self.gamma = float(gamma)
+        self.register_buffer("target", None, persistent=False)
+        self.register_buffer("level", None, persistent=False)
+
+    def reset(self, alpha: Tensor, device: torch.device) -> None:
+        self.target = torch.as_tensor(alpha, dtype=torch.float64, device=device).clone()
+        self.level = self.target.clone()
+
+    def threshold(self, t: int, distribution: ScoreDistribution) -> Tensor:
+        return distribution.threshold(self.level)
+
+    def observe(
+        self,
+        t: int,
+        scores: Tensor,
+        thresholds: Tensor,
+        released: Tensor,
+        levels: Tensor | None = None,
+    ) -> None:
+        miss = (scores > thresholds).to(torch.float64)
+        stepped = self.level + self.gamma * (self.target - miss)
+        self.level = torch.where(released, stepped, self.level)
+
+
+__all__ = ["ACIUpdate"]

tscp/online/base.py

@@ -0,0 +1,67 @@
+"""The online-update interface (axis 5)."""
+
+import torch
+from torch import Tensor
+
+from tscp.quantile import ScoreDistribution
+from tscp.utils import Component
+
+
+class OnlineUpdate(Component):
+    """How the threshold reacts to its own misses.
+
+    The composer answers the test queries in time order, numbered ``t = 0, 1, ...``.
+    Before answering a query it passes every outcome that has become known
+    (:meth:`observe`); then it asks for the query's threshold (:meth:`threshold`),
+    handing over the conformal distribution of its score. Updates that move a level (the
+    ACI family) read the distribution's threshold at their level, several levels at once
+    if they run experts; updates that move the threshold itself (conformal PID, SAOCP)
+    may adjust or ignore it.
+
+    Outcomes arrive late: the outcome of step ``h`` of a query is known ``h`` base steps
+    after its origin, so later queries are often answered before earlier outcomes
+    arrive. An update that keeps state per query, such as expert thresholds, keys it by
+    ``t``.
+    """
+
+    #: Whether :meth:`observe` needs the realized level of each outcome (DtACI does).
+    needs_levels: bool = False
+
+    def reset(self, alpha: Tensor, device: torch.device) -> None:
+        """Start a new pass with target miscoverage ``alpha``, ``[*kept]`` in float64."""
+        raise NotImplementedError
+
+    def threshold(self, t: int, distribution: ScoreDistribution) -> Tensor:
+        """Return the threshold ``[*kept]`` for query ``t``.
+
+        Args:
+            t: Position of the query, counting from 0 in time order.
+            distribution: The conformal distribution of query ``t``'s score, of batch
+                shape ``[*kept]``.
+        """
+        raise NotImplementedError
+
+    def observe(
+        self,
+        t: int,
+        scores: Tensor,
+        thresholds: Tensor,
+        released: Tensor,
+        levels: Tensor | None = None,
+    ) -> None:
+        """Take an outcome of query ``t``.
+
+        Args:
+            t: Position of the query the outcome belongs to.
+            scores: ``[*kept]`` scores of its residuals.
+            thresholds: ``[*kept]`` thresholds it was answered with.
+            released: ``[*kept]`` bool, True for the elements whose outcome is known now.
+                Per-step outcomes arrive one step at a time; missing targets never arrive.
+            levels: ``[*kept]`` realized levels, ``distribution.level(scores)`` under the
+                distribution query ``t`` was answered with, if :attr:`needs_levels`;
+                otherwise None.
+        """
+        raise NotImplementedError
+
+
+__all__ = ["OnlineUpdate"]

tscp/quantile/__init__.py

@@ -0,0 +1,16 @@
+"""Axis 4, quantile model: how weighted scores become a threshold."""
+
+from .base import Pool, QuantileModel
+from .distribution import TOLERANCE, EmpiricalDistribution, ScoreDistribution
+from .weighted import WeightedQuantile, uniform_quantile, weighted_quantile
+
+__all__ = [
+    "TOLERANCE",
+    "EmpiricalDistribution",
+    "Pool",
+    "QuantileModel",
+    "ScoreDistribution",
+    "WeightedQuantile",
+    "uniform_quantile",
+    "weighted_quantile",
+]

tscp/quantile/base.py

@@ -0,0 +1,73 @@
+"""The quantile-model interface (axis 4)."""
+
+import torch
+from torch import Tensor
+
+from tscp.data import ResidualArchive
+from tscp.quantile.distribution import ScoreDistribution
+from tscp.utils import Component
+
+
+class Pool:
+    """The scores a predictor calibrates with, with the rows they come from.
+
+    The composer builds one per call to ``predict``. It sorts the scores once, however
+    many queries read them.
+
+    Args:
+        scores: ``[P, *kept]`` pool scores.
+        rows: The ``P`` pool rows, a row subset of the archive, for quantile models and
+            weights that use features of the rows.
+        valid: ``[P, *kept]`` bool, False where a score is undefined; None if all are.
+    """
+
+    def __init__(self, scores: Tensor, rows: ResidualArchive, valid: Tensor | None = None):
+        if scores.shape[0] != len(rows):
+            raise ValueError("the pool needs one row per score")
+        if valid is not None and valid.shape != scores.shape:
+            raise ValueError("valid must have the shape of the scores")
+        self.scores = scores
+        self.rows = rows
+        self.valid = valid
+        self._sorted: tuple[Tensor, Tensor] | None = None
+
+    def __len__(self) -> int:
+        return self.scores.shape[0]
+
+    @property
+    def kept(self) -> tuple[int, ...]:
+        """Sizes of the kept dimensions."""
+        return tuple(self.scores.shape[1:])
+
+    def sorted(self) -> tuple[Tensor, Tensor]:
+        """The scores sorted along the pool dimension, and the indices that sort them."""
+        if self._sorted is None:
+            values, indices = torch.sort(self.scores, dim=0)
+            self._sorted = (values, indices)
+        return self._sorted
+
+
+class QuantileModel(Component):
+    """How weighted scores become a threshold.
+
+    :meth:`forward` takes the calibration pool, the weights of its scores for each query
+    and the query rows, and returns a :class:`~tscp.quantile.ScoreDistribution` from which
+    the composer and online updates read thresholds at any level.
+
+    ``weights=None`` means every pool score counts the same for every query; the composer
+    passes None only when every pool score is defined and available. Otherwise weights
+    are ``[Q, P]`` or broadcastable to ``[Q, P, *kept]``, nonnegative, and already zero for
+    scores that are undefined or not yet known at the query.
+
+    A model that regresses scores on features of the rows, such as quantile regression on
+    recent residuals, reads them from ``pool.rows`` and ``query`` and their timelines.
+    """
+
+    def forward(
+        self, pool: Pool, weights: Tensor | None, query: ResidualArchive
+    ) -> ScoreDistribution:
+        """Return the distribution of the score at each of the ``Q`` query rows."""
+        raise NotImplementedError
+
+
+__all__ = ["Pool", "QuantileModel"]

tscp/quantile/distribution.py

@@ -0,0 +1,170 @@
+"""Conformal distributions of scores: a threshold at any level, and the level of any score."""
+
+import torch
+from torch import Tensor
+
+from tscp.typing import Alpha
+
+TOLERANCE = 1e-9
+"""Levels within this of a cumulative weight count as reached, so float rounding in
+``1 - alpha`` or in cumulative sums cannot shift a threshold by one score."""
+
+
+class ScoreDistribution:
+    """The conformal distribution of a score, for a batch of queries and kept elements.
+
+    Quantile models (axis 4) return one. :meth:`threshold` gives the threshold at any
+    miscoverage level, several levels at once if asked; :meth:`level` gives the largest
+    level at which a score is still covered, its conformal p-value. Online updates use
+    both: the ACI family moves the level it asks a threshold for, and DtACI weighs its
+    experts by the level of each realized score.
+
+    The batch shape is ``[Q, *kept]`` for one distribution per query, or ``[*kept]`` when
+    every query shares one.
+    """
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        raise NotImplementedError
+
+    def threshold(self, alpha: Alpha) -> Tensor:
+        """Return the threshold at miscoverage ``alpha``.
+
+        ``alpha`` broadcasts against :attr:`batch_shape` from the right and may add leading
+        dimensions to ask for several levels at once: ``alpha`` of shape ``[L, *batch]``
+        gives thresholds ``[L, *batch]``. A level at or below 0 gives ``+inf`` (an
+        unbounded region), and one at or above 1 gives ``-inf`` (an empty region).
+        """
+        raise NotImplementedError
+
+    def level(self, scores: Tensor) -> Tensor:
+        """Return the largest miscoverage level whose threshold still covers ``scores``.
+
+        ``scores`` broadcast against :attr:`batch_shape`. The result, in float64, is the
+        conformal p-value of each score (up to the rounding tolerance of
+        :meth:`threshold`): 1 below every pool score, and the test point's share of the
+        weight above them all.
+        """
+        raise NotImplementedError
+
+    def select(self, index: int) -> "ScoreDistribution":
+        """Return the distribution of one query, of batch shape ``[*kept]``."""
+        raise NotImplementedError
+
+
+class EmpiricalDistribution(ScoreDistribution):
+    """The weighted empirical distribution of pool scores, with the test point at ``+inf``.
+
+    With weights ``w_i`` on pool scores ``s_i`` and the test point's weight ``w``, the
+    threshold at ``alpha`` is the smallest pool score ``s`` with
+    ``sum(w_i for s_i <= s) / (sum(w_i) + w) >= 1 - alpha`` (Tibshirani et al., NeurIPS
+    2019; Barber et al., Annals of Statistics 2023), or ``+inf`` when the pool never
+    reaches ``1 - alpha``. With equal weights it is the ``ceil((1 - alpha)(n + w))``-th
+    smallest of the ``n`` pool scores.
+
+    Args:
+        ordered: ``[P, *kept]`` pool scores, sorted along the pool dimension.
+        cumulative: None for equal weights. Otherwise the cumulative weights of the sorted
+            scores divided by the total weight plus ``test_weight``: ``[Q, P, *kept]``, or
+            ``[P, *kept]`` for one query.
+        test_weight: Weight of the test point; with equal weights, relative to 1 per score.
+    """
+
+    def __init__(self, ordered: Tensor, cumulative: Tensor | None = None, test_weight=1.0):
+        self.ordered = ordered
+        self.cumulative = cumulative
+        self.test_weight = float(test_weight)
+
+    @classmethod
+    def weighted(
+        cls, ordered: Tensor, order: Tensor, weights: Tensor, test_weight: float = 1.0
+    ) -> "EmpiricalDistribution":
+        """Build a distribution per query from sorted scores and unsorted weights.
+
+        Args:
+            ordered: ``[P, *kept]`` pool scores, sorted along the pool dimension.
+            order: ``[P, *kept]`` indices that sort the scores, as ``torch.sort`` returns.
+            weights: Nonnegative weights ``[Q, P]``, or broadcastable to ``[Q, P, *kept]``.
+            test_weight: Weight of the test point.
+        """
+        num_scores, kept = ordered.shape[0], tuple(ordered.shape[1:])
+        if weights.ndim < 2 or weights.shape[1] != num_scores:
+            raise ValueError("weights must have shape [Q, P] or [Q, P, *kept] with P pool scores")
+        weights = weights.reshape(*weights.shape, *([1] * (2 + len(kept) - weights.ndim)))
+        full = (weights.shape[0], num_scores, *kept)
+        part = weights.to(torch.float64).broadcast_to(full)
+        part = torch.gather(part, 1, order.unsqueeze(0).expand(full))
+        total = part.sum(dim=1, keepdim=True) + test_weight
+        return cls(ordered, torch.cumsum(part, dim=1) / total, test_weight)
+
+    @property
+    def _pool_dim(self) -> int:
+        """Position of the pool dimension, counted from the right."""
+        return -self.ordered.ndim
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        if self.cumulative is None:
+            return self.ordered.shape[1:]
+        shape = list(self.cumulative.shape)
+        del shape[self._pool_dim]
+        return torch.Size(shape)
+
+    def threshold(self, alpha: Alpha) -> Tensor:
+        level = 1 - torch.as_tensor(alpha, dtype=torch.float64, device=self.ordered.device)
+        if self.cumulative is None:
+            return self._uniform_threshold(level)
+        shape = torch.broadcast_shapes(level.shape, self.batch_shape)
+        level = level.broadcast_to(shape)
+        pool_dim, num_scores = self._pool_dim, self.ordered.shape[0]
+        # Cumulative weights never decrease, so the first index that reaches the level is
+        # the number of indices that do not; it equals P when the level is never reached.
+        unreached = self.cumulative < (level - TOLERANCE).unsqueeze(pool_dim)
+        first = unreached.sum(dim=pool_dim)
+        index = first.clamp(max=num_scores - 1).unsqueeze(pool_dim)
+        values = torch.gather(self.ordered.expand(unreached.shape), pool_dim, index)
+        values = torch.where(first >= num_scores, torch.inf, values.squeeze(pool_dim).double())
+        return _extremes(values, level).to(self.ordered.dtype)
+
+    def _uniform_threshold(self, level: Tensor) -> Tensor:
+        num_scores, kept = self.ordered.shape[0], tuple(self.ordered.shape[1:])
+        shape = torch.broadcast_shapes(level.shape, kept)
+        level = level.broadcast_to(shape)
+        rank = torch.ceil((level - TOLERANCE) * (num_scores + self.test_weight))
+        index = (rank - 1).clamp(0, num_scores - 1).to(torch.int64).reshape(-1, *kept)
+        source = self.ordered.unsqueeze(0).expand(index.shape[0], num_scores, *kept)
+        values = torch.gather(source, 1, index.unsqueeze(1)).squeeze(1).reshape(shape)
+        values = torch.where(rank > num_scores, torch.inf, values.to(torch.float64))
+        return _extremes(values, level).to(self.ordered.dtype)
+
+    def level(self, scores: Tensor) -> Tensor:
+        scores = torch.as_tensor(scores, device=self.ordered.device)
+        pool_dim, num_scores = self._pool_dim, self.ordered.shape[0]
+        shape = torch.broadcast_shapes(scores.shape, self.batch_shape)
+        scores = scores.broadcast_to(shape)
+        below = (self.ordered < scores.unsqueeze(pool_dim)).sum(dim=pool_dim)
+        if self.cumulative is None:
+            return 1 - below.to(torch.float64) / (num_scores + self.test_weight)
+        full = list(shape)
+        full.insert(len(full) + pool_dim + 1, num_scores)
+        index = (below - 1).clamp(min=0).unsqueeze(pool_dim)
+        mass = torch.gather(self.cumulative.expand(full), pool_dim, index).squeeze(pool_dim)
+        return 1 - torch.where(below > 0, mass, 0.0)
+
+    def select(self, index: int) -> "EmpiricalDistribution":
+        if self.cumulative is None or self.cumulative.ndim == self.ordered.ndim:
+            return self
+        return EmpiricalDistribution(self.ordered, self.cumulative[index], self.test_weight)
+
+
+def _extremes(values: Tensor, level: Tensor) -> Tensor:
+    """Apply the conventions at the ends: ``alpha <= 0`` covers all, ``alpha >= 1`` nothing.
+
+    Without them, a test weight of 0 would give the largest pool score at ``alpha = 0``,
+    and the long-run guarantee of the ACI family relies on an unbounded region there.
+    """
+    values = torch.where(level >= 1, torch.inf, values)
+    return torch.where(level <= 0, -torch.inf, values)
+
+
+__all__ = ["TOLERANCE", "EmpiricalDistribution", "ScoreDistribution"]