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

[hugo2046]

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.

[adityaanikam]

Verified both halves precisely. TSDataSampler.init does self.data = data.swaplevel().sort_index(), which physically sorts rows by the post-swap index, instrument first per the data_index docstring. get_index() then does self.data_index.swaplevel(), which only swaps the labels inside each tuple back to (datetime, instrument), it never touches row order. So the returned index reads like day by day pairs but the actual sequence is still grouped by instrument, exactly matching your repro output.

Your replacement for DailyBatchSampler looks correct, grouping by the actual datetime values in positions rather than assuming contiguous blocks sidesteps the whole issue. Given this also affects get_index()'s own docstring claim about day by day sampling being supported, is the fix meant to stay scoped to DailyBatchSampler, or does get_index() itself need a doc fix or a real reordering so other callers relying on that docstring aren't quietly wrong too?