# 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](#417-train--validation--test-with-two-splits) 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](./4.2._Standardization_and_Normalization.md), 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
```rust,ignore
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>);
```
| `x` | `Array2<f64>` | Feature matrix, shape `(n_samples, n_features)`, taken by value |
| `y` | `Array1<A>` | Labels, length `n_samples`. The element type `A` is generic |
| `test_size` | `Option<f64>` | Fraction of samples for the test set. `None` means `0.3` |
| `random_state` | `Option<u64>` | Seed for the shuffle. `None` defers to the global seed or entropy |
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](./4.3._Label_Encoding.md).
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
```rust
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.
```rust,ignore
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](../Chapter-01/1.4._Your_First_End_to_End_Model.md) for a
complete pipeline and [classification metrics](../Chapter-05/5.2._Classification_Metrics.md) 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.
```rust
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`](../Chapter-03/3.1._The_Sequential_Model.md) 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.
```rust
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](../Chapter-07/7.1._Reproducibility_and_Random_Seeds.md) covers
the full rules. One edge case catches people specifically with splits:
```rust
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.
```rust
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](../Chapter-01/1.6._Error_Handling.md). The table below lists every failure
case.
| `n_samples == 0` | `Error::EmptyInput` | Payload `"dataset"` |
| `x.nrows() != y.len()` | `Error::DimensionMismatch { expected, found }` | `expected` is the row count, `found` the label count |
| `test_size <= 0.0` or `>= 1.0` | `Error::InvalidParameter { name, reason }` | `name` is `"test_size"`. Both bounds are exclusive |
| `n_samples == 1` (plain split) | `Error::InvalidInput` | Cannot form both a train and a test set from one row |
| A class with fewer than 2 samples (stratified) | `Error::InvalidInput` | Every class must land on both sides |
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.
```rust
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.