stagereg 0.1.0__tar.gz

stagereg-0.1.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2026, Intelligent System Prognostics for Operations & Maintenance Research Group
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

stagereg-0.1.0/PKG-INFO

@@ -0,0 +1,115 @@
+Metadata-Version: 2.4
+Name: stagereg
+Version: 0.1.0
+Summary: Stage-oriented probabilistic regression
+Author: Giorgio Tosti Balducci, Intelligent System Prognostics for Operations & Maintenance Research Group
+Author-email: Giorgio Tosti Balducci <g.b.l.tostibalducci@tudelft.nl>
+License-Expression: BSD-3-Clause
+License-File: LICENSE
+Classifier: Development Status :: 3 - Alpha
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3 :: Only
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Dist: numpy>=2.0
+Requires-Dist: flax>=0.11 ; extra == 'jax'
+Requires-Dist: jax>=0.6 ; extra == 'jax'
+Requires-Dist: optax>=0.2.4 ; extra == 'jax'
+Requires-Dist: matplotlib>=3.9 ; extra == 'plot'
+Requires-Python: >=3.12
+Project-URL: Repository, https://github.com/GroupiSP/stagereg
+Project-URL: Documentation, https://groupisp.github.io/stagereg/
+Project-URL: Issues, https://github.com/GroupiSP/stagereg/issues
+Provides-Extra: jax
+Provides-Extra: plot
+Description-Content-Type: text/markdown
+
+# stagereg
+
+[![PyPI version](https://img.shields.io/pypi/v/stagereg)](https://pypi.org/project/stagereg/)
+[![Python versions](https://img.shields.io/pypi/pyversions/stagereg)](https://pypi.org/project/stagereg/)
+[![License: BSD-3-Clause](https://img.shields.io/badge/license-BSD--3--Clause-blue)](https://github.com/GroupiSP/stagereg/blob/main/LICENSE)
+
+`stagereg` is a library for stage-oriented probabilistic regression: training
+models that predict a distribution over the target, not just a point estimate.
+Training runs as a sequence of explicit, checkpointed stages, such as fitting a
+mean model and then a separate variance model on its residuals. A core of
+backend-neutral contracts is implemented by an optional JAX backend built on
+Flax NNX and Optax.
+
+**Documentation:** <https://groupisp.github.io/stagereg/>
+
+## Install
+
+<!-- --8<-- [start:install] -->
+Install with the JAX backend (JAX, Flax and Optax):
+
+```bash
+pip install "stagereg[jax]"
+```
+
+The optional `plot` extra also installs matplotlib, for plotting results:
+
+```bash
+pip install "stagereg[jax,plot]"
+```
+
+With uv, use the equivalent `uv add` form:
+
+```bash
+uv add "stagereg[jax]"
+uv add "stagereg[jax,plot]"
+```
+
+Drop the `[jax]` extra for the backend-neutral contracts alone, which depend
+only on NumPy.
+
+The documentation tracks the `main` branch, with no per-release docs before
+1.0, so it may describe features newer than your installed version.
+<!-- --8<-- [end:install] -->
+
+## Minimal example
+
+<!-- --8<-- [start:minimal-example] -->
+Fit a linear model to `y = 3x + 2` for a fixed number of epochs with the JAX
+backend:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from stagereg.core.types import Batch
+from stagereg.jax import create_optimizer, initialize_training_state, run_supervised
+
+
+def squared_error(model, inputs, targets, sample_weight, key, training):
+    return jnp.mean(jnp.square(model(inputs) - targets))
+
+
+inputs = jnp.linspace(-1.0, 1.0, 32).reshape(-1, 1)
+targets = 3.0 * inputs + 2.0
+
+
+def loader(*, split, epoch):
+    return [Batch(inputs=inputs, targets=targets)]
+
+
+model = nnx.Linear(1, 1, rngs=nnx.Rngs(0))
+optimizer = create_optimizer(model, optax.sgd(learning_rate=0.1))
+state = initialize_training_state(model, optimizer, rng_key=jax.random.key(0))
+
+result = run_supervised(
+    model=model,
+    optimizer=optimizer,
+    train_loader=loader,
+    loss=squared_error,
+    state=state,
+    epochs=3,
+)
+print(f"final training loss: {result.loss:.4f}")
+```
+<!-- --8<-- [end:minimal-example] -->

stagereg-0.1.0/README.md

@@ -0,0 +1,87 @@
+# stagereg
+
+[![PyPI version](https://img.shields.io/pypi/v/stagereg)](https://pypi.org/project/stagereg/)
+[![Python versions](https://img.shields.io/pypi/pyversions/stagereg)](https://pypi.org/project/stagereg/)
+[![License: BSD-3-Clause](https://img.shields.io/badge/license-BSD--3--Clause-blue)](https://github.com/GroupiSP/stagereg/blob/main/LICENSE)
+
+`stagereg` is a library for stage-oriented probabilistic regression: training
+models that predict a distribution over the target, not just a point estimate.
+Training runs as a sequence of explicit, checkpointed stages, such as fitting a
+mean model and then a separate variance model on its residuals. A core of
+backend-neutral contracts is implemented by an optional JAX backend built on
+Flax NNX and Optax.
+
+**Documentation:** <https://groupisp.github.io/stagereg/>
+
+## Install
+
+<!-- --8<-- [start:install] -->
+Install with the JAX backend (JAX, Flax and Optax):
+
+```bash
+pip install "stagereg[jax]"
+```
+
+The optional `plot` extra also installs matplotlib, for plotting results:
+
+```bash
+pip install "stagereg[jax,plot]"
+```
+
+With uv, use the equivalent `uv add` form:
+
+```bash
+uv add "stagereg[jax]"
+uv add "stagereg[jax,plot]"
+```
+
+Drop the `[jax]` extra for the backend-neutral contracts alone, which depend
+only on NumPy.
+
+The documentation tracks the `main` branch, with no per-release docs before
+1.0, so it may describe features newer than your installed version.
+<!-- --8<-- [end:install] -->
+
+## Minimal example
+
+<!-- --8<-- [start:minimal-example] -->
+Fit a linear model to `y = 3x + 2` for a fixed number of epochs with the JAX
+backend:
+
+```python
+import jax
+import jax.numpy as jnp
+import optax
+from flax import nnx
+
+from stagereg.core.types import Batch
+from stagereg.jax import create_optimizer, initialize_training_state, run_supervised
+
+
+def squared_error(model, inputs, targets, sample_weight, key, training):
+    return jnp.mean(jnp.square(model(inputs) - targets))
+
+
+inputs = jnp.linspace(-1.0, 1.0, 32).reshape(-1, 1)
+targets = 3.0 * inputs + 2.0
+
+
+def loader(*, split, epoch):
+    return [Batch(inputs=inputs, targets=targets)]
+
+
+model = nnx.Linear(1, 1, rngs=nnx.Rngs(0))
+optimizer = create_optimizer(model, optax.sgd(learning_rate=0.1))
+state = initialize_training_state(model, optimizer, rng_key=jax.random.key(0))
+
+result = run_supervised(
+    model=model,
+    optimizer=optimizer,
+    train_loader=loader,
+    loss=squared_error,
+    state=state,
+    epochs=3,
+)
+print(f"final training loss: {result.loss:.4f}")
+```
+<!-- --8<-- [end:minimal-example] -->

stagereg-0.1.0/pyproject.toml

@@ -0,0 +1,106 @@
+[project]
+name = "stagereg"
+version = "0.1.0"
+description = "Stage-oriented probabilistic regression"
+readme = "README.md"
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+requires-python = ">=3.12"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = ["numpy>=2.0"]
+
+[[project.authors]]
+name = "Giorgio Tosti Balducci"
+email = "g.b.l.tostibalducci@tudelft.nl"
+
+[[project.authors]]
+name = "Intelligent System Prognostics for Operations & Maintenance Research Group"
+
+[project.optional-dependencies]
+jax = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+]
+plot = ["matplotlib>=3.9"]
+
+[project.urls]
+Repository = "https://github.com/GroupiSP/stagereg"
+Documentation = "https://groupisp.github.io/stagereg/"
+Issues = "https://github.com/GroupiSP/stagereg/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.3,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+source-include = ["tests/**"]
+source-exclude = [
+    "**/__pycache__",
+    "**/*.pyc",
+]
+
+[tool.docformatter]
+recursive = true
+black = true
+style = "google"
+
+[tool.pytest.ini_options]
+testpaths = [
+    "tests",
+    "docs_site",
+    "README.md",
+]
+addopts = [
+    "--markdown-docs",
+    "--markdown-docs-syntax=superfences",
+]
+
+[tool.pyright]
+typeCheckingMode = "basic"
+reportUnusedVariable = false
+
+[tool.ruff.lint.per-file-ignores]
+"examples/jax/cmapss/*.py" = ["E402"]
+"tests/jax/test_cmapss_*.py" = ["E402"]
+"examples/jax/tracking/*.py" = ["E402"]
+
+[dependency-groups]
+dev = [
+    "docformatter>=1.7.7",
+    "hypothesis>=6.140.4",
+    "pre-commit>=4.5.1",
+    "pytest>=9.0.2",
+    "pytest-markdown-docs>=0.9.2",
+    "pyyaml>=6.0",
+    "ruff>=0.14.9",
+]
+docs = [
+    "mkdocs-material>=9.7.7",
+    "mkdocstrings[python]>=1.0.6",
+]
+example-cmapss = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+    "pandas>=2.2",
+    "scikit-learn>=1.9.0",
+    "seaborn==0.13.2",
+    "matplotlib>=3.11.0",
+]
+example-tracking = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+    "matplotlib>=3.11.0",
+    "tensorboard>=2.20",
+    "tensorboardX>=2.6.4",
+]

stagereg-0.1.0/pyproject.toml.orig

@@ -0,0 +1,101 @@
+[project]
+name = "stagereg"
+version = "0.1.0"
+description = "Stage-oriented probabilistic regression"
+readme = "README.md"
+authors = [
+    { name = "Giorgio Tosti Balducci", email = "g.b.l.tostibalducci@tudelft.nl" },
+    { name = "Intelligent System Prognostics for Operations & Maintenance Research Group" },
+]
+license = "BSD-3-Clause"
+license-files = ["LICENSE"]
+requires-python = ">=3.12"
+classifiers = [
+    "Development Status :: 3 - Alpha",
+    "Intended Audience :: Science/Research",
+    "Programming Language :: Python :: 3",
+    "Programming Language :: Python :: 3 :: Only",
+    "Programming Language :: Python :: 3.12",
+    "Programming Language :: Python :: 3.13",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy>=2.0",
+]
+
+[project.optional-dependencies]
+jax = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+]
+plot = [
+    "matplotlib>=3.9",
+]
+
+[project.urls]
+Repository = "https://github.com/GroupiSP/stagereg"
+Documentation = "https://groupisp.github.io/stagereg/"
+Issues = "https://github.com/GroupiSP/stagereg/issues"
+
+[build-system]
+requires = ["uv_build>=0.12.3,<0.13.0"]
+build-backend = "uv_build"
+
+[tool.uv.build-backend]
+# The sdist also carries the tests, so downstream packagers can check the build. The examples
+# stay in the repository, so the tests that exercise them need a repository checkout.
+source-include = ["tests/**"]
+source-exclude = ["**/__pycache__", "**/*.pyc"]
+
+[dependency-groups]
+dev = [
+    "docformatter>=1.7.7",
+    "hypothesis>=6.140.4",
+    "pre-commit>=4.5.1",
+    "pytest>=9.0.2",
+    "pytest-markdown-docs>=0.9.2",
+    "pyyaml>=6.0",
+    "ruff>=0.14.9",
+]
+docs = [
+    "mkdocs-material>=9.7.7",
+    "mkdocstrings[python]>=1.0.6",
+]
+
+example-cmapss = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+    "pandas>=2.2",
+    "scikit-learn>=1.9.0",
+    "seaborn==0.13.2",
+    "matplotlib>=3.11.0"
+]
+
+example-tracking = [
+    "flax>=0.11",
+    "jax>=0.6",
+    "optax>=0.2.4",
+    "matplotlib>=3.11.0",
+    "tensorboard>=2.20",
+    "tensorboardX>=2.6.4",
+]
+
+[tool.docformatter]
+recursive = true
+black = true
+style = "google"
+
+[tool.pytest.ini_options]
+testpaths = ["tests", "docs_site", "README.md"]
+addopts = ["--markdown-docs", "--markdown-docs-syntax=superfences"]
+
+[tool.pyright]
+typeCheckingMode = "basic"
+reportUnusedVariable = false
+
+[tool.ruff.lint.per-file-ignores]
+"examples/jax/cmapss/*.py" = ["E402"]
+"tests/jax/test_cmapss_*.py" = ["E402"]
+"examples/jax/tracking/*.py" = ["E402"]

stagereg-0.1.0/src/stagereg/__init__.py

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

stagereg-0.1.0/src/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-0.1.0/src/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