Repository: koaning/scikit-lego (python)

Issue:
[FEATURE] `TimeGapSplit` with stride 
In the current implementation of `TimeGapSplit` the argument `validation_duration` is used as stride as well.

This leads to the impossibility of having overlapping validation periods. What about the possibility of adding a stride parameter to indicate the time shift for the training period?

The concrete use case is as follows: I would like to validate a model that retrains every week but forecast for multiple weeks ahead.

Happy to give it a try in the next following days if there is interest in adding this feature.

Pull request:
feat: Add `stride_duration` in `TimeGapSplit`
# Description

Introduces stride in `TimeGapSplit`

Closes #569 

## Type of change

- [ ] Bug fix (non-breaking change which fixes an issue)
- [X] New feature (non-breaking change which adds functionality)
- [ ] Breaking change (fix or feature that would cause existing functionality to not work as expected)


## Checklist:

- [X] My code follows the style guidelines (ruff)
- [X] I have commented my code, particularly in hard-to-understand areas
- [X] I have made corresponding changes to the documentation (also to the readme.md)
- [X] I have added tests that prove my fix is effective or that my feature works
- [ ] I have added tests to check whether the new feature adheres to the sklearn convention
- [X] New and existing unit tests pass locally with my changes


Patch:
diff --git a/sklego/model_selection.py b/sklego/model_selection.py
index 8e5a053..a2dad17 100644
--- a/sklego/model_selection.py
+++ b/sklego/model_selection.py
@@ -30,11 +30,11 @@ class TimeGapSplit:
     Each validation fold doesn't overlap. The entire `window` moves by 1 `valid_duration` until there is not enough
     data.
 
-    If this would lead to more splits then specified with `n_splits`, the `window` moves by `valid_duration` times the
+    If this would lead to more splits than specified with `n_splits`, the `window` moves by `valid_duration` times the
     fraction of possible splits and requested splits:
 
-    - `n_possible_splits = (total_length - train_duration-gap_duration) // valid_duration`
-    - `time_shift = valid_duration * n_possible_splits / n_slits`
+    - `n_possible_splits = (total_length - train_duration - gap_duration) // valid_duration`
+    - `time_shift = valid_duration * n_possible_splits / n_splits`
 
     so the CV spans the whole dataset.
 
@@ -49,19 +49,17 @@ class TimeGapSplit:
     date_serie : Series
         Series with the date, that should have all the indices of X used in the split() method.
         If the Series is not pandas-like (for example, if it's a Polars Series, which does not have
-        an index) then it must the same same length as the `X` and `y` objects passed to `split`.
+        an index) then it must the same length as the `X` and `y` objects passed to `split`.
     valid_duration : datetime.timedelta
         Retraining period.
+    stride_duration : datetime.timedelta | None
+        The time shift for the training period. If `None`, then it fallbacks to `valid_duration` value.
     train_duration : datetime.timedelta | None, default=None
         Historical training data.
     gap_duration : datetime.timedelta, default=timedelta(0)
-        Forward looking window of the target. The period of the forward looking window necessary to create your target
-        variable.
-
-        This period is dropped at the end of your training folds due to lack of recent data.
-
-        In production you would have not been able to create the target for that period, and you would have drop it from
-        the training data.
+        Forward-looking window of the target. The period of the forward-looking window necessary to create your target
+        variable. This period is dropped at the end of your training folds due to lack of recent data. In production you
+        would have not been able to create the target for that period, and you would have to drop it from the training data.
     n_splits : int | None, default=None
         Number of splits.
     window : Literal["rolling", "expanding"], default="rolling"
@@ -90,21 +88,34 @@ class TimeGapSplit:
         self,
         date_serie,
         valid_duration,
+        stride_duration=None,
         train_duration=None,
         gap_duration=timedelta(0),
         n_splits=None,
         window="rolling",
     ):
+        # If stride length is not defined, set it equal to the length validation set
+        if stride_duration is None:
+            stride_duration = valid_duration
+
+        if stride_duration <= timedelta(milliseconds=0):
+            msg = f"`stride_duration` should be a positive timedelta, found {stride_duration}"
+            raise ValueError(msg)
+
         if (train_duration is None) and (n_splits is None):
             raise ValueError("Either train_duration or n_splits have to be defined")
 
         if (train_duration is not None) and (train_duration <= gap_duration):
             raise ValueError("gap_duration is longer than train_duration, it should be shorter.")
 
+        if (train_duration is not None) and (train_duration <= stride_duration):
+            raise ValueError("stride_duration is longer than train_duration, it should be shorter.")
+
         self.date_serie = nw.from_native(date_serie, series_only=True).alias("__date__")
         self.train_duration = train_duration
         self.valid_duration = valid_duration
         self.gap_duration = gap_duration
+        self.stride_duration = stride_duration
         self.n_splits = n_splits
         self.window = window
 
@@ -149,7 +160,7 @@ class TimeGapSplit:
 
         if len(X) != len(X_index_df):
             raise AssertionError(
-                "X and X_index_df are not the same length, " "there must be some index missing in 'self.date_serie'"
+                "X and X_index_df are not the same length, there must be some index missing in 'self.date_serie'"
             )
 
         date_min = X_index_df["__date__"].min()
@@ -157,30 +168,35 @@ class TimeGapSplit:
         date_length = X_index_df["__date__"].max() - X_index_df["__date__"].min()
 
         if (self.train_duration is None) and (self.n_splits is not None):
-            self.train_duration = date_length - (self.gap_duration + self.valid_duration * self.n_splits)
+            self.train_duration = date_length - (
+                self.gap_duration + self.valid_duration + (self.n_splits - 1) * self.stride_duration
+            )
 
         if (self.train_duration is not None) and (self.train_duration <= self.gap_duration):
             raise ValueError("gap_duration is longer than train_duration, it should be shorter.")
 
-        n_split_max = (date_length - self.train_duration - self.gap_duration) / self.valid_duration
+        n_split_max = (
+            1 + (date_length - self.train_duration - self.gap_duration - self.valid_duration) / self.stride_duration
+        )
+
         if self.n_splits:
             if n_split_max < self.n_splits:
                 raise ValueError(
                     (
                         "Number of folds requested = {1} are greater"
-                        " than maximum  ={0} possible without"
-                        " overlapping validation sets."
+                        " than maximum  ={0} possible"
+                        " based on the given values."
                     ).format(n_split_max, self.n_splits)
                 )
 
         current_date = date_min
         start_date = date_min
-        # if the n_splits is smaller than what would usually be done for train val and gap duration,
-        # the next fold is slightly further in time than just valid_duration
+        # if the n_splits is smaller than what would usually be done for train, validation, stride and gap duration,
+        # the next fold is slightly further in time than just stride_duration
         if self.n_splits is not None:
-            time_shift = self.valid_duration * n_split_max / self.n_splits
+            time_shift = self.stride_duration * n_split_max / self.n_splits
         else:
-            time_shift = self.valid_duration
+            time_shift = self.stride_duration
         while True:
             if current_date + self.train_duration + time_shift + self.gap_duration > date_max:
                 break
@@ -384,8 +400,7 @@ class GroupTimeSeriesSplit(_BaseKFold):
     def __init__(self, n_splits):
         if not isinstance(n_splits, numbers.Integral):
             raise ValueError(
-                "The number of folds must be of Integral type. "
-                "%s of type %s was passed." % (n_splits, type(n_splits))
+                "The number of folds must be of Integral type. %s of type %s was passed." % (n_splits, type(n_splits))
             )
         n_splits = int(n_splits)
 
@@ -442,7 +457,7 @@ class GroupTimeSeriesSplit(_BaseKFold):
         n_groups = np.unique(groups).shape[0]
         if self.n_splits >= n_groups:
             raise ValueError(
-                ("n_splits({0}) must be less than the amount" " of unique groups({1}).").format(self.n_splits, n_groups)
+                ("n_splits({0}) must be less than the amount of unique groups({1}).").format(self.n_splits, n_groups)
             )
         return list(self._iter_test_indices(X, y, groups))
 
@@ -481,9 +496,7 @@ class GroupTimeSeriesSplit(_BaseKFold):
         """
         unique_groups = len(set(groups))
         warning = (
-            "Finding the optimal split points"
-            " with {0} unique groups and n_splits at {1}"
-            " can take several minutes."
+            "Finding the optimal split points with {0} unique groups and n_splits at {1} can take several minutes."
         ).format(unique_groups, self.n_splits)
         if self.n_splits == 4 and unique_groups > 250:
             warn(


Test patch:
diff --git a/tests/test_model_selection/test_timegapsplit.py b/tests/test_model_selection/test_timegapsplit.py
index 9be0a10..1184936 100644
--- a/tests/test_model_selection/test_timegapsplit.py
+++ b/tests/test_model_selection/test_timegapsplit.py
@@ -1,5 +1,4 @@
-import datetime
-from datetime import timedelta
+from datetime import datetime, timedelta
 
 import numpy as np
 import pandas as pd
@@ -24,36 +23,45 @@ X_train = train[["A", "B", "C"]]
 y_train = train["y"]
 
 
-def test_timegapsplit():
+@pytest.mark.parametrize(
+    ("stride_duration", "expected"),
+    [
+        (
+            3,
+            [
+                (np.array([0, 1, 2, 3, 4]), np.array([5, 6, 7])),
+                (np.array([3, 4, 5, 6, 7]), np.array([8, 9, 10])),
+                (np.array([6, 7, 8, 9, 10]), np.array([11, 12, 13])),
+                (np.array([9, 10, 11, 12, 13]), np.array([14, 15, 16])),
+                (np.array([12, 13, 14, 15, 16]), np.array([17, 18, 19])),
+                (np.array([15, 16, 17, 18, 19]), np.array([20, 21, 22])),
+            ],
+        ),
+        (
+            2,
+            [
+                (np.array([0, 1, 2, 3, 4]), np.array([5, 6, 7])),
+                (np.array([2, 3, 4, 5, 6]), np.array([7, 8, 9])),
+                (np.array([4, 5, 6, 7, 8]), np.array([9, 10, 11])),
+                (np.array([6, 7, 8, 9, 10]), np.array([11, 12, 13])),
+                (np.array([8, 9, 10, 11, 12]), np.array([13, 14, 15])),
+                (np.array([10, 11, 12, 13, 14]), np.array([15, 16, 17])),
+                (np.array([12, 13, 14, 15, 16]), np.array([17, 18, 19])),
+                (np.array([14, 15, 16, 17, 18]), np.array([19, 20, 21])),
+                (np.array([16, 17, 18, 19, 20]), np.array([21, 22, 23])),
+            ],
+        ),
+    ],
+)
+def test_timegapsplit(stride_duration, expected):
     cv = TimeGapSplit(
         date_serie=df["date"],
         train_duration=timedelta(days=5),
         valid_duration=timedelta(days=3),
+        stride_duration=timedelta(days=stride_duration),
         gap_duration=timedelta(days=0),
     )
 
-    for i, indices in enumerate(cv.split(X_train, y_train)):
-        train_mindate = df.loc[X_train.iloc[indices[0]].index]["date"].min()
-        train_maxdate = df.loc[X_train.iloc[indices[0]].index]["date"].max()
-        valid_mindate = df.loc[X_train.iloc[indices[1]].index]["date"].min()
-        valid_maxdate = df.loc[X_train.iloc[indices[1]].index]["date"].max()
-
-        assert train_mindate <= train_maxdate <= valid_mindate <= valid_maxdate
-
-    # regression testing, check if output changes of the last fold
-    assert train_mindate == datetime.datetime.strptime("2018-01-16", "%Y-%m-%d")
-    assert train_maxdate == datetime.datetime.strptime("2018-01-20", "%Y-%m-%d")
-    assert valid_mindate == datetime.datetime.strptime("2018-01-21", "%Y-%m-%d")
-    assert valid_maxdate == datetime.datetime.strptime("2018-01-23", "%Y-%m-%d")
-
-    expected = [
-        (np.array([0, 1, 2, 3, 4]), np.array([5, 6, 7])),
-        (np.array([3, 4, 5, 6, 7]), np.array([8, 9, 10])),
-        (np.array([6, 7, 8, 9, 10]), np.array([11, 12, 13])),
-        (np.array([9, 10, 11, 12, 13]), np.array([14, 15, 16])),
-        (np.array([12, 13, 14, 15, 16]), np.array([17, 18, 19])),
-        (np.array([15, 16, 17, 18, 19]), np.array([20, 21, 22])),
-    ]
     for result_indices, expected_indices in zip(list(cv.split(X_train, y_train)), expected):
         np.testing.assert_array_equal(result_indices[0], expected_indices[0])
         np.testing.assert_array_equal(result_indices[1], expected_indices[1])
@@ -66,16 +74,9 @@ def test_timegapsplit():
         date_serie=pl.from_pandas(date_serie),
         train_duration=timedelta(days=5),
         valid_duration=timedelta(days=3),
+        stride_duration=timedelta(days=stride_duration),
         gap_duration=timedelta(days=0),
     )
-    expected = [
-        (np.array([0, 1, 2, 3, 4]), np.array([5, 6, 7])),
-        (np.array([3, 4, 5, 6, 7]), np.array([8, 9, 10])),
-        (np.array([6, 7, 8, 9, 10]), np.array([11, 12, 13])),
-        (np.array([9, 10, 11, 12, 13]), np.array([14, 15, 16])),
-        (np.array([12, 13, 14, 15, 16]), np.array([17, 18, 19])),
-        (np.array([15, 16, 17, 18, 19]), np.array([20, 21, 22])),
-    ]
     for result_indices, expected_indices in zip(
         list(cv.split(pl.from_pandas(X_train), pl.from_pandas(y_train))), expected
     ):
@@ -162,6 +163,50 @@ def test_timegapsplit_with_a_gap():
         assert valid_mindate - train_maxdate >= gap_duration
 
 
+def test_timegapsplit_stride_is_zero():
+    with pytest.raises(ValueError):
+        _ = TimeGapSplit(
+            date_serie=df["date"],
+            train_duration=timedelta(days=5),
+            stride_duration=timedelta(days=0),
+            valid_duration=timedelta(days=3),
+            gap_duration=timedelta(days=5),
+            n_splits=None,
+        )
+
+
+def test_timegapsplit_stride_longer_than_train():
+    with pytest.raises(ValueError):
+        _ = TimeGapSplit(
+            date_serie=df["date"],
+            train_duration=timedelta(days=5),
+            stride_duration=timedelta(days=6),
+            valid_duration=timedelta(days=3),
+            gap_duration=timedelta(days=5),
+            n_splits=None,
+        )
+
+
+def test_timegapsplit_with_stride():
+    stride_duration = timedelta(days=2)
+    cv_gap = TimeGapSplit(
+        date_serie=df["date"],
+        train_duration=timedelta(days=5),
+        stride_duration=stride_duration,
+        valid_duration=timedelta(days=3),
+        gap_duration=timedelta(days=1),
+    )
+
+    for i, indices in enumerate(cv_gap.split(X_train, y_train)):
+        train_mindate = df.loc[X_train.iloc[indices[0]].index]["date"].min()
+        train_maxdate = df.loc[X_train.iloc[indices[0]].index]["date"].max()
+        valid_mindate = df.loc[X_train.iloc[indices[1]].index]["date"].min()
+        valid_maxdate = df.loc[X_train.iloc[indices[1]].index]["date"].max()
+
+        assert train_mindate <= train_maxdate <= valid_mindate <= valid_maxdate
+        assert valid_mindate - train_maxdate >= stride_duration
+
+
 def test_timegapsplit_with_gridsearch():
     cv = TimeGapSplit(
         date_serie=df["date"],
@@ -194,46 +239,46 @@ def test_timegapsplit_summary():
 
     expected_data = {
         "Start date": [
-            datetime.datetime(2018, 1, 1, 0, 0),
-            datetime.datetime(2018, 1, 6, 0, 0),
-            datetime.datetime(2018, 1, 4, 0, 0),
-            datetime.datetime(2018, 1, 9, 0, 0),
-            datetime.datetime(2018, 1, 7, 0, 0),
-            datetime.datetime(2018, 1, 12, 0, 0),
-            datetime.datetime(2018, 1, 10, 0, 0),
-            datetime.datetime(2018, 1, 15, 0, 0),
-            datetime.datetime(2018, 1, 13, 0, 0),
-            datetime.datetime(2018, 1, 18, 0, 0),
-            datetime.datetime(2018, 1, 16, 0, 0),
-            datetime.datetime(2018, 1, 21, 0, 0),
+            datetime(2018, 1, 1, 0, 0),
+            datetime(2018, 1, 6, 0, 0),
+            datetime(2018, 1, 4, 0, 0),
+            datetime(2018, 1, 9, 0, 0),
+            datetime(2018, 1, 7, 0, 0),
+            datetime(2018, 1, 12, 0, 0),
+            datetime(2018, 1, 10, 0, 0),
+            datetime(2018, 1, 15, 0, 0),
+            datetime(2018, 1, 13, 0, 0),
+            datetime(2018, 1, 18, 0, 0),
+            datetime(2018, 1, 16, 0, 0),
+            datetime(2018, 1, 21, 0, 0),
         ],
         "End date": [
-            datetime.datetime(2018, 1, 5, 0, 0),
-            datetime.datetime(2018, 1, 8, 0, 0),
-            datetime.datetime(2018, 1, 8, 0, 0),
-            datetime.datetime(2018, 1, 11, 0, 0),
-            datetime.datetime(2018, 1, 11, 0, 0),
-            datetime.datetime(2018, 1, 14, 0, 0),
-            datetime.datetime(2018, 1, 14, 0, 0),
-            datetime.datetime(2018, 1, 17, 0, 0),
-            datetime.datetime(2018, 1, 17, 0, 0),
-            datetime.datetime(2018, 1, 20, 0, 0),
-            datetime.datetime(2018, 1, 20, 0, 0),
-            datetime.datetime(2018, 1, 23, 0, 0),
+            datetime(2018, 1, 5, 0, 0),
+            datetime(2018, 1, 8, 0, 0),
+            datetime(2018, 1, 8, 0, 0),
+            datetime(2018, 1, 11, 0, 0),
+            datetime(2018, 1, 11, 0, 0),
+            datetime(2018, 1, 14, 0, 0),
+            datetime(2018, 1, 14, 0, 0),
+            datetime(2018, 1, 17, 0, 0),
+            datetime(2018, 1, 17, 0, 0),
+            datetime(2018, 1, 20, 0, 0),
+            datetime(2018, 1, 20, 0, 0),
+            datetime(2018, 1, 23, 0, 0),
         ],
         "Period": [
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
-            datetime.timedelta(days=4),
-            datetime.timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
+            timedelta(days=4),
+            timedelta(days=2),
         ],
         "Unique days": [5, 3, 5, 3, 5, 3, 5, 3, 5, 3, 5, 3],
         "nbr samples": [5, 3, 5, 3, 5, 3, 5, 3, 5, 3, 5, 3],