stagereg 0.1.0__py3-none-any.whl

stagereg/__init__.py

@@ -0,0 +1 @@
+"""Probabilistic regression with explicit, composable training stages."""

stagereg/core/__init__.py

@@ -0,0 +1,149 @@
+"""Backend-neutral contracts and utilities for probabilistic regression."""
+
+from stagereg.core.checkpoints import (
+    Checkpoint,
+    CheckpointStore,
+    InMemoryCheckpointStore,
+)
+from stagereg.core.distributions import (
+    DistributionHead,
+    DistributionLoss,
+    Likelihood,
+    Loss,
+    PredictionLoss,
+    PredictiveDistribution,
+)
+from stagereg.core.early_stopping import (
+    EarlyStopper,
+    EarlyStoppingDecision,
+    EarlyStoppingState,
+    OptimizationMode,
+)
+from stagereg.core.losses import (
+    BetaNLLLoss,
+    GaussianNLLLoss,
+    NegativeLogLikelihoodLoss,
+    SquaredErrorLoss,
+    add_epsilon,
+)
+from stagereg.core.metric_registry import (
+    ContinuousRankedProbabilityScore,
+    EpochMetric,
+    EpochPredictionData,
+    EvaluationGrid,
+    IntervalCoverage,
+    MetricRequirements,
+    NegativeLogLikelihood,
+    PointContinuousRankedProbabilityScore,
+    PredictionInterval,
+    RootMeanSquaredError,
+    SampleContinuousRankedProbabilityScore,
+    WeightedSpread,
+)
+from stagereg.core.metrics import (
+    cdf,
+    coverage,
+    crps,
+    point_crps,
+    rmse,
+    sample_crps,
+    wsu,
+)
+from stagereg.core.naming import (
+    MetricTag,
+    Split,
+    flatten_parameters,
+    metric_tag,
+    parse_metric_tag,
+)
+from stagereg.core.protocols import (
+    Dataset,
+    LoaderFactory,
+    Optimizer,
+    Step,
+    ValidationStrategy,
+)
+from stagereg.core.stages import StageState, TrainingStage, validate_transition
+from stagereg.core.tracking import (
+    Decision,
+    EventSink,
+    ExperimentTracker,
+    TrackerEventSink,
+    TrainingEvent,
+)
+from stagereg.core.types import (
+    Array,
+    Batch,
+    CheckpointRef,
+    ParameterRole,
+    PyTree,
+    StageResult,
+    TrainingState,
+    ValidationResult,
+)
+
+__all__ = [
+    "Array",
+    "Batch",
+    "BetaNLLLoss",
+    "Checkpoint",
+    "CheckpointRef",
+    "CheckpointStore",
+    "ContinuousRankedProbabilityScore",
+    "Dataset",
+    "Decision",
+    "DistributionHead",
+    "DistributionLoss",
+    "EarlyStopper",
+    "EarlyStoppingDecision",
+    "EarlyStoppingState",
+    "EpochMetric",
+    "EpochPredictionData",
+    "EvaluationGrid",
+    "EventSink",
+    "ExperimentTracker",
+    "GaussianNLLLoss",
+    "InMemoryCheckpointStore",
+    "IntervalCoverage",
+    "Likelihood",
+    "LoaderFactory",
+    "Loss",
+    "MetricRequirements",
+    "MetricTag",
+    "NegativeLogLikelihood",
+    "NegativeLogLikelihoodLoss",
+    "OptimizationMode",
+    "Optimizer",
+    "ParameterRole",
+    "PointContinuousRankedProbabilityScore",
+    "PredictionInterval",
+    "PredictionLoss",
+    "PredictiveDistribution",
+    "PyTree",
+    "RootMeanSquaredError",
+    "SampleContinuousRankedProbabilityScore",
+    "Split",
+    "SquaredErrorLoss",
+    "StageResult",
+    "StageState",
+    "Step",
+    "TrackerEventSink",
+    "TrainingEvent",
+    "TrainingStage",
+    "TrainingState",
+    "ValidationResult",
+    "ValidationStrategy",
+    "WeightedSpread",
+    "add_epsilon",
+    "cdf",
+    "coverage",
+    "crps",
+    "flatten_parameters",
+    "metric_tag",
+    "parse_metric_tag",
+    "point_crps",
+    "rmse",
+    "sample_crps",
+    "validate_transition",
+    "wsu",
+]

stagereg/core/checkpoints.py

@@ -0,0 +1,100 @@
+"""Checkpoint value objects and persistence protocol."""
+
+from __future__ import annotations
+
+from collections.abc import Mapping
+from dataclasses import dataclass, field
+from typing import Any, Protocol
+
+from stagereg.core.types import TrainingState
+
+
+@dataclass(frozen=True)
+class Checkpoint:
+    """The complete state necessary to resume staged training.
+
+    Attributes:
+        state: Snapshot of the [`TrainingState`][stagereg.core.TrainingState] at
+            the time the checkpoint was saved.
+        epoch: Epoch of the active stage at which the checkpoint was saved.
+        parameters: Snapshot of the model's parameters, in the backend's own
+            format. Defaults to ``None``.
+        optimizer_state: Snapshot of the optimizer's state, in the backend's own
+            format. Defaults to ``None``.
+        rng_state: Random key to resume training with. It mirrors the
+            ``rng_state`` of the embedded ``state``. Defaults to ``None``.
+        early_stopping_state: State of the early stopper, such as an
+            [`EarlyStoppingState`][stagereg.core.EarlyStoppingState], so that
+            patience and the best value carry over on resume. Defaults to
+            ``None``.
+        metadata: Information about the checkpoint, such as the stage that saved
+            it. Defaults to an empty mapping.
+    """
+
+    state: TrainingState
+    epoch: int
+    parameters: Any | None = None
+    optimizer_state: Any | None = None
+    rng_state: Any | None = None
+    early_stopping_state: Any | None = None
+    metadata: Mapping[str, Any] = field(default_factory=dict)
+
+
+class CheckpointStore(Protocol):
+    """Persists checkpoints under opaque keys."""
+
+    def save(self, key: str, checkpoint: Checkpoint) -> None: ...
+
+    def load(self, key: str) -> Checkpoint: ...
+
+    def exists(self, key: str) -> bool: ...
+
+
+class InMemoryCheckpointStore:
+    """A [`CheckpointStore`][stagereg.core.CheckpointStore] kept in process memory.
+
+    Checkpoints live only for the lifetime of the store, which makes it the natural
+    choice for examples, tests, and short runs that only need the best checkpoint of the
+    current process restored at the end. Saving under an existing key replaces the
+    checkpoint held there.
+    """
+
+    def __init__(self) -> None:
+        self._checkpoints: dict[str, Checkpoint] = {}
+
+    def save(self, key: str, checkpoint: Checkpoint) -> None:
+        """Store a checkpoint under a key, replacing any checkpoint there.
+
+        Args:
+            key: The opaque key to store the checkpoint under.
+            checkpoint: The checkpoint to store.
+
+        Returns:
+            None.
+        """
+        self._checkpoints[key] = checkpoint
+
+    def load(self, key: str) -> Checkpoint:
+        """Return the checkpoint stored under a key.
+
+        Args:
+            key: The key the checkpoint was saved under.
+
+        Returns:
+            The stored checkpoint.
+
+        Raises:
+            KeyError: If no checkpoint has been saved under ``key``.
+        """
+        return self._checkpoints[key]
+
+    def exists(self, key: str) -> bool:
+        """Report whether a checkpoint is stored under a key.
+
+        Args:
+            key: The key to look up.
+
+        Returns:
+            ``True`` if a checkpoint is stored under ``key``.
+        """
+        return key in self._checkpoints

stagereg/core/distributions.py

@@ -0,0 +1,60 @@
+"""Protocols for distribution-valued model predictions and objectives."""
+
+from __future__ import annotations
+
+from typing import Protocol
+
+from stagereg.core.types import Array, Batch
+
+
+class PredictiveDistribution(Protocol):
+    """A probability distribution produced by a regression model.
+
+    This protocol is the backend-neutral model/loss boundary. Epoch metrics use backend
+    adapters to materialize host arrays rather than retaining distribution objects,
+    which intentionally have no concatenation or indexing contract.
+    """
+
+    @property
+    def batch_shape(self) -> tuple[int, ...]: ...
+
+    @property
+    def event_shape(self) -> tuple[int, ...]: ...
+
+    def log_prob(self, targets: Array) -> Array: ...
+
+    def sample(self, key: Array, sample_shape: tuple[int, ...] = ()) -> Array: ...
+
+    def mean(self) -> Array: ...
+
+    def variance(self) -> Array: ...
+
+
+class DistributionHead(Protocol):
+    """Produces a predictive distribution from model features."""
+
+    def __call__(self, features: Array) -> PredictiveDistribution: ...
+
+
+class Likelihood(Protocol):
+    """Evaluates a target under a predictive distribution."""
+
+    def log_prob(self, prediction: PredictiveDistribution, targets: Array) -> Array: ...
+
+
+class DistributionLoss(Protocol):
+    """Computes an unreduced loss from a predictive distribution."""
+
+    def per_example(
+        self, prediction: PredictiveDistribution, batch: Batch
+    ) -> Array: ...
+
+
+class PredictionLoss(Protocol):
+    """Computes an unreduced loss from deterministic array predictions."""
+
+    def per_example(self, prediction: Array, batch: Batch) -> Array: ...
+
+
+# Backward-compatible alias for the original distribution-loss protocol.
+Loss = DistributionLoss

stagereg/core/early_stopping.py

@@ -0,0 +1,252 @@
+"""Backend-neutral early-stopping policies."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from enum import StrEnum
+from math import isfinite
+
+from stagereg.core.naming import Split
+
+
+class OptimizationMode(StrEnum):
+    """The direction in which a monitored metric is considered improved."""
+
+    MIN = "min"
+    MAX = "max"
+
+
+@dataclass(frozen=True)
+class EarlyStoppingState:
+    """Checkpointable state for a configured early-stopping policy.
+
+    Attributes:
+        metric: Name of the monitored metric.
+        mode: Direction in which the metric is considered improved.
+        min_delta: Minimum change required to qualify as an improvement.
+        patience: Number of consecutive non-improving epochs tolerated
+            before stopping.
+        source: Split the monitored metric is measured on.
+        best_score: Best value observed so far, or ``None`` if no value
+            has been observed yet.
+        best_epoch: Epoch at which ``best_score`` was observed, or
+            ``None`` if no value has been observed yet.
+        non_improving_epochs: Number of consecutive epochs without an
+            improvement.
+        stopped: Whether patience has been exhausted and training should
+            stop.
+    """
+
+    metric: str
+    mode: OptimizationMode
+    min_delta: float
+    patience: int
+    source: Split = Split.VALIDATION
+    best_score: float | None = None
+    best_epoch: int | None = None
+    non_improving_epochs: int = 0
+    stopped: bool = False
+
+    def __post_init__(self) -> None:
+        """Validate checkpointable early-stopping state.
+
+        Raises:
+            ValueError: If configuration or observation fields cannot represent
+                a state produced by [`EarlyStopper`][stagereg.core.EarlyStopper].
+        """
+        if not self.metric:
+            raise ValueError("metric must be a non-empty string.")
+        if self.patience < 0:
+            raise ValueError("patience must be non-negative.")
+        if not isfinite(self.min_delta) or self.min_delta < 0:
+            raise ValueError("min_delta must be finite and non-negative.")
+        if self.best_score is not None and not isfinite(self.best_score):
+            raise ValueError("best_score must be finite when provided.")
+        if self.best_epoch is not None and self.best_epoch < 0:
+            raise ValueError("best_epoch must be non-negative when provided.")
+        if (self.best_score is None) != (self.best_epoch is None):
+            raise ValueError("best_score and best_epoch must be provided together.")
+        if self.non_improving_epochs < 0:
+            raise ValueError("non_improving_epochs must be non-negative.")
+        if self.best_score is None and self.non_improving_epochs != 0:
+            raise ValueError(
+                "non_improving_epochs must be zero before the first observation."
+            )
+        expected_stopped = self.non_improving_epochs > self.patience
+        if self.stopped is not expected_stopped:
+            raise ValueError(
+                "stopped must indicate whether non-improving epochs exceed patience."
+            )
+
+
+@dataclass(frozen=True)
+class EarlyStoppingDecision:
+    """The result of observing one monitored metric value.
+
+    Attributes:
+        state: Updated early-stopping state after the observation.
+        improved: Whether the observed value improved on the best score.
+    """
+
+    state: EarlyStoppingState
+    improved: bool
+
+    @property
+    def should_stop(self) -> bool:
+        """Whether the configured patience has been exhausted."""
+        return self.state.stopped
+
+
+class EarlyStopper:
+    """Tracks a scalar metric and decides when training should stop."""
+
+    def __init__(
+        self,
+        *,
+        metric: str,
+        mode: OptimizationMode | str,
+        patience: int,
+        min_delta: float = 0.0,
+        source: Split | str = Split.VALIDATION,
+        state: EarlyStoppingState | None = None,
+    ) -> None:
+        """Initialize an early stopper.
+
+        Args:
+            metric: Name of the metric to monitor.
+            mode: Direction in which the metric is considered improved,
+                either ``"min"`` or ``"max"``.
+            patience: Number of consecutive non-improving epochs
+                tolerated before stopping.
+            min_delta: Minimum change required to qualify as an
+                improvement. Defaults to ``0.0``.
+            source: Split the monitored metric is measured on, either
+                ``"train"`` or ``"validation"``. Defaults to
+                ``Split.VALIDATION``.
+            state: Previously checkpointed state to resume from. Must
+                match the configuration passed to this constructor.
+
+        Raises:
+            ValueError: If ``metric`` is empty, ``patience`` is
+                negative, ``min_delta`` is not a finite non-negative
+                number, ``mode`` or ``source`` is not a recognised value,
+                or ``state`` does not match the given configuration.
+        """
+        if not metric:
+            raise ValueError("metric must be a non-empty string.")
+        if patience < 0:
+            raise ValueError("patience must be non-negative.")
+        if not isfinite(min_delta) or min_delta < 0:
+            raise ValueError("min_delta must be finite and non-negative.")
+
+        normalized_mode = OptimizationMode(mode)
+        normalized_source = Split(source)
+        self._state = state or EarlyStoppingState(
+            metric=metric,
+            mode=normalized_mode,
+            min_delta=min_delta,
+            patience=patience,
+            source=normalized_source,
+        )
+        if (
+            self._state.metric != metric
+            or self._state.mode != normalized_mode
+            or self._state.min_delta != min_delta
+            or self._state.patience != patience
+            or self._state.source != normalized_source
+        ):
+            raise ValueError("state does not match the early-stopper configuration.")
+
+    @property
+    def state(self) -> EarlyStoppingState:
+        """Return the current checkpointable stopping state."""
+        return self._state
+
+    def expects_validation(self) -> bool:
+        """Whether the monitored metric must come from validation.
+
+        Returns:
+            ``True`` if the metric configured for this stopper is
+            measured on the validation split (i.e. requires a
+            [`ValidationStrategy`][stagereg.core.ValidationStrategy] to be
+            supplied to the training loop), ``False`` if it is measured
+            on the training split.
+        """
+        return self._state.source is Split.VALIDATION
+
+    def monitored_metric_name(self) -> str:
+        """Return the name of the metric this stopper monitors.
+
+        Returns:
+            The metric name passed to the constructor, used by callers
+            to look up the corresponding value in per-epoch metrics.
+        """
+        return self._state.metric
+
+    def observe(self, value: float, *, epoch: int) -> EarlyStoppingDecision:
+        """Record a metric observed after ``epoch`` and return its decision.
+
+        Args:
+            value: The monitored metric value observed for this epoch.
+            epoch: The epoch at which ``value`` was observed.
+
+        Returns:
+            An [`EarlyStoppingDecision`][stagereg.core.EarlyStoppingDecision]
+            with the updated state and whether ``value`` improved on the best
+            score so far.
+
+        Raises:
+            ValueError: If ``epoch`` is negative or ``value`` is not
+                finite.
+            RuntimeError: If early stopping has already been triggered.
+        """
+        if epoch < 0:
+            raise ValueError("epoch must be non-negative.")
+        if not isfinite(value):
+            raise ValueError("observed metric value must be finite.")
+        if self._state.stopped:
+            raise RuntimeError("cannot observe a metric after early stopping.")
+
+        improved = self._is_improvement(value)
+        if improved:
+            self._state = EarlyStoppingState(
+                metric=self._state.metric,
+                mode=self._state.mode,
+                min_delta=self._state.min_delta,
+                patience=self._state.patience,
+                source=self._state.source,
+                best_score=value,
+                best_epoch=epoch,
+            )
+        else:
+            non_improving_epochs = self._state.non_improving_epochs + 1
+            self._state = EarlyStoppingState(
+                metric=self._state.metric,
+                mode=self._state.mode,
+                min_delta=self._state.min_delta,
+                patience=self._state.patience,
+                source=self._state.source,
+                best_score=self._state.best_score,
+                best_epoch=self._state.best_epoch,
+                non_improving_epochs=non_improving_epochs,
+                stopped=non_improving_epochs > self._state.patience,
+            )
+        return EarlyStoppingDecision(state=self._state, improved=improved)
+
+    def _is_improvement(self, value: float) -> bool:
+        """Determine whether ``value`` improves on the current best score.
+
+        Args:
+            value: The candidate metric value.
+
+        Returns:
+            ``True`` if ``value`` is an improvement over the current best
+            score given the configured mode and ``min_delta``.
+        """
+        if self._state.best_score is None:
+            return True
+        match self._state.mode:
+            case OptimizationMode.MIN:
+                return value < self._state.best_score - self._state.min_delta
+            case OptimizationMode.MAX:
+                return value > self._state.best_score + self._state.min_delta

stagereg/core/losses.py

@@ -0,0 +1,130 @@
+"""Backend-neutral per-example objectives for probabilistic regression."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable
+from dataclasses import dataclass, field
+
+from stagereg.core.distributions import PredictiveDistribution
+from stagereg.core.types import Array, Batch
+
+
+def _identity(value: Array) -> Array:
+    """Return ``value`` unchanged; the default no-op gradient-stopping hook."""
+    return value
+
+
+@dataclass(frozen=True)
+class NegativeLogLikelihoodLoss:
+    """Configurable negative log-likelihood for predictive distributions.
+
+    Attributes:
+        beta: Variance-reweighting exponent in ``[0, 1]``. Zero recovers
+            ordinary negative log-likelihood.
+        stop_gradient: Callable excluding the variance weight from gradient
+            propagation. Defaults to the identity, which lets gradients flow
+            through the weight. With JAX and ``beta > 0``, pass
+            ``jax.lax.stop_gradient``; otherwise a different objective from
+            beta-NLL is trained.
+        target_transform: Callable applied to targets before evaluating the
+            predictive distribution. Defaults to the identity transform.
+    """
+
+    beta: float = 0.0
+    stop_gradient: Callable[[Array], Array] = field(default=_identity)
+    target_transform: Callable[[Array], Array] = field(default=_identity)
+
+    def __post_init__(self) -> None:
+        """Validate the variance-reweighting exponent.
+
+        Raises:
+            ValueError: If ``beta`` is outside ``[0, 1]``.
+        """
+        if not 0.0 <= self.beta <= 1.0:
+            raise ValueError("beta must be within [0, 1].")
+
+    def per_example(self, prediction: PredictiveDistribution, batch: Batch) -> Array:
+        """Compute the optionally transformed and beta-weighted NLL.
+
+        Args:
+            prediction: The predictive distribution produced for ``batch``.
+            batch: Batch whose targets are scored under ``prediction``.
+
+        Returns:
+            An array of per-example negative log-likelihood values.
+
+        Raises:
+            ValueError: If ``batch.targets`` is missing.
+        """
+        if batch.targets is None:
+            raise ValueError("batch.targets must be provided.")
+        targets = self.target_transform(batch.targets)
+        nll = -prediction.log_prob(targets)
+        if self.beta == 0.0:
+            return nll
+        weight = self.stop_gradient(prediction.variance()) ** self.beta
+        return nll * weight
+
+
+@dataclass(frozen=True)
+class SquaredErrorLoss:
+    """Per-example squared error for deterministic predictions."""
+
+    def per_example(self, prediction: Array, batch: Batch) -> Array:
+        """Compute squared prediction errors.
+
+        Args:
+            prediction: Deterministic model predictions.
+            batch: Batch containing regression targets.
+
+        Returns:
+            Elementwise squared prediction errors.
+
+        Raises:
+            ValueError: If ``batch.targets`` is missing.
+        """
+        if batch.targets is None:
+            raise ValueError("batch.targets must be provided.")
+        if prediction.shape != batch.targets.shape:
+            raise ValueError("prediction and batch.targets must have matching shapes.")
+        return (prediction - batch.targets) ** 2
+
+
+def add_epsilon(epsilon: float = 1e-12) -> Callable[[Array], Array]:
+    """Create a target transform that adds a positive stabilization offset.
+
+    Args:
+        epsilon: Positive finite offset added to each target.
+
+    Returns:
+        A callable adding ``epsilon`` to an array.
+
+    Raises:
+        ValueError: If ``epsilon`` is not positive and finite.
+    """
+    if not math.isfinite(epsilon) or epsilon <= 0.0:
+        raise ValueError("epsilon must be positive and finite.")
+
+    def transform(targets: Array) -> Array:
+        return targets + epsilon
+
+    return transform
+
+
+# Compatibility aliases for the public pre-consolidation names. Both preserve
+# their original constructor behavior while the canonical API uses
+# NegativeLogLikelihoodLoss.
+GaussianNLLLoss = NegativeLogLikelihoodLoss
+"""Alias of [`NegativeLogLikelihoodLoss`][stagereg.core.NegativeLogLikelihoodLoss].
+
+Kept for compatibility with the pre-consolidation name; prefer the canonical name.
+"""
+
+BetaNLLLoss = NegativeLogLikelihoodLoss
+"""Alias of [`NegativeLogLikelihoodLoss`][stagereg.core.NegativeLogLikelihoodLoss].
+
+Kept for compatibility with the pre-consolidation name; prefer the canonical name. The
+beta-NLL weighting is the `beta` argument of
+[`NegativeLogLikelihoodLoss`][stagereg.core.NegativeLogLikelihoodLoss].
+"""