Skip to main content

train_test_split

Function train_test_split 

Source
pub fn train_test_split(
    n_samples: usize,
    test_ratio: f64,
    seed: u64,
) -> Result<IndexSplit, DatasetError>
Expand description

Split 0..n_samples into shuffled train and test index lists.

The test set gets round(n_samples * test_ratio) rows, clamped so that neither side is empty whenever there are at least two samples. The train set gets the rest. Both lists are in shuffled order, so a dataset stored grouped by class (the common case) does not produce a train set missing a class.

To keep each class’s proportion intact, use stratified_split instead.

§Parameters

  • n_samples - Total number of samples to split.
  • test_ratio - Fraction of samples to place in the test set, in 0.0..=1.0.
  • seed - Seed for the internal generator. The same seed always yields the same split.

§Returns

  • IndexSplit - The (train, test) row indices. Together they are a permutation of 0..n_samples, and they never overlap.

§Errors

  • DatasetError::ValidationError - Returns this when n_samples is 0, or when test_ratio is not a finite value in 0.0..=1.0.

§Example

use dataset_ml::preprocessing::train_test_split;

let (train, test) = train_test_split(150, 0.2, 42).unwrap();
assert_eq!(train.len(), 120);
assert_eq!(test.len(), 30);

// The two sides are disjoint and cover everything.
let mut all: Vec<usize> = train.iter().chain(test.iter()).copied().collect();
all.sort_unstable();
assert_eq!(all, (0..150).collect::<Vec<_>>());