Merge pull request #103 from m9o8/ENH-spatio-temporal-splits

Feature: Flexible Spatio-Temporal Splits and Dynamic Multi-Group Visualization

Author: 4Freye

README.md

@@ -48,6 +48,30 @@ for train_idx, test_idx in splits:
     print("Test:"); display(panel_data.loc[test_idx])
 ```
 
+### Spatio-Temporal Cross-Validation
+
+panelsplit can also handle combined spatio-temporal holdouts by factoring in entity hierarchies (e.g., states or cities) to prevent cluster-level leakage. You can simultaneously validate on unobserved time periods *and* structurally unobserved groups:
+
+```python
+from sklearn.model_selection import StratifiedGroupKFold
+
+# Create spatial splits that evaluate cluster-level combinations robustly:
+panel_split = PanelSplit(
+    periods=panel_data.year,
+    n_splits=2,
+    groups=panel_data["country_id"],
+    group_splitter=StratifiedGroupKFold(n_splits=3) # Use any valid Scikit-Learn group methodology!
+)
+
+# You can also pass arbitrarily nested multi-column groups!
+# PanelSplit will internally flatten them into a single composite group identifier for KFold slicing.
+# e.g., groups = panel_data[["country_id", "city_id"]]
+
+# Lazy Evaluation securely propagates X and y through the StratifiedGroupKFold!
+splits = panel_split.split(X=panel_data, y=panel_data["y"])
+# Yields 6 total sub-splits (2 temporal cuts x 3 spatial stratified holds)!
+```
+
 For more examples and detailed usage instructions, refer to the [examples](examples) directory in this repository. Also feel free to check out [an introductory article on panelsplit](https://towardsdatascience.com/how-to-cross-validate-your-panel-data-in-python-9ad981ddd043).
 
 ## Background

examples/An introduction to PanelSplit.ipynb

@@ -1,16 +1,17 @@
 import warnings
-from typing import Optional, Union, TYPE_CHECKING, Any
-from numpy.typing import NDArray
-from narwhals.typing import IntoDataFrame, IntoSeries
-from .utils.typing import ArrayLike, CVIndices
+from typing import TYPE_CHECKING, Any, Optional, Union
 
 import narwhals as nw
 import numpy as np
-from sklearn.model_selection import TimeSeriesSplit
+from narwhals.typing import IntoDataFrame, IntoSeries
+from numpy.typing import NDArray
+from sklearn.model_selection import GroupKFold, TimeSeriesSplit
 
+from .utils.typing import ArrayLike, CVIndices
 from .utils.validation import (
     _safe_indexing,
     _to_numpy_array,
+    check_groups,
     check_labels,
     check_periods,
     get_index_or_col_from_df,
@@ -64,6 +65,11 @@ class PanelSplit:
     include_train_in_test : bool
         Whether to include all training sets in their respective test sets. If set to
         True, overrides ``include_first_train_in_test``. Default is False.
+    groups : Optional[Any]
+        A 1D/2D array or DataFrame of spatial groupings/IDs for implementing spatio-temporal holdouts.
+        If provided, tests will simultaneously cross-validate over spatial nested structures using GroupKFold. Default is None.
+    group_splitter : Optional[Any]
+        A scikit-learn compatible splitter (e.g., `StratifiedGroupKFold(n_splits=3)`) used to build spatial splits natively. Default is `GroupKFold(n_splits=2)`.
 
     Attributes
     ----------
@@ -101,6 +107,8 @@ def __init__(
         max_train_size: Optional[int] = None,
         include_first_train_in_test: bool = False,
         include_train_in_test: bool = False,
+        groups: Optional[Any] = None,
+        group_splitter: Optional[Any] = None,
     ) -> None:
         periods = check_periods(periods)
 
@@ -130,11 +138,42 @@ def __init__(
             self._include_first_train_in_test = include_first_train_in_test
         else:
             self._include_first_train_in_test = True
+
         self._u_periods_cv = self._split_unique_periods(indices, unique_periods_array)
         self._periods = _to_numpy_array(periods)
         self._snapshots = _to_numpy_array(snapshots) if snapshots is not None else None
+
+        self._groups = check_groups(groups) if groups is not None else None
+
+        if self._groups is not None:
+            if group_splitter is None:
+                self._group_splitter = GroupKFold(n_splits=2)
+            else:
+                self._group_splitter = group_splitter
+            if len(self._groups) != len(self._periods):
+                raise ValueError(
+                    f"groups size ({len(self._groups)}) does not match periods size ({len(self._periods)})"
+                )
+        else:
+            self._group_splitter = None
+
         self.n_splits = n_splits
-        self.train_test_splits = self._gen_splits()
+        if self._groups is not None:
+            self.n_splits = n_splits * self._group_splitter.get_n_splits()  # type: ignore[union-attr]
+
+        self._temporal_splits = self._gen_splits()
+
+        self.train_test_splits = self._temporal_splits
+        if self._groups is not None:
+            try:
+                self.train_test_splits = self._compute_spatio_temporal_splits(
+                    X=None, y=None
+                )
+            except Exception as e:
+                warnings.warn(
+                    f"Could not cleanly pre-generate spatial splits in __init__: {e}. Passing X and y to split() natively at runtime."
+                )
+                self.train_test_splits = []
 
     def _split_unique_periods(self, indices: Any, unique_periods: NDArray) -> CVIndices:
         """
@@ -200,6 +239,30 @@ def _gen_splits(self) -> CVIndices:
 
         return train_test_splits
 
+    def _compute_spatio_temporal_splits(
+        self,
+        X: Optional[ArrayLike] = None,
+        y: Optional[ArrayLike] = None,
+    ) -> CVIndices:
+        """
+        Intersect internal time cuts with lazy spatial matrices natively.
+        """
+        spatio_temporal_splits = []
+        dummy_X = np.zeros(len(self._periods)) if X is None else X
+
+        for train_indices, test_indices in self._temporal_splits:
+            for sp_train, sp_test in self._group_splitter.split(
+                dummy_X, y, groups=self._groups
+            ):
+                final_train_indices = np.intersect1d(
+                    train_indices, sp_train, assume_unique=True
+                )
+                final_test_indices = np.intersect1d(
+                    test_indices, sp_test, assume_unique=True
+                )
+                spatio_temporal_splits.append((final_train_indices, final_test_indices))
+        return spatio_temporal_splits
+
     def split(
         self,
         X: Optional[ArrayLike] = None,
@@ -212,9 +275,9 @@ def split(
         Parameters
         ----------
         X : Optional[ArrayLike]
-            Ignored; included for compatibility.
+            Data matrix to base splits on (used by logic like StratifiedGroupKFold).
         y : Optional[ArrayLike]
-            Ignored; included for compatibility.
+            Target variables to base splits on (used by logic like StratifiedGroupKFold).
         groups : Optional[np.ndarray]
             Ignored; included for compatibility.
 
@@ -241,7 +304,18 @@ def split(
         ...     print("Train:", train, "Test:", test)
         Train: [0 1] Test: [2]
         """
-        return self.train_test_splits
+        if self._groups is None:
+            return self.train_test_splits  # type: ignore[return-value]
+
+        if X is not None or y is not None:
+            self.train_test_splits = self._compute_spatio_temporal_splits(X=X, y=y)
+
+        if not self.train_test_splits:
+            raise ValueError(
+                "train_test_splits is uncomputed. Your selected group_splitter requires passing X and y explicitly to the .split() method to calculate strata boundaries."
+            )
+
+        return self.train_test_splits  # type: ignore[return-value]
 
     def get_n_splits(
         self,

panelsplit/cross_validation.py

@@ -1,33 +1,49 @@
+from typing import Optional, Tuple, Union
+
 import matplotlib.pyplot as plt
+import numpy as np
+
 from .cross_validation import PanelSplit
-from typing import Tuple, Optional
+from .utils.typing import ArrayLike
 
 
 def plot_splits(
-    panel_split: PanelSplit, show: bool = True
-) -> Optional[Tuple[plt.Figure, plt.Axes]]:
+    panel_split: PanelSplit,
+    X: Optional[ArrayLike] = None,
+    y: Optional[ArrayLike] = None,
+    n_groups: int = 2,
+    show: bool = True,
+) -> Optional[Tuple[plt.Figure, Union[plt.Axes, np.ndarray]]]:
     """
     Visualize time series cross-validation splits using a scatter plot.
 
     Each split is plotted on a separate horizontal line: blue markers represent training indices
     and red markers represent test indices.
 
+    If the PanelSplit instance uses groups for spatio-temporal holdouts, this will create
+    `n_groups` subplots, each visualizing the train/test periods for an individual, randomly
+    sampled or sequential subset of the unique groups.
+
     Parameters
     ----------
     panel_split : PanelSplit
         An instance of PanelSplit containing the cross-validation splits.
-        It must have an attribute `_u_periods_cv`, which should be an iterable of tuples,
-        each in the form `(train_index, test_index)`. Both `train_index` and `test_index`
-        are array-like collections of period indices.
+    X : ArrayLike, optional
+        Features dataset. Needed if the group_splitter relies on X for boundary calculation.
+    y : ArrayLike, optional
+        Target dataset. Needed if the group_splitter relies on y for boundary calculation.
+    n_groups : int, default=2
+        The number of subgroups to plot side-by-side if groups are used.
     show : bool, default=True
         If True, the plot is immediately displayed using `plt.show()`.
-        If False, the function returns the matplotlib Figure and Axes objects for further customization.
+        If False, the function returns the matplotlib Figure and Axes objects.
 
     Returns
     -------
-    Optional[Tuple[plt.Figure, plt.Axes]]
+    Optional[Tuple[plt.Figure, Union[plt.Axes, np.ndarray]]]
         If `show` is False, returns a tuple `(fig, ax)` where `fig` is the matplotlib Figure
-        and `ax` is the Axes object. If `show` is True, the plot is displayed and the function returns None.
+        and `ax` is the Axes object (or an array of Axes objects if groups are used).
+        If `show` is True, the plot is displayed and the function returns None.
 
     Examples
     --------
@@ -36,16 +52,15 @@ def plot_splits(
     >>> import matplotlib.pyplot as plt
     >>> periods = np.array([1, 2, 3, 4, 5, 6])
     >>> ps = PanelSplit(periods, n_splits=3)
-    >>> # Display the plot immediately
     >>> plot_splits(ps)
-    >>> # Or return the Figure and Axes for customization
-    >>> fig, ax = plot_splits(ps, show=False)
-    >>> ax.set_title("A custom plot of cross-validation splits")
-    >>> plt.show()
     """
+
+    if panel_split._groups is not None:
+        return _plot_group_subplots(panel_split, X=X, y=y, n_groups=n_groups, show=show)
+
     split_output = panel_split._u_periods_cv
     splits = len(split_output)
-    fig, ax = plt.subplots()
+    fig, ax = plt.subplots(figsize=(10, 5))
 
     for i, (train_index, test_index) in enumerate(split_output):
         ax.scatter(train_index, [i] * len(train_index), color="blue", marker=".", s=50)
@@ -54,13 +69,100 @@ def plot_splits(
     ax.set_xlabel("Periods")
     ax.set_ylabel("Split")
     ax.set_title("Cross-validation splits")
-    ax.set_yticks(range(splits))  # Set the number of ticks on the y-axis
-    ax.set_yticklabels(
-        [f"{i}" for i in range(splits)]
-    )  # Set custom labels for the y-axis
+    ax.set_yticks(range(splits))
+    ax.set_yticklabels([f"{i}" for i in range(splits)])
 
     if show:
         plt.show()
         return None
     else:
         return fig, ax
+
+
+def _plot_group_subplots(
+    panel_split: PanelSplit,
+    X: Optional[ArrayLike],
+    y: Optional[ArrayLike],
+    n_groups: int,
+    show: bool,
+) -> Optional[Tuple[plt.Figure, Union[plt.Axes, np.ndarray]]]:
+    unique_groups = np.unique(np.asarray(panel_split._groups))
+    selected_groups = unique_groups[:n_groups]
+    actual_groups = len(selected_groups)
+
+    fig, axes = plt.subplots(
+        ncols=actual_groups,
+        sharey=True,
+        sharex=True,
+        figsize=(min(16, 6 * actual_groups), 6),
+    )
+
+    if actual_groups == 1:
+        axes = np.array([axes])
+
+    splits = panel_split.split(X=X, y=y)
+    n_total_splits = len(splits)
+    n_temporal_splits = (
+        len(panel_split._temporal_splits)
+        if hasattr(panel_split, "_temporal_splits")
+        else n_total_splits
+    )
+    n_spatial_splits = (
+        n_total_splits // n_temporal_splits if n_temporal_splits > 0 else 1
+    )
+
+    for ax_idx, group in enumerate(selected_groups):
+        ax = axes[ax_idx]
+        group_mask = panel_split._groups == group
+
+        group_indices = np.where(group_mask)[0]
+        for i, (train_indices, test_indices) in enumerate(splits):
+            group_train_indices = np.intersect1d(
+                train_indices, group_indices, assume_unique=True
+            )
+            group_test_indices = np.intersect1d(
+                test_indices, group_indices, assume_unique=True
+            )
+
+            train_periods = panel_split._periods[group_train_indices]
+            test_periods = panel_split._periods[group_test_indices]
+
+            temporal_idx = i // n_spatial_splits
+
+            if len(train_periods) > 0:
+                ax.scatter(
+                    train_periods,
+                    [temporal_idx] * len(train_periods),
+                    color="blue",
+                    marker=".",
+                    s=50,
+                )
+            if len(test_periods) > 0:
+                ax.scatter(
+                    test_periods,
+                    [temporal_idx] * len(test_periods),
+                    color="red",
+                    marker=".",
+                    s=50,
+                )
+
+        ax.set_xlabel("Periods")
+        if ax_idx == 0:
+            ax.set_ylabel("Split")
+        ax.set_title(f"Cross-validation splits: {group}")
+        ax.set_yticks(range(n_temporal_splits))
+        ax.set_yticklabels([f"{j}" for j in range(n_temporal_splits)])
+
+    total_groups = len(unique_groups)
+    remaining_groups = total_groups - actual_groups
+    if remaining_groups > 0:
+        fig.suptitle(
+            f"Cross-validation splits: {remaining_groups} other groups not plotted"
+        )
+
+    plt.tight_layout()
+    if show:
+        plt.show()
+        return None
+    else:
+        return fig, (axes if actual_groups > 1 else axes[0])