Quick Start¶

import torch
import torch.nn as nn
from torchregress.losses import WeightedMSELoss, WeightedHuberLoss
from torchregress.metrics import mean_absolute_error, r2_score, rmse

model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 1))

# Huber is a robust default for mild outliers; delta=1.0 transitions
# between quadratic and linear at |r| = 1.
loss_fn = WeightedHuberLoss(delta=1.0)

optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)
for x, y, mask, w in dataloader:
    pred = model(x)
    loss = loss_fn(pred, y, mask=mask, weights=w)
    loss.backward(); optimizer.step(); optimizer.zero_grad()

# Evaluation: report RMSE + MAE + R² as secondary summaries
with torch.no_grad():
    y_pred = model(x_test)
    print(f"RMSE: {rmse(y_pred, y_test):.4f}  "
          f"MAE: {mean_absolute_error(y_pred, y_test):.4f}  "
          f"R²: {r2_score(y_pred, y_test):.4f}")

from torchregress.losses import GaussianNLLLoss
from torchregress.metrics import crps_gaussian, gaussian_nll

# Head: [mean, log_var] per example
model = nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 2))

loss_fn = GaussianNLLLoss()

for x, y in dataloader:
    out = model(x)                      # [B, 2]
    loss = loss_fn(out, y)              # -log p(y | μ(x), σ²(x))
    loss.backward(); optimizer.step(); optimizer.zero_grad()

# Inference + density-aware evaluation
with torch.no_grad():
    out = model(x_test)
    mu, logvar = out[:, 0], out[:, 1]
    var = torch.exp(logvar)
    std = torch.sqrt(var)
    lower, upper = mu - 1.96 * std, mu + 1.96 * std
    crps = crps_gaussian(mu, y_test, std)          # proper scoring rule
    gnll = gaussian_nll(mu, y_test, var)           # NLL on test set

from torchregress.losses import (
    WeightedHuberLoss, CauchyLoss, TukeyBiweightLoss, AdaptiveRobustLoss
)

# Mild outliers (5–10%): Huber is a safe default
loss_fn = WeightedHuberLoss(delta=1.0)

# Moderate outliers (10–25%): Cauchy is more aggressive
loss_fn = CauchyLoss(c=1.0)

# Severe contamination (>25%): Tukey rejects outliers entirely
loss_fn = TukeyBiweightLoss(c=4.685)

# Unknown noise regime: learn the shape parameter
loss_fn = AdaptiveRobustLoss()

# IMPORTANT: when the loss has learnable parameters (AdaptiveRobust,
# Barron), pass them to the optimizer:
optimizer = torch.optim.Adam([
    {"params": model.parameters(), "lr": 1e-3},
    {"params": loss_fn.parameters(), "lr": 1e-2},
])

from torchregress.losses import SplitConformal, CQR

# Split conformal: residuals |y - ŷ| on a held-out calibration set
cp = SplitConformal(alpha=0.1)                       # target 90% coverage
cp.calibrate(y_pred_cal, y_cal)
lower, upper = cp.predict_interval(y_pred_test)

# Conformalized Quantile Regression (CQR): adapt to heteroscedastic noise
cqr = CQR(alpha=0.1)
cqr.calibrate(quantile_preds_cal, y_cal)             # quantile predictions
lower, upper = cqr.predict_interval(quantile_preds_test)

from torchregress.ensemble import BaseEnsembleModel
from torchregress.metrics import uncertainty_decomposition

# A heteroscedastic head that outputs [mean, log_var]
def heteroscedastic_head():
    return nn.Sequential(nn.Linear(10, 64), nn.ReLU(), nn.Linear(64, 2))

ens = BaseEnsembleModel(base_model=heteroscedastic_head(),
                   ensemble_size=5, base_seed=0)
ens.fit(train_loader, optimizer_factory=torch.optim.Adam, lr=1e-3)

# Per-member predictions
mu_per_member = torch.stack([m(x_test) for m in ens.models], dim=0)  # [M, B]
var_per_member = ...  # from each member's (mean, log_var) head

# Decomposition
decomp = uncertainty_decomposition(mu_per_member, var_per_member)
# decomp["epistemic_uncertainty"], decomp["aleatoric_uncertainty"], decomp["total_uncertainty"]

from torchregress.test_time import (
    BayesianLinearHead, ScoreCDFReweighter, WeightedSplitConformalAdapter
)

# 1. Closed-form Bayesian last layer — adapts to a few labeled or
#    unlabeled examples without backprop
head = BayesianLinearHead(in_features=64, out_features=1, noise_variance=0.1)
head.fit(features_train, y_train)              # one-shot
head.partial_fit(features_stream, y_stream)   # streaming with forgetting

# 2. OT-style conformal reweighting under non-exchangeable shift
rw = ScoreCDFReweighter().fit(cal_scores, target_scores)
adapter = WeightedSplitConformalAdapter(alpha=0.1).calibrate(cal_scores, rw.weights_)
mask = adapter.predict_from_test_scores(test_scores)

# 3. Full pipeline: prior transport + feature alignment + conformal
from torchregress.test_time import (
    ShiftFactoredPredictiveTransport, ShiftFactoredTransportConfig
)
transport = ShiftFactoredPredictiveTransport(ShiftFactoredTransportConfig())
transport.fit_source(source_pb, source_targets=src_y, source_inputs=src_x)
adapted_pb = transport.adapt_unlabeled_target(target_pb, target_inputs=tgt_x)

from torchregress.causal import dr_ate, causal_overlap_report

ate = dr_ate(x, t, y,
             outcome_model=lambda: GradientBoostingRegressor(n_estimators=50),
             propensity_model=lambda: GradientBoostingClassifier(n_estimators=50),
             folds=2, alpha=0.05, seed=42)
print(f"ATE = {ate['estimate']:.3f}  CI = [{ate['ci_lower']:.3f}, {ate['ci_upper']:.3f}]")

# Always inspect overlap diagnostics — a low ESS in either arm means
# the estimate is dominated by a few influential points
print(f"min ESS: {ate['diagnostics']['min_group_ess']:.1f}  "
      f"overlap rate: {ate['diagnostics']['overlap_rate']:.3f}")