fix(model): group DailyBatchSampler batches by datetime positions

qlib/contrib/model/pytorch_gats_ts.py

@@ -26,19 +26,19 @@
 class DailyBatchSampler(Sampler):
     def __init__(self, data_source):
         self.data_source = data_source
-        # calculate number of samples in each batch
-        self.daily_count = (
-            pd.Series(index=self.data_source.get_index()).groupby("datetime", group_keys=False).size().values
-        )
-        self.daily_index = np.roll(np.cumsum(self.daily_count), 1)  # calculate begin index of each batch
-        self.daily_index[0] = 0
+        # ``TSDataSampler.get_index()`` swaps the level labels back to (datetime,
+        # instrument) but keeps the instrument-major row order, so the rows of a
+        # trading day are not contiguous. Group the row positions by the "datetime"
+        # level instead of assuming contiguous blocks.
+        index = self.data_source.get_index()
+        positions = pd.Series(np.arange(len(index)), index=index.get_level_values("datetime"))
+        self.batches = [g.to_numpy() for _, g in positions.groupby(level=0, sort=True)]
 
     def __iter__(self):
-        for idx, count in zip(self.daily_index, self.daily_count):
-            yield np.arange(idx, idx + count)
+        yield from self.batches
 
     def __len__(self):
-        return len(self.data_source)
+        return len(self.batches)
 
 
 class GATs(Model):
@@ -332,7 +332,7 @@ def predict(self, dataset):
 
             preds.append(pred)
 
-        return pd.Series(np.concatenate(preds), index=dl_test.get_index())
+        return pd.Series(np.concatenate(preds), index=dl_test.get_index()[np.concatenate(sampler_test.batches)])
 
 
 class GATModel(nn.Module):

qlib/data/dataset/__init__.py

@@ -478,6 +478,12 @@ def get_index(self):
         """
         Get the pandas index of the data, it will be useful in following scenarios
         - Special sampler will be used (e.g. user want to sample day by day)
+
+        NOTE: the level labels are swapped back to (datetime, instrument), but the
+        rows are NOT reordered: they keep the instrument-major order described in
+        "Indices design" above. Samplers that need day-by-day batches should group
+        the row positions by the "datetime" level instead of assuming the rows of
+        each day are contiguous.
         """
         return self.data_index.swaplevel()  # to align the order of multiple index of original data received by __init__
 

tests/test_daily_batch_sampler.py

@@ -0,0 +1,37 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import unittest
+
+import numpy as np
+import pandas as pd
+
+from qlib.contrib.model.pytorch_gats_ts import DailyBatchSampler
+from qlib.data.dataset import TSDataSampler
+
+
+class TestDailyBatchSampler(unittest.TestCase):
+    def setUp(self):
+        dates = pd.date_range("2020-01-01", periods=5, freq="B")
+        instruments = ["A", "B", "C"]
+        index = pd.MultiIndex.from_product([dates, instruments], names=["datetime", "instrument"])
+        data = pd.DataFrame({"feature": np.arange(len(index)), "label": np.arange(len(index))}, index=index)
+        self.data_source = TSDataSampler(data, dates[0], dates[-1], step_len=2)
+
+    def test_each_batch_is_one_day(self):
+        sampler = DailyBatchSampler(self.data_source)
+        index = self.data_source.get_index()
+        for batch in sampler:
+            batch_index = index[batch]
+            self.assertEqual(batch_index.get_level_values("datetime").nunique(), 1)
+            self.assertEqual(batch_index.get_level_values("instrument").nunique(), 3)
+
+    def test_batches_cover_all_samples_once(self):
+        sampler = DailyBatchSampler(self.data_source)
+        positions = np.concatenate(list(sampler))
+        self.assertEqual(sorted(positions.tolist()), list(range(len(self.data_source.get_index()))))
+        self.assertEqual(len(sampler), 5)
+
+
+if __name__ == "__main__":
+    unittest.main()