torchregress

A PyTorch library for regression, uncertainty quantification, and robust estimation

torchregress is a PyTorch library of regression losses, metrics, and calibration tools for problems where you need more than a single point prediction — uncertainty, robustness, and messy real-world data included.

It plugs into standard training loops and supports missing data and sample weights out of the box.

Features & Capabilities

torchregress is built for regression problems that are messy. In plain terms, it helps you:

• Predict distributions — Train models that output means, spreads, quantiles, or full predictive distributions when a single point estimate is not enough.
• Measure and separate uncertainty — Distinguish irreducible data noise from model uncertainty, and evaluate whether predicted ranges are trustworthy.
• Stay robust on dirty data — Handle missing values, sample weights, outliers, noisy inputs and labels, censored targets, and imbalanced or rare outcomes without bolting on ad hoc fixes.
• Calibrate and guarantee coverage — Turn a trained model into well-calibrated intervals or conformal prediction sets with stated coverage properties.
• Adapt to distribution shifts — Lightweight test-time tools when deployment data differs from what you trained on.
• Evaluate the right metrics — Metrics for point error, interval quality, distributional accuracy, and calibration — not just average squared error.

Every loss and metric drops into a normal PyTorch training loop. For method names, API details, and worked examples, see the documentation.

Metrics are tested against reference implementations (scoringrules, properscoring, scipy, scikit-learn), and before 0.3.0 the whole library went through an audit in which every defect found got a regression test.

Installation

From PyPI

pip install torchregress

Optional extras:

pip install "torchregress[flows]"   # normalizing flows (zuko)
pip install "torchregress[viz]"     # plotting helpers (matplotlib)

torchregress needs Python 3.12 or newer (tested on 3.12, 3.13 and 3.14) and PyTorch 2.13 or newer. Install the PyTorch build for your hardware first (see pytorch.org) if you need CUDA.

Dev / Source Setup (pixi)

This project uses pixi for development.

git clone https://github.com/astroai/torchregress.git
cd torchregress
pixi install

pixi run test        # pytest + coverage
pixi run lint        # ruff check + format --check
pixi run typecheck   # ty
pixi run docs        # zensical build --strict
pixi run ci          # lint + typecheck + test + docs

Onboarding Recipes

1. Parametric Aleatoric UQ (Heteroscedastic Gaussian)

Train a network that outputs both a mean $\mu(x)$ and log-variance $\log\sigma^2(x)$ using a proper scoring rule (Gaussian NLL), and evaluate with the Continuous Ranked Probability Score (CRPS).

import torch
import torch.nn as nn
from torchregress.losses import GaussianNLLLoss
from torchregress.metrics import crps_gaussian

# Model outputs [mean, log_variance]
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 2))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
loss_fn = GaussianNLLLoss()

# Training loop
for x, y in train_loader:
    optimizer.zero_grad()
    pred = model(x)  # Shape: [Batch, 2]
    loss = loss_fn(pred, y)
    loss.backward()
    optimizer.step()

# Evaluation
model.eval()
with torch.no_grad():
    out = model(x_val)
    mu, logvar = out[:, 0], out[:, 1]
    std = torch.exp(0.5 * logvar)
    val_crps = crps_gaussian(mu, std, y_val)
    print(f"Validation CRPS: {val_crps:.4f}")

2. Non-Parametric Quantile Regression

Estimate multiple target quantiles simultaneously with crossover penalties to prevent quantile crossing.

import torch
import torch.nn as nn
from torchregress.losses import MultiQuantileLoss, QuantileCrossoverLoss

# Predict 3 quantiles: 10%, 50%, 90%
quantiles = [0.1, 0.5, 0.9]
model = nn.Sequential(nn.Linear(10, 32), nn.ReLU(), nn.Linear(32, 3))
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

pinball_loss = MultiQuantileLoss(quantiles=quantiles)
crossover_penalty = QuantileCrossoverLoss(quantiles=quantiles)

# Training loop
for x, y in train_loader:
    optimizer.zero_grad()
    pred = model(x)  # Shape: [Batch, 3]
    loss = pinball_loss(pred, y) + 0.1 * crossover_penalty(pred, y)
    loss.backward()
    optimizer.step()

3. Distribution-Free Conformal Prediction

Wrap a pre-trained regression model to output interval predictions with guaranteed $90%$ coverage.

import torch
from torchregress.losses import SplitConformal

# Assume base_model is a pre-trained point estimator model(x) -> y_pred
base_model.eval()

# Calibrate conformal thresholds
conformal = SplitConformal(alpha=0.1)
x_cal = torch.randn(200, 10)  # Calibration features
y_cal = torch.randn(200, 1)   # Calibration targets
with torch.no_grad():
    y_cal_pred = base_model(x_cal)
conformal.calibrate(y_cal_pred, y_cal)

# Predict 90% prediction intervals on new test targets
x_test = torch.randn(10, 10)
with torch.no_grad():
    y_test_pred = base_model(x_test)
lower, upper = conformal.predict_interval(y_test_pred)

Development & Contribution

We welcome contributions! See CONTRIBUTING.md for setup, quality checks, tests, and documentation guidelines. Release notes are in the CHANGELOG.

License

This project is licensed under the MIT License - see the LICENSE file for details.