From 396378bb84a4c69e133c97c36433de712d14f859 Mon Sep 17 00:00:00 2001 From: Akshay Kumar Date: Thu, 20 Aug 2026 23:25:02 +0530 Subject: [PATCH] fix(model): resolve cross-day batching issue in DailyBatchSampler --- qlib/contrib/model/pytorch_gats_ts.py | 16 ++++++---------- 1 file changed, 6 insertions(+), 10 deletions(-) diff --git a/qlib/contrib/model/pytorch_gats_ts.py b/qlib/contrib/model/pytorch_gats_ts.py index 09f0ac08b25..a57f22fab38 100644 --- a/qlib/contrib/model/pytorch_gats_ts.py +++ b/qlib/contrib/model/pytorch_gats_ts.py @@ -26,19 +26,15 @@ 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 = 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 +328,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):