Dataset (#4)

* Dataset is an iterable of DataEntry

* test_forecast_multivariate

* formatting

* offset can also be W-MON

* fix type

* fourier_time_features_from_frequency_str for weekly data

* randomly shuffle dataset for each worker

* filedataset is not scriptable

* read file randomly

* list and file datasets shuffle the time series in train

* do not shuffle time series in multivariate grouper

* fix tests

* formatting

* formatting

pts/dataset/artificial.py

@@ -72,7 +72,7 @@ class ArtificialDataset:
         return TrainDatasets(
             metadata=self.metadata,
             train=ListDataset(self.train, self.freq),
-            test=ListDataset(self.test, self.freq),
+            test=ListDataset(self.test, self.freq, is_train=False),
         )
 
 
@@ -697,7 +697,7 @@ class RecipeDataset(ArtificialDataset):
         return TrainDatasets(
             metadata=metadata,
             train=ListDataset(train_data, metadata.freq),
-            test=ListDataset(test_data, metadata.freq),
+            test=ListDataset(test_data, metadata.freq, is_train=False),
         )
 
 
@@ -776,6 +776,7 @@ def constant_dataset() -> Tuple[DatasetInfo, Dataset, Dataset]:
             for i in range(10)
         ],
         freq=metadata.freq,
+        is_train=False
     )
 
     info = DatasetInfo(

pts/dataset/common.py

@@ -4,8 +4,12 @@ from typing import Any, Dict, Iterable, NamedTuple, Sized, List, Optional, Itera
 import pandas as pd
 from pydantic import BaseModel
 
+# Dictionary used for data flowing through the transformations.
 DataEntry = Dict[str, Any]
 
+# A Dataset is an iterable of DataEntry.
+Dataset = Iterable[DataEntry]
+
 
 class SourceContext(NamedTuple):
     source: str
@@ -37,16 +41,6 @@ class FieldName:
     FORECAST_START = "forecast_start"
 
 
-class Dataset(Sized, Iterable[DataEntry], ABC):
-    @abstractmethod
-    def __iter__(self) -> Iterator[DataEntry]:
-        pass
-
-    @abstractmethod
-    def __len__(self):
-        pass
-
-
 class CategoricalFeatureInfo(BaseModel):
     name: str
     cardinality: str
@@ -78,6 +72,7 @@ class TrainDatasets(NamedTuple):
     train: Dataset
     test: Optional[Dataset] = None
 
+
 class DateConstants:
     """
     Default constants for specific dates.

pts/dataset/file_dataset.py

@@ -3,6 +3,7 @@ from pathlib import Path
 from typing import NamedTuple
 from typing import Iterator, List
 import glob
+import random
 
 import rapidjson as json
 
@@ -36,12 +37,17 @@ class JsonLinesFile:
         JSON Lines file.
     """
 
-    def __init__(self, path) -> None:
+    def __init__(self, path: Path, is_train: bool = True) -> None:
         self.path = path
+        self.is_train = is_train
 
     def __iter__(self):
         with open(self.path) as jsonl_file:
-            for line_number, raw in enumerate(jsonl_file, start=1):
+            lines = jsonl_file.read().splitlines()
+            if self.is_train:
+                random.shuffle(lines)
+
+            for line_number, raw in enumerate(lines, start=1):
                 span = Span(path=self.path, line=line_number)
                 try:
                     yield Line(json.loads(raw), span=span)
@@ -74,7 +80,10 @@ class FileDataset(Dataset):
         Whether to accept only univariate target time series.
     """
 
-    def __init__(self, path: Path, freq: str, one_dim_target: bool = True,) -> None:
+    def __init__(
+        self, path: Path, freq: str, one_dim_target: bool = True, is_train: bool = True
+    ) -> None:
+        self.is_train = is_train
         self.path = path
         self.process = ProcessDataEntry(freq, one_dim_target=one_dim_target)
         if not self.files():
@@ -82,7 +91,7 @@ class FileDataset(Dataset):
 
     def __iter__(self) -> Iterator[DataEntry]:
         for path in self.files():
-            for line in JsonLinesFile(path):
+            for line in JsonLinesFile(path, self.is_train):
                 data = self.process(line.content)
                 data["source"] = SourceContext(
                     source=line.span.path, row=line.span.line

pts/dataset/list_dataset.py

@@ -1,4 +1,5 @@
 from typing import Iterable
+import random
 
 from .common import DataEntry, Dataset, SourceContext
 from .process import ProcessDataEntry
@@ -6,10 +7,17 @@ from .process import ProcessDataEntry
 
 class ListDataset(Dataset):
     def __init__(
-        self, data_iter: Iterable[DataEntry], freq: str, one_dim_target: bool = True
+        self,
+        data_iter: Iterable[DataEntry],
+        freq: str,
+        one_dim_target: bool = True,
+        is_train: bool = True,
     ) -> None:
         process = ProcessDataEntry(freq, one_dim_target)
         self.list_data = [process(data) for data in data_iter]
+        if is_train:
+            random.shuffle(self.list_data)
+
 
     def __iter__(self):
         source_name = "list_data"

pts/dataset/multivariate_grouper.py

@@ -21,6 +21,7 @@ from typing import Callable, Optional
 from .common import DataEntry, Dataset, FieldName, DateConstants
 from .list_dataset import ListDataset
 
+
 class MultivariateGrouper:
     """
     The MultivariateGrouper takes a univariate dataset and groups it into a
@@ -121,7 +122,9 @@ class MultivariateGrouper:
         grouped_data[FieldName.START] = self.first_timestamp
         grouped_data[FieldName.FEAT_STATIC_CAT] = [0]
 
-        return ListDataset([grouped_data], freq=self.frequency, one_dim_target=False)
+        return ListDataset(
+            [grouped_data], freq=self.frequency, one_dim_target=False, is_train=False
+        )
 
     def _prepare_test_data(self, dataset: Dataset) -> ListDataset:
         logging.info("group test time-series to datasets")
@@ -142,7 +145,9 @@ class MultivariateGrouper:
             grouped_data[FieldName.FEAT_STATIC_CAT] = [0]
             all_entries.append(grouped_data)
 
-        return ListDataset(all_entries, freq=self.frequency, one_dim_target=False)
+        return ListDataset(
+            all_entries, freq=self.frequency, one_dim_target=False, is_train=False
+        )
 
     def _align_data_entry(self, data: DataEntry) -> np.array:
         ts = self.to_ts(data)
@@ -204,4 +209,4 @@ class MultivariateGrouper:
                 freq=data[FieldName.START].freq,
             ),
         )
- 
+

pts/dataset/repository/_gp_copula_2019.py

@@ -111,7 +111,7 @@ def generate_gp_copula_dataset(dataset_path: Path, dataset_name: str):
 
     download_dataset(dataset_path.parent, ds_info)
     save_metadata(dataset_path, ds_info)
-    save_dataset(dataset_path / "train" , ds_info)
+    save_dataset(dataset_path / "train", ds_info)
     save_dataset(dataset_path / "test", ds_info)
     clean_up_dataset(dataset_path, ds_info)
 
@@ -137,7 +137,9 @@ def save_metadata(dataset_path: Path, ds_info: GPCopulaDataset):
 
 
 def save_dataset(dataset_path: Path, ds_info: GPCopulaDataset):
-    dataset = list(FileDataset(dataset_path / "*.json", freq=ds_info.freq))
+    dataset = list(
+        FileDataset(dataset_path / "*.json", freq=ds_info.freq, is_train=False)
+    )
     shutil.rmtree(dataset_path)
     train_file = dataset_path / "data.json"
     save_to_file(

pts/dataset/transformed_iterable_dataset.py

@@ -1,5 +1,6 @@
 import itertools
 from typing import Dict, Iterable, Iterator, Optional
+import random
 
 import numpy as np
 import torch
@@ -8,10 +9,12 @@ from pts.transform import Transformation
 
 from .common import DataEntry, Dataset
 
+
 class TransformedIterableDataset(torch.utils.data.IterableDataset):
     def __init__(
         self, dataset: Dataset, is_train: bool, transform: Transformation
     ) -> None:
+        super().__init__()
         self.dataset = dataset
         self.transform = transform
         self.is_train = is_train
@@ -28,7 +31,7 @@ class TransformedIterableDataset(torch.utils.data.IterableDataset):
                 for x in itertools.chain([first], collection):
                     yield x
 
-    def __iter__(self) -> Dict[str, np.ndarray]:
+    def __iter__(self) -> Iterator[Dict[str, np.ndarray]]:
         if self._cur_iter is None:
             self._cur_iter = self.transform(
                 self._iterate_forever(self.dataset), is_train=self.is_train
@@ -43,4 +46,4 @@ class TransformedIterableDataset(torch.utils.data.IterableDataset):
             }
 
     # def __len__(self) -> int:
-    #     return len(self.dataset)
+    #     return len(self.dataset)

pts/dataset/utils.py

@@ -60,7 +60,7 @@ def load_datasets(metadata, train, test) -> TrainDatasets:
     """
     meta = MetaData.parse_file(metadata)
     train_ds = FileDataset(train, meta.freq)
-    test_ds = FileDataset(test, meta.freq) if test else None
+    test_ds = FileDataset(test, meta.freq, is_train=False) if test else None
 
     return TrainDatasets(metadata=meta, train=train_ds, test=test_ds)
 

pts/feature/lag.py

@@ -85,7 +85,7 @@ def get_lags_for_frequency(
 
     if offset.name == "M":
         lags = _make_lags_for_month(offset.n)
-    elif offset.name == "W-SUN":
+    elif offset.name == "W-SUN" or offset.name == "W-MON":
         lags = _make_lags_for_week(offset.n)
     elif offset.name == "D":
         lags = _make_lags_for_day(offset.n) + _make_lags_for_week(offset.n / 7.0)

pts/feature/time_feature.py

@@ -172,7 +172,8 @@ def fourier_time_features_from_frequency_str(freq_str: str) -> List[TimeFeature]
 
     features = {
         "M": ["weekofyear"],
-        "W": ["daysinmonth", "weekofyear"],
+        "W-SUN": ["daysinmonth", "weekofyear"],
+        "W-MON": ["daysinmonth", "weekofyear"],
         "D": ["dayofweek"],
         "B": ["dayofweek", "dayofyear"],
         "H": ["hour", "dayofweek"],

pts/transform/transform.py

@@ -1,5 +1,5 @@
 from abc import ABC, abstractmethod
-from typing import Callable, Iterator, List
+from typing import Callable, Iterator, Iterable, List
 from functools import reduce
 
 
@@ -12,11 +12,11 @@ MAX_IDLE_TRANSFORMS = 100
 class Transformation(ABC):
     @abstractmethod
     def __call__(
-        self, data_it: Iterator[DataEntry], is_train: bool
+        self, data_it: Iterable[DataEntry], is_train: bool
     ) -> Iterator[DataEntry]:
         pass
 
-    def estimate(self, data_it: Iterator[DataEntry]) -> Iterator[DataEntry]:
+    def estimate(self, data_it: Iterable[DataEntry]) -> Iterator[DataEntry]:
         return data_it  # default is to pass through without estimation
 
     def chain(self, other: "Transformation") -> "Chain":
@@ -41,7 +41,7 @@ class Chain(Transformation):
                 self.transformations.append(transformation)
 
     def __call__(
-        self, data_it: Iterator[DataEntry], is_train: bool
+        self, data_it: Iterable[DataEntry], is_train: bool
     ) -> Iterator[DataEntry]:
         tmp = data_it
         for t in self.transformations:
@@ -54,7 +54,7 @@ class Chain(Transformation):
 
 class Identity(Transformation):
     def __call__(
-        self, data_it: Iterator[DataEntry], is_train: bool
+        self, data_it: Iterable[DataEntry], is_train: bool
     ) -> Iterator[DataEntry]:
         return data_it
 
@@ -64,7 +64,7 @@ class MapTransformation(Transformation):
     Base class for Transformations that returns exactly one result per input in the stream.
     """
 
-    def __call__(self, data_it: Iterator[DataEntry], is_train: bool) -> Iterator:
+    def __call__(self, data_it: Iterable[DataEntry], is_train: bool) -> Iterator:
         for data_entry in data_it:
             try:
                 yield self.map_transform(data_entry.copy(), is_train)
@@ -110,7 +110,7 @@ class FlatMapTransformation(Transformation):
     elements from the input stream.
     """
 
-    def __call__(self, data_it: Iterator[DataEntry], is_train: bool) -> Iterator:
+    def __call__(self, data_it: Iterable[DataEntry], is_train: bool) -> Iterator:
         num_idle_transforms = 0
         for data_entry in data_it:
             num_idle_transforms += 1

test/dataset/test_multivariate_grouper.py

@@ -43,12 +43,7 @@ UNIVARIATE_TS = [
 
 MULTIVARIATE_TS = [
     [{"start": "2014-09-07", "target": [[1, 2, 3, 4], [5, 6, 7, 8]]}],
-    [
-        {
-            "start": "2014-09-07",
-            "target": [[1, 2, 3, 4, 2.5], [6.5, 5, 6, 7, 8]],
-        }
-    ],
+    [{"start": "2014-09-07", "target": [[1, 2, 3, 4, 2.5], [6.5, 5, 6, 7, 8]],}],
     [{"start": "2014-09-07", "target": [[1, 2, 3, 4], [0, 0, 0, 0]]}],
     [
         {
@@ -72,21 +67,17 @@ TRAIN_FILL_RULE = [np.mean, np.mean, np.mean, np.mean, lambda x: 0.0]
 def test_multivariate_grouper_train(
     univariate_ts, multivariate_ts, train_fill_rule
 ) -> None:
-    univariate_ds = ListDataset(univariate_ts, freq="1D")
+    univariate_ds = ListDataset(univariate_ts, freq="1D", is_train=False)
     multivariate_ds = ListDataset(
-        multivariate_ts, freq="1D", one_dim_target=False
+        multivariate_ts, freq="1D", one_dim_target=False, is_train=False
     )
 
     grouper = MultivariateGrouper(train_fill_rule=train_fill_rule)
     assert (
-        list(grouper(univariate_ds))[0]["target"]
-        == list(multivariate_ds)[0]["target"]
+        list(grouper(univariate_ds))[0]["target"] == list(multivariate_ds)[0]["target"]
     ).all()
 
-    assert (
-        list(grouper(univariate_ds))[0]["start"]
-        == list(multivariate_ds)[0]["start"]
-    )
+    assert list(grouper(univariate_ds))[0]["start"] == list(multivariate_ds)[0]["start"]
 
 
 UNIVARIATE_TS_TEST = [
@@ -121,30 +112,21 @@ MAX_TARGET_DIM = [2, 1]
 
 @pytest.mark.parametrize(
     "univariate_ts, multivariate_ts, test_fill_rule, max_target_dim",
-    zip(
-        UNIVARIATE_TS_TEST,
-        MULTIVARIATE_TS_TEST,
-        TEST_FILL_RULE,
-        MAX_TARGET_DIM,
-    ),
+    zip(UNIVARIATE_TS_TEST, MULTIVARIATE_TS_TEST, TEST_FILL_RULE, MAX_TARGET_DIM,),
 )
 def test_multivariate_grouper_test(
     univariate_ts, multivariate_ts, test_fill_rule, max_target_dim
 ) -> None:
-    univariate_ds = ListDataset(univariate_ts, freq="1D")
+    univariate_ds = ListDataset(univariate_ts, freq="1D", is_train=False)
     multivariate_ds = ListDataset(
-        multivariate_ts, freq="1D", one_dim_target=False
+        multivariate_ts, freq="1D", one_dim_target=False, is_train=False
     )
 
     grouper = MultivariateGrouper(
-        test_fill_rule=test_fill_rule,
-        num_test_dates=2,
-        max_target_dim=max_target_dim,
+        test_fill_rule=test_fill_rule, num_test_dates=2, max_target_dim=max_target_dim,
     )
 
-    for grouped_data, multivariate_data in zip(
-        grouper(univariate_ds), multivariate_ds
-    ):
+    for grouped_data, multivariate_data in zip(grouper(univariate_ds), multivariate_ds):
         assert (grouped_data["target"] == multivariate_data["target"]).all()
 
         assert grouped_data["start"] == multivariate_data["start"]

test/model/test_forecast.py

@@ -92,3 +92,33 @@ def test_DistributionForecast():
     assert forecast.prediction_length == pred_length
     assert len(forecast.index) == pred_length
     assert forecast.index[0] == pd.Timestamp(START_DATE)
+
+
+@pytest.mark.parametrize(
+    "forecast, exp_index",
+    [
+        (
+            SampleForecast(
+                samples=np.random.normal(size=(100, 7, 3)),
+                start_date=pd.Timestamp("2020-01-01 00:00:00"),
+                freq="1D",
+            ),
+            pd.date_range(
+                start=pd.Timestamp("2020-01-01 00:00:00"), freq="1D", periods=7,
+            ),
+        ),
+        (
+            DistributionForecast(
+                Uniform(low=torch.zeros(size=(5, 2)), high=torch.ones(size=(5, 2)),),
+                start_date=pd.Timestamp("2020-01-01 00:00:00"),
+                freq="W",
+            ),
+            pd.date_range(
+                start=pd.Timestamp("2020-01-01 00:00:00"), freq="W", periods=5,
+            ),
+        ),
+    ],
+)
+def test_forecast_multivariate(forecast, exp_index):
+    assert forecast.prediction_length == len(exp_index)
+    assert np.all(forecast.index == exp_index)