diff --git a/alpha/ml_signal.py b/alpha/ml_signal.py index af464bd..dcdc889 100644 --- a/alpha/ml_signal.py +++ b/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 ) diff --git a/tests/test_alpha/test_ml_signal.py b/tests/test_alpha/test_ml_signal.py index 5d9f874..e57fff6 100644 --- a/tests/test_alpha/test_ml_signal.py +++ b/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_features(self, sample_factors): 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 @@ def test_predict(self, sample_factors, sample_forward_returns): 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 @@ def test_dataclass(self): 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"