4.1. Train-Test Split

Before you fit a model, you split your data into a training set and a test set. The module rustyml::utils::train_test_split gives you 2 functions for this. train_test_split makes a plain random partition. train_test_split_stratified makes a partition that keeps each class in the same proportion on both sides. Both functions are deterministic under a fixed seed. Both take ownership of your arrays, so you can pass the returned partitions straight to a model. This page explains when to use each function, the exact contract each one follows, and the common mistakes people make.

4.1.1. Why hold out data at all

A model measured only on its own training data gives a false sense of quality. The model can memorize the training rows. Examples include a decision tree grown deep enough, a k-nearest-neighbors classifier with k = 1, and an over-parameterized network. Such a model reports near-perfect accuracy that collapses on new data. The metric that matters is generalization: performance on data from the same distribution that the model never saw during fitting. The only honest way to estimate generalization is to set some data aside before training and keep the fitting process away from it.

That is the entire job of a train/test split. The training set is what fit learns from. The test set stands in for future data. Use it exactly once, at the end, to score the finished model. If you tune anything against the test set, such as a hyperparameter, a threshold, or a feature choice, it stops being held out. It starts to leak into your model, and your reported score drifts back toward the optimistic training-set number. When you need to tune settings, carve out a third slice, a validation set, and keep the test set sealed. Section 4.1.7 shows this pattern.

One related leak is common enough to flag here. Any statistic you compute over the whole dataset before splitting has already seen the test rows. Examples include a feature mean and standard deviation for standardization, a min or max value for normalization, and a label vocabulary. Fit those transforms on the training partition only, then apply them to the test partition. Split first, transform second.

4.1.2. The train_test_split signature

pub fn train_test_split<A: Clone>(
    x: Array2<f64>,
    y: Array1<A>,
    test_size: Option<f64>,
    random_state: Option<u64>,
) -> Result<TrainTestSplit<A>, Error>;

pub fn train_test_split_stratified<A: Clone + Eq + Hash>(
    x: Array2<f64>,
    y: Array1<A>,
    test_size: Option<f64>,
    random_state: Option<u64>,
) -> Result<TrainTestSplit<A>, Error>;

pub type TrainTestSplit<A> = (Array2<f64>, Array2<f64>, Array1<A>, Array1<A>);

The label type is generic. Plain train_test_split requires only A: Clone, so i32, usize, f64, and &str labels all work. The crate tests exercise i32 and &str labels directly. Stratification groups rows by class, so it tightens the bound to A: Clone + Eq + Hash. Integers and string slices satisfy that bound. Raw f64 labels do not, because floats do not implement Eq or Hash. This is a good reason to encode class labels as integers before you stratify. See label encoding.

The return value is a 4-tuple in the order (x_train, x_test, y_train, y_test). Both feature matrices come first, then both label vectors. This matches the order of scikit-learn’s X_train, X_test, y_train, y_test, so code ported from Python keeps the same order. Note 3 differences from scikit-learn. The default test_size here is 0.3, not scikit-learn’s 0.25. There is no shuffle flag, because the split always shuffles. For time-series data that needs contiguous, order-preserving slices, do not use this function. Slice the arrays yourself instead. Stratification is a separate function, not a stratify= argument.

You can reach these functions 3 ways. Use the fully qualified path rustyml::utils::train_test_split::train_test_split. Use the flattened path rustyml::utils::{train_test_split, train_test_split_stratified}. Or use the prelude with use rustyml::prelude::*;. The examples on this page use the fully qualified module path.

4.1.3. A basic split

use ndarray::{Array1, Array2};
use rustyml::utils::train_test_split::train_test_split;

fn main() {
    let x = Array2::from_shape_vec(
        (10, 2),
        vec![
            0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 10.0, 11.0, 12.0, 13.0, 14.0,
            15.0, 16.0, 17.0, 18.0, 19.0,
        ],
    )
    .unwrap();
    let y = Array1::from(vec![0, 1, 0, 1, 0, 1, 0, 1, 0, 1]);

    let (x_train, x_test, y_train, y_test) =
        train_test_split(x, y, Some(0.3), Some(42)).unwrap();

    // round(10 * 0.3) = 3 test rows, the remaining 7 are training rows.
    assert_eq!(x_train.nrows(), 7);
    assert_eq!(x_test.nrows(), 3);
    assert_eq!(y_train.len(), 7);
    assert_eq!(y_test.len(), 3);
    println!("train {} / test {}", x_train.nrows(), x_test.nrows());
}

The test set size is round(n_samples * test_size). For 10 samples at 0.3, this gives 3 test rows and 7 training rows. Rows stay aligned: x[i] and y[i] always land in the same partition. Every output row matches exactly one input row. No row is duplicated across the two sides, and no row is dropped. So x_train.nrows() + x_test.nrows() always equals n_samples.

x comes back as Array2<f64> and y comes back as Array1<A>. The partitions feed directly into a model’s fit and predict methods. Most estimators take the feature matrix and labels by reference.

model.fit(&x_train, &y_train)?;
let predictions = model.predict(&x_test)?;
// score `predictions` against `y_test` with a metric from Chapter 5.

See your first end-to-end model for a complete pipeline and classification metrics for scoring the held-out predictions.

4.1.4. Why shuffling matters

Real datasets are rarely stored in random order. Exports are often sorted by label, by timestamp, or by collection batch. Suppose you took the first 70% of an iris-style file as training data. You might train on 2 species and test on a third species the model never saw. That is a guaranteed failure, and it says nothing about the model. Shuffling before slicing breaks this structure, so both partitions become representative samples of the same distribution.

train_test_split always shuffles the row indices before it splits. You never need to pre-sort your data. You cannot turn shuffling off. The only setting you control is the seed. This makes the function unsuitable for problems where order carries meaning. One example is forecasting, where the test set must come strictly after the training set in time. For those problems, slice the arrays by hand instead of using this function.

4.1.5. Stratification and imbalanced classes

Random shuffling gives each row an equal chance of landing in the test set. It does not guarantee that a rare class appears on both sides. Consider a fraud-detection dataset with a negative-to-positive ratio of 8 to 1. A plain split can, by chance, put every positive example into training and leave the test set with none. The test set then cannot measure fraud detection at all, because it holds no positive examples to score against. The smaller the minority class is relative to test_size, the more likely this outcome becomes.

train_test_split_stratified fixes this by splitting each class independently. It groups the row indices by label, in first-appearance order, so the result stays deterministic for a given seed. It shuffles within each group and applies test_size to each group separately. It clamps the result so every class keeps at least 1 sample on each side. This preserves the per-class proportions of the input in both partitions.

use ndarray::{Array1, Array2};
use rustyml::utils::train_test_split::train_test_split_stratified;

fn main() {
    // 8 samples of class 0 and 2 samples of class 1 (a 4:1 imbalance).
    let x = Array2::from_shape_fn((10, 1), |(i, _)| i as f64);
    let mut labels = vec![0i32; 8];
    labels.extend(vec![1i32; 2]);
    let y = Array1::from(labels);

    let (_x_train, _x_test, y_train, y_test) =
        train_test_split_stratified(x, y, Some(0.3), Some(42)).unwrap();

    let count = |a: &Array1<i32>, c: i32| a.iter().filter(|&&l| l == c).count();

    // The minority class survives on both sides. This is guaranteed, not luck.
    assert!(count(&y_train, 1) >= 1, "class 1 must remain in train");
    assert!(count(&y_test, 1) >= 1, "class 1 must remain in test");
    println!(
        "test set: {} of class 0, {} of class 1",
        count(&y_test, 0),
        count(&y_test, 1)
    );
}

Here, the majority class contributes round(8 * 0.3) = 2 test rows. The minority class contributes round(2 * 0.3) = 1 test row, which the clamp also holds inside the allowed range of [1, class_size - 1]. Both classes appear on both sides, every time, for every seed. On a balanced dataset, stratification simply reproduces the requested ratio for each class. For example, 6 samples per class at test_size = 0.5 gives 3 test rows and 3 train rows for each class.

One structural difference matters here. Stratification concatenates each class’s test slice and each class’s train slice in turn. So the returned rows are ordered in class blocks: all of class 0, then all of class 1, and so on. This differs from the global shuffle that plain train_test_split produces. Within a class, the order is shuffled, but the classes themselves are not interleaved. This block order has no effect on an estimator that shuffles internally, such as the Sequential network, which shuffles minibatches each epoch. If you feed the labels to an order-sensitive process that does not reshuffle, keep the class blocks in mind.

Use stratification when the label is categorical and the classes are uneven. For a balanced regression target or roughly even classes, the plain split works well and is simpler to use.

4.1.6. Reproducible splits: random_state vs the global seed

A stable test set matters more than it might seem. If the split changes on every run, your reported accuracy varies for reasons that have nothing to do with the model. You can no longer tell a real improvement from split noise. Fix the seed.

The direct control is random_state. Passing Some(seed) makes the shuffle fully reproducible and independent of everything else. An explicit seed uses exactly that value and never touches the crate’s global seed stream.

use ndarray::{Array1, Array2};
use rustyml::utils::train_test_split::train_test_split;

fn main() {
    let x = Array2::from_shape_fn((20, 3), |(i, j)| (i + j) as f64);
    let y = Array1::from_iter(0..20i32);

    let a = train_test_split(x.clone(), y.clone(), Some(0.25), Some(42)).unwrap();
    let b = train_test_split(x, y, Some(0.25), Some(42)).unwrap();

    // Same seed -> byte-identical partitions.
    assert_eq!(a.0, b.0); // x_train
    assert_eq!(a.1, b.1); // x_test
    assert_eq!(a.2, b.2); // y_train
    assert_eq!(a.3, b.3); // y_test
    println!("reproduced a split of {} training rows", a.0.nrows());
}

If random_state is None, the split falls back to the thread-local global seed set by rustyml::set_global_seed. Reproducibility and random seeds covers the full rules. One edge case catches people specifically with splits:

use ndarray::{Array1, Array2};
use rustyml::set_global_seed;
use rustyml::utils::train_test_split::train_test_split;

fn main() {
    // Fix the whole run's randomness up front.
    set_global_seed(123);

    let x = Array2::from_shape_fn((12, 2), |(i, j)| (i + j) as f64);
    let y = Array1::from_iter(0..12i32);

    // random_state = None derives its seed from the global stream set above.
    let (x_train, x_test, _, _) = train_test_split(x, y, None, None).unwrap();
    println!("train {} / test {}", x_train.nrows(), x_test.nrows());
}

A single set_global_seed call makes the whole program run reproducible, as long as you construct the randomized components in the same order each run. Each unseeded consumer draws a fresh sub-seed from the global stream, in construction order. The global stream advances on every draw. If you call train_test_split(.., None, None) twice under one set_global_seed, the two calls receive different sub-seeds. So the two calls produce different splits. To get the same split every time, seed it explicitly with Some(seed). This is the recommended approach for anything you re-run and compare. Alternatively, reset the global seed before the call. An explicit Some seed never consumes the global stream. So you can safely pin your split with Some(42) and still let your model draw its weights from a program-wide set_global_seed. Neither call disturbs the other.

4.1.7. Train / validation / test with two splits

There is no dedicated 3-way split function. You build one by calling train_test_split twice. First, peel off the sealed test set. Then split what remains into training and validation sets. Each call takes its arrays by value and returns owned arrays, so the remainder flows straight into the second call with no cloning.

use ndarray::{Array1, Array2};
use rustyml::utils::train_test_split::train_test_split;

fn main() {
    let x = Array2::from_shape_fn((20, 2), |(i, j)| (i + j) as f64);
    let y = Array1::from_iter(0..20i32);

    // 1) Peel off the test set: 20% of 20 -> 4 rows, 16 remain.
    let (x_rest, x_test, y_rest, y_test) =
        train_test_split(x, y, Some(0.2), Some(42)).unwrap();

    // 2) Split the remaining 16 into train and validation: 25% -> 4 val, 12 train.
    let (x_train, x_val, y_train, y_val) =
        train_test_split(x_rest, y_rest, Some(0.25), Some(7)).unwrap();

    assert_eq!(x_train.nrows(), 12);
    assert_eq!(x_val.nrows(), 4);
    assert_eq!(x_test.nrows(), 4);
    println!(
        "train {} / val {} / test {}",
        y_train.len(),
        y_val.len(),
        y_test.len()
    );
}

Watch the arithmetic here. The second test_size is a fraction of the remaining rows, not of the original count. Peeling off 20% and then taking 25% of the rest gives a 60/20/20 split of the whole dataset, not 55/25/20. Seed both calls, with the same value or different values, since the 2 calls are independent. This makes the 3-way partition reproducible end to end. When the target is a categorical class, use train_test_split_stratified for both stages. This keeps the class balance intact through both cuts.

4.1.8. Edge cases and error handling

Both functions validate their inputs first. Both return rustyml::error::Error, so failures are typed values you can match on, rather than panics. See error handling. The table below lists every failure case.

The test_size bounds are strictly exclusive. 0.0 and 1.0 are both rejected, and so are negative values and values above 1.0. Either endpoint would leave one partition empty. Inside the valid range, the computed test count is clamped to [1, n_samples - 1]. So no legal test_size ever produces an empty side. For example, with 10 samples, test_size = 0.99 rounds to 10 and clamps to 9 test rows, leaving 1 row for training. With test_size = 0.01, the count rounds to 0 and clamps up to 1 test row. A dataset of 2 samples is a special case, handled before any rounding. That split is always 1 train row and 1 test row, regardless of test_size.

Stratification adds 1 hard requirement: every class needs at least 2 samples, 1 for each side. A singleton class raises Error::InvalidInput. It is not silently dropped. scikit-learn raises the same kind of failure for its least-populated class. The plain split has no such rule, so a class can legally appear on only 1 side. That is exactly the risk stratification removes.

use ndarray::{Array1, Array2};
use rustyml::error::Error;
use rustyml::utils::train_test_split::{train_test_split, train_test_split_stratified};

fn main() {
    // test_size must lie strictly inside (0, 1).
    let x = Array2::from_shape_fn((5, 2), |(i, j)| (i + j) as f64);
    let y = Array1::from_iter(0..5i32);
    match train_test_split(x, y, Some(1.0), Some(42)) {
        Err(Error::InvalidParameter { name, .. }) => {
            assert_eq!(name, "test_size");
            println!("rejected test_size = 1.0 on parameter `{name}`");
        }
        other => panic!("expected InvalidParameter, got {other:?}"),
    }

    // A stratified split needs >= 2 samples in every class.
    let x2 = Array2::from_shape_fn((5, 1), |(i, _)| i as f64);
    let y2 = Array1::from(vec![0i32, 0, 1, 1, 2]); // class 2 is a singleton
    match train_test_split_stratified(x2, y2, Some(0.3), Some(42)) {
        Err(Error::InvalidInput(msg)) => {
            println!("stratified split rejected a singleton class: {msg}");
        }
        other => panic!("expected InvalidInput, got {other:?}"),
    }
}

Error is #[non_exhaustive], so a match over it needs a trailing arm. Here, the other => arm also checks that the expected variant came back. In production code, return the error with ? instead of using panic. This lets the error propagate up to the code that handles failures for the whole pipeline.