From fec626887180486eca7c77060356ed90009eec1a Mon Sep 17 00:00:00 2001 From: Yimin Du <5088553+initial-d@users.noreply.github.com> Date: Tue, 27 Oct 2026 10:18:41 +0800 Subject: [PATCH 1/2] Guard correlation loss gradients for constant batches --- CHANGELOG.md | 11 ++++++ src/mlquant/models/losses.py | 27 ++++++++++++--- tests/test_losses.py | 67 ++++++++++++++++++++++++++++++++++++ 3 files changed, 100 insertions(+), 5 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 1c3eaee..eeaa5c3 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,17 @@ ## Unreleased +### Correlation-loss gradients + +- `ICLoss` and `RankICLoss` return zero loss and zero gradients for constant + prediction/label batches and singleton batches. Previously the reported loss + could be zero while backward produced NaNs from `sqrt(0)`, poisoning optimizer + state. Undefined batches contribute no learning signal; nonconstant batches + retain the previous formula and gradients. Inputs must be aligned and finite. +- Empty/broadcast shapes are rejected. Sign-aware `gamma` and soft-rank + `temperature` must be finite and positive. This is training correctness, + not predictive-quality or investment evidence. + ### Full-pipeline validation: execution coverage - Reports now show joint factor-mask coverage, per-split training/prediction diff --git a/src/mlquant/models/losses.py b/src/mlquant/models/losses.py index daa0430..7a6a4a1 100644 --- a/src/mlquant/models/losses.py +++ b/src/mlquant/models/losses.py @@ -20,6 +20,8 @@ """ from __future__ import annotations +import math + import torch import torch.nn as nn @@ -40,8 +42,8 @@ class AdjMSELoss(nn.Module): def __init__(self, gamma: float = 0.1) -> None: super().__init__() - if gamma <= 0: - raise ValueError("gamma must be positive") + if not math.isfinite(gamma) or gamma <= 0: + raise ValueError("gamma must be finite and positive") self.gamma = gamma def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: @@ -54,14 +56,27 @@ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: def _pearson(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor: + if x.shape != y.shape or not x.numel(): + raise ValueError("correlation inputs must have the same nonempty shape") xm = x - x.mean() ym = y - y.mean() - denom = (xm.pow(2).sum().sqrt() * ym.pow(2).sum().sqrt()).clamp_min(1e-12) - return (xm * ym).sum() / denom + xss, yss = xm.pow(2).sum(), ym.pow(2).sum() + defined = (xss > 0) & (yss > 0) + # sqrt'(0) is infinite: clamping the product after sqrt cannot protect + # backward. Replace zero-variance operands before taking either root. + sx = torch.where(defined, xss, torch.ones_like(xss)).sqrt() + sy = torch.where(defined, yss, torch.ones_like(yss)).sqrt() + denom = (sx * sy).clamp_min(1e-12) + numerator = (xm * ym).sum() + return torch.where(defined, numerator / denom, numerator * 0.0) class ICLoss(nn.Module): - """Negative Pearson IC across a batch (e.g. one cross-section).""" + """Negative Pearson IC across a batch (e.g. one cross-section). + + A constant prediction or label batch has undefined correlation and returns + zero loss/gradients. Pass aligned, finite tensors from one cross-section. + """ def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: return -_pearson(pred, target) @@ -77,6 +92,8 @@ class RankICLoss(nn.Module): def __init__(self, temperature: float = 1.0) -> None: super().__init__() + if not math.isfinite(temperature) or temperature <= 0: + raise ValueError("temperature must be finite and positive") self.temperature = temperature def forward(self, pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor: diff --git a/tests/test_losses.py b/tests/test_losses.py index 3b3e436..3134909 100644 --- a/tests/test_losses.py +++ b/tests/test_losses.py @@ -2,6 +2,7 @@ from __future__ import annotations import torch +import pytest from mlquant.models.losses import AdjMSELoss, ICLoss, RankICLoss @@ -27,3 +28,69 @@ def test_rank_ic_loss_negative_correlation_sign(): loss = RankICLoss(temperature=0.05)(x, y) # Loss is -(rank corr) ≈ -(-1) = +1 assert loss.item() > 0.5 + + +@pytest.mark.parametrize("loss_type", [ICLoss, RankICLoss]) +@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32, torch.float64]) +@pytest.mark.parametrize("case", ["prediction", "label", "both", "singleton"]) +def test_constant_correlation_batches_have_zero_finite_gradients(loss_type, dtype, case): + p = [1.] if case == "singleton" else ([1.] * 4 if case in ("prediction", "both") else [1., 2., 3., 4.]) + t = [2.] if case == "singleton" else ([2.] * 4 if case in ("label", "both") else [4., 1., 3., 2.]) + pred = torch.tensor(p, dtype=dtype, requires_grad=True) + target = torch.tensor(t, dtype=dtype, requires_grad=True) + loss = loss_type()(pred, target) + loss.backward() + assert loss.item() == 0 + assert torch.equal(pred.grad, torch.zeros_like(pred)) + assert torch.equal(target.grad, torch.zeros_like(target)) + + +@pytest.mark.parametrize("loss_type", [ICLoss, RankICLoss]) +def test_nondegenerate_correlation_values_and_gradients_match_previous_formula(loss_type): + pred = torch.tensor([.7, -1.2, .3, 2.4], dtype=torch.float64, requires_grad=True) + target = torch.tensor([.1, -.8, 1.5, .9], dtype=torch.float64, requires_grad=True) + loss_fn = loss_type() + x, y = (pred, target) if loss_type is ICLoss else (loss_fn._soft_rank(pred), loss_fn._soft_rank(target)) + a, b = x - x.mean(), y - y.mean() + reference = -(a * b).sum() / (a.square().sum().sqrt() * b.square().sum().sqrt()).clamp_min(1e-12) + actual = loss_fn(pred, target) + expected_grad = torch.autograd.grad(reference, (pred, target), retain_graph=True) + actual_grad = torch.autograd.grad(actual, (pred, target)) + torch.testing.assert_close(actual, reference, atol=0, rtol=0) + for before, after in zip(expected_grad, actual_grad): + torch.testing.assert_close(after, before, atol=1e-14, rtol=1e-14) + + +@pytest.mark.parametrize("bad", [0., -1., float("nan"), float("inf"), -float("inf")]) +@pytest.mark.parametrize("loss_type,option", [(AdjMSELoss, "gamma"), (RankICLoss, "temperature")]) +def test_loss_hyperparameters_are_finite_positive(loss_type, option, bad): + with pytest.raises(ValueError, match="finite and positive"): + loss_type(**{option: bad}) + + +@pytest.mark.parametrize("loss_type", [ICLoss, RankICLoss]) +def test_correlation_losses_reject_empty_or_broadcast_shapes(loss_type): + with pytest.raises(ValueError, match="same nonempty shape"): + loss_type()(torch.empty(0), torch.empty(0)) + with pytest.raises(ValueError, match="same nonempty shape"): + loss_type()(torch.ones(3, 1), torch.ones(3)) + + +@pytest.mark.parametrize("loss_type", [ICLoss, RankICLoss]) +def test_constant_batch_does_not_poison_the_next_optimizer_step(loss_type): + weights = torch.nn.Parameter(torch.tensor([1., 1., 1.])) + optimizer = torch.optim.Adam([weights], lr=.01) + loss_fn = loss_type() + for target in [torch.tensor([1., 2., 3.]), torch.ones(3)]: + optimizer.zero_grad() + loss_fn(weights, target).backward() + optimizer.step() + assert weights.isfinite().all() + # An unrelated informative step must not inherit NaN Adam moments. + optimizer.zero_grad() + (weights - torch.tensor([0., 2., 3.])).square().mean().backward() + optimizer.step() + assert weights.isfinite().all() + assert not torch.equal(weights.detach(), torch.ones(3)) + assert optimizer.state[weights]["exp_avg"].isfinite().all() + assert optimizer.state[weights]["exp_avg_sq"].isfinite().all() From c955c1a8b937cdf78a33b1ed7cedb3ddf16e0a33 Mon Sep 17 00:00:00 2001 From: Yimin Du <5088553+initial-d@users.noreply.github.com> Date: Wed, 28 Oct 2026 10:45:21 +0800 Subject: [PATCH 2/2] Train correlation objectives on complete daily cross-sections --- CHANGELOG.md | 10 ++ docs/backtest_assumptions.md | 15 ++- src/mlquant/training/__init__.py | 3 +- src/mlquant/training/samplers.py | 36 ++++++++ src/mlquant/training/trainer.py | 79 ++++++++++++---- tests/test_cross_section_training.py | 132 +++++++++++++++++++++++++++ tests/test_trainer_split.py | 2 +- 7 files changed, 254 insertions(+), 23 deletions(-) create mode 100644 src/mlquant/training/samplers.py create mode 100644 tests/test_cross_section_training.py diff --git a/CHANGELOG.md b/CHANGELOG.md index eeaa5c3..a3d09af 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -4,6 +4,16 @@ ### Correlation-loss gradients +- The trainer automatically uses one complete observed date per batch for + `ICLoss`/`RankICLoss`, shuffling dates rather than stocks across dates. Daily + losses are averaged equally over dates. `batch_mode="date"` supports custom + daily objectives; ordinary regression retains flat sample batching. +- Time validation now holds out complete latest dates, with no date shared + between training and validation. `val_split` is a date fraction in this mode. + Explicit random daily splits also preserve whole dates. Flat training no longer + drops a dataset smaller than `batch_size`. These changes can alter historical + model diagnostics; no archived performance result is revalidated here. + - `ICLoss` and `RankICLoss` return zero loss and zero gradients for constant prediction/label batches and singleton batches. Previously the reported loss could be zero while backward produced NaNs from `sqrt(0)`, poisoning optimizer diff --git a/docs/backtest_assumptions.md b/docs/backtest_assumptions.md index 19b659e..15909a0 100644 --- a/docs/backtest_assumptions.md +++ b/docs/backtest_assumptions.md @@ -50,13 +50,26 @@ A publishable result should document the train, validation, and test split, plus The CLI pipeline splits by date. `train_frac` (default `0.6`) sets the first out-of-sample date index; the model trains only on earlier dates, its -validation set is the most recent slice of those training samples +validation set is the most recent slice of complete training dates (`val_split_mode: time`), and the backtest scores only the out-of-sample window. `summary.json` records both windows under `evaluation`. The one-day label of the last training date uses the close of the first out-of-sample date, which is already known when the first out-of-sample signal is formed, so no purge gap is needed for the default one-day horizon; longer horizons need one. +`Trainer` holds out whole latest observed dates for `val_split_mode="time"`; +`val_split` therefore denotes a fraction of dates, not individual stock-days. +The default `batch_mode="auto"` uses full daily cross-sections for `ICLoss` and +`RankICLoss`, with one optimizer step per date and equally weighted daily losses. +`batch_size` is ignored for daily batching: memory depends on the largest daily +cross-section, and soft RankIC still requires quadratic pairwise rank memory. +Use `batch_mode="date"` for a custom daily objective. Ordinary regression uses +flat sample batches and sample-weighted epoch loss; it retains the final partial +batch. Explicit `val_split_mode="random"` splits whole dates in daily mode and +individual samples in flat mode. Random splits are reproducibility options, +not chronological out-of-sample evidence. Constant/singleton IC batches contribute +zero loss and zero gradients; they are retained, not counted as informative IC. + This is a single split, not a walk-forward evaluation. Demo metrics come from synthetic data and are expected to hover around zero or below after costs. `scripts/public_data_validation.py` runs rolling walk-forward splits. diff --git a/src/mlquant/training/__init__.py b/src/mlquant/training/__init__.py index 701b9aa..ac7f484 100644 --- a/src/mlquant/training/__init__.py +++ b/src/mlquant/training/__init__.py @@ -2,5 +2,6 @@ from .dataset import FactorDataset from .augment import gbm_augment from .trainer import Trainer, TrainConfig +from .samplers import CrossSectionBatchSampler -__all__ = ["FactorDataset", "gbm_augment", "Trainer", "TrainConfig"] +__all__ = ["FactorDataset", "gbm_augment", "Trainer", "TrainConfig", "CrossSectionBatchSampler"] diff --git a/src/mlquant/training/samplers.py b/src/mlquant/training/samplers.py new file mode 100644 index 0000000..1192771 --- /dev/null +++ b/src/mlquant/training/samplers.py @@ -0,0 +1,36 @@ +"""Batch complete daily cross-sections without mixing their correlation targets.""" +from __future__ import annotations + +import torch +from torch.utils.data import Sampler + + +def date_boundaries(dates: torch.Tensor) -> list[int]: + """Start/end offsets for finite, date-ordered samples, including the sentinel.""" + if (dates.ndim != 1 or not torch.isfinite(dates).all() + or (dates[1:] < dates[:-1]).any()): + raise ValueError("sample dates must be a finite nondecreasing vector") + if not dates.numel(): + return [0] + starts = (dates[1:] != dates[:-1]).nonzero().flatten().add(1).tolist() + return [0, *starts, dates.numel()] + + +class CrossSectionBatchSampler(Sampler[list[int]]): + """One full observed date per batch; optionally shuffle the order of dates. + + Missing stocks stay missing. No padding, cross-date mixing, within-day + splitting or dropped small batches. Memory for boundaries grows with dates. + """ + + def __init__(self, dates: torch.Tensor, *, shuffle: bool = False) -> None: + self.boundaries = date_boundaries(dates) + self.shuffle = shuffle + + def __len__(self) -> int: + return len(self.boundaries) - 1 + + def __iter__(self): + order = torch.randperm(len(self)).tolist() if self.shuffle else range(len(self)) + for day in order: + yield list(range(self.boundaries[day], self.boundaries[day + 1])) diff --git a/src/mlquant/training/trainer.py b/src/mlquant/training/trainer.py index d66ff02..6e9bdbb 100644 --- a/src/mlquant/training/trainer.py +++ b/src/mlquant/training/trainer.py @@ -8,6 +8,7 @@ from __future__ import annotations from dataclasses import dataclass, field +import math from pathlib import Path from typing import Callable, Iterable @@ -15,7 +16,9 @@ import torch.nn as nn from torch.utils.data import DataLoader, Subset +from ..models.losses import ICLoss, RankICLoss from .dataset import FactorDataset +from .samplers import CrossSectionBatchSampler, date_boundaries @dataclass @@ -29,11 +32,14 @@ class TrainConfig: log_every: int = 50 save_dir: str = "checkpoints" val_split: float = 0.1 - # "time" holds out the most recent samples (FactorDataset is ordered by - # date), so validation never sees days that precede training days. - # "random" restores the old shuffled split for backward compatibility. + # "time" holds out the most recent complete dates. "random" splits dates + # for daily batches, or individual samples for legacy regression batches. val_split_mode: str = "time" extra: dict = field(default_factory=dict) + # "auto" uses complete daily batches for IC/RankIC, flat batches otherwise. + # "date" also supports custom cross-sectional objectives. In date mode, + # batch_size is ignored and losses are averaged equally over observed dates. + batch_mode: str = "auto" class Trainer: @@ -56,23 +62,53 @@ def __init__( # Public API # ------------------------------------------------------------------ def fit(self, dataset: FactorDataset) -> dict: - n_val = max(1, int(len(dataset) * self.cfg.val_split)) - n_train = len(dataset) - n_val - if self.cfg.val_split_mode == "time": - train_ds = Subset(dataset, range(n_train)) - val_ds = Subset(dataset, range(n_train, len(dataset))) - elif self.cfg.val_split_mode == "random": + if not math.isfinite(self.cfg.val_split) or not 0 < self.cfg.val_split < 1: + raise ValueError("val_split must be finite and between zero and one") + if self.cfg.epochs < 1 or self.cfg.batch_size < 1: + raise ValueError("epochs and batch_size must be positive") + if self.cfg.val_split_mode not in ("time", "random"): + raise ValueError("val_split_mode must be 'time' or 'random'") + if self.cfg.batch_mode not in ("auto", "samples", "date"): + raise ValueError("batch_mode must be 'auto', 'samples' or 'date'") + correlation = isinstance(self.loss_fn, (ICLoss, RankICLoss)) + if correlation and self.cfg.batch_mode == "samples": + raise ValueError("IC/RankIC require complete date batches") + daily = self.cfg.batch_mode == "date" or (self.cfg.batch_mode == "auto" and correlation) + if len(dataset) < 2 or dataset.dates.shape != (len(dataset),): + raise ValueError("training requires at least two samples with aligned dates") + if self.cfg.val_split_mode == "time" or daily: + boundaries = date_boundaries(dataset.dates) + n_dates = len(boundaries) - 1 + if n_dates < 2: + raise ValueError("date splitting requires at least two observed dates") + n_val_dates = max(1, int(n_dates * self.cfg.val_split)) + if self.cfg.val_split_mode == "time": + split = boundaries[n_dates - n_val_dates] + train_ds = Subset(dataset, range(split)) + val_ds = Subset(dataset, range(split, len(dataset))) + else: + validation_days = set(torch.randperm( + n_dates, generator=torch.Generator().manual_seed(0) + )[:n_val_dates].tolist()) + train_indices, val_indices = [], [] + for day in range(n_dates): + target = val_indices if day in validation_days else train_indices + target.extend(range(boundaries[day], boundaries[day + 1])) + train_ds, val_ds = Subset(dataset, train_indices), Subset(dataset, val_indices) + else: + n_val = max(1, int(len(dataset) * self.cfg.val_split)) train_ds, val_ds = torch.utils.data.random_split( - dataset, [n_train, n_val], + dataset, [len(dataset) - n_val, n_val], generator=torch.Generator().manual_seed(0), ) + if daily: + train_loader = DataLoader(train_ds, batch_sampler=CrossSectionBatchSampler( + dataset.dates[train_ds.indices], shuffle=True)) + val_loader = DataLoader(val_ds, batch_sampler=CrossSectionBatchSampler( + dataset.dates[val_ds.indices])) else: - raise ValueError( - f"val_split_mode must be 'time' or 'random', got {self.cfg.val_split_mode!r}" - ) - train_loader = DataLoader(train_ds, batch_size=self.cfg.batch_size, - shuffle=True, drop_last=True) - val_loader = DataLoader(val_ds, batch_size=self.cfg.batch_size) + train_loader = DataLoader(train_ds, batch_size=self.cfg.batch_size, shuffle=True) + val_loader = DataLoader(val_ds, batch_size=self.cfg.batch_size) history: dict = {"train_loss": [], "val_loss": []} save_dir = Path(self.cfg.save_dir) @@ -103,18 +139,21 @@ def predict(self, features: torch.Tensor) -> torch.Tensor: def _run_epoch(self, loader: Iterable, *, train: bool) -> float: self.model.train(train) total, n = 0.0, 0 + daily = isinstance(getattr(loader, "batch_sampler", None), CrossSectionBatchSampler) for x, y in loader: x = x.to(self.cfg.device) y = y.to(self.cfg.device) if train: self.optim.zero_grad(set_to_none=True) - preds = self.model(x) - loss = self.loss_fn(preds, y) + with torch.set_grad_enabled(train): + preds = self.model(x) + loss = self.loss_fn(preds, y) if train: loss.backward() if self.cfg.grad_clip: nn.utils.clip_grad_norm_(self.model.parameters(), self.cfg.grad_clip) self.optim.step() - total += loss.item() * x.shape[0] - n += x.shape[0] + weight = 1 if daily else x.shape[0] + total += loss.item() * weight + n += weight return total / max(1, n) diff --git a/tests/test_cross_section_training.py b/tests/test_cross_section_training.py new file mode 100644 index 0000000..f2af409 --- /dev/null +++ b/tests/test_cross_section_training.py @@ -0,0 +1,132 @@ +"""Daily objectives must not become pooled cross-date correlations.""" +import pytest +import torch +from torch.utils.data import DataLoader + +from mlquant.models.losses import ICLoss, RankICLoss +from mlquant.training import CrossSectionBatchSampler, FactorDataset, TrainConfig, Trainer + + +def _ragged_dataset(): + factors = torch.zeros(5, 5, 2) + mask = torch.zeros(5, 5, dtype=torch.bool) + labels = torch.zeros(5, 5) + # Intersection with the next day's legacy mask gives counts 2, 4, 1, 1. + for day, size in enumerate([2, 5, 4, 1, 5]): + factors[day, :, 0] = day + factors[day, :, 1] = torch.arange(5) + labels[day] = 100 * day + torch.arange(5) + mask[day, :size] = True + return FactorDataset(factors, mask, labels) + + +def test_sampler_keeps_complete_ragged_dates_and_all_samples(): + dates = torch.tensor([0, 0, 3, 3, 3, 9]) + assert list(CrossSectionBatchSampler(dates)) == [[0, 1], [2, 3, 4], [5]] + shuffled = list(CrossSectionBatchSampler(dates, shuffle=True)) + assert sorted(i for batch in shuffled for i in batch) == list(range(6)) + assert all(dates[batch].unique().numel() == 1 for batch in shuffled) + assert len(CrossSectionBatchSampler(torch.empty(0))) == 0 + for invalid in (torch.tensor([2, 1]), torch.tensor([0., float("nan")]), torch.ones(2, 1)): + with pytest.raises(ValueError, match="nondecreasing vector"): + CrossSectionBatchSampler(invalid) + + +@pytest.mark.parametrize("loss_type", [ICLoss, RankICLoss]) +@pytest.mark.parametrize("split_mode", ["time", "random"]) +def test_fit_automatically_uses_full_dates_for_correlation_losses(tmp_path, loss_type, split_mode): + seen = {True: [], False: []} + model = torch.nn.Linear(2, 1) + + class ScalarModel(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = model + + def forward(self, x): + return self.linear(x).squeeze(-1) + + class RecordLoss(loss_type): + def forward(self, pred, target): + days = target.floor_divide(100) + assert days.unique().numel() == 1 + seen[model.training].append((int(days[0]), len(target))) + assert torch.is_grad_enabled() == model.training + return super().forward(pred, target) + + cfg = TrainConfig(epochs=1, batch_size=1, val_split=.5, + val_split_mode=split_mode, save_dir=str(tmp_path)) + history = Trainer(ScalarModel(), RecordLoss(), cfg).fit(_ragged_dataset()) + assert dict(seen[True] + seen[False]) == {0: 2, 1: 4, 2: 1, 3: 1} + train_days, val_days = {d for d, _ in seen[True]}, {d for d, _ in seen[False]} + assert not train_days & val_days + if split_mode == "time": + assert max(train_days) < min(val_days) + assert all(torch.isfinite(torch.tensor(history[key])).all() for key in history) + assert (tmp_path / "best.pt").exists() + + +def test_daily_ic_avoids_the_spurious_positive_pooled_correlation(tmp_path): + # Each day has reversed stock ordering, but between-day level offsets make + # pooled correlation positive. This fixture is an objective check, not alpha. + x = torch.tensor([[0.], [1.], [2.], [100.], [101.], [102.]], dtype=torch.float32) + y = torch.tensor([2., 1., 0., 102., 101., 100.]) + dates = torch.tensor([0, 0, 0, 1, 1, 1]) + model = torch.nn.Linear(1, 1, bias=False) + with torch.no_grad(): + model.weight.fill_(1) + class Flatten(torch.nn.Module): + def __init__(self): + super().__init__() + self.linear = model + + def forward(self, values): + return self.linear(values).squeeze(-1) + + trainer = Trainer(Flatten(), ICLoss(), TrainConfig(save_dir=str(tmp_path))) + loader = DataLoader(torch.utils.data.TensorDataset(x, y), + batch_sampler=CrossSectionBatchSampler(dates)) + assert trainer._run_epoch(loader, train=False) == pytest.approx(1.) + assert ICLoss()(x.flatten(), y).item() < -.99 + + +def test_daily_validation_weights_dates_equally_and_disables_gradients(tmp_path): + x = torch.zeros(6, 1) + y = torch.tensor([1., 3., 3., 3., 3., 3.]) + model = torch.nn.Linear(1, 1, bias=False) + loader = DataLoader(torch.utils.data.TensorDataset(x, y), + batch_sampler=CrossSectionBatchSampler(torch.tensor([0, 1, 1, 1, 1, 1]))) + trainer = Trainer(model, lambda p, t: ((p.flatten() - t) ** 2).mean(), + TrainConfig(save_dir=str(tmp_path))) + assert trainer._run_epoch(loader, train=False) == pytest.approx(5.) + + +def test_small_regression_dataset_is_not_dropped(tmp_path): + ds = _ragged_dataset() + model = torch.nn.Sequential(torch.nn.Linear(2, 1), torch.nn.Flatten(0)) + before = model[0].weight.detach().clone() + trainer = Trainer(model, torch.nn.MSELoss(), + TrainConfig(epochs=1, batch_size=1024, save_dir=str(tmp_path))) + history = trainer.fit(ds) + assert history["train_loss"][0] > 0 + assert not torch.equal(model[0].weight.detach(), before) + + +@pytest.mark.parametrize("overrides,match", [ + ({"val_split": 0}, "val_split"), ({"val_split": 1}, "val_split"), + ({"val_split": float("nan")}, "val_split"), ({"batch_mode": "unknown"}, "batch_mode"), + ({"batch_mode": "samples"}, "complete date"), ({"batch_size": 0}, "positive"), +]) +def test_invalid_correlation_training_config_is_rejected(tmp_path, overrides, match): + cfg = TrainConfig(save_dir=str(tmp_path), **overrides) + trainer = Trainer(torch.nn.Linear(2, 1), ICLoss(), cfg) + with pytest.raises(ValueError, match=match): + trainer.fit(_ragged_dataset()) + assert not (tmp_path / "last.pt").exists() + + +def test_time_split_rejects_a_single_observed_date(tmp_path): + ds = FactorDataset(torch.ones(2, 3, 1), torch.ones(2, 3, dtype=torch.bool), torch.ones(2, 3)) + trainer = Trainer(torch.nn.Linear(1, 1), torch.nn.MSELoss(), TrainConfig(save_dir=str(tmp_path))) + with pytest.raises(ValueError, match="two observed dates"): + trainer.fit(ds) diff --git a/tests/test_trainer_split.py b/tests/test_trainer_split.py index 6584f89..3c5e2cb 100644 --- a/tests/test_trainer_split.py +++ b/tests/test_trainer_split.py @@ -38,7 +38,7 @@ def spy(subset, *args, **kwargs): last_train_date = ds.dates[seen["train"]].max() first_val_date = ds.dates[seen["val"]].min() - assert last_train_date <= first_val_date + assert last_train_date < first_val_date def test_unknown_split_mode_is_rejected(tmp_path):