Skip to main content

stratified_split

Function stratified_split 

Source
pub fn stratified_split<T: Hash + Eq>(
    labels: &[T],
    test_ratio: f64,
    seed: u64,
) -> Result<IndexSplit, DatasetError>
Expand description

Split into train and test index lists that preserve each class’s proportion.

Like train_test_split, but this function draws the split within each class rather than over the whole dataset. A class that holds 10% of the samples then holds about 10% of the train set and 10% of the test set. This matters for the imbalanced loaders: sms_spam (13% spam), covtype (its rarest cover type is under 0.5%), and kddcup99. An unstratified split can omit a rare class from the test set completely.

Every class with at least two members contributes at least one row to each side. A class with a single member contributes it to the train set.

§Parameters

  • labels - The per-sample class labels, of any comparable, hashable type (&str, String, u8, char, and so on). This covers every label type this crate produces.
  • test_ratio - Fraction of each class 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, each in shuffled order. Together they are a permutation of 0..labels.len().

§Errors

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

§Example

use dataset_ml::preprocessing::stratified_split;

// Nine samples of class "a", one of class "b".
let labels = ["a", "a", "a", "a", "a", "a", "a", "a", "a", "b"];
let (train, test) = stratified_split(&labels, 0.5, 7).unwrap();

// The lone "b" cannot be in both sides, so it stays in the train set.
assert!(train.contains(&9));
assert_eq!(train.len() + test.len(), 10);