src/tsfm_bench/data/dataset.py

from __future__ import annotations

import math
import os
from enum import Enum
from functools import cached_property
from pathlib import Path
from typing import Iterable, Iterator

from dotenv import load_dotenv
from gluonts.dataset import DataEntry
from gluonts.dataset.common import ProcessDataEntry
from gluonts.dataset.split import TestData, TrainingDataset, split
from gluonts.itertools import Map
from gluonts.time_feature import norm_freq_str
from gluonts.transform import Transformation
from pandas.tseries.frequencies import to_offset
from toolz import compose

from tsfm_bench.data.interface import TimeSeriesDataSource, TimeSeriesRecord
from tsfm_bench.data.registry import load_data_source

TEST_SPLIT = 0.1
MAX_WINDOW = 20

M4_PRED_LENGTH_MAP = {
    "A": 6,
    "Q": 8,
    "M": 18,
    "W": 13,
    "D": 14,
    "H": 48,
}

PRED_LENGTH_MAP = {
    "M": 12,
    "W": 8,
    "D": 30,
    "H": 48,
    "T": 48,
    "S": 60,
}


class Term(Enum):
    SHORT = "short"
    MEDIUM = "medium"
    LONG = "long"

    @property
    def multiplier(self) -> int:
        if self == Term.SHORT:
            return 1
        if self == Term.MEDIUM:
            return 10
        return 15


def itemize_start(data_entry: DataEntry) -> DataEntry:
    start = data_entry["start"]
    if hasattr(start, "item"):
        data_entry["start"] = start.item()
    return data_entry


def maybe_reconvert_freq(freq: str) -> str:
    deprecated_map = {
        "Y": "A",
        "YE": "A",
        "QE": "Q",
        "ME": "M",
        "h": "H",
        "min": "T",
        "s": "S",
        "us": "U",
    }
    return deprecated_map.get(freq, freq)


class MultivariateToUnivariate(Transformation):
    def __init__(self, field: str):
        self.field = field

    def __call__(
        self, data_it: Iterable[DataEntry], is_train: bool = False
    ) -> Iterator[DataEntry]:
        for data_entry in data_it:
            item_id = data_entry["item_id"]
            val_ls = list(data_entry[self.field])
            for dim_id, val in enumerate(val_ls):
                univariate_entry = data_entry.copy()
                univariate_entry[self.field] = val
                univariate_entry["item_id"] = f"{item_id}_dim{dim_id}"
                yield univariate_entry


def _record_to_entry(record: TimeSeriesRecord) -> dict:
    import pandas as pd

    return {
        "item_id": record.item_id,
        "start": pd.Period(record.start),
        "target": record.target,
        "freq": record.freq,
    }


class Dataset:
    """GluonTS-compatible dataset backed by a streaming data source."""

    def __init__(
        self,
        name: str,
        term: Term | str = Term.SHORT,
        to_univariate: bool = False,
        source: TimeSeriesDataSource | None = None,
        config_path: str | Path | None = None,
    ):
        load_dotenv()
        if source is None:
            env_config = os.getenv("TSFM_BENCH_CONFIG")
            source = load_data_source(config_path or env_config)

        self.source = source
        self.metadata = source.get_metadata(name)
        self.name = name
        self.term = Term(term)

        records = list(source.stream(name))
        self._raw_entries = [_record_to_entry(record) for record in records]

        first_target = self._raw_entries[0]["target"]
        target_dim = first_target.shape[0] if getattr(first_target, "ndim", 1) > 1 else 1

        process = ProcessDataEntry(
            self.freq,
            one_dim_target=target_dim == 1,
        )
        self.gluonts_dataset = Map(compose(process, itemize_start), self._raw_entries)
        if to_univariate:
            self.gluonts_dataset = MultivariateToUnivariate("target").apply(
                self.gluonts_dataset
            )

    @cached_property
    def prediction_length(self) -> int:
        freq = norm_freq_str(to_offset(self.freq).name)
        freq = maybe_reconvert_freq(freq)
        pred_len = PRED_LENGTH_MAP.get(freq, 48)
        return self.term.multiplier * pred_len

    @cached_property
    def freq(self) -> str:
        return self.metadata.frequency

    @cached_property
    def target_dim(self) -> int:
        target = self._raw_entries[0]["target"]
        return target.shape[0] if getattr(target, "ndim", 1) > 1 else 1

    @cached_property
    def windows(self) -> int:
        w = math.ceil(TEST_SPLIT * self._min_series_length / self.prediction_length)
        return min(max(1, w), MAX_WINDOW)

    @cached_property
    def _min_series_length(self) -> int:
        lengths = []
        for entry in self._raw_entries:
            target = entry["target"]
            lengths.append(target.shape[-1] if getattr(target, "ndim", 1) > 1 else len(target))
        return min(lengths)

    @property
    def training_dataset(self) -> TrainingDataset:
        training_dataset, _ = split(
            self.gluonts_dataset, offset=-self.prediction_length * (self.windows + 1)
        )
        return training_dataset

    @property
    def validation_dataset(self) -> TrainingDataset:
        validation_dataset, _ = split(
            self.gluonts_dataset, offset=-self.prediction_length * self.windows
        )
        return validation_dataset

    @property
    def test_data(self) -> TestData:
        _, test_template = split(
            self.gluonts_dataset, offset=-self.prediction_length * self.windows
        )
        return test_template.generate_instances(
            prediction_length=self.prediction_length,
            windows=self.windows,
            distance=self.prediction_length,
        )