#2319·qlib

DailyBatchSampler in pytorch_gats_ts.py yields cross-day batches: TSDataSampler.get_index() swaps index labels but not row order

Author: hugo2046Created Aug 16, 2026Updated Sep 7, 2026

Bug description

DailyBatchSampler (qlib/contrib/model/pytorch_gats_ts.py) assumes rows belonging to the same trading day are contiguous in the data source: it computes per-day counts via groupby("datetime").size(), converts them to start offsets with cumsum, and slices contiguous ranges np.arange(idx, idx + count).

However, TSDataSampler stores data instrument-major — qlib/data/dataset/__init__.py builds it as self.data = data.swaplevel().sort_index() (and the data_index docstring itself says index order <instrument, datetime>). get_index() then returns self.data_index.swaplevel(), which swaps the labels of the MultiIndex back to (datetime, instrument) but does not reorder rows.

As a result each "daily" batch actually contains one instrument across many consecutive days, not one day's cross-section. Any model trained with this sampler (e.g. the GATs benchmark, and code copying this sampler) silently trains its graph/attention over "one stock's history" instead of "one day's cross-section".

Minimal reproduction

import numpy as np
import pandas as pd
from qlib.data.dataset import TSDataSampler

dates = pd.date_range("2020-01-01", periods=5, freq="B")
insts = ["A", "B", "C"]
idx = pd.MultiIndex.from_product([dates, insts], names=["datetime", "instrument"])
df = pd.DataFrame({"f": range(len(idx)), "label": range(len(idx))}, index=idx)

s = TSDataSampler(df, dates[0], dates[-1], step_len=2)
print(s.get_index()[:6].tolist())
# [(2020-01-01, A), (2020-01-02, A), (2020-01-03, A), (2020-01-06, A), (2020-01-07, A), (2020-01-01, B)]
# -> labels say (datetime, instrument), but rows are instrument-major.

# DailyBatchSampler's first "day" batch is rows 0..2 (3 instruments expected),
# which are actually instrument A on three different days.

Note the docstring of get_index() explicitly advertises the day-by-day use case: "Special sampler will be used (e.g. user want to sample day by day)".

Expected behavior

Either:

  1. get_index() documents that row order is instrument-major (and DailyBatchSampler is fixed to not assume contiguity), or
  2. DailyBatchSampler groups actual row positions by datetime instead of assuming contiguous blocks.

Suggested fix (drop-in for DailyBatchSampler)

class DailyBatchSampler(Sampler):
    def __init__(self, data_source):
        self.data_source = data_source
        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):
        yield from self.batches

    def __len__(self):
        return len(self.batches)

When using this, downstream index alignment for predictions must follow the sampler's iteration order, e.g. dl_test.get_index()[np.concatenate(self.batches)].

Environment

Verified against current main: qlib/data/dataset/__init__.py (get_index, and data.swaplevel().sort_index() in TSDataSampler.__init__) and qlib/contrib/model/pytorch_gats_ts.py (DailyBatchSampler) are unchanged. Found while implementing a cross-sectional model whose per-day-batch unit test ("each batch contains exactly one datetime") failed against the GATs-style sampler.