fix(ml): split cross-validation on dates, not on flattened rows (BUG-23)

alpha/ml_signal.py

@@ -149,21 +149,18 @@ def __init__(
         self.embargo = embargo
         self.mode = mode
 
-    def split(self, n_samples: int):
-        """Generate train/test index pairs with purge and embargo.
+    def _folds(self, n_units: int):
+        """Yield (train_start, train_end, test_start, test_end) in unit terms.
 
-        Args:
-            n_samples: Total number of samples.
-
-        Yields:
-            (train_indices, test_indices) tuples.
+        Split out from `split()` so the same arithmetic can be applied to a
+        date axis or to a row axis.
         """
         min_size = self.train_size + self.gap + self.test_size
-        if n_samples < min_size:
-            raise ValueError(f"Need at least {min_size} samples, got {n_samples}")
+        if n_units < min_size:
+            raise ValueError(f"Need at least {min_size} units, got {n_units}")
 
         # Step size accounts for test_size + embargo (next fold's train starts after embargo)
-        available = n_samples - self.train_size - self.gap
+        available = n_units - self.train_size - self.gap
         step = max(1, (available - self.test_size) // max(self.n_splits - 1, 1))
 
         for i in range(self.n_splits):
@@ -174,15 +171,57 @@ def split(self, n_samples: int):
 
             train_end = self.train_size + i * step
             test_start = train_end + self.gap  # purge gap
-            test_end = min(test_start + self.test_size, n_samples)
+            test_end = min(test_start + self.test_size, n_units)
 
-            if train_end >= n_samples or test_start >= n_samples:
+            if train_end >= n_units or test_start >= n_units:
                 break
 
-            train_idx = list(range(train_start, min(train_end, n_samples)))
+            yield train_start, min(train_end, n_units), test_start, test_end
+
+    def split(self, n_samples: int):
+        """Generate train/test index pairs with purge and embargo.
+
+        Args:
+            n_samples: Total number of samples.
+
+        Yields:
+            (train_indices, test_indices) tuples.
+        """
+        for train_start, train_end, test_start, test_end in self._folds(n_samples):
+            train_idx = list(range(train_start, train_end))
             test_idx = list(range(test_start, test_end))
+            if train_idx and test_idx:
+                yield train_idx, test_idx
+
+    def split_by_dates(self, row_dates):
+        """Split on the date axis when each row is a (date, asset) pair.
+
+        `split()` counts in rows. On a cross-sectional panel that makes
+        `train_size=504` about twelve trading days rather than two years, and a
+        ten-row purge gap removes less than a single cross-section -- so train
+        and test end up sharing dates, and the labels of those dates leak
+        across the split.
+
+        This applies the identical fold arithmetic to the unique date axis and
+        returns row indices, so callers do not change.
 
-            if len(train_idx) > 0 and len(test_idx) > 0:
+        Args:
+            row_dates: sequence of dates, one per row, in row order.
+
+        Yields:
+            (train_row_indices, test_row_indices) tuples.
+        """
+        dates = pd.Index(pd.unique(pd.Index(row_dates)))
+        rows_by_date: dict = {}
+        for i, d in enumerate(row_dates):
+            rows_by_date.setdefault(d, []).append(i)
+
+        for train_start, train_end, test_start, test_end in self._folds(len(dates)):
+            train_dates = dates[train_start:train_end]
+            test_dates = dates[test_start:test_end]
+            train_idx = [i for d in train_dates for i in rows_by_date.get(d, ())]
+            test_idx = [i for d in test_dates for i in rows_by_date.get(d, ())]
+            if train_idx and test_idx:
                 yield train_idx, test_idx
 
 
@@ -298,6 +337,12 @@ def _prepare_dataset(
 
         Flattens all dates and assets into a single feature matrix.
         Each row = (date, asset) pair.
+
+        Returns:
+            (X, y, feature_names, row_dates) -- `row_dates` carries the date of
+            each row, so a splitter can divide on the date axis instead of
+            counting rows. Row counts per date vary (NaN assets are dropped),
+            so this cannot be reconstructed by the caller.
         """
         feature_names = sorted(factors.keys())[:self.config.top_n_features]
         first_factor = factors[feature_names[0]]
@@ -306,6 +351,7 @@ def _prepare_dataset(
 
         X_list = []
         y_list = []
+        dates_list = []
 
         for i, date in enumerate(dates):
             actual_idx = start_idx + i
@@ -332,11 +378,17 @@ def _prepare_dataset(
             if len(X_valid) > 0:
                 X_list.append(X_valid)
                 y_list.append(y_valid)
+                dates_list.extend([date] * len(X_valid))
 
         if not X_list:
-            return np.array([]), np.array([]), feature_names
+            return np.array([]), np.array([]), feature_names, pd.Index([])
 
-        return np.vstack(X_list), np.concatenate(y_list), feature_names
+        return (
+            np.vstack(X_list),
+            np.concatenate(y_list),
+            feature_names,
+            pd.Index(dates_list),
+        )
 
     def train(
         self,
@@ -356,7 +408,7 @@ def train(
         n_dates = len(first_factor.index)
 
         # Prepare full dataset for CV
-        X_full, y_full, feature_names = self._prepare_dataset(
+        X_full, y_full, feature_names, row_dates = self._prepare_dataset(
             factors, forward_returns, 0, n_dates
         )
 
@@ -374,7 +426,9 @@ def train(
         )
 
         cv_results = []
-        for fold, (train_idx, test_idx) in enumerate(cv.split(len(X_full))):
+        # Split on dates, not rows: X_full has one row per (date, asset), so a
+        # row-based split would put the same trading day in both train and test.
+        for fold, (train_idx, test_idx) in enumerate(cv.split_by_dates(row_dates)):
             X_train, y_train = X_full[train_idx], y_full[train_idx]
             X_test, y_test = X_full[test_idx], y_full[test_idx]
 
@@ -498,7 +552,7 @@ def get_shap_values(self, factors: dict[str, pd.DataFrame], n_samples: int = 500
 
         first_factor = list(factors.values())[0]
         n_dates = len(first_factor.index)
-        X, _, feature_names = self._prepare_dataset(factors, pd.DataFrame(), 0, n_dates)
+        X, _, feature_names, _ = self._prepare_dataset(factors, pd.DataFrame(), 0, n_dates)
 
         if len(X) == 0:
             return {"error": "no data"}
@@ -550,7 +604,7 @@ def generate(
             if i - last_train_idx >= self.config.retrain_frequency or model is None:
                 # Prepare training data: all dates before i
                 train_start = max(0, i - self.config.train_window)
-                X_train, y_train, feat_names = self._prepare_dataset(
+                X_train, y_train, feat_names, _ = self._prepare_dataset(
                     factors, forward_returns, train_start, i
                 )
 

tests/test_alpha/test_ml_signal.py

@@ -15,9 +15,16 @@
 
 @pytest.fixture
 def sample_factors():
-    """Create sample factor data for testing."""
+    """Create sample factor data for testing.
+
+    `n_dates` has to satisfy the configured training window in *days*, not in
+    rows. The default `train_window` is 504 (two years), plus the purge gap and
+    the test window, so a 300-day panel cannot support it -- the previous 300
+    only "worked" because the splitter counted flattened (date, asset) rows and
+    so trained on about ten days of data while believing it had 504.
+    """
     np.random.seed(42)
-    n_dates, n_assets = 300, 50
+    n_dates, n_assets = 800, 50
     dates = pd.bdate_range("2022-01-01", periods=n_dates)
     assets = [f"stock_{i:03d}" for i in range(n_assets)]
 
@@ -96,12 +103,14 @@
 
     def test_prepare_dataset(self, sample_factors, sample_forward_returns):
         gen = MLSignalGenerator()
-        X, y, names = gen._prepare_dataset(
-            sample_factors, sample_forward_returns, 0, 300
+        X, y, names, row_dates = gen._prepare_dataset(
+            sample_factors, sample_forward_returns, 0, 200
         )
         assert len(X) > 0
         assert len(y) == len(X)
         assert len(names) == 5
+        # One date per row, so a splitter can divide on the date axis.
+        assert len(row_dates) == len(X)
 
     @pytest.mark.skipif(not HAS_LGB, reason="lightgbm not installed")
     def test_train_lightgbm(self, sample_factors, sample_forward_returns):
@@ -127,8 +136,10 @@
         gen.train(sample_factors, sample_forward_returns)
         signal = gen.predict(sample_factors)
         assert isinstance(signal, pd.DataFrame)
-        assert signal.shape[0] == 300
-        assert signal.shape[1] == 50
+        # Derive from the fixture rather than hardcoding, so growing the panel
+        # does not silently invalidate this assertion.
+        assert signal.shape[0] == sample_factors["momentum"].shape[0]
+        assert signal.shape[1] == sample_factors["momentum"].shape[1]
         # Cross-sectional rank should be in [-0.5, 0.5]
         for col in signal.columns:
             valid = signal[col].dropna()
@@ -196,3 +207,78 @@
         assert perf.date == "2024-01-01"
         assert perf.test_ic == 0.03
         assert perf.test_icir == 1.5
+
+
+class TestTimeSeriesCVDates:
+    """Regression for BUG-23.
+
+    `split()` counts rows, but `_prepare_dataset` returns one row per
+    (date, asset) pair. Counting rows made `train_size=504` about ten trading
+    days rather than two years, and a ten-row purge gap less than a single
+    cross-section -- so train and test shared dates and the labels of those
+    dates leaked across the split.
+    """
+
+    def test_train_and_test_are_disjoint_on_dates(
+        self, sample_factors, sample_forward_returns
+    ):
+        gen = MLSignalGenerator()
+        n_dates = len(sample_factors["momentum"].index)
+        _, _, _, row_dates = gen._prepare_dataset(
+            sample_factors, sample_forward_returns, 0, n_dates
+        )
+        cv = TimeSeriesCV(n_splits=2, train_size=252, test_size=63, gap=10)
+
+        folds = list(cv.split_by_dates(row_dates))
+        assert folds, "expected at least one fold"
+
+        for train_idx, test_idx in folds:
+            train_dates = set(row_dates[train_idx])
+            test_dates = set(row_dates[test_idx])
+            overlap = train_dates & test_dates
+            assert not overlap, f"train and test share {len(overlap)} dates"
+            assert max(train_dates) < min(test_dates), "train must precede test"
+
+    def test_train_size_is_counted_in_days(self, sample_factors, sample_forward_returns):
+        """504 must mean 504 trading days, not 504 flattened rows."""
+        gen = MLSignalGenerator()
+        n_dates = len(sample_factors["momentum"].index)
+        _, _, _, row_dates = gen._prepare_dataset(
+            sample_factors, sample_forward_returns, 0, n_dates
+        )
+        cv = TimeSeriesCV(n_splits=1, train_size=252, test_size=63, gap=10)
+        train_idx, _ = next(iter(cv.split_by_dates(row_dates)))
+        assert len(set(row_dates[train_idx])) == 252
+
+    @pytest.mark.skipif(not HAS_LGB, reason="lightgbm not installed")
+    def test_train_splits_on_dates_not_rows(
+        self, sample_factors, sample_forward_returns, monkeypatch
+    ):
+        """The guard the two tests above cannot provide.
+
+        They exercise `split_by_dates` directly, so they would still pass if
+        `train()` went back to calling `split(len(X_full))`. This asserts which
+        one the caller actually uses.
+        """
+        import quant_platform.alpha.ml_signal as ml
+
+        calls = {"by_dates": 0, "by_rows": 0}
+        orig_by_dates = ml.TimeSeriesCV.split_by_dates
+        orig_by_rows = ml.TimeSeriesCV.split
+
+        def spy_by_dates(self, row_dates):
+            calls["by_dates"] += 1
+            return orig_by_dates(self, row_dates)
+
+        def spy_by_rows(self, n_samples):
+            calls["by_rows"] += 1
+            return orig_by_rows(self, n_samples)
+
+        monkeypatch.setattr(ml.TimeSeriesCV, "split_by_dates", spy_by_dates)
+        monkeypatch.setattr(ml.TimeSeriesCV, "split", spy_by_rows)
+
+        cfg = MLSignalConfig(model_type="lightgbm", n_splits=2)
+        MLSignalGenerator(config=cfg).train(sample_factors, sample_forward_returns)
+
+        assert calls["by_dates"] > 0, "train() did not split on dates"
+        assert calls["by_rows"] == 0, "train() fell back to counting rows"