microsoft / microsoft/qlib

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

Open
#2,319 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

Dominant language
Python
Stars
48.7k
Forks
7.7k
PR merge metrics
No merged PRs in 30d

Description

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.

Contributor guide

No contributing guide indexed for this repository

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start with qlib/data/dataset/init.py, especially TSDataSampler.init and get_index(), then inspect DailyBatchSampler in qlib/contrib/model/pytorch_gats_ts.py. Reproduce the ordering shown in the issue and run the existing per-day-batch unit test; done means each batch contains one datetime and prediction index alignment follows the sampler iteration order.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, pandas, python
Domain
data, machine-learning, testing-qa
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Active
Clarity
Clearly specified
Newbie friendliness
74/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.