tscp 0.1.0.dev0__py3-none-any.whl

tscp/methods/engine.py

@@ -0,0 +1,128 @@
+"""The online event loop.
+
+An online method answers queries in time order and learns from outcomes as they become
+known. :class:`OnlineEngine` turns one outcome into the release mask an update observes;
+:func:`replay` runs it over the test queries of an archive, releasing each outcome once
+its target time has passed, by the rule of :mod:`tscp.data.availability`. The engine
+itself does not depend on archives: code that answers live queries can drive it with
+outcomes as they arrive.
+"""
+
+from collections.abc import Iterable
+
+import torch
+from torch import Tensor
+
+from tscp.online import OnlineUpdate
+from tscp.quantile import ScoreDistribution
+
+
+def group_selectors(
+    kept: tuple[int, ...], step_kept: bool, groups: int, device: torch.device
+) -> Tensor:
+    """``[G, *kept]`` bool: the kept elements whose outcome arrives with each group.
+
+    With steps kept, group ``g`` is horizon step ``g`` (the first kept dimension); with
+    a score over the horizon there is one group, released after the last step.
+    """
+    if not step_kept:
+        return torch.ones((1, *kept), dtype=torch.bool, device=device)
+    selectors = torch.zeros((groups, *kept), dtype=torch.bool, device=device)
+    for group in range(groups):
+        selectors[group, group] = True
+    return selectors
+
+
+class OnlineEngine:
+    """Feeds an online update its queries and outcomes.
+
+    Args:
+        update: The update to drive; :meth:`OnlineUpdate.reset` is called here.
+        levels: ``[*kept]`` target miscoverage of each kept element.
+        step_kept: Whether the score keeps the step dimension, so that each step's
+            outcome arrives separately.
+        num_groups: Number of release groups: the number of steps if they are kept, else 1.
+    """
+
+    def __init__(self, update: OnlineUpdate, levels: Tensor, step_kept: bool, num_groups: int):
+        self.update = update
+        self.selectors = group_selectors(tuple(levels.shape), step_kept, num_groups, levels.device)
+        update.reset(levels, levels.device)
+
+    def threshold(self, t: int, distribution: ScoreDistribution) -> Tensor:
+        """The threshold of query ``t``."""
+        return self.update.threshold(t, distribution)
+
+    def release(
+        self,
+        t: int,
+        group: int,
+        scores: Tensor,
+        thresholds: Tensor,
+        valid: Tensor | None,
+        levels: Tensor | None,
+    ) -> None:
+        """Pass the outcome of release group ``group`` of query ``t`` to the update."""
+        released = self.selectors[group]
+        if valid is not None:
+            released = released & valid
+        self.update.observe(t, scores, thresholds, released, levels)
+
+
+def replay(
+    engine: OnlineEngine,
+    origins: Tensor,
+    offsets: Tensor,
+    chunks: Iterable[tuple[int, int, ScoreDistribution]],
+    scores: Tensor,
+    valid: Tensor | None,
+    out: Tensor,
+) -> Tensor:
+    """Answer every query in time order, releasing each outcome once its target time passes.
+
+    At each query origin, every outcome whose target time has come is released first;
+    then the query is answered. Outcomes that would arrive at the same origin as a query
+    are known before it is answered.
+
+    Args:
+        engine: The engine to drive.
+        origins: ``[Q]`` query origins.
+        offsets: ``[G]`` steps after which each release group becomes known.
+        chunks: ``(start, stop, distribution)`` for consecutive chunks of queries, in order;
+            ``distribution`` covers queries ``start .. stop - 1``.
+        scores: ``[Q, *kept]`` scores of the queries' residuals, released once known.
+        valid: ``[Q, *kept]`` bool, or None if every score is defined.
+        out: ``[Q, *kept]`` tensor that receives the thresholds.
+
+    Returns:
+        ``out``.
+    """
+    origin_list, offset_list = origins.tolist(), offsets.tolist()
+    levels = None
+    if engine.update.needs_levels:
+        levels = torch.empty(out.shape, dtype=torch.float64, device=out.device)
+    pending = [0] * len(offset_list)  # first query whose outcome in group g is unreleased
+    for start, stop, distribution in chunks:
+        for u in range(start, stop):
+            for group, offset in enumerate(offset_list):
+                while pending[group] < u and origin_list[pending[group]] + offset <= origin_list[u]:
+                    t = pending[group]
+                    engine.release(
+                        t,
+                        group,
+                        scores[t],
+                        out[t],
+                        None if valid is None else valid[t],
+                        None if levels is None else levels[t],
+                    )
+                    pending[group] += 1
+            current = distribution.select(u - start)
+            out[u] = engine.threshold(u, current)
+            if levels is not None:
+                # Stored now, while the distribution is at hand; the update sees it only
+                # when the outcome is released.
+                levels[u] = current.level(scores[u])
+    return out
+
+
+__all__ = ["OnlineEngine", "group_selectors", "replay"]

tscp/methods/io.py

@@ -0,0 +1,106 @@
+"""Fitted methods in self-contained ``.npz`` files.
+
+A file holds a method's config and its fitted state, ``state_dict()``, with a JSON header
+that records the format, the schema version, the tscp version, the config and the
+fingerprints of the state and of the archive the method was fitted on. Files are written
+without pickle and never overwritten; loading checks the state against its fingerprint.
+
+Calibration state is not saved, because calibration is cheap: load, then ``calibrate``.
+"""
+
+import json
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from tscp import __version__
+from tscp.data import ResidualArchive
+from tscp.utils import Component, array_fingerprint, from_config
+
+FORMAT = "tscp.method"
+SCHEMA_VERSION = 1
+_PREFIX = "state."
+
+
+def save_method(method: Component, path: str | Path) -> None:
+    """Write a fitted method, or any fitted component, to a new ``.npz`` file.
+
+    Raises:
+        FileExistsError: If ``path`` exists; files are never overwritten.
+        ValueError: If the config cannot be written as JSON.
+    """
+    state = {name: value.detach().cpu() for name, value in method.state_dict().items()}
+    arrays = {_PREFIX + name: value.numpy() for name, value in state.items()}
+    header = {
+        "format": FORMAT,
+        "schema_version": SCHEMA_VERSION,
+        "tscp_version": __version__,
+        "config": method.to_config(),
+        "state_fingerprint": array_fingerprint(state),
+        "fit_archive": getattr(method, "fit_fingerprint", None),
+    }
+    try:
+        text = json.dumps(header, sort_keys=True, allow_nan=False)
+    except (TypeError, ValueError) as error:
+        raise ValueError(f"the method's config cannot be written as JSON: {error}") from error
+    arrays["header"] = np.asarray(text)
+    with Path(path).open("xb") as handle:
+        np.savez_compressed(handle, **arrays)
+
+
+def load_method(path: str | Path, archive: ResidualArchive | None = None) -> Component:
+    """Read a method written by :func:`save_method`, fitted and ready to calibrate.
+
+    The method is rebuilt from its config and its fitted state is loaded. A component
+    whose fitted state lives outside torch sets ``refit_on_load``; it is fitted again on
+    ``archive.split("fit")``, so ``archive`` must then be the archive the method was
+    fitted on, which its fingerprint checks.
+
+    Raises:
+        ValueError: If the file is not a tscp method, has another schema version, holds
+            state that does not match its fingerprint, or needs an archive to refit that is
+            missing or is not the one the method was fitted on.
+    """
+    with np.load(path, allow_pickle=False) as data:
+        if "header" not in data.files:
+            raise ValueError(f"{path} is not a tscp method: it has no header")
+        header = json.loads(str(data["header"].item()))
+        if header.get("format") != FORMAT:
+            raise ValueError(f"{path} is not a tscp method")
+        if header.get("schema_version") != SCHEMA_VERSION:
+            raise ValueError(
+                f"{path} has schema version {header.get('schema_version')}; "
+                f"this tscp reads version {SCHEMA_VERSION}"
+            )
+        unexpected = [
+            name for name in data.files if name != "header" and not name.startswith(_PREFIX)
+        ]
+        if unexpected:
+            raise ValueError(f"{path} holds unexpected arrays: {unexpected}")
+        state = {
+            name.removeprefix(_PREFIX): torch.from_numpy(data[name])
+            for name in data.files
+            if name.startswith(_PREFIX)
+        }
+    if array_fingerprint(state) != header["state_fingerprint"]:
+        raise ValueError(f"{path}: the fitted state does not match the fingerprint in the header")
+    method = from_config(header["config"])
+    method.load_state_dict(state)
+    refit = [m for m in method.modules() if isinstance(m, Component) and m.refit_on_load]
+    if refit:
+        if archive is None:
+            raise ValueError(
+                "this method refits components on load: pass the archive it was fitted on"
+            )
+        if archive.fingerprint() != header["fit_archive"]:
+            raise ValueError("the archive is not the one the method was fitted on")
+        plan = archive.splits
+        validation = archive.split("validation") if "validation" in plan else None
+        for component in refit:
+            component.fit(archive.split("fit"), validation)
+    method.fit_fingerprint = header["fit_archive"]
+    return method
+
+
+__all__ = ["FORMAT", "SCHEMA_VERSION", "load_method", "save_method"]

tscp/methods/nexcp.py

@@ -0,0 +1,42 @@
+"""Nonexchangeable conformal prediction."""
+
+from tscp.methods.base import ConformalPredictor
+from tscp.registry import register
+from tscp.scores import AbsoluteScore, Score
+from tscp.sources import ExpandingSource, Source
+from tscp.weights import ExponentialWeights
+
+
+@register("method", "nexcp")
+class NexCP(ConformalPredictor):
+    """Nonexchangeable conformal prediction (Barber, Candès, Ramdas & Tibshirani, 2023).
+
+    Weighs each past residual by ``decay ** age``, so recent errors count more, and adds
+    each test residual to the pool once it is known. The coverage loss under
+    nonexchangeability is bounded by the weighted total-variation distance terms of the
+    paper. With ``decay=1`` it is split conformal over the same pool.
+
+    Args:
+        decay: Weight ratio per base time step, in (0, 1].
+        score: Defaults to :class:`~tscp.scores.AbsoluteScore`.
+        source: Defaults to :class:`~tscp.sources.ExpandingSource`;
+            :class:`~tscp.sources.HoldoutSource` keeps the calibration split fixed.
+        alpha: Target miscoverage level.
+    """
+
+    def __init__(
+        self,
+        decay: float = 0.99,
+        score: Score | None = None,
+        source: Source | None = None,
+        alpha: float = 0.1,
+    ):
+        super().__init__(
+            score=score if score is not None else AbsoluteScore(),
+            source=source if source is not None else ExpandingSource(),
+            weights=ExponentialWeights(decay),
+            alpha=alpha,
+        )
+
+
+__all__ = ["NexCP"]

tscp/methods/split.py

@@ -0,0 +1,30 @@
+"""Split conformal prediction."""
+
+from tscp.methods.base import ConformalPredictor
+from tscp.registry import register
+from tscp.scores import AbsoluteScore, Score
+from tscp.sources import HoldoutSource
+
+
+@register("method", "split")
+class SplitCP(ConformalPredictor):
+    """Split conformal prediction (Vovk et al., 2005; Lei et al., JASA 2018).
+
+    Scores the calibration split, weighs every score equally and takes the
+    ``ceil((1 - alpha)(n + 1))``-th smallest as the threshold. Coverage is at least
+    ``1 - alpha`` in finite samples when calibration and test residuals are exchangeable.
+
+    Args:
+        score: Defaults to :class:`~tscp.scores.AbsoluteScore`.
+        alpha: Target miscoverage level.
+    """
+
+    def __init__(self, score: Score | None = None, alpha: float = 0.1):
+        super().__init__(
+            score=score if score is not None else AbsoluteScore(),
+            source=HoldoutSource(),
+            alpha=alpha,
+        )
+
+
+__all__ = ["SplitCP"]

tscp/metrics/__init__.py

@@ -0,0 +1,30 @@
+"""Metrics by output type, plus the default report.
+
+Region size is always per-context log-volume, summarized per dimension; the mean of raw
+volumes is never reported, because the widest regions dominate it.
+"""
+
+from .conditional import grouped_coverage, rolling_coverage, worst_window_coverage
+from .interval import (
+    coverage,
+    coverage_gap,
+    log_width,
+    mean_width,
+    observed_regions,
+    winkler_score,
+)
+from .report import Report, evaluate
+
+__all__ = [
+    "Report",
+    "coverage",
+    "coverage_gap",
+    "evaluate",
+    "grouped_coverage",
+    "log_width",
+    "mean_width",
+    "observed_regions",
+    "rolling_coverage",
+    "winkler_score",
+    "worst_window_coverage",
+]

tscp/metrics/conditional.py

@@ -0,0 +1,66 @@
+"""Coverage by group and over rolling windows."""
+
+import torch
+from torch import Tensor
+
+
+def _hits_and_counts(covered: Tensor, observed: Tensor | None) -> tuple[Tensor, Tensor]:
+    counts = torch.ones_like(covered, dtype=torch.float64)
+    if observed is not None:
+        counts = observed.to(torch.float64)
+    return covered.to(torch.float64) * counts, counts
+
+
+def rolling_coverage(covered: Tensor, window: int, observed: Tensor | None = None) -> Tensor:
+    """Coverage over each run of ``window`` consecutive query origins, per element.
+
+    Args:
+        covered: ``[Q, ...]`` bool, as returned by ``region.contains``.
+        window: Number of consecutive origins per window.
+        observed: Optional mask shaped like ``covered``; unobserved targets are skipped.
+
+    Returns:
+        ``[Q - window + 1, ...]``; NaN for a window without observed targets.
+    """
+    if not 1 <= window <= covered.shape[0]:
+        raise ValueError("window must lie between 1 and the number of origins")
+    hits, counts = _hits_and_counts(covered, observed)
+    zero = torch.zeros_like(hits[:1])
+    hits = torch.cat([zero, hits.cumsum(0)])
+    counts = torch.cat([zero, counts.cumsum(0)])
+    return (hits[window:] - hits[:-window]) / (counts[window:] - counts[:-window])
+
+
+def worst_window_coverage(covered: Tensor, window: int, observed: Tensor | None = None) -> Tensor:
+    """Lowest rolling coverage over every window and element."""
+    rolling = rolling_coverage(covered, window, observed)
+    return torch.where(rolling.isnan(), torch.inf, rolling).min()
+
+
+def grouped_coverage(
+    covered: Tensor, groups: Tensor, observed: Tensor | None = None
+) -> tuple[Tensor, Tensor]:
+    """Coverage within each group of query origins, such as the hour of day.
+
+    Args:
+        covered: ``[Q, ...]`` bool.
+        groups: ``[Q]`` integer group label of each origin.
+        observed: Optional mask shaped like ``covered``.
+
+    Returns:
+        The sorted group labels and the coverage in each.
+    """
+    if groups.shape != covered.shape[:1]:
+        raise ValueError("groups must hold one label per query origin")
+    hits, counts = _hits_and_counts(covered, observed)
+    hits = hits.reshape(covered.shape[0], -1).sum(1)
+    counts = counts.reshape(covered.shape[0], -1).sum(1)
+    labels, inverse = torch.unique(groups, return_inverse=True)
+    group_hits = torch.zeros(len(labels), dtype=torch.float64, device=covered.device)
+    group_counts = torch.zeros_like(group_hits)
+    group_hits.index_add_(0, inverse, hits)
+    group_counts.index_add_(0, inverse, counts)
+    return labels, group_hits / group_counts
+
+
+__all__ = ["grouped_coverage", "rolling_coverage", "worst_window_coverage"]

tscp/metrics/interval.py

@@ -0,0 +1,75 @@
+"""Coverage and size of regions; interval-specific scores."""
+
+import torch
+from torch import Tensor
+
+from tscp.regions import Interval, Region, move_event_dims
+
+
+def _mean(values: Tensor, observed: Tensor | None) -> Tensor:
+    values = values.to(torch.float64)
+    if observed is None:
+        return values.mean()
+    return values[observed].mean()
+
+
+def observed_regions(region: Region, observed: Tensor | None) -> Tensor | None:
+    """Return which regions have all their targets observed, or None if every target is.
+
+    Args:
+        region: The regions.
+        observed: ``[Q, H, N]`` mask of observed targets (or the region's batch shape for a
+            region per element), or None.
+
+    Returns:
+        A bool tensor of the region's batch shape, or None.
+    """
+    if observed is None or bool(observed.all()):
+        return None
+    if not region.event_dims:
+        return observed
+    moved = move_event_dims(observed, region.event_dims)
+    return moved.reshape(*region.batch_shape, -1).all(dim=-1)
+
+
+def coverage(region: Region, residuals: Tensor, observed: Tensor | None = None) -> Tensor:
+    """Fraction of regions that cover their targets, among those with every target observed."""
+    return _mean(region.contains(residuals), observed_regions(region, observed))
+
+
+def coverage_gap(
+    region: Region, residuals: Tensor, alpha: float, observed: Tensor | None = None
+) -> Tensor:
+    """Coverage minus its target ``1 - alpha``; negative means under-coverage."""
+    return coverage(region, residuals, observed) - (1 - alpha)
+
+
+def log_width(interval: Interval) -> Tensor:
+    """Log-width of each interval: its log-volume in one dimension."""
+    return interval.log_volume()
+
+
+def mean_width(interval: Interval, observed: Tensor | None = None) -> Tensor:
+    """Arithmetic mean width.
+
+    Only for comparing with published one-step tables, which report it. Report size as
+    log-width otherwise: a mean of raw widths is dominated by the widest.
+    """
+    return _mean(interval.width(), observed)
+
+
+def winkler_score(interval: Interval, residuals: Tensor, alpha: float) -> Tensor:
+    """Winkler (interval) score of each interval: width plus ``2/alpha`` times the miss."""
+    below = (interval.lower - residuals).clamp(min=0)
+    above = (residuals - interval.upper).clamp(min=0)
+    return (interval.upper - interval.lower) + (2 / alpha) * (below + above)
+
+
+__all__ = [
+    "coverage",
+    "coverage_gap",
+    "log_width",
+    "mean_width",
+    "observed_regions",
+    "winkler_score",
+]

tscp/metrics/report.py

@@ -0,0 +1,135 @@
+"""The default evaluation report."""
+
+import dataclasses
+import math
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+from torch import Tensor
+
+from tscp.data import ResidualArchive
+from tscp.metrics.conditional import grouped_coverage, worst_window_coverage
+from tscp.metrics.interval import observed_regions, winkler_score
+from tscp.regions import RESIDUAL_DIMS, Interval, Region
+
+
+@dataclass(frozen=True)
+class Report:
+    """What :func:`evaluate` measures.
+
+    Attributes:
+        alpha: The target miscoverage level.
+        num_regions: Regions evaluated: those whose targets are all observed.
+        coverage: Fraction of evaluated regions that cover their targets.
+        coverage_gap: ``coverage - (1 - alpha)``; negative means under-coverage.
+        log_volume_per_dim: Mean log-volume per dimension (for an interval, the
+            log-width), over bounded, nonempty regions only.
+        infinite_fraction: Fraction of evaluated regions that are unbounded.
+        empty_fraction: Fraction of evaluated regions that are empty or of zero volume.
+        winkler: Mean Winkler score of bounded, nonempty intervals; None for other regions.
+        coverage_by_step: Coverage at each horizon step, for regions made per step.
+        log_width_by_step: Mean log-width at each horizon step, for intervals made per step.
+        joint_coverage: Fraction of origins (and series, unless covered jointly) whose
+            targets are covered at every step of the horizon.
+        worst_window_coverage: Lowest coverage over windows of consecutive origins, if a
+            window was given.
+        grouped_coverage: Coverage per group label, if groups were given.
+    """
+
+    alpha: float
+    num_regions: int
+    coverage: float
+    coverage_gap: float
+    log_volume_per_dim: float
+    infinite_fraction: float
+    empty_fraction: float
+    winkler: float | None = None
+    coverage_by_step: tuple[float, ...] | None = None
+    log_width_by_step: tuple[float, ...] | None = None
+    joint_coverage: float | None = None
+    worst_window_coverage: float | None = None
+    grouped_coverage: dict[int, float] | None = None
+
+    def to_dict(self) -> dict[str, Any]:
+        """The report as a JSON-serializable dict."""
+        return dataclasses.asdict(self)
+
+
+def _mean(values: Tensor, mask: Tensor) -> float:
+    return float(values[mask].to(torch.float64).mean()) if mask.any() else math.nan
+
+
+def evaluate(
+    region: Region,
+    archive: ResidualArchive,
+    alpha: float,
+    *,
+    window: int | None = None,
+    groups: Tensor | None = None,
+) -> Report:
+    """Score regions against the residuals they were made for.
+
+    Works for any region type. A region that covers some dimensions jointly counts as
+    covering only if all its targets lie inside, and it is evaluated only if all its
+    targets are observed. Size is the per-context log-volume divided by the region's
+    dimension, averaged over bounded, nonempty regions only; unbounded and empty regions
+    are counted, never averaged in.
+
+    Args:
+        region: Regions from ``predict``, one per row of ``archive`` and element not
+            covered jointly.
+        archive: The rows the regions were made for, usually ``archive.split("test")``.
+        alpha: The target miscoverage level.
+        window: If given, also report the worst coverage over windows of this many origins.
+        groups: If given, ``[Q]`` labels; also report coverage per label.
+    """
+    residuals = archive.residuals
+    sizes = dict(zip(RESIDUAL_DIMS, residuals.shape[1:], strict=True))
+    kept = [name for name in RESIDUAL_DIMS if name not in region.event_dims]
+    expected = (len(archive), *(sizes[name] for name in kept))
+    if tuple(region.batch_shape) != expected:
+        raise ValueError(
+            f"regions must have batch shape {expected}: one per row and element not covered "
+            f"jointly, not {tuple(region.batch_shape)}"
+        )
+    covered = region.contains(residuals)
+    valid = observed_regions(region, archive.observed)
+    if valid is None:
+        valid = torch.ones_like(covered)
+    log_volume = region.log_volume()
+    proper = valid & torch.isfinite(log_volume)
+    coverage = _mean(covered, valid)
+    fields: dict[str, Any] = {
+        "alpha": float(alpha),
+        "num_regions": int(valid.sum()),
+        "coverage": coverage,
+        "coverage_gap": coverage - (1 - alpha),
+        "log_volume_per_dim": _mean(log_volume / region.dimension, proper),
+        "infinite_fraction": _mean(torch.isposinf(log_volume), valid),
+        "empty_fraction": _mean(torch.isneginf(log_volume), valid),
+    }
+    if isinstance(region, Interval):
+        fields["winkler"] = _mean(winkler_score(region, residuals, alpha), proper)
+    if "step" in kept:
+        step = 1  # the step dimension of the batch
+        fields["coverage_by_step"] = tuple(
+            _mean(covered.select(step, h), valid.select(step, h)) for h in range(sizes["step"])
+        )
+        if isinstance(region, Interval):
+            fields["log_width_by_step"] = tuple(
+                _mean(log_volume.select(step, h), proper.select(step, h))
+                for h in range(sizes["step"])
+            )
+        fields["joint_coverage"] = _mean(covered.all(dim=step), valid.all(dim=step))
+    else:
+        fields["joint_coverage"] = coverage
+    if window is not None:
+        fields["worst_window_coverage"] = float(worst_window_coverage(covered, window, valid))
+    if groups is not None:
+        labels, values = grouped_coverage(covered, groups, valid)
+        fields["grouped_coverage"] = dict(zip(labels.tolist(), values.tolist(), strict=True))
+    return Report(**fields)
+
+
+__all__ = ["Report", "evaluate"]

tscp/nn/__init__.py

@@ -0,0 +1,13 @@
+"""Helpers for learned components: losses, row batches and a reproducible training loop.
+
+A learned component is a :class:`~tscp.utils.Component` with torch parameters, fitted in
+its ``fit(train, validation)``: it batches the fit split with :class:`RowBatches` and
+fits itself with :func:`train`, which stops early on the validation split. Its fitted
+parameters live in ``state_dict()``, so a saved method keeps them.
+"""
+
+from .data import RowBatches
+from .losses import interval_score, pinball_loss
+from .training import TrainingResult, seeded, train
+
+__all__ = ["RowBatches", "TrainingResult", "interval_score", "pinball_loss", "seeded", "train"]