Fix daily batching for time-series samplers

qlib/contrib/model/pytorch_gats_ts.py

@@ -26,19 +26,20 @@
 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
+        index = self.data_source.get_index()
+        positions = pd.Series(np.arange(len(index)), index=index)
+        self.daily_index = [group.values for _, group in positions.groupby(level="datetime", 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.daily_index
 
     def __len__(self):
-        return len(self.data_source)
+        return len(self.daily_index)
+
+    def get_index(self):
+        """Return the data source index in sampler iteration order."""
+        positions = np.concatenate(self.daily_index)
+        return self.data_source.get_index()[positions]
 
 
 class GATs(Model):
@@ -332,7 +333,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=sampler_test.get_index())
 
 
 class GATModel(nn.Module):

tests/model/test_pytorch_gats_ts.py

@@ -0,0 +1,25 @@
+import pandas as pd
+
+from qlib.contrib.model.pytorch_gats_ts import DailyBatchSampler
+from qlib.data.dataset import TSDataSampler
+
+
+def test_daily_batch_sampler_groups_actual_positions_by_datetime():
+    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": range(len(index)), "label": range(len(index))}, index=index)
+    data_source = TSDataSampler(data, dates[0], dates[-1], step_len=2)
+
+    sampler = DailyBatchSampler(data_source)
+    source_index = data_source.get_index()
+    batches = list(sampler)
+
+    assert len(sampler) == len(dates)
+    for date, batch in zip(dates, batches):
+        batch_index = source_index[batch]
+        assert batch_index.get_level_values("datetime").unique().tolist() == [date]
+        assert batch_index.get_level_values("instrument").tolist() == instruments
+
+    expected_index = pd.MultiIndex.from_product([dates, instruments], names=["datetime", "instrument"])
+    assert sampler.get_index().equals(expected_index)