[ENH] Add train/test split strategy for v2 datamodules

pytorch_forecasting/data/data_module/_encoder_decoder_data_module.py

@@ -135,6 +135,7 @@ def __init__(
         batch_size: int = 32,
         num_workers: int = 0,
         train_val_test_split: tuple = (0.7, 0.15, 0.15),
+        train_val_test_split_strategy: str = "random",
     ):
         self.time_series_dataset = time_series_dataset
         self.max_encoder_length = max_encoder_length
@@ -153,6 +154,7 @@ def __init__(
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.train_val_test_split = train_val_test_split
+        self.train_val_test_split_strategy = train_val_test_split_strategy
 
         warn(
             "EncoderDecoderTimeSeriesDataModule is part of an experimental "
@@ -964,7 +966,14 @@ def _ensure_split(self):
             return
 
         total_series = len(self.time_series_dataset)
-        self._split_indices = torch.randperm(total_series)
+        if self.train_val_test_split_strategy == "random":
+            self._split_indices = torch.randperm(total_series)
+        elif self.train_val_test_split_strategy == "sequential":
+            self._split_indices = torch.arange(total_series)
+        else:
+            raise ValueError(
+                f"Unknown split strategy: {self.train_val_test_split_strategy}"
+            )
 
         self._train_size = int(self.train_val_test_split[0] * total_series)
         self._val_size = int(self.train_val_test_split[1] * total_series)

pytorch_forecasting/data/data_module/_tslib_data_module.py

@@ -309,6 +309,7 @@ def __init__(
         batch_size: int = 32,
         num_workers: int = 0,
         train_val_test_split: tuple[float, float, float] = (0.7, 0.15, 0.15),
+        train_val_test_split_strategy: str = "random",
         collate_fn: Callable | None = None,
         **kwargs,
     ) -> None:
@@ -323,6 +324,7 @@ def __init__(
         self.batch_size = batch_size
         self.num_workers = num_workers
         self.train_val_test_split = train_val_test_split
+        self.train_val_test_split_strategy = train_val_test_split_strategy
         self.collate_fn = (
             collate_fn if collate_fn is not None else self.__class__.collate_fn
         )  # noqa: E501
@@ -683,8 +685,6 @@ def setup(self, stage: str | None = None) -> None:
             If None, the data module will be setup for training.
         """
 
-        # TODO: Add support for temporal/random/group splits.
-        # Currently, it only supports random splits.
         # Handle the case where the dataset is empty.
 
         total_series = len(self.time_series_dataset)
@@ -695,9 +695,17 @@ def setup(self, stage: str | None = None) -> None:
                 "Please provide a non-empty dataset."
             )
 
+        if self.train_val_test_split_strategy == "random":
+            self._indices = torch.randperm(total_series)
+        elif self.train_val_test_split_strategy == "sequential":
+            self._indices = torch.arange(total_series)
+        else:
+            raise ValueError(
+                f"Unknown split strategy: {self.train_val_test_split_strategy}"
+            )
+
         # this is a very rudimentary way to handle the splits when
         # the dataset is of size equal to 1 or 2.
-        self._indices = torch.randperm(total_series)
         if total_series == 1:
             self._train_indices = self._indices
             self._val_indices = self._indices

pytorch_forecasting/data/tests/test_tslib_data_module.py

@@ -409,6 +409,40 @@ def test_dataloader_pipeline(tslib_data_module):
     assert y_batch.shape[1] == tslib_data_module.prediction_length
 
 
+def test_sequential_split_strategy(sample_timeseries_data):
+    """Test the TslibDataModule with sequential split strategy."""
+
+    dm_seq = TslibDataModule(
+        time_series_dataset=sample_timeseries_data,
+        context_length=8,
+        prediction_length=4,
+        batch_size=2,
+        train_val_test_split=(0.6, 0.2, 0.2),
+        train_val_test_split_strategy="sequential",
+    )
+
+    dm_seq.setup(stage="fit")
+
+    total_series = len(sample_timeseries_data)
+    expected_train = int(total_series * 0.6)
+    expected_val = int(total_series * 0.2)
+
+    # Check if indices are purely sequential
+    import torch
+
+    assert torch.all(
+        dm_seq._train_indices == torch.arange(0, expected_train)
+    ), "Train indices should be sequential."
+    assert torch.all(
+        dm_seq._val_indices
+        == torch.arange(expected_train, expected_train + expected_val)
+    ), "Val indices should be sequential."
+    assert torch.all(
+        dm_seq._test_indices
+        == torch.arange(expected_train + expected_val, total_series)
+    ), "Test indices should be sequential."
+
+
 def test_different_split_ratios(sample_timeseries_data):
     """Test the TslibDataModule with different train/val/test split ratios."""
 
@@ -528,5 +562,5 @@ def test_multivariate_target():
     x, y = dm.train_dataset[0]
 
     assert (
-        y.shape[-1] == 2
-    ), "Target should have two dimensions for n_features for multivariate target."
+        isinstance(y, list) and len(y) == 2
+    ), "Target should be a list of two tensors for multivariate target."