Data Splitters

This module provides utilities for splitting datasets into train/validation/test sets.

Summary

Splitter Classes

random.RandomSplitter	Random splitting strategy for datasets.
coloring.ColoringSplitter	Coloring-based splitting strategy for distribution shift experiments.

Class Documentation

class RandomSplitter(val_size: int | float = 0.1, test_size: int | float = 0.2)

Bases: Splitter

Random splitting strategy for datasets.

Randomly divides a dataset into train, validation, and test splits. Ensures reproducibility when numpy’s random seed is set externally before calling fit().

The splitting is done in the following order: 1. Test (if test_size > 0) 2. Validation (if val_size > 0) 3. Training (remaining samples)

Parameters:

• val_size (Union[int, float], optional) – Size of validation set. If float, represents fraction of dataset. If int, represents absolute number of samples. Defaults to 0.1.

• test_size (Union[int, float], optional) – Size of test set. If float, represents fraction of dataset. If int, represents absolute number of samples. Defaults to 0.2.

Example

>>> # 70% train, 10% val, 20% test
>>> splitter = RandomSplitter(val_size=0.1, test_size=0.2)
>>> splitter.fit(dataset)
>>> print(f"Train: {splitter.train_len}, Val: {splitter.val_len}, Test: {splitter.test_len}")
Train: 700, Val: 100, Test: 200

fit(dataset: ConceptDataset) → None

Randomly split the dataset into train/val/test sets.

Creates a random permutation of dataset indices and divides them according to specified split sizes. Sets the _fitted flag to True upon completion.

Parameters:

dataset – The ConceptDataset to split.

Raises:

ValueError – If split sizes exceed dataset size.

class ColoringSplitter(root: str, seed: int | None = None, val_size: int | float = 0.1, test_size: int | float = 0.2)

Bases: Splitter

Coloring-based splitting strategy for distribution shift experiments.

Divides a dataset into train/val/test splits based on a pre-computed coloring scheme stored in a JSON file. This ensures that training and validation sets contain samples with ‘training’ coloring, while test sets contain samples with ‘test’ coloring.

This is useful for: - Out-of-distribution (OOD) evaluation - Domain adaptation experiments - Controlled distribution shift scenarios

Note: Assumes the dataset is already shuffled and that a coloring file exists at {root}/coloring_mode_seed_{seed}.json

Parameters:

• root (str) – Root directory containing the coloring mode JSON file.

• seed (int, optional) – Random seed used to identify the coloring file. Defaults to None.

• val_size (Union[int, float], optional) – Validation set size (from ‘training’ colored samples). Defaults to 0.1.

• test_size (Union[int, float], optional) – Test set size (from ‘test’ colored samples). Defaults to 0.2.

Example

>>> # Create a coloring file first: coloring_mode_seed_42.json
>>> # Format: {"0": "training", "1": "training", "2": "test", ...}
>>>
>>> splitter = ColoringSplitter(
...     root='data/my_dataset',
...     seed=42,
...     val_size=0.1,
...     test_size=0.2
... )
>>> splitter.fit(dataset)
>>> # Train/val from 'training' samples, test from 'test' samples

fit(dataset: ConceptDataset) → None

Split dataset based on coloring scheme from JSON file.

Loads the coloring mode file and divides indices into ‘training’ and ‘test’ groups. Then allocates samples from each group to the appropriate splits (train/val from ‘training’, test from ‘test’).

Parameters:

dataset – The ConceptDataset to split.

Raises:

ValueError – If coloring file doesn’t exist, or if there aren’t enough samples of a particular coloring mode to satisfy the requested splits.