tscp 0.1.0.dev0__py3-none-any.whl

tscp/quantile/weighted.py

@@ -0,0 +1,103 @@
+"""The conformal quantile of (weighted) scores."""
+
+import math
+
+import torch
+from torch import Tensor
+
+from tscp.data import ResidualArchive
+from tscp.quantile.base import Pool, QuantileModel
+from tscp.quantile.distribution import EmpiricalDistribution
+from tscp.registry import register
+from tscp.typing import Alpha
+
+_MAX_ELEMENTS = 2**23
+
+
+def uniform_quantile(scores: Tensor, alpha: Alpha, test_weight: float = 1.0) -> Tensor:
+    """Conformal quantile of equally weighted scores.
+
+    The threshold is the ``ceil((1 - alpha) * (n + w))``-th smallest of the ``n`` pool
+    scores, where ``w`` is the test point's weight (1 for split conformal); ``+inf`` when
+    that rank exceeds ``n`` or ``alpha <= 0``, and ``-inf`` when ``alpha >= 1``.
+
+    Args:
+        scores: ``[P, *kept]`` pool scores.
+        alpha: Miscoverage level; a float, or a tensor that broadcasts against ``kept``
+            from the right (for example ``[*kept]`` or ``[Q, *kept]``).
+        test_weight: Weight of the test point, relative to 1 per pool score.
+
+    Returns:
+        Thresholds of shape ``broadcast(alpha.shape, kept)``, in the dtype of ``scores``.
+    """
+    ordered = torch.sort(scores, dim=0).values
+    return EmpiricalDistribution(ordered, None, test_weight).threshold(alpha)
+
+
+def weighted_quantile(
+    scores: Tensor, weights: Tensor, alpha: Alpha, test_weight: float = 1.0
+) -> Tensor:
+    """Conformal quantile of weighted scores.
+
+    For each query the threshold is the smallest pool score ``s`` with
+    ``sum(w_i for s_i <= s) / (sum(w_i) + test_weight) >= 1 - alpha``: the
+    ``(1 - alpha)``-quantile of the weighted pool with the test point's weight at
+    ``+inf`` (Tibshirani et al., NeurIPS 2019; Barber et al., Annals of Statistics 2023).
+    It is ``+inf`` when the pool never reaches ``1 - alpha`` or ``alpha <= 0``, and
+    ``-inf`` when ``alpha >= 1``. With all weights 1 it equals :func:`uniform_quantile`.
+
+    Args:
+        scores: ``[P, *kept]`` pool scores.
+        weights: Nonnegative weights ``[Q, P]``, or broadcastable to ``[Q, P, *kept]``.
+        alpha: Miscoverage level; a float, or a tensor broadcastable to ``[Q, *kept]``.
+        test_weight: Weight of the test point.
+
+    Returns:
+        Thresholds ``[Q, *kept]``, in the dtype of ``scores``.
+    """
+    num_scores, kept = scores.shape[0], tuple(scores.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")
+    num_queries = weights.shape[0]
+    alpha = torch.as_tensor(alpha, dtype=torch.float64, device=scores.device)
+    alpha = alpha.broadcast_to((num_queries, *kept))
+    ordered, order = torch.sort(scores, dim=0)
+    out = torch.empty((num_queries, *kept), dtype=scores.dtype, device=scores.device)
+    chunk = max(1, _MAX_ELEMENTS // max(1, num_scores * math.prod(kept)))
+    for start in range(0, num_queries, chunk):
+        stop = min(start + chunk, num_queries)
+        part = EmpiricalDistribution.weighted(ordered, order, weights[start:stop], test_weight)
+        out[start:stop] = part.threshold(alpha[start:stop])
+    return out
+
+
+@register("quantile", "weighted")
+class WeightedQuantile(QuantileModel):
+    """Empirical conformal quantile of (weighted) scores, with the test point at ``+inf``.
+
+    Returns an :class:`~tscp.quantile.EmpiricalDistribution`: one shared by every query
+    when ``weights`` is None, one per query otherwise. Either way the pool is sorted once.
+
+    Args:
+        test_weight: Weight of the test point relative to the pool weights; 1 for split
+            conformal and NexCP.
+    """
+
+    def __init__(self, test_weight: float = 1.0):
+        super().__init__()
+        if not test_weight >= 0:
+            raise ValueError("test_weight must be nonnegative")
+        self.test_weight = float(test_weight)
+
+    def forward(
+        self, pool: Pool, weights: Tensor | None, query: ResidualArchive
+    ) -> EmpiricalDistribution:
+        ordered, order = pool.sorted()
+        if weights is None:
+            return EmpiricalDistribution(ordered, None, self.test_weight)
+        if (weights < 0).any() or weights.isnan().any():
+            raise ValueError("weights must be nonnegative numbers")
+        return EmpiricalDistribution.weighted(ordered, order, weights, self.test_weight)
+
+
+__all__ = ["WeightedQuantile", "uniform_quantile", "weighted_quantile"]

tscp/regions/__init__.py

@@ -0,0 +1,15 @@
+"""Axis 6, output: what the user gets, with a membership test and a size."""
+
+from .base import RESIDUAL_DIMS, Region, check_event_dims, move_event_dims, restore_event_dims
+from .box import Box
+from .interval import Interval
+
+__all__ = [
+    "RESIDUAL_DIMS",
+    "Box",
+    "Interval",
+    "Region",
+    "check_event_dims",
+    "move_event_dims",
+    "restore_event_dims",
+]

tscp/regions/base.py

@@ -0,0 +1,139 @@
+"""The region interface (axis 6, output)."""
+
+import math
+
+import torch
+from torch import Tensor
+
+RESIDUAL_DIMS = ("step", "series")
+"""Names of the dimensions of ``[H, N]``, in the order residuals store them."""
+
+
+def check_event_dims(event_dims) -> tuple[str, ...]:
+    """Return ``event_dims`` as a tuple in ``[H, N]`` order, checking the names."""
+    event_dims = tuple(event_dims)
+    unknown = sorted(set(event_dims) - set(RESIDUAL_DIMS))
+    if unknown or len(set(event_dims)) != len(event_dims):
+        raise ValueError(f"event_dims must name distinct dimensions of {RESIDUAL_DIMS}")
+    return tuple(name for name in RESIDUAL_DIMS if name in event_dims)
+
+
+def move_event_dims(values: Tensor, event_dims: tuple[str, ...]) -> Tensor:
+    """Reorder ``[..., H, N]`` values to ``[..., *kept, *event]``.
+
+    The dimensions named in ``event_dims`` move to the end and the others stay in front,
+    each group in ``[H, N]`` order. This is how a region that covers some dimensions
+    jointly lines residuals or forecasts up with its bounds.
+    """
+    if values.ndim < 2:
+        raise ValueError("expected values of shape [..., H, N]")
+    lead = values.ndim - 2
+    kept = [lead + i for i, name in enumerate(RESIDUAL_DIMS) if name not in event_dims]
+    event = [lead + i for i, name in enumerate(RESIDUAL_DIMS) if name in event_dims]
+    return values.permute(*range(lead), *kept, *event)
+
+
+def restore_event_dims(values: Tensor, event_dims: tuple[str, ...]) -> Tensor:
+    """Invert :func:`move_event_dims`: reorder ``[..., *kept, *event]`` back to ``[..., H, N]``."""
+    if values.ndim < 2:
+        raise ValueError("expected values of shape [..., *kept, *event] with two residual dims")
+    order = [name for name in RESIDUAL_DIMS if name not in event_dims]
+    order += [name for name in RESIDUAL_DIMS if name in event_dims]
+    lead = values.ndim - 2
+    return values.permute(*range(lead), *(lead + order.index(name) for name in RESIDUAL_DIMS))
+
+
+class Region:
+    """A batch of prediction regions, one per query origin and kept element.
+
+    Regions live in residual coordinates: a target ``y`` with point forecast ``f`` is
+    covered when ``y - f`` lies in the region. Coverage and volume do not change under
+    this shift, so evaluation never needs the forecasts; :meth:`to_target` maps a region
+    back to target units.
+
+    A region covers the dimensions of ``[H, N]`` in :attr:`event_dims` jointly and makes
+    one region per element of the others. :meth:`contains` always takes the residuals of
+    the query rows, ``[Q, H, N]``, and returns one answer per region, of shape
+    :attr:`batch_shape`.
+
+    A subclass stores its regions in tensors whose leading dimensions are the batch
+    dimensions, returns them from :meth:`tensors`, and can be rebuilt from them by its
+    constructor; indexing and device moves then work without more code. A subclass with
+    settings that are not tensors overrides :meth:`_rebuild`.
+    """
+
+    #: Dimensions of ``[H, N]`` that each region covers jointly.
+    event_dims: tuple[str, ...] = ()
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        """``[Q, *kept]``: one region per query origin and element not covered jointly."""
+        raise NotImplementedError
+
+    @property
+    def dimension(self) -> int:
+        """Number of coordinates each region covers jointly (1 for an interval)."""
+        raise NotImplementedError
+
+    def tensors(self) -> dict[str, Tensor]:
+        """The tensors that define the regions, by constructor argument name."""
+        raise NotImplementedError
+
+    def contains(self, residuals: Tensor) -> Tensor:
+        """Return a bool tensor of shape :attr:`batch_shape`, True where covered.
+
+        Args:
+            residuals: ``[Q, H, N]`` residuals (or any points in the region's coordinates).
+        """
+        raise NotImplementedError
+
+    def log_volume(self) -> Tensor:
+        """Return the log-volume of each region.
+
+        ``+inf`` means unbounded; ``-inf`` means empty or of zero volume.
+        """
+        raise NotImplementedError
+
+    def to_target(self, forecasts: Tensor) -> "Region":
+        """Return the regions shifted by the ``[Q, H, N]`` point forecasts, in target units."""
+        raise NotImplementedError
+
+    def _rebuild(self, tensors: dict[str, Tensor]) -> "Region":
+        """Return a region of the same type and settings, made of other tensors."""
+        return type(self)(**tensors)
+
+    def __getitem__(self, index) -> "Region":
+        """Select regions along the batch dimensions, as ``region[3]`` or ``region[mask]``."""
+        items = index if isinstance(index, tuple) else (index,)
+        if any(item is Ellipsis for item in items):
+            raise IndexError("regions do not support Ellipsis in indices")
+        if len(items) > len(self.batch_shape):
+            raise IndexError("too many indices: regions are indexed along batch dimensions")
+        return self._rebuild({name: value[index] for name, value in self.tensors().items()})
+
+    def to(self, device: torch.device | str) -> "Region":
+        """Return the regions on ``device``."""
+        return self._rebuild({name: value.to(device) for name, value in self.tensors().items()})
+
+    def __len__(self) -> int:
+        return self.batch_shape[0]
+
+    def __repr__(self) -> str:
+        parts = [f"batch_shape={tuple(self.batch_shape)}"]
+        if self.event_dims:
+            parts.append(f"event_dims={self.event_dims}")
+        return f"{type(self).__name__}({', '.join(parts)})"
+
+
+def event_size(shape: torch.Size, event_ndim: int) -> int:
+    """Number of coordinates in the last ``event_ndim`` dimensions of ``shape``."""
+    return math.prod(shape[len(shape) - event_ndim :]) if event_ndim else 1
+
+
+__all__ = [
+    "RESIDUAL_DIMS",
+    "Region",
+    "check_event_dims",
+    "move_event_dims",
+    "restore_event_dims",
+]

tscp/regions/box.py

@@ -0,0 +1,82 @@
+"""Boxes: an interval per coordinate, covered jointly."""
+
+import torch
+from torch import Tensor
+
+from tscp.regions.base import Region, check_event_dims, event_size, move_event_dims
+
+
+class Box(Region):
+    """Axis-aligned boxes over the dimensions in ``event_dims``.
+
+    A box covers a point when every coordinate lies within its bounds. Bounds have shape
+    ``[*batch, *event]``: the batch dimensions (the query origin, then whichever of step
+    and series is not covered jointly), then the dimensions in ``event_dims``, both in
+    ``[H, N]`` order. For example, a box per series over the whole horizon has bounds
+    ``[Q, N, H]`` and ``event_dims=("step",)``.
+
+    The log-volume is the sum of the coordinates' log-widths. A box with an empty or
+    zero-width coordinate has log-volume ``-inf``; otherwise any infinite coordinate makes
+    it ``+inf``.
+
+    Args:
+        lower: Lower bounds ``[*batch, *event]``, in residual coordinates.
+        upper: Upper bounds, the same shape as ``lower``.
+        event_dims: The dimensions of ``[H, N]`` each box covers jointly.
+    """
+
+    def __init__(self, lower: Tensor, upper: Tensor, event_dims=("step",)):
+        event_dims = check_event_dims(event_dims)
+        if not event_dims:
+            raise ValueError("a box covers at least one dimension jointly; use an Interval")
+        lower = torch.as_tensor(lower)
+        upper = torch.as_tensor(upper, device=lower.device)
+        if lower.shape != upper.shape:
+            raise ValueError("lower and upper must have the same shape")
+        if not (lower.is_floating_point() and upper.is_floating_point()):
+            raise ValueError("box bounds must be floating tensors")
+        if lower.ndim < len(event_dims):
+            raise ValueError("bounds must end with one dimension per entry of event_dims")
+        self.lower = lower
+        self.upper = upper
+        self.event_dims = event_dims
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        return self.lower.shape[: self.lower.ndim - len(self.event_dims)]
+
+    @property
+    def dimension(self) -> int:
+        return event_size(self.lower.shape, len(self.event_dims))
+
+    def tensors(self) -> dict[str, Tensor]:
+        return {"lower": self.lower, "upper": self.upper}
+
+    def _rebuild(self, tensors: dict[str, Tensor]) -> "Box":
+        return Box(tensors["lower"], tensors["upper"], self.event_dims)
+
+    def _flat(self, values: Tensor) -> Tensor:
+        """``[*batch, *event]`` values with the event dimensions flattened into one."""
+        return values.reshape(*self.batch_shape, self.dimension)
+
+    def is_empty(self) -> Tensor:
+        """True where some coordinate has ``lower > upper``."""
+        return self._flat(self.upper < self.lower).any(dim=-1)
+
+    def contains(self, residuals: Tensor) -> Tensor:
+        points = move_event_dims(residuals, self.event_dims)
+        inside = (self.lower <= points) & (points <= self.upper)
+        return self._flat(inside).all(dim=-1)
+
+    def log_volume(self) -> Tensor:
+        width = self._flat(self.upper - self.lower)
+        log_width = torch.log(width.clamp(min=0))  # -inf for empty or zero-width coordinates
+        degenerate = torch.isneginf(log_width).any(dim=-1)
+        return torch.where(degenerate, -torch.inf, log_width.sum(dim=-1))
+
+    def to_target(self, forecasts: Tensor) -> "Box":
+        shift = move_event_dims(forecasts, self.event_dims)
+        return Box(self.lower + shift, self.upper + shift, self.event_dims)
+
+
+__all__ = ["Box"]

tscp/regions/interval.py

@@ -0,0 +1,62 @@
+"""Intervals, one per element."""
+
+import torch
+from torch import Tensor
+
+from tscp.regions.base import Region
+
+
+class Interval(Region):
+    """Closed intervals ``[lower, upper]``, one per element of the batch.
+
+    Bounds may be infinite: ``(-inf, inf)`` is the whole line, and ``lower > upper``
+    (for instance ``(inf, -inf)``) is empty.
+
+    Args:
+        lower: Lower bounds, in residual coordinates.
+        upper: Upper bounds, the same shape as ``lower``.
+    """
+
+    event_dims = ()
+
+    def __init__(self, lower: Tensor, upper: Tensor):
+        lower = torch.as_tensor(lower)
+        upper = torch.as_tensor(upper, device=lower.device)
+        if lower.shape != upper.shape:
+            raise ValueError("lower and upper must have the same shape")
+        if not (lower.is_floating_point() and upper.is_floating_point()):
+            raise ValueError("interval bounds must be floating tensors")
+        self.lower = lower
+        self.upper = upper
+
+    @property
+    def batch_shape(self) -> torch.Size:
+        return self.lower.shape
+
+    @property
+    def dimension(self) -> int:
+        return 1
+
+    def tensors(self) -> dict[str, Tensor]:
+        return {"lower": self.lower, "upper": self.upper}
+
+    def is_empty(self) -> Tensor:
+        """True where ``lower > upper``."""
+        return self.upper < self.lower
+
+    def width(self) -> Tensor:
+        """``upper - lower``, with empty intervals at 0."""
+        return torch.where(self.is_empty(), 0.0, self.upper - self.lower)
+
+    def contains(self, residuals: Tensor) -> Tensor:
+        return (self.lower <= residuals) & (residuals <= self.upper)
+
+    def log_volume(self) -> Tensor:
+        width = self.upper - self.lower
+        return torch.where(self.is_empty(), -torch.inf, torch.log(width.clamp(min=0)))
+
+    def to_target(self, forecasts: Tensor) -> "Interval":
+        return Interval(self.lower + forecasts, self.upper + forecasts)
+
+
+__all__ = ["Interval"]

tscp/registry.py

@@ -0,0 +1,179 @@
+"""Short names for components and methods, and plugins that add more.
+
+Every built-in component and preset is registered under a kind and a short name, so that
+a config or a spec can name it::
+
+    from tscp import registry
+
+    weights = registry.build("weights", "exponential", decay=0.99)
+
+A package that adds components registers them the same way, with :func:`register`, and
+declares an entry point in the group ``tscp.plugins`` whose value is the module that does
+the registering. tscp imports those modules when a name or a class is not found, so a
+spec or config that names a plugin's component works once the plugin is installed.
+"""
+
+import importlib
+import importlib.metadata
+import inspect
+from collections.abc import Callable
+from typing import Any
+
+from tscp.utils.config import qualified_name
+
+KINDS = ("score", "source", "weights", "quantile", "update", "features", "method")
+"""Kinds of registered classes: the component axes, feature maps, and methods."""
+
+ENTRY_POINT_GROUP = "tscp.plugins"
+"""The entry-point group of plugins: each entry point names a module to import."""
+
+_ENTRIES: dict[str, dict[str, type]] = {kind: {} for kind in KINDS}
+_BUILTINS = (
+    "tscp.scores",
+    "tscp.sources",
+    "tscp.weights",
+    "tscp.quantile",
+    "tscp.online",
+    "tscp.features",
+    "tscp.methods",
+)
+_LOADED_PLUGINS: set[str] = set()
+
+
+def _check_kind(kind: str) -> None:
+    if kind not in KINDS:
+        raise ValueError(f"unknown kind {kind!r}; the kinds are {KINDS}")
+
+
+def register(kind: str, name: str, *, replace: bool = False) -> Callable[[type], type]:
+    """Class decorator that makes a class buildable as ``build(kind, name, ...)``.
+
+    Names are case-insensitive. Registering a name that another class holds raises a
+    ValueError, so two packages never silently override each other or tscp's own
+    components; pass ``replace=True`` to take a name on purpose. Registering the same
+    class again, as when its module is reloaded, does nothing.
+
+    Args:
+        kind: One of :data:`KINDS`.
+        name: The short name.
+        replace: Take the name even if another class holds it.
+    """
+    _check_kind(kind)
+    key = name.lower()
+
+    def decorator(cls: type) -> type:
+        current = _ENTRIES[kind].get(key)
+        same = current is not None and qualified_name(current) == qualified_name(cls)
+        if current is not None and not same and not replace:
+            raise ValueError(
+                f"{kind} {key!r} is already registered to {qualified_name(current)}; "
+                "pass replace=True to take the name"
+            )
+        _ENTRIES[kind][key] = cls
+        return cls
+
+    return decorator
+
+
+def entries(kind: str) -> dict[str, type]:
+    """The live name-to-class table of a kind.
+
+    It fills as modules register; unlike :func:`lookup`, reading it imports nothing.
+    """
+    _check_kind(kind)
+    return _ENTRIES[kind]
+
+
+def _load_builtins() -> None:
+    for module in _BUILTINS:
+        importlib.import_module(module)
+
+
+def load_plugins() -> list[str]:
+    """Import the module of every installed plugin not imported yet.
+
+    Returns:
+        The names of the entry points imported by this call.
+
+    Raises:
+        ImportError: If a plugin's module fails to import; the message names the plugin.
+    """
+    loaded = []
+    for entry_point in importlib.metadata.entry_points(group=ENTRY_POINT_GROUP):
+        if entry_point.name in _LOADED_PLUGINS:
+            continue
+        try:
+            entry_point.load()
+        except Exception as error:
+            raise ImportError(
+                f"the tscp plugin {entry_point.name!r} ({entry_point.value}) failed to load: "
+                f"{error}"
+            ) from error
+        _LOADED_PLUGINS.add(entry_point.name)
+        loaded.append(entry_point.name)
+    return loaded
+
+
+def lookup(kind: str, name: str) -> type:
+    """Return the class registered as ``name`` for ``kind``.
+
+    The built-in components are imported first; plugins are imported if the name is
+    still unknown.
+
+    Raises:
+        ValueError: If no built-in or installed plugin registers the name.
+    """
+    _check_kind(kind)
+    _load_builtins()
+    key = name.lower()
+    if key not in _ENTRIES[kind]:
+        load_plugins()
+    if key not in _ENTRIES[kind]:
+        raise ValueError(f"unknown {kind} {name!r}; registered: {sorted(_ENTRIES[kind])}")
+    return _ENTRIES[kind][key]
+
+
+def build(kind: str, name: str, *, ignore_unknown: bool = False, **kwargs: Any) -> Any:
+    """Build the class registered as ``name`` for ``kind`` with keyword arguments.
+
+    Args:
+        kind: One of :data:`KINDS`.
+        name: A registered name, in any case.
+        ignore_unknown: Drop keyword arguments the class does not take, which helps when
+            one config sweeps several classes. By default they raise.
+        **kwargs: Arguments of the class.
+
+    Raises:
+        ValueError: If the name is not registered.
+        TypeError: If an argument is unknown and ``ignore_unknown`` is False.
+    """
+    cls = lookup(kind, name)
+    parameters = inspect.signature(cls.__init__).parameters
+    takes_any = any(p.kind is inspect.Parameter.VAR_KEYWORD for p in parameters.values())
+    accepted = [p for p in parameters if p != "self"]
+    unknown = [] if takes_any else sorted(set(kwargs) - set(accepted))
+    if unknown:
+        if not ignore_unknown:
+            raise TypeError(f"{cls.__name__} does not take {unknown}; it takes {accepted}")
+        kwargs = {key: value for key, value in kwargs.items() if key in accepted}
+    return cls(**kwargs)
+
+
+def names(kind: str) -> list[str]:
+    """The registered names of a kind, built-ins and installed plugins included."""
+    _check_kind(kind)
+    _load_builtins()
+    load_plugins()
+    return sorted(_ENTRIES[kind])
+
+
+__all__ = [
+    "ENTRY_POINT_GROUP",
+    "KINDS",
+    "build",
+    "entries",
+    "load_plugins",
+    "lookup",
+    "names",
+    "register",
+]

tscp/scores/__init__.py

@@ -0,0 +1,9 @@
+"""Axis 2, score: how an error becomes a number."""
+
+from .absolute import AbsoluteScore
+from .base import Score
+from .cqr import CQRScore
+from .normalized import NormalizedScore
+from .signed import SignedScore
+
+__all__ = ["AbsoluteScore", "CQRScore", "NormalizedScore", "Score", "SignedScore"]

tscp/scores/absolute.py

@@ -0,0 +1,22 @@
+"""The absolute residual."""
+
+from torch import Tensor
+
+from tscp.data import ResidualArchive
+from tscp.regions import Interval
+from tscp.registry import register
+from tscp.scores.base import Score
+
+
+@register("score", "absolute")
+class AbsoluteScore(Score):
+    """``|r|``, per element; a threshold ``q`` gives the symmetric interval ``[-q, q]``."""
+
+    def forward(self, rows: ResidualArchive) -> Tensor:
+        return rows.residuals.abs()
+
+    def region(self, threshold: Tensor, rows: ResidualArchive) -> Interval:
+        return Interval(-threshold, threshold)
+
+
+__all__ = ["AbsoluteScore"]