Support purged time validation and record split provenance

CHANGELOG.md

@@ -2,6 +2,18 @@
 
 ## Unreleased
 
+### Purged date validation
+
+- `TrainConfig.purge_dates` removes complete observed signal dates immediately
+  before time validation. The default is zero; callers using multi-day labels
+  can select a gap from their label horizon and information timing. Random splits
+  reject a nonzero gap rather than claiming chronological isolation.
+- `Trainer.split_info` and checkpoint `split.json` record effective batch/split
+  modes, sample/date counts and date ranges for training, validation and purged
+  partitions. The gap leaves validation dates unchanged and never sends purged
+  samples to either loader. This is protocol support, not market-performance
+  evidence or automatic detection of label overlap.
+
 ### Correlation-loss gradients
 
 - The trainer automatically uses one complete observed date per batch for

docs/backtest_assumptions.md

@@ -57,6 +57,20 @@ 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.
 
+Set `train.purge_dates` (default `0`) for an explicit gap immediately before
+the trainer's time-validation boundary. It counts complete observed signal dates
+in `FactorDataset`, not calendar days, and keeps the latest validation dates fixed.
+For an h-period label known at close of t+h, a gap of at least h-1 signal periods
+keeps the last training endpoint at or before the first validation signal close;
+use h if it must be strictly earlier. Sparse observations can make an observed-date
+gap more conservative. Select the gap from the actual label horizon, data timing
+and research protocol: the trainer does not infer that information. This gap
+controls the internal training/validation split; the CLI's outer out-of-sample
+boundary and separate walk-forward scripts need their own horizon-aware policy.
+Nonzero gaps require `val_split_mode="time"` and at least one remaining training
+date. `Trainer.split_info` and `save_dir/split.json` record all three partitions'
+sample/date counts and numeric date-ID ranges so the chosen split can be inspected.
+
 `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

src/mlquant/training/trainer.py

@@ -9,6 +9,7 @@
 
 from dataclasses import dataclass, field
 import math
+import json
 from pathlib import Path
 from typing import Callable, Iterable
 
@@ -40,6 +41,10 @@ class TrainConfig:
     # "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"
+    # Remove this many complete observed signal dates immediately before time
+    # validation. Choose from label horizon/information timing; zero keeps the
+    # default one-day close-to-close protocol. Random splits cannot use a gap.
+    purge_dates: int = 0
 
 
 class Trainer:
@@ -52,6 +57,7 @@ def __init__(
         self.model = model.to(cfg.device)
         self.loss_fn = loss_fn
         self.cfg = cfg
+        self.split_info: dict = {}
         self.optim = torch.optim.AdamW(
             model.parameters(),
             lr=cfg.lr,
@@ -70,12 +76,19 @@ def fit(self, dataset: FactorDataset) -> dict:
             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'")
+        if type(self.cfg.purge_dates) is not int or self.cfg.purge_dates < 0:
+            raise ValueError("purge_dates must be a nonnegative integer")
+        if self.cfg.purge_dates and self.cfg.val_split_mode != "time":
+            raise ValueError("purge_dates requires val_split_mode='time'")
         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 not torch.isfinite(dataset.dates).all():
+            raise ValueError("sample dates must be finite")
+        purged_indices = range(0)
         if self.cfg.val_split_mode == "time" or daily:
             boundaries = date_boundaries(dataset.dates)
             n_dates = len(boundaries) - 1
@@ -84,8 +97,13 @@ def fit(self, dataset: FactorDataset) -> dict:
             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))
+                n_train_dates = n_dates - n_val_dates - self.cfg.purge_dates
+                if n_train_dates < 1:
+                    raise ValueError("purge_dates leaves no training dates")
+                train_end = boundaries[n_train_dates]
+                train_ds = Subset(dataset, range(train_end))
                 val_ds = Subset(dataset, range(split, len(dataset)))
+                purged_indices = range(train_end, split)
             else:
                 validation_days = set(torch.randperm(
                     n_dates, generator=torch.Generator().manual_seed(0)
@@ -111,8 +129,25 @@ def fit(self, dataset: FactorDataset) -> dict:
             val_loader = DataLoader(val_ds, batch_size=self.cfg.batch_size)
 
         history: dict = {"train_loss": [], "val_loss": []}
+        partitions = {"train": train_ds.indices, "validation": val_ds.indices,
+                      "purged": purged_indices}
+        self.split_info = {
+            "schema_version": 1, "val_split_mode": self.cfg.val_split_mode,
+            "batch_mode": "date" if daily else "samples", "val_split": float(self.cfg.val_split),
+            "purge_dates": self.cfg.purge_dates, "purge_unit": "observed_signal_dates",
+            "partitions": {},
+        }
+        for name, indices in partitions.items():
+            dates = dataset.dates[indices]
+            self.split_info["partitions"][name] = {
+                "samples": len(indices), "observed_dates": dates.unique().numel(),
+                "date_range": [dates.min().item(), dates.max().item()] if dates.numel() else None,
+            }
         save_dir = Path(self.cfg.save_dir)
         save_dir.mkdir(parents=True, exist_ok=True)
+        (save_dir / "split.json").write_text(
+            json.dumps(self.split_info, indent=2, sort_keys=True, allow_nan=False) + "\n"
+        )
         best_val = float("inf")
 
         for epoch in range(self.cfg.epochs):

tests/test_purged_training.py

@@ -0,0 +1,84 @@
+"""Label-horizon-aware time splitting and inspectable partition provenance."""
+import json
+
+import pytest
+import torch
+
+from mlquant.models.losses import ICLoss
+from mlquant.training import FactorDataset, TrainConfig, Trainer
+
+
+def _dataset():
+    # Four-period labels, only where both endpoints exist. Keep NaN for the
+    # unavailable tail because the legacy dataset has no separate label mask.
+    close = 10 + torch.arange(18).float()[:, None] + torch.arange(3).float()[None, :]
+    labels = torch.full_like(close, float("nan"))
+    labels[:-4] = close[4:] / close[:-4] - 1
+    return FactorDataset(close[..., None], torch.ones_like(close, dtype=torch.bool), labels)
+
+
+@pytest.mark.parametrize("objective", ["mse", "ic"])
+@pytest.mark.parametrize("purge", [0, 3, 4])
+def test_gap_removes_full_dates_keeps_validation_fixed_and_records_counts(tmp_path, monkeypatch, objective, purge):
+    ds = _dataset()
+    seen = {}
+    real_loader = torch.utils.data.DataLoader
+
+    def record(subset, *args, **kwargs):
+        shuffled = kwargs.get("shuffle", getattr(kwargs.get("batch_sampler"), "shuffle", False))
+        seen["train" if shuffled else "validation"] = list(subset.indices)
+        return real_loader(subset, *args, **kwargs)
+
+    monkeypatch.setattr("mlquant.training.trainer.DataLoader", record)
+    model = torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Flatten(0))
+    criterion = torch.nn.MSELoss() if objective == "mse" else ICLoss()
+    trainer = Trainer(model, criterion, TrainConfig(epochs=1, val_split=.25,
+                                                  purge_dates=purge, save_dir=str(tmp_path)))
+    history = trainer.fit(ds)
+    train_dates = set(ds.dates[seen["train"]].tolist())
+    validation_dates = set(ds.dates[seen["validation"]].tolist())
+    assert train_dates == set(range(11 - purge))
+    assert validation_dates == {11, 12, 13}
+    excluded = set(range(11 - purge, 11))
+    assert excluded.isdisjoint(train_dates | validation_dates)
+    if purge == 3:
+        assert max(train_dates) + 4 == min(validation_dates)
+    elif purge == 4:
+        assert max(train_dates) + 4 < min(validation_dates)
+    info = json.loads((tmp_path / "split.json").read_text())
+    assert info == trainer.split_info
+    assert info["purge_unit"] == "observed_signal_dates"
+    assert info["partitions"]["validation"] == {"samples": 9, "observed_dates": 3, "date_range": [11, 13]}
+    assert info["partitions"]["purged"]["samples"] == 3 * purge
+    assert info["partitions"]["purged"]["observed_dates"] == purge
+    assert sum(p["samples"] for p in info["partitions"].values()) == len(ds)
+    assert set(history) == {"train_loss", "val_loss"}
+
+
+def test_sparse_date_ids_use_observed_dates_and_never_split_a_date(tmp_path):
+    ds = _dataset()
+    ds.dates = ds.dates * 5 + 100
+    trainer = Trainer(torch.nn.Sequential(torch.nn.Linear(1, 1), torch.nn.Flatten(0)),
+                      torch.nn.MSELoss(), TrainConfig(epochs=1, val_split=.25, purge_dates=2,
+                                                      save_dir=str(tmp_path)))
+    trainer.fit(ds)
+    info = trainer.split_info["partitions"]
+    assert info["train"]["date_range"] == [100, 140]
+    assert info["purged"] == {"samples": 6, "observed_dates": 2, "date_range": [145, 150]}
+    assert info["validation"]["date_range"] == [155, 165]
+
+
+@pytest.mark.parametrize("overrides,match", [
+    ({"purge_dates": -1}, "nonnegative integer"), ({"purge_dates": 1.5}, "nonnegative integer"),
+    ({"purge_dates": True}, "nonnegative integer"), ({"purge_dates": 20}, "no training dates"),
+    ({"purge_dates": 1, "val_split_mode": "random"}, "requires val_split_mode"),
+])
+def test_invalid_gap_preserves_existing_checkpoint_and_creates_no_report(tmp_path, overrides, match):
+    old = tmp_path / "last.pt"
+    old.write_bytes(b"existing checkpoint")
+    trainer = Trainer(torch.nn.Linear(1, 1), torch.nn.MSELoss(),
+                      TrainConfig(epochs=1, save_dir=str(tmp_path), **overrides))
+    with pytest.raises(ValueError, match=match):
+        trainer.fit(_dataset())
+    assert old.read_bytes() == b"existing checkpoint"
+    assert not (tmp_path / "split.json").exists()