Skip to content

[ENH] Add more train_test_split strategies #2087

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
Muhammad-Rebaal wants to merge 39 commits into
base: main
Choose a base branch
from
+1,043 −56
Open

[ENH] Add more train_test_split strategies #2087

Show file tree
Hide file tree
Changes from 35 commits
Commits
Show all changes
39 commits
Select commit Hold shift + click to select a range
12c5e57
feat:Add more train_test_split strategies
Muhammad-Rebaal Feb 23, 2026
d4e48b3
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 2, 2026
455472c
Merge branch 'main' of https://github.com/Muhammad-Rebaal/pytorch-for…
Muhammad-Rebaal Mar 2, 2026
1758ec3
Merge branch 'train_test_split' of https://github.com/Muhammad-Rebaal…
Muhammad-Rebaal Mar 2, 2026
29734d7
fix: Failed PyTest
Muhammad-Rebaal Mar 2, 2026
263ae38
fix: ruff formatting
Muhammad-Rebaal Mar 2, 2026
d089be4
fix: failure of test_univariate_forecast
Muhammad-Rebaal Mar 2, 2026
90d345a
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 2, 2026
af59b4c
fix: Module not found error
Muhammad-Rebaal Mar 3, 2026
35bf4c3
Merge branch 'train_test_split' of https://github.com/Muhammad-Rebaal…
Muhammad-Rebaal Mar 3, 2026
aae1f96
fix: formatting issue
Muhammad-Rebaal Mar 3, 2026
9248866
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 9, 2026
ab5b4d0
fix:Reverse the documentation related issues
Muhammad-Rebaal Mar 9, 2026
b623d59
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 10, 2026
573d1fb
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 16, 2026
2d0456c
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 21, 2026
7e5c849
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 25, 2026
7e3ae57
fix: Remove out of scope logic of categorical_encoder
Muhammad-Rebaal Mar 30, 2026
67ba8ee
Merge branch 'main' into train_test_split
Muhammad-Rebaal Mar 30, 2026
2295183
Merge branch 'main' into train_test_split
Muhammad-Rebaal Apr 3, 2026
d482898
Merge branch 'main' of https://github.com/Muhammad-Rebaal/pytorch-for…
Muhammad-Rebaal Jun 25, 2026
3be0d0a
Merge branch 'train_test_split' of https://github.com/Muhammad-Rebaal…
Muhammad-Rebaal Jun 25, 2026
3694122
Merge branch 'main' into train_test_split
Muhammad-Rebaal Jun 30, 2026
b6ef270
feat: Implementation added to avoid data leakage for the temporal split
Muhammad-Rebaal Jul 13, 2026
8e22de7
fix: Added the accidently removed docstring
Muhammad-Rebaal Jul 13, 2026
ac4a762
Merge branch 'main' of https://github.com/Muhammad-Rebaal/pytorch-for…
Muhammad-Rebaal Jul 16, 2026
54ff0ab
Fix: Code Quality
Muhammad-Rebaal Jul 16, 2026
5017d01
Merge branch 'main' into train_test_split
Muhammad-Rebaal Jul 29, 2026
2ee716e
Merge branch 'main' into train_test_split
Muhammad-Rebaal Aug 2, 2026
a4a47b4
Merge branch 'main' of https://github.com/Muhammad-Rebaal/pytorch-for…
Muhammad-Rebaal Aug 3, 2026
84a5bca
Merge branch 'main' into train_test_split
Muhammad-Rebaal Aug 3, 2026
d71376d
Merge branch 'train_test_split' of https://github.com/Muhammad-Rebaal…
Muhammad-Rebaal Aug 3, 2026
a6f3a24
[ENH] Dynamic Cutoff is implemented
Muhammad-Rebaal Aug 3, 2026
54b5de8
train test split can handle timestamps & raise warning when the user …
Muhammad-Rebaal Aug 5, 2026
a87d7e4
Merge branch 'main' into train_test_split
Muhammad-Rebaal Aug 7, 2026
5144273
[ENH] Added Implementation for group-time-split
Muhammad-Rebaal Aug 7, 2026
f68b441
Merge branch 'train_test_split' of https://github.com/Muhammad-Rebaal…
Muhammad-Rebaal Aug 7, 2026
15543a9
fix code quality
Muhammad-Rebaal Aug 7, 2026
ab5d862
Merge branch 'main' into train_test_split
Muhammad-Rebaal Aug 11, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
149 changes: 126 additions & 23 deletions pytorch_forecasting/data/data_module/_encoder_decoder_data_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -135,6 +135,8 @@ def __init__(
batch_size: int = 32,
num_workers: int = 0,
train_val_test_split: tuple = (0.7, 0.15, 0.15),
split_strategy: str = "random",
temporal_cutoffs: dict[str, float] | None = None,
):
self.time_series_dataset = time_series_dataset
self.max_encoder_length = max_encoder_length
Expand All @@ -153,6 +155,8 @@ def __init__(
self.batch_size = batch_size
self.num_workers = num_workers
self.train_val_test_split = train_val_test_split
self.split_strategy = split_strategy
self.temporal_cutoffs = temporal_cutoffs

warn(
"EncoderDecoderTimeSeriesDataModule is part of an experimental "
Expand Down Expand Up @@ -959,21 +963,40 @@ def _resolve_target_normalizer(self, train_indices: torch.Tensor) -> None:
self._target_normalizer = ScalerAdapter(normalizer)

def _ensure_split(self):

@phoeenniixx phoeenniixx Aug 12, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

please add all the split statergies here

"""Compute train/val/test indices once and cache them."""
if hasattr(self, "_split_indices"):
"""
Compute train/val/test indices once and cache them,
respecting split_strategy.
"""
if hasattr(self, "_split_done"):
return

from pytorch_forecasting.data.splitters import (
random_series_split,
stratified_series_split,
)

total_series = len(self.time_series_dataset)
self._split_indices = torch.randperm(total_series)

self._train_size = int(self.train_val_test_split[0] * total_series)
self._val_size = int(self.train_val_test_split[1] * total_series)
if self.split_strategy in ["random", "group"]:
self._train_indices, self._val_indices, self._test_indices = (
random_series_split(total_series, self.train_val_test_split)
)
elif self.split_strategy == "stratified":
self._train_indices, self._val_indices, self._test_indices = (
stratified_series_split(
self.time_series_dataset,
target_idx=0,
train_val_test_split=self.train_val_test_split,
)
)
elif self.split_strategy == "temporal":
self._train_indices = torch.arange(total_series)
self._val_indices = torch.arange(total_series)
self._test_indices = torch.arange(total_series)
else:
raise ValueError(f"Unknown split_strategy: {self.split_strategy}")

self._train_indices = self._split_indices[: self._train_size]
self._val_indices = self._split_indices[
self._train_size : self._train_size + self._val_size
]
self._test_indices = self._split_indices[self._train_size + self._val_size :]
self._split_done = True

def _make_dataset(self, indices: torch.Tensor):
"""Preprocess a set of series indices into a windowed Dataset.
Expand Down Expand Up @@ -1017,23 +1040,103 @@ def setup(self, stage: str | None = None):
if not self._feature_scalers_fitted:
self._fit_scalers(self._train_indices)
if not hasattr(self, "train_dataset") or not hasattr(self, "val_dataset"):
self._train_preprocessed, self.train_windows, self.train_dataset = (
self._make_dataset(self._train_indices)
)
self._val_preprocessed, self.val_windows, self.val_dataset = (
self._make_dataset(self._val_indices)
)
if self.split_strategy == "temporal":
# Build all windows, then split them by timestamp
all_windows = self._create_windows(self._train_indices)
series_timestamps = {}
for idx in self._train_indices:
series_idx = (
idx.item() if isinstance(idx, torch.Tensor) else idx
)
sample = self.time_series_dataset[series_idx]
series_timestamps[series_idx] = sample["t"]

from pytorch_forecasting.data.splitters import temporal_window_split

t_win, v_win, te_win = temporal_window_split(
all_windows,
self.train_val_test_split,
series_timestamps,
self.temporal_cutoffs,
)
self.train_windows, self.val_windows, self.test_windows = (
t_win,
v_win,
te_win,
)

# Preprocess ALL series (train, val, test share the same series)
all_indices = self._train_indices
preprocessed = {
idx.item(): self._preprocess_data(idx.item())
for idx in all_indices
}
self._train_preprocessed = preprocessed
self._val_preprocessed = preprocessed

self.train_dataset = self._ProcessedEncoderDecoderDataset(
self,
self.train_windows,
preprocessed,
self.add_relative_time_idx,
)
self.val_dataset = self._ProcessedEncoderDecoderDataset(
self, self.val_windows, preprocessed, self.add_relative_time_idx
)
else:
self._train_preprocessed, self.train_windows, self.train_dataset = (
self._make_dataset(self._train_indices)
)
self._val_preprocessed, self.val_windows, self.val_dataset = (
self._make_dataset(self._val_indices)
)

elif stage == "test":
if not hasattr(self, "test_dataset"):
self._test_preprocessed, self.test_windows, self.test_dataset = (
self._make_dataset(self._test_indices)
)
if self.split_strategy == "temporal":
if not hasattr(self, "test_windows"):
# Recompute temporal test windows if fit wasn't called first
total_series = len(self.time_series_dataset)
all_windows = self._create_windows(torch.arange(total_series))
series_timestamps = {}
for idx in range(total_series):
sample = self.time_series_dataset[idx]
series_timestamps[idx] = sample["t"]

from pytorch_forecasting.data.splitters import (
temporal_window_split,
)

_, _, self.test_windows = temporal_window_split(
all_windows,
self.train_val_test_split,
series_timestamps,
self.temporal_cutoffs,
)

preprocessed = {
idx: self._preprocess_data(idx)
for idx in {w[0] for w in self.test_windows}
}
self.test_dataset = self._ProcessedEncoderDecoderDataset(
self,
self.test_windows,
preprocessed,
self.add_relative_time_idx,
)
else:
self._test_preprocessed, self.test_windows, self.test_dataset = (
self._make_dataset(self._test_indices)
)

elif stage == "predict":
predict_indices = torch.arange(len(self.time_series_dataset))
self._predict_preprocessed, self.predict_windows, self.predict_dataset = (
self._make_dataset(predict_indices)
)
if not hasattr(self, "predict_dataset"):
predict_indices = torch.arange(len(self.time_series_dataset))
(
self._predict_preprocessed,
self.predict_windows,
self.predict_dataset,
) = self._make_dataset(predict_indices)

def train_dataloader(self):
return DataLoader(
Expand Down
101 changes: 73 additions & 28 deletions pytorch_forecasting/data/data_module/_tslib_data_module.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -309,6 +309,8 @@ def __init__(
batch_size: int = 32,
num_workers: int = 0,
train_val_test_split: tuple[float, float, float] = (0.7, 0.15, 0.15),
split_strategy: str = "random",
temporal_cutoffs: dict[str, float] | None = None,
collate_fn: Callable | None = None,
**kwargs,
) -> None:
Expand All @@ -323,6 +325,8 @@ def __init__(
self.batch_size = batch_size
self.num_workers = num_workers
self.train_val_test_split = train_val_test_split
self.split_strategy = split_strategy
self.temporal_cutoffs = temporal_cutoffs
self.collate_fn = (
collate_fn if collate_fn is not None else self.__class__.collate_fn
) # noqa: E501
Expand Down Expand Up @@ -695,51 +699,92 @@ def setup(self, stage: str | None = None) -> None:
"Please provide a non-empty dataset."
)

# 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
self._test_indices = self._indices
elif total_series == 2:
self._train_indices = self._indices[0:1]
self._val_indices = self._indices[1:2]
self._test_indices = self._indices[1:2]
else:
self._train_size = int(self.train_val_test_split[0] * total_series)
self._val_size = int(self.train_val_test_split[1] * total_series)

self._train_indices = self._indices[: self._train_size]
self._val_indices = self._indices[
self._train_size : self._train_size + self._val_size
]
from pytorch_forecasting.data.splitters import (
random_series_split,
stratified_series_split,
temporal_window_split,
)

self._test_indices = self._indices[
self._train_size + self._val_size : total_series
]
if self.split_strategy in ["random", "group"]:
self._train_indices, self._val_indices, self._test_indices = (
random_series_split(total_series, self.train_val_test_split)
)
elif self.split_strategy == "stratified":
self._train_indices, self._val_indices, self._test_indices = (
stratified_series_split(
self.time_series_dataset,
target_idx=0,
train_val_test_split=self.train_val_test_split,
)
)
elif self.split_strategy == "temporal":
self._train_indices = torch.arange(total_series)
self._val_indices = torch.arange(total_series)
self._test_indices = torch.arange(total_series)
else:
raise ValueError(f"Unknown split_strategy: {self.split_strategy}")

if stage == "fit" or stage is None:
if not hasattr(self, "_train_dataset") or not hasattr(self, "_val_dataset"):
self._train_windows = self._create_windows(self._train_indices)
self._val_windows = self._create_windows(self._val_indices)
if not hasattr(self, "_train_windows") or not hasattr(self, "_val_windows"):
if self.split_strategy == "temporal":

@agobbifbk agobbifbk Jun 22, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Can you provide a visual example (or a snippet of code) with let say:
1 - 3 groups
2- group 1 has 2 series with variable length
3- group 2 has 1 series with a given length
4- group 3 has 4 series with variable length
So we can visually see how the split will apply for each group? THX

@Muhammad-Rebaal Muhammad-Rebaal Jun 30, 2026

Copy link
Copy Markdown
Member Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Diagram :

Untitled Diagram-2026-06-30T13-24-27

Code :

def temporal_window_split(
    windows: list[tuple[int, int, int, int]],
    train_val_test_split: tuple[float, float, float],
) -> tuple[
    list[tuple[int, int, int, int]],
    list[tuple[int, int, int, int]],
    list[tuple[int, int, int, int]],
]:
 
    # Group windows by series_idx
    series_windows = {}
    for w in windows:
        s_idx = w[0]
        if s_idx not in series_windows:
            series_windows[s_idx] = []
        series_windows[s_idx].append(w)

    train_windows, val_windows, test_windows = [], [], []

    for s_idx, sw in series_windows.items():
        # Ensure windows are sorted by time (start_idx: w[1])
        sw.sort(key=lambda x: x[1])
        total_w = len(sw)

        train_end = int(np.round(train_val_test_split[0] * total_w))
        if train_end == 0 and train_val_test_split[0] > 0 and total_w > 0:
            train_end = 1

        val_end = train_end + int(np.round(train_val_test_split[1] * total_w))
        if val_end > total_w:
            val_end = total_w

        train_windows.extend(sw[:train_end])
        val_windows.extend(sw[train_end:val_end])
        test_windows.extend(sw[val_end:])

    return train_windows, val_windows, test_windows

For the Temporal split, the filtering happens at the Window level. The DataModule assigns ALL series indices to Train, Val, and Test. It generates ALL possible windows for the entire dataset, and then passes that massive list of windows to temporal_window_split(). This function groups the windows by their series_idx, sorts them chronologically, and slices the window arrays based on our split ratios (e.g., the first 70% of windows for Series A go to Train).

@agobbifbk agobbifbk Jul 1, 2026

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Nice for the representation, it is much easier now to understand :-)
I have a point to raise: if we call it temporal_split and the two series of group 1 are consecutive (for example yearly csv data) this procedure will break the data leakage paradigm, see for example https://arxiv.org/html/2412.11376v1

To avoid information leakage during the evaluation phase in Section 4.3, each dataset is chronologically split into training, validation, and test sets with a ratio of 6:2:2

We can still have this splitting methodology, but I won't call it temporal_split because it is not a real temporal split. You should look at the timestamp and decide which samples goes in the correct place.
Any thoughts @phoeenniixx ?

all_windows = self._create_windows(self._train_indices)

series_timestamps = {}
for idx in self._train_indices:
series_idx = (
idx.item() if isinstance(idx, torch.Tensor) else idx
)
sample = self.time_series_dataset[series_idx]
series_timestamps[series_idx] = sample["t"]

t_win, v_win, te_win = temporal_window_split(
all_windows,
self.train_val_test_split,
series_timestamps,
self.temporal_cutoffs,
)

self._train_windows, self._val_windows, self._test_windows = (
t_win,
v_win,
te_win,
)
else:
self._train_windows = self._create_windows(self._train_indices)
self._val_windows = self._create_windows(self._val_indices)

self.train_dataset = _TslibDataset(
dataset=self.time_series_dataset,
data_module=self,
windows=self._train_windows,
add_relative_time_idx=self.add_relative_time_idx,
)

self.val_dataset = _TslibDataset(
dataset=self.time_series_dataset,
data_module=self,
windows=self._val_windows,
add_relative_time_idx=self.add_relative_time_idx,
)

elif stage == "test":
if not hasattr(self, "_test_dataset"):
self._test_windows = self._create_windows(self._test_indices)
if not hasattr(self, "_test_windows") or self.test_dataset is None:
if self.split_strategy == "temporal":
all_windows = self._create_windows(torch.arange(total_series))

series_timestamps = {}
for idx in range(total_series):
sample = self.time_series_dataset[idx]
series_timestamps[idx] = sample["t"]

_, _, self._test_windows = temporal_window_split(
all_windows,
self.train_val_test_split,
series_timestamps,
self.temporal_cutoffs,
)

else:
self._test_windows = self._create_windows(self._test_indices)

self.test_dataset = _TslibDataset(
dataset=self.time_series_dataset,
Expand Down
Loading
Loading