model_selection_rs/splitters/

time_series_split.rs

//! Time-aware splitting: rolling-origin / expanding-window.

use super::CvSplitter;
use crate::error::{ModelSelectionError, Result};

/// Time-series cross-validation (rolling-origin evaluation).
///
/// Successive splits grow their training window into the past while the test set
/// is always the *next* chronological chunk. Order is never shuffled: for every
/// split, every training index is strictly earlier than every test index (with
/// an optional [`gap`](TimeSeriesSplit::with_gap) between them).
///
/// **Input convention:** samples are assumed to be in chronological row order.
/// This splitter works purely on positions, not on explicit timestamps — sort
/// your data by time before using it. (Positional order keeps the API simple and
/// matches scikit-learn's `TimeSeriesSplit`; a timestamp-aware variant is
/// intentionally out of scope.)
///
/// Configuration mirrors scikit-learn:
/// * `n_splits` — number of train/test splits.
/// * [`max_train_size`](TimeSeriesSplit::with_max_train_size) — cap the training
///   window to a fixed size (a rolling window) instead of an ever-expanding one.
/// * [`gap`](TimeSeriesSplit::with_gap) — drop this many samples between the end
///   of train and the start of test, modelling a real-world delay before an
///   outcome/label is known.
/// * [`test_size`](TimeSeriesSplit::with_test_size) — samples per test set
///   (defaults to `n_samples / (n_splits + 1)`).
///
/// ```
/// use model_selection_rs::splitters::{CvSplitter, TimeSeriesSplit};
///
/// let tss = TimeSeriesSplit::new(3).unwrap();
/// for (train, test) in tss.split(12).unwrap() {
///     // train is always entirely before test
///     assert!(train.iter().max() < test.iter().min());
/// }
/// ```
#[derive(Debug, Clone)]
pub struct TimeSeriesSplit {
    n_splits: usize,
    max_train_size: Option<usize>,
    gap: usize,
    test_size: Option<usize>,
}

impl TimeSeriesSplit {
    /// Create a `TimeSeriesSplit` with `n_splits` expanding-window splits.
    ///
    /// # Errors
    ///
    /// Returns [`ModelSelectionError::InvalidSplitCount`] if `n_splits < 1`.
    pub fn new(n_splits: usize) -> Result<Self> {
        if n_splits < 1 {
            return Err(ModelSelectionError::InvalidSplitCount {
                msg: format!("n_splits must be >= 1, got {n_splits}"),
            });
        }
        Ok(Self {
            n_splits,
            max_train_size: None,
            gap: 0,
            test_size: None,
        })
    }

    /// Cap the training window to `max_train_size` samples (rolling window).
    #[must_use]
    pub fn with_max_train_size(mut self, max_train_size: usize) -> Self {
        self.max_train_size = Some(max_train_size);
        self
    }

    /// Insert a `gap` of dropped samples between train and test.
    #[must_use]
    pub fn with_gap(mut self, gap: usize) -> Self {
        self.gap = gap;
        self
    }

    /// Fix the number of samples in each test set.
    #[must_use]
    pub fn with_test_size(mut self, test_size: usize) -> Self {
        self.test_size = Some(test_size);
        self
    }
}

impl CvSplitter for TimeSeriesSplit {
    fn split(&self, n_samples: usize) -> Result<Vec<(Vec<usize>, Vec<usize>)>> {
        let n_folds = self.n_splits;
        let test_size = self.test_size.unwrap_or_else(|| n_samples / (n_folds + 1));

        if test_size == 0 {
            return Err(ModelSelectionError::InsufficientTrainWindow {
                msg: format!(
                    "computed test_size is 0 for n_samples={n_samples}, n_splits={n_folds}; \
                     supply more samples or a smaller n_splits"
                ),
            });
        }

        // The first test set starts here; each subsequent one is `test_size`
        // later. We need room for every training window plus the gap.
        let total_test = n_folds.checked_mul(test_size).ok_or_else(|| {
            ModelSelectionError::InvalidSplitCount {
                msg: "n_splits * test_size overflowed".to_string(),
            }
        })?;
        if total_test >= n_samples {
            return Err(ModelSelectionError::InsufficientTrainWindow {
                msg: format!(
                    "n_splits({n_folds}) * test_size({test_size}) = {total_test} leaves no room \
                     for a training set in n_samples={n_samples}"
                ),
            });
        }

        let first_test_start = n_samples - total_test;
        let mut splits = Vec::with_capacity(n_folds);
        for fold in 0..n_folds {
            let test_start = first_test_start + fold * test_size;
            let test_end = test_start + test_size;

            // Train ends `gap` samples before the test set begins.
            let train_end = test_start.checked_sub(self.gap).ok_or_else(|| {
                ModelSelectionError::InsufficientTrainWindow {
                    msg: format!(
                        "gap({}) is larger than the available history before fold {fold}",
                        self.gap
                    ),
                }
            })?;
            if train_end == 0 {
                return Err(ModelSelectionError::InsufficientTrainWindow {
                    msg: format!(
                        "fold {fold} has an empty training window (gap={}, test_size={test_size})",
                        self.gap
                    ),
                });
            }
            let train_start = match self.max_train_size {
                Some(mts) => train_end.saturating_sub(mts),
                None => 0,
            };

            let train: Vec<usize> = (train_start..train_end).collect();
            let test: Vec<usize> = (test_start..test_end).collect();
            splits.push((train, test));
        }
        Ok(splits)
    }

    fn n_splits(&self) -> usize {
        self.n_splits
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    /// The property that matters most: train strictly precedes test everywhere.
    fn assert_chronological(splits: &[(Vec<usize>, Vec<usize>)], gap: usize) {
        for (train, test) in splits {
            let max_train = *train.iter().max().unwrap();
            let min_test = *test.iter().min().unwrap();
            assert!(max_train < min_test, "train index >= test index");
            assert!(min_test - max_train > gap, "gap not respected");
        }
    }

    #[test]
    fn expanding_window_is_chronological() {
        let tss = TimeSeriesSplit::new(4).unwrap();
        let splits = tss.split(20).unwrap();
        assert_eq!(splits.len(), 4);
        assert_chronological(&splits, 0);
        // Training set grows monotonically.
        let train_lens: Vec<usize> = splits.iter().map(|(tr, _)| tr.len()).collect();
        assert!(train_lens.windows(2).all(|w| w[0] < w[1]));
    }

    #[test]
    fn fixed_window_caps_train_size() {
        let tss = TimeSeriesSplit::new(3).unwrap().with_max_train_size(4);
        let splits = tss.split(20).unwrap();
        assert_chronological(&splits, 0);
        assert!(splits.iter().all(|(tr, _)| tr.len() <= 4));
    }

    #[test]
    fn gap_is_respected() {
        let tss = TimeSeriesSplit::new(3).unwrap().with_gap(2);
        let splits = tss.split(30).unwrap();
        assert_chronological(&splits, 2);
    }

    #[test]
    fn errors_when_not_enough_samples() {
        let tss = TimeSeriesSplit::new(10).unwrap();
        assert!(matches!(
            tss.split(5),
            Err(ModelSelectionError::InsufficientTrainWindow { .. })
        ));
    }

    #[test]
    fn custom_test_size() {
        let tss = TimeSeriesSplit::new(3).unwrap().with_test_size(2);
        let splits = tss.split(20).unwrap();
        assert!(splits.iter().all(|(_, te)| te.len() == 2));
        assert_chronological(&splits, 0);
    }
}