API reference

Auto-generated from the source docstrings. Everything below is also importable directly from the top-level splytters namespace (e.g. from splytters import cluster_split).

Adversarial splitters

adversarial

Adversarial splitting algorithms that minimize train-test similarity.

These methods create "hard" evaluation sets where test samples are dissimilar from training samples, testing model generalization.

cluster_split

cluster_split(embeddings: ArrayLike, train_size: float | int = 0.7, method: str = 'kmeans', n_clusters: int = 10, random_state: int = 42, *, strategy: str = 'size', y: ArrayLike | None = None, **cluster_kwargs: Any) -> tuple[np.ndarray, np.ndarray]

Split dataset by assigning entire clusters to train or test.

Prevents 'cluster leakage' where similar samples end up on both sides. The embeddings are clustered (method), then whole clusters are assigned to train/test according to strategy.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
method	str	clustering algorithm, 'kmeans' or 'dbscan'	'kmeans'
n_clusters	int	number of clusters (kmeans only)	10
random_state	int	for reproducibility	42
strategy	str	cluster -> train/test assignment policy: - "size" (default): greedily fill train with the largest clusters until the target ratio is met. The original, target-ratio-driven behavior; DBSCAN noise points go to test. - "centroid": rank clusters by distance from the global centroid, nearest -> train, farthest -> test (adversarial). This is what :func:centroid_adversarial_split delegates to. - "subset_sum": select a subset of clusters whose pooled per-class counts best match a class-balanced target test set. Requires y. - "closest": seed the most isolated cluster and grow it by nearest-neighbor into a single coherent test "pocket" (adversarial).	'size'
y	ArrayLike | None	class labels of shape (n_samples,). Required for strategy="subset_sum"; ignored by the other strategies.	None
**cluster_kwargs	Any	passed to the clustering algorithm	{}

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Raises:

Type	Description
ValueError	on an unknown method or strategy, or if strategy="subset_sum" is used without a valid y.

References

The "subset_sum" and "closest" strategies implement SUBSET-SUM-SPLIT and CLOSEST-SPLIT from Züfle, Dankers & Titov (2023), "Latent Feature-based Data Splits to Improve Generalisation Evaluation," GenBench Workshop @ EMNLP, pp. 112-129. https://aclanthology.org/2023.genbench-1.9 . They cluster a model's hidden representations and assign whole clusters to test to expose latent-space "blind spots"; their analysis finds that the resulting difficulty does not correlate with surface-level properties (e.g. length), complementing the surface-based :mod:splytters.sorters.

The label-balanced idea behind "subset_sum" traces to the earlier ClusterDataSplit of Wecker, Friedrich & Adel (2020), "ClusterDataSplit: Exploring Challenging Clustering-Based Data Splits for Model Performance Evaluation," Eval4NLP @ COLING, which clusters data into splits that differ lexically from train while keeping the label distribution fixed. https://aclanthology.org/2020.eval4nlp-1.15 . See :func:cluster_kfold for their cross-validation variant.

Napoli & White (2025), "Clustering-Based Validation Splits for Model Selection under Domain Shift," TMLR, ground this family theoretically: maximizing the MMD between the two sets is equivalent to kernel k-means clustering (k=2), and they replace ClusterDataSplit's greedy class-balancing with a linear program (with convergence guarantees). https://openreview.net/forum?id=Q692C0WtiD . See :func:mmd_maximized_split for the MMD-maximizing objective.

Source code in splytters/adversarial.py

def cluster_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    method: str = "kmeans",
    n_clusters: int = 10,
    random_state: int = 42,
    *,
    strategy: str = "size",
    y: ArrayLike | None = None,
    **cluster_kwargs: Any,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Split dataset by assigning entire clusters to train or test.

    Prevents 'cluster leakage' where similar samples end up on both sides. The
    embeddings are clustered (``method``), then whole clusters are assigned to
    train/test according to ``strategy``.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        method: clustering algorithm, 'kmeans' or 'dbscan'
        n_clusters: number of clusters (kmeans only)
        random_state: for reproducibility
        strategy: cluster -> train/test assignment policy:
            - ``"size"`` (default): greedily fill train with the largest clusters
              until the target ratio is met. The original, target-ratio-driven
              behavior; DBSCAN noise points go to test.
            - ``"centroid"``: rank clusters by distance from the global centroid,
              nearest -> train, farthest -> test (adversarial). This is what
              :func:`centroid_adversarial_split` delegates to.
            - ``"subset_sum"``: select a subset of clusters whose pooled per-class
              counts best match a class-balanced target test set. Requires ``y``.
            - ``"closest"``: seed the most isolated cluster and grow it by
              nearest-neighbor into a single coherent test "pocket" (adversarial).
        y: class labels of shape (n_samples,). Required for ``strategy="subset_sum"``;
            ignored by the other strategies.
        **cluster_kwargs: passed to the clustering algorithm

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set

    Raises:
        ValueError: on an unknown ``method`` or ``strategy``, or if
            ``strategy="subset_sum"`` is used without a valid ``y``.

    References:
        The ``"subset_sum"`` and ``"closest"`` strategies implement SUBSET-SUM-SPLIT
        and CLOSEST-SPLIT from Züfle, Dankers & Titov (2023), "Latent Feature-based
        Data Splits to Improve Generalisation Evaluation," GenBench Workshop @ EMNLP,
        pp. 112-129. https://aclanthology.org/2023.genbench-1.9 . They cluster a
        model's hidden representations and assign whole clusters to test to expose
        latent-space "blind spots"; their analysis finds that the resulting
        difficulty does not correlate with surface-level properties (e.g. length),
        complementing the surface-based :mod:`splytters.sorters`.

        The label-balanced idea behind ``"subset_sum"`` traces to the earlier
        ClusterDataSplit of Wecker, Friedrich & Adel (2020), "ClusterDataSplit:
        Exploring Challenging Clustering-Based Data Splits for Model Performance
        Evaluation," Eval4NLP @ COLING, which clusters data into splits that
        differ lexically from train while keeping the label distribution fixed.
        https://aclanthology.org/2020.eval4nlp-1.15 . See :func:`cluster_kfold`
        for their cross-validation variant.

        Napoli & White (2025), "Clustering-Based Validation Splits for Model
        Selection under Domain Shift," TMLR, ground this family theoretically:
        maximizing the MMD between the two sets is equivalent to kernel k-means
        clustering (k=2), and they replace ClusterDataSplit's greedy
        class-balancing with a linear program (with convergence guarantees).
        https://openreview.net/forum?id=Q692C0WtiD . See
        :func:`mmd_maximized_split` for the MMD-maximizing objective.
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    valid_strategies = {"size", "centroid", "subset_sum", "closest"}
    if strategy not in valid_strategies:
        raise ValueError(
            f"Unknown strategy: {strategy!r}. Choose from {sorted(valid_strategies)}"
        )

    n_samples = len(embeddings)

    if strategy == "subset_sum":
        if y is None:
            raise ValueError("strategy='subset_sum' requires class labels `y`")
        y = np.asarray(y)
        if len(y) != n_samples:
            raise ValueError(
                f"y has length {len(y)} but embeddings has {n_samples} rows"
            )

    if method == "kmeans":
        clusterer = KMeans(
            n_clusters=n_clusters,
            random_state=random_state,
            n_init="auto",
            **cluster_kwargs,
        )
    elif method == "dbscan":
        clusterer = DBSCAN(**cluster_kwargs)
    else:
        raise ValueError(f"Unknown clustering method: {method}")

    labels = clusterer.fit_predict(embeddings)
    cluster_to_indices: dict[int, list[int]] = defaultdict(list)
    for idx, label in enumerate(labels):
        cluster_to_indices[int(label)].append(idx)

    target_train = resolve_n_train(n_samples, train_size)
    target_test = n_samples - target_train

    if strategy == "size":
        train, test = _assign_by_size(cluster_to_indices, target_train)
    elif strategy == "centroid":
        train, test = _assign_by_centroid(embeddings, cluster_to_indices, target_train)
    elif strategy == "closest":
        train, test = _assign_closest(embeddings, cluster_to_indices, target_test)
    else:  # subset_sum
        train, test = _assign_subset_sum(cluster_to_indices, y, target_test, n_samples)

    return as_index_array(train), as_index_array(test)

centroid_adversarial_split

centroid_adversarial_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_clusters: int = 10, random_state: int = 42, **cluster_kwargs: Any) -> tuple[np.ndarray, np.ndarray]

Adversarial cluster split: assign clusters nearest to global centroid to train, furthest clusters to test.

Combines centroid-distance ranking with cluster-based splitting to maximize train-test dissimilarity while keeping similar samples together.

Thin wrapper over cluster_split(strategy="centroid").

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_clusters	int	number of clusters	10
random_state	int	for reproducibility	42
**cluster_kwargs	Any	passed to KMeans	{}

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def centroid_adversarial_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_clusters: int = 10,
    random_state: int = 42,
    **cluster_kwargs: Any,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial cluster split: assign clusters nearest to global centroid
    to train, furthest clusters to test.

    Combines centroid-distance ranking with cluster-based splitting to
    maximize train-test dissimilarity while keeping similar samples together.

    Thin wrapper over ``cluster_split(strategy="centroid")``.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_clusters: number of clusters
        random_state: for reproducibility
        **cluster_kwargs: passed to KMeans

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    return cluster_split(
        embeddings,
        train_size=train_size,
        method="kmeans",
        n_clusters=n_clusters,
        random_state=random_state,
        strategy="centroid",
        **cluster_kwargs,
    )

cluster_kfold

cluster_kfold(embeddings: ArrayLike, y: ArrayLike, n_folds: int = 5, method: str = 'kmeans', n_clusters: int | None = None, random_state: int = 42, **cluster_kwargs: Any) -> np.ndarray

Partition data into challenging, label-balanced cross-validation folds.

Clusters the embeddings, then assigns whole clusters to n_folds folds so that each fold is lexically coherent -- similar samples share a fold, making every fold differ from the data trained on for it -- while its label distribution stays close to the global one. This yields a "challenging" CV: each held-out fold is a cluster-based blind spot rather than a random sample.

The result is a per-sample fold id, usable directly with :class:sklearn.model_selection.PredefinedSplit::

from sklearn.model_selection import PredefinedSplit, cross_validate
folds = cluster_kfold(embeddings, y, n_folds=5)
cross_validate(model, X, y, cv=PredefinedSplit(folds))

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
y	ArrayLike	class labels of shape (n_samples,); used to balance folds by label	required
n_folds	int	number of cross-validation folds (default 5)	5
method	str	clustering algorithm, 'kmeans' or 'dbscan'	'kmeans'
n_clusters	int | None	number of clusters (kmeans only). Defaults to max(10, n_folds * 3) -- comfortably above n_folds so every fold can receive whole clusters.	None
random_state	int	for reproducibility	42
**cluster_kwargs	Any	passed to the clustering algorithm	{}

Returns:

Name	Type	Description
fold_ids	ndarray	integer ndarray of shape (n_samples,), values in
	ndarray	[0, n_folds). For fold k the CV test set is fold_ids == k
	ndarray	and the training set is the rest.

Raises:

Type	Description
ValueError	on an unknown method, an y/embeddings length mismatch, an out-of-range n_folds, or fewer clusters than folds.

References

Implements the challenging clustering-based cross-validation of Wecker, Friedrich & Adel (2020), "ClusterDataSplit: Exploring Challenging Clustering-Based Data Splits for Model Performance Evaluation," Eval4NLP @ COLING -- folds that are lexically distinct from train while preserving label balance. https://aclanthology.org/2020.eval4nlp-1.15

Source code in splytters/adversarial.py

def cluster_kfold(
    embeddings: ArrayLike,
    y: ArrayLike,
    n_folds: int = 5,
    method: str = "kmeans",
    n_clusters: int | None = None,
    random_state: int = 42,
    **cluster_kwargs: Any,
) -> np.ndarray:
    """Partition data into challenging, label-balanced cross-validation folds.

    Clusters the embeddings, then assigns whole clusters to ``n_folds`` folds so
    that each fold is lexically coherent -- similar samples share a fold, making
    every fold differ from the data trained on for it -- while its label
    distribution stays close to the global one. This yields a "challenging" CV:
    each held-out fold is a cluster-based blind spot rather than a random sample.

    The result is a per-sample fold id, usable directly with
    :class:`sklearn.model_selection.PredefinedSplit`::

        from sklearn.model_selection import PredefinedSplit, cross_validate
        folds = cluster_kfold(embeddings, y, n_folds=5)
        cross_validate(model, X, y, cv=PredefinedSplit(folds))

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        y: class labels of shape (n_samples,); used to balance folds by label
        n_folds: number of cross-validation folds (default 5)
        method: clustering algorithm, 'kmeans' or 'dbscan'
        n_clusters: number of clusters (kmeans only). Defaults to
            ``max(10, n_folds * 3)`` -- comfortably above ``n_folds`` so every
            fold can receive whole clusters.
        random_state: for reproducibility
        **cluster_kwargs: passed to the clustering algorithm

    Returns:
        fold_ids: integer ndarray of shape (n_samples,), values in
        ``[0, n_folds)``. For fold ``k`` the CV test set is ``fold_ids == k``
        and the training set is the rest.

    Raises:
        ValueError: on an unknown ``method``, an ``y``/embeddings length
            mismatch, an out-of-range ``n_folds``, or fewer clusters than folds.

    References:
        Implements the challenging clustering-based cross-validation of Wecker,
        Friedrich & Adel (2020), "ClusterDataSplit: Exploring Challenging
        Clustering-Based Data Splits for Model Performance Evaluation," Eval4NLP
        @ COLING -- folds that are lexically distinct from train while preserving
        label balance. https://aclanthology.org/2020.eval4nlp-1.15
    """
    embeddings = validate_split_inputs(embeddings, 0.5)  # 0.5: unused placeholder
    y = np.asarray(y)
    n_samples = len(embeddings)

    if len(y) != n_samples:
        raise ValueError(f"y has length {len(y)} but embeddings has {n_samples} rows")
    if not 2 <= n_folds <= n_samples:
        raise ValueError(f"n_folds must be in [2, {n_samples}], got {n_folds}")

    if n_clusters is None:
        n_clusters = max(10, n_folds * 3)
    n_clusters = min(n_clusters, n_samples)
    if n_clusters < n_folds:
        raise ValueError(
            f"need at least n_folds={n_folds} clusters, got n_clusters={n_clusters}"
        )

    if method == "kmeans":
        clusterer = KMeans(
            n_clusters=n_clusters,
            random_state=random_state,
            n_init="auto",
            **cluster_kwargs,
        )
    elif method == "dbscan":
        clusterer = DBSCAN(**cluster_kwargs)
    else:
        raise ValueError(f"Unknown clustering method: {method}")

    labels = clusterer.fit_predict(embeddings)
    cluster_to_indices: dict[int, list[int]] = defaultdict(list)
    for idx, label in enumerate(labels):
        cluster_to_indices[int(label)].append(idx)

    # Folds are whole clusters, so we need at least n_folds distinct clusters.
    # KMeans is guarded by the n_clusters check above, but DBSCAN chooses its
    # own cluster count and can return fewer (e.g. one dense cluster plus a
    # noise group), which would silently leave CV folds empty. Fail loudly.
    n_found = len(cluster_to_indices)
    if n_found < n_folds:
        raise ValueError(
            f"clustering produced only {n_found} cluster(s), fewer than "
            f"n_folds={n_folds}, which would leave empty folds. Adjust the "
            f"clustering (e.g. DBSCAN eps/min_samples) or reduce n_folds."
        )

    classes = np.unique(y)
    class_idx = {c: k for k, c in enumerate(classes)}

    # Per-cluster class-count vectors.
    clusters = []
    for idxs in cluster_to_indices.values():
        vec = np.zeros(len(classes))
        labs, counts = np.unique(y[idxs], return_counts=True)
        for c, cnt in zip(labs, counts, strict=True):
            vec[class_idx[c]] = cnt
        clusters.append((idxs, vec))

    # Target per-fold class counts: the global label distribution, split n_folds ways.
    global_counts = np.array([np.sum(y == c) for c in classes], dtype=float)
    target = global_counts / n_folds

    # Place the largest clusters first into the fold that overshoots the per-fold
    # target least (so under-filled and empty folds are preferred, keeping labels
    # balanced); ties go to the currently-smallest fold to balance fold sizes.
    fold_counts = [np.zeros(len(classes)) for _ in range(n_folds)]
    fold_ids = np.empty(n_samples, dtype=np.intp)
    clusters.sort(key=lambda t: len(t[0]), reverse=True)
    for idxs, vec in clusters:
        k = min(
            range(n_folds),
            key=lambda f: (
                float(np.maximum(fold_counts[f] + vec - target, 0).sum()),
                float(fold_counts[f].sum()),
            ),
        )
        fold_counts[k] = fold_counts[k] + vec
        fold_ids[idxs] = k

    return fold_ids

minority_split

minority_split(embeddings: ArrayLike, y: ArrayLike, n_clusters: int = 10, method: str = 'kmeans', random_state: int = 42, **cluster_kwargs: Any) -> tuple[np.ndarray, np.ndarray]

Bias-amplified split via per-cluster minority labels.

Clusters the embeddings, then for each cluster routes instances whose label is the cluster's majority label to train (the "biased" majority a shortcut-learning model would exploit) and instances with any minority label to test (the "anti-biased" examples that defy the local pattern), yielding a hard, bias-amplified evaluation set.

Unlike the size-driven splitters, the train/test sizes are determined by the data's bias structure (test = the minority-label instances), so this function takes no train_size. It inspects per-cluster label counts only; it never references global class frequencies.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
y	ArrayLike	class labels of shape (n_samples,)	required
n_clusters	int	number of clusters (kmeans only)	10
method	str	clustering algorithm, 'kmeans' or 'dbscan'	'kmeans'
random_state	int	for reproducibility	42
**cluster_kwargs	Any	passed to the clustering algorithm	{}

Returns:

Name	Type	Description
train_indices	ndarray	majority-label ("biased") instances, pooled over clusters
test_indices	ndarray	minority-label ("anti-biased") instances, pooled over clusters

Raises:

Type	Description
ValueError	on an unknown method, a y/embeddings length mismatch, or if no minority examples exist (every cluster is label-pure -- try a different n_clusters or embeddings).

Warns:

Type	Description
UserWarning	if the test set is degenerately small (under 5% of the samples). The clusters are then nearly label-pure, so the split has too few minority examples to be informative -- prefer a different n_clusters or a train_size-driven splitter such as :func:class_boundary_split.

References

Implements the "minority examples" notion of bias from Reif & Schwartz (2023), "Fighting Bias with Bias: Promoting Model Robustness by Amplifying Dataset Biases," Findings of ACL, pp. 13169-13189 (https://aclanthology.org/2023.findings-acl.833): cluster the data, treat each cluster's majority label as biased (train) and its minority labels as anti-biased (test). Their other two notions -- dataset cartography (training dynamics) and partial-input models -- are model-in-the-loop / task-specific and out of scope for a static splitter.

Source code in splytters/adversarial.py

def minority_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    n_clusters: int = 10,
    method: str = "kmeans",
    random_state: int = 42,
    **cluster_kwargs: Any,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Bias-amplified split via per-cluster minority labels.

    Clusters the embeddings, then for each cluster routes instances whose label
    is the cluster's *majority* label to train (the "biased" majority a
    shortcut-learning model would exploit) and instances with any *minority*
    label to test (the "anti-biased" examples that defy the local pattern),
    yielding a hard, bias-amplified evaluation set.

    Unlike the size-driven splitters, the train/test sizes are determined by the
    data's bias structure (test = the minority-label instances), so this function
    takes no ``train_size``. It inspects per-cluster label counts only; it never
    references global class frequencies.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        y: class labels of shape (n_samples,)
        n_clusters: number of clusters (kmeans only)
        method: clustering algorithm, 'kmeans' or 'dbscan'
        random_state: for reproducibility
        **cluster_kwargs: passed to the clustering algorithm

    Returns:
        train_indices: majority-label ("biased") instances, pooled over clusters
        test_indices: minority-label ("anti-biased") instances, pooled over clusters

    Raises:
        ValueError: on an unknown ``method``, a ``y``/embeddings length mismatch,
            or if no minority examples exist (every cluster is label-pure -- try
            a different ``n_clusters`` or embeddings).

    Warns:
        UserWarning: if the test set is degenerately small (under 5% of the
            samples). The clusters are then nearly label-pure, so the split has
            too few minority examples to be informative -- prefer a different
            ``n_clusters`` or a ``train_size``-driven splitter such as
            :func:`class_boundary_split`.

    References:
        Implements the "minority examples" notion of bias from Reif & Schwartz
        (2023), "Fighting Bias with Bias: Promoting Model Robustness by Amplifying
        Dataset Biases," Findings of ACL, pp. 13169-13189
        (https://aclanthology.org/2023.findings-acl.833): cluster the data, treat
        each cluster's majority label as biased (train) and its minority labels as
        anti-biased (test). Their other two notions -- dataset cartography
        (training dynamics) and partial-input models -- are model-in-the-loop /
        task-specific and out of scope for a static splitter.
    """
    embeddings = validate_split_inputs(embeddings, 0.5)  # 0.5: unused placeholder
    y = np.asarray(y)
    n_samples = len(embeddings)
    if len(y) != n_samples:
        raise ValueError(f"y has length {len(y)} but embeddings has {n_samples} rows")

    if method == "kmeans":
        clusterer = KMeans(
            n_clusters=n_clusters,
            random_state=random_state,
            n_init="auto",
            **cluster_kwargs,
        )
    elif method == "dbscan":
        clusterer = DBSCAN(**cluster_kwargs)
    else:
        raise ValueError(f"Unknown clustering method: {method}")

    labels = clusterer.fit_predict(embeddings)
    cluster_to_indices: dict[int, list[int]] = defaultdict(list)
    for idx, label in enumerate(labels):
        cluster_to_indices[int(label)].append(idx)

    train: list[int] = []
    test: list[int] = []
    for idxs in cluster_to_indices.values():
        idx_arr = np.asarray(idxs)
        cls_labels = y[idx_arr]
        vals, counts = np.unique(cls_labels, return_counts=True)
        majority = vals[np.argmax(counts)]  # ties -> lowest label (deterministic)
        is_majority = cls_labels == majority
        train.extend(idx_arr[is_majority].tolist())
        test.extend(idx_arr[~is_majority].tolist())

    if not test:
        raise ValueError(
            "no minority examples found (every cluster is label-pure); "
            "try a different n_clusters or embeddings"
        )

    test_fraction = len(test) / n_samples
    if test_fraction < _MINORITY_DEGENERATE_FRACTION:
        warnings.warn(
            f"minority_split produced a degenerate test set: {len(test)} of "
            f"{n_samples} samples ({test_fraction:.1%}). The clusters are nearly "
            "label-pure, so almost no 'minority' examples exist and this split is "
            "likely uninformative. Try a different n_clusters, or a splitter with "
            "an explicit train_size such as class_boundary_split.",
            stacklevel=2,
        )

    return as_index_array(sorted(train)), as_index_array(sorted(test))

minority_grow_split

minority_grow_split(embeddings: ArrayLike, y: ArrayLike, train_size: float | int = 0.7, n_clusters: int = 10, method: str = 'kmeans', metric: str = 'euclidean', random_state: int = 42, stratify: str = 'none', **cluster_kwargs: Any) -> tuple[np.ndarray, np.ndarray]

Region-grown variant of :func:minority_split with a target test size.

Seeds the test set with the per-cluster minority-label ("anti-biased") instances found by :func:minority_split, then greedily grows it: at each step the not-yet-test sample closest to the current test set (single-link distance to its nearest test member) is moved to test, until the test set reaches the train_size-implied target. This keeps the hard, bias-defying flavor of the minority seed while letting you dial the test fraction -- unlike :func:minority_split, whose size is dictated by the data and is often degenerately small.

Class balance is controlled by stratify (mirrors cluster_split's strategy string convention):

• "none" (default): growth is purely geometric and ignores y, so the test set is generally not class-balanced -- it expands into one region and inherits that region's dominant class (same un-balanced behavior as :func:minority_split).
• "global": one shared proximity frontier, but each class is capped at its data-proportional share of the test set; a class drops out of the frontier once full. The test set stays one contiguous region whose label mix tracks the data.
• "per_class": each class grows its own nearest-to-seed neighborhood up to its quota, independently. The test set is the union of these per-class regions (so not a single contiguous blob), but every class contributes its geometrically-tightest hard samples. A class with no seed member is anchored at its sample nearest the overall seed.

For both stratified modes the quotas are data-proportional and the minority seed is always kept, so a class whose seed already over-fills its quota may still exceed its share.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
y	ArrayLike	class labels of shape (n_samples,)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set; sets the target test size (n_samples - n_train).	0.7
n_clusters	int	number of clusters for the minority seed (kmeans only)	10
method	str	clustering algorithm for the seed, 'kmeans' or 'dbscan'	'kmeans'
metric	str	distance metric (scipy.spatial.distance.cdist) for growing	'euclidean'
random_state	int	for reproducibility of the clustering seed	42
stratify	str	class-balance policy for growth, one of "none" (default, proximity only), "global" (shared frontier with per-class quotas), or "per_class" (independent per-class neighborhoods). See above.	'none'
**cluster_kwargs	Any	passed to the seed clustering algorithm	{}

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for the training set
test_indices	ndarray	ndarray of indices for the test set (the minority seed plus its grown neighborhood). If the seed already meets or exceeds the target size, it is returned as-is (test may exceed the target). Under a stratified mode the test set may fall short of the target if every class hits its quota first.

Raises:

Type	Description
ValueError	on an unknown method or stratify, a y/embeddings length mismatch, or if the seed is empty (every cluster is label-pure -- see :func:minority_split).

Warns:

Type	Description
UserWarning	if the resulting test set contains no samples for one or more classes (those classes then can't be evaluated). Prefer a stratified mode or a larger train_size.

References

Extends the bias-amplified minority seed of Reif & Schwartz (2023); see :func:minority_split. The proximity growth is the single-linkage region growth also used by cluster_split(strategy="closest").

Source code in splytters/adversarial.py

def minority_grow_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    train_size: float | int = 0.7,
    n_clusters: int = 10,
    method: str = "kmeans",
    metric: str = "euclidean",
    random_state: int = 42,
    stratify: str = "none",
    **cluster_kwargs: Any,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Region-grown variant of :func:`minority_split` with a target test size.

    Seeds the test set with the per-cluster minority-label ("anti-biased")
    instances found by :func:`minority_split`, then greedily grows it: at each
    step the not-yet-test sample *closest* to the current test set (single-link
    distance to its nearest test member) is moved to test, until the test set
    reaches the ``train_size``-implied target. This keeps the hard, bias-defying
    flavor of the minority seed while letting you dial the test fraction --
    unlike :func:`minority_split`, whose size is dictated by the data and is
    often degenerately small.

    Class balance is controlled by ``stratify`` (mirrors ``cluster_split``'s
    ``strategy`` string convention):

    - ``"none"`` (default): growth is purely geometric and ignores ``y``, so the
      test set is generally *not* class-balanced -- it expands into one region and
      inherits that region's dominant class (same un-balanced behavior as
      :func:`minority_split`).
    - ``"global"``: one shared proximity frontier, but each class is capped at its
      data-proportional share of the test set; a class drops out of the frontier
      once full. The test set stays one contiguous region whose label mix tracks
      the data.
    - ``"per_class"``: each class grows its *own* nearest-to-seed neighborhood up
      to its quota, independently. The test set is the union of these per-class
      regions (so not a single contiguous blob), but every class contributes its
      geometrically-tightest hard samples. A class with no seed member is anchored
      at its sample nearest the overall seed.

    For both stratified modes the quotas are data-proportional and the minority
    seed is always kept, so a class whose seed already over-fills its quota may
    still exceed its share.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        y: class labels of shape (n_samples,)
        train_size: fraction in (0, 1) or absolute count for the training set;
            sets the target test size (``n_samples - n_train``).
        n_clusters: number of clusters for the minority seed (kmeans only)
        method: clustering algorithm for the seed, 'kmeans' or 'dbscan'
        metric: distance metric (``scipy.spatial.distance.cdist``) for growing
        random_state: for reproducibility of the clustering seed
        stratify: class-balance policy for growth, one of ``"none"`` (default,
            proximity only), ``"global"`` (shared frontier with per-class quotas),
            or ``"per_class"`` (independent per-class neighborhoods). See above.
        **cluster_kwargs: passed to the seed clustering algorithm

    Returns:
        train_indices: ndarray of indices for the training set
        test_indices: ndarray of indices for the test set (the minority seed plus
            its grown neighborhood). If the seed already meets or exceeds the
            target size, it is returned as-is (test may exceed the target). Under a
            stratified mode the test set may fall short of the target if every
            class hits its quota first.

    Raises:
        ValueError: on an unknown ``method`` or ``stratify``, a ``y``/embeddings
            length mismatch, or if the seed is empty (every cluster is label-pure
            -- see :func:`minority_split`).

    Warns:
        UserWarning: if the resulting test set contains no samples for one or more
            classes (those classes then can't be evaluated). Prefer a stratified
            mode or a larger ``train_size``.

    References:
        Extends the bias-amplified minority seed of Reif & Schwartz (2023); see
        :func:`minority_split`. The proximity growth is the single-linkage region
        growth also used by ``cluster_split(strategy="closest")``.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    y = np.asarray(y)
    n_samples = len(embeddings)
    if len(y) != n_samples:
        raise ValueError(f"y has length {len(y)} but embeddings has {n_samples} rows")
    if stratify not in _STRATIFY_MODES:
        raise ValueError(
            f"Unknown stratify: {stratify!r}. Choose from {list(_STRATIFY_MODES)}"
        )

    # Seed with minority_split; suppress its degenerate-size warning since the
    # whole point here is to grow that seed to a usable size.
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        _, seed = minority_split(
            embeddings,
            y,
            n_clusters=n_clusters,
            method=method,
            random_state=random_state,
            **cluster_kwargs,
        )

    test_mask = np.zeros(n_samples, dtype=bool)
    test_mask[seed] = True

    target_n_test = n_samples - resolve_n_train(n_samples, train_size)
    target_n_test = max(1, min(target_n_test, n_samples - 1))

    if test_mask.sum() < target_n_test:
        if stratify == "per_class":
            _grow_per_class(embeddings, y, test_mask, target_n_test, metric)
        else:
            _grow_global(
                embeddings, y, test_mask, target_n_test, metric,
                stratify=stratify == "global",
            )

    test_idx = np.flatnonzero(test_mask)
    train_idx = np.flatnonzero(~test_mask)

    # A class wholly absent from the test set can't be evaluated. Possible here
    # because the minority seed skips majority-everywhere classes and proximity
    # growth need not reach them (stratify quotas mitigate but don't guarantee
    # when a class's proportional quota rounds to 0).
    missing = np.setdiff1d(np.unique(y), y[test_idx])
    if missing.size:
        warnings.warn(
            f"minority_grow_split produced a test set with no samples for "
            f"class(es) {missing.tolist()}, so those classes cannot be evaluated. "
            "Try stratify='global' or 'per_class', a larger train_size (bigger "
            "test set), or a different n_clusters.",
            stacklevel=2,
        )

    return as_index_array(train_idx), as_index_array(test_idx)

class_boundary_split

class_boundary_split(embeddings: ArrayLike, y: ArrayLike, train_size: float | int = 0.7, metric: str = 'euclidean', reference: str = 'centroids', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Class-stratified adversarial split on per-class decision boundaries.

Within each class, ranks samples by how close they sit to other classes and routes the closest (most confusable, near-boundary) ones to the test set until that class's test quota is filled. The remaining, more class-typical samples go to train. Because the quota is applied per class, the test set stays label-balanced (stratified) while still being hard: it concentrates the examples a model is most likely to confuse across categories.

"Closeness to other classes" is measured by reference:

• "centroids" (default): distance to the nearest other class centroid. Cheap -- O(n*K) for K classes.
• "samples": nearest-enemy distance, i.e. distance to the single closest sample belonging to any other class. Sharper boundaries, but O(n^2).

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
y	ArrayLike	class labels of shape (n_samples,)	required
train_size	float | int	fraction in (0, 1) or absolute count, applied per class to size each class's train portion. A fraction is recommended; an absolute count is clamped to each class's size.	0.7
metric	str	distance metric passed to scipy.spatial.distance.cdist	'euclidean'
reference	str	"centroids" or "samples" (see above)	'centroids'
random_state	int	accepted for API consistency; this split is deterministic	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of class-typical (interior) sample indices
test_indices	ndarray	ndarray of near-boundary, cross-class-confusable indices

Raises:

Type	Description
ValueError	on a y/embeddings length mismatch, an unknown reference, or fewer than two distinct classes (no "other class" to measure distance to).

References

A label-stratified, per-class variant of the adversarial-distance idea of Søgaard, Ebert, Bastings & Filippova (2021), "We Need to Talk About Random Splits," EACL (https://aclanthology.org/2021.eacl-main.156): hard splits push dissimilar samples into test. Here "dissimilar" is measured toward other classes rather than toward the training set, yielding a boundary-focused test set; contrast :func:distance_adversarial_split (unsupervised, distance from the global centroid).

Source code in splytters/adversarial.py

def class_boundary_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    reference: str = "centroids",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Class-stratified adversarial split on per-class decision boundaries.

    Within each class, ranks samples by how close they sit to *other* classes
    and routes the closest (most confusable, near-boundary) ones to the test set
    until that class's test quota is filled. The remaining, more class-typical
    samples go to train. Because the quota is applied per class, the test set
    stays label-balanced (stratified) while still being hard: it concentrates the
    examples a model is most likely to confuse across categories.

    "Closeness to other classes" is measured by ``reference``:

    - ``"centroids"`` (default): distance to the nearest *other* class centroid.
      Cheap -- O(n*K) for K classes.
    - ``"samples"``: nearest-enemy distance, i.e. distance to the single closest
      sample belonging to any other class. Sharper boundaries, but O(n^2).

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        y: class labels of shape (n_samples,)
        train_size: fraction in (0, 1) or absolute count, applied *per class* to
            size each class's train portion. A fraction is recommended; an
            absolute count is clamped to each class's size.
        metric: distance metric passed to ``scipy.spatial.distance.cdist``
        reference: ``"centroids"`` or ``"samples"`` (see above)
        random_state: accepted for API consistency; this split is deterministic

    Returns:
        train_indices: ndarray of class-typical (interior) sample indices
        test_indices: ndarray of near-boundary, cross-class-confusable indices

    Raises:
        ValueError: on a ``y``/embeddings length mismatch, an unknown
            ``reference``, or fewer than two distinct classes (no "other class"
            to measure distance to).

    References:
        A label-stratified, per-class variant of the adversarial-distance idea of
        Søgaard, Ebert, Bastings & Filippova (2021), "We Need to Talk About
        Random Splits," EACL (https://aclanthology.org/2021.eacl-main.156): hard
        splits push dissimilar samples into test. Here "dissimilar" is measured
        toward *other classes* rather than toward the training set, yielding a
        boundary-focused test set; contrast :func:`distance_adversarial_split`
        (unsupervised, distance from the global centroid).
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    y = np.asarray(y)
    n_samples = len(embeddings)
    if len(y) != n_samples:
        raise ValueError(f"y has length {len(y)} but embeddings has {n_samples} rows")
    if reference not in ("centroids", "samples"):
        raise ValueError(f"reference must be 'centroids' or 'samples', got {reference!r}")

    classes = np.unique(y)
    if len(classes) < 2:
        raise ValueError(
            f"need at least 2 distinct classes for a boundary split, got {len(classes)}"
        )

    train: list[int] = []
    test: list[int] = []
    for k in classes:
        members = np.flatnonzero(y == k)
        others = np.flatnonzero(y != k)
        n_train_k = min(resolve_n_train(len(members), train_size), len(members))
        n_test_k = len(members) - n_train_k
        if n_test_k <= 0:
            train.extend(members.tolist())
            continue
        dist = _other_class_distance(embeddings, members, others, y, reference, metric)
        # Closest to other classes first -> test; stable sort keeps ties deterministic.
        order = np.argsort(dist, kind="stable")
        test.extend(members[order[:n_test_k]].tolist())
        train.extend(members[order[n_test_k:]].tolist())

    return as_index_array(sorted(train)), as_index_array(sorted(test))

decision_boundary_split

decision_boundary_split(embeddings: ArrayLike, y: ArrayLike, train_size: float | int = 0.7, *, model: str = 'linear_svc', score: str = 'confidence', stratify: str = 'per_class', cv: int = 5, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Supervised adversarial split on a learned decision boundary.

Fits a fast linear classifier under cross-validation and routes the samples a real model finds hardest -- those closest to the learned decision boundary -- to the test set, leaving the confident, class-interior samples for train. The learned-boundary counterpart to :func:class_boundary_split (which instead measures geometric distance to other classes).

Hardness is scored out of fold via :class:~sklearn.model_selection.StratifiedKFold: each sample's margin comes from a model that did not train on it, so the score reflects genuine difficulty rather than in-sample memorization.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim).	required
y	ArrayLike	class labels of shape (n_samples,).	required
train_size	float | int	fraction in (0, 1) or absolute count. With stratify="per_class" it sizes each class's train portion (so the test set stays label-balanced); with "global" it sizes the whole train set.	0.7
model	str	surrogate classifier -- "linear_svc" (default, :class:~sklearn.svm.LinearSVC), "logistic" (:class:~sklearn.linear_model.LogisticRegression), or "rbf_svc" (kernel SVM, :class:~sklearn.svm.SVC with an RBF kernel and gamma="scale"). The linear models are fast; "rbf_svc" captures non-linear boundaries -- and finds genuinely harder splits on non-linearly-separable embeddings -- but costs O(n^2)+ per fold, so reserve it for smaller n. Each is standardized via a per-fold :class:~sklearn.preprocessing.StandardScaler.	'linear_svc'
score	str	hardness measure. "confidence" (default) uses the top-1 minus top-2 margin and works for both models; "entropy" uses predictive entropy and requires probabilities (model="logistic" only).	'confidence'
stratify	str	"per_class" (default; per-class test quota -> label-balanced test) or "global" (purest-hard; may unbalance classes).	'per_class'
cv	int	number of cross-validation folds, clamped to the smallest class count. Each fold supplies the out-of-fold margins for its held-out samples.	5
random_state	int	seeds the fold shuffling and the classifier; the split is deterministic (ties broken by index).	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of confident, class-interior sample indices.
test_indices	ndarray	ndarray of near-boundary, model-confusable indices.

Raises:

Type	Description
ValueError	on a y/embeddings length mismatch, fewer than two distinct classes, an unknown model/score/stratify, an entropy score with linear_svc, or a class with fewer than two samples (too small to hold out).

References

Uncertainty / margin sampling from active learning (Lewis & Gale, 1994; Scheffer et al., 2001; Settles, 2009). Like :func:class_boundary_split, a label-aware variant of the hard-split motivation of Soegaard et al. (2021), "We Need to Talk About Random Splits" (https://aclanthology.org/2021.eacl-main.156), but using a learned boundary rather than embedding geometry.

Source code in splytters/adversarial.py

def decision_boundary_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    train_size: float | int = 0.7,
    *,
    model: str = "linear_svc",
    score: str = "confidence",
    stratify: str = "per_class",
    cv: int = 5,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Supervised adversarial split on a learned decision boundary.

    Fits a fast linear classifier under cross-validation and routes the samples
    a real model finds hardest -- those closest to the learned decision boundary
    -- to the test set, leaving the confident, class-interior samples for train.
    The learned-boundary counterpart to :func:`class_boundary_split` (which
    instead measures geometric distance to other classes).

    Hardness is scored *out of fold* via
    :class:`~sklearn.model_selection.StratifiedKFold`: each sample's margin comes
    from a model that did not train on it, so the score reflects genuine
    difficulty rather than in-sample memorization.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim).
        y: class labels of shape (n_samples,).
        train_size: fraction in (0, 1) or absolute count. With
            ``stratify="per_class"`` it sizes each class's train portion (so the
            test set stays label-balanced); with ``"global"`` it sizes the whole
            train set.
        model: surrogate classifier -- ``"linear_svc"`` (default,
            :class:`~sklearn.svm.LinearSVC`), ``"logistic"``
            (:class:`~sklearn.linear_model.LogisticRegression`), or ``"rbf_svc"``
            (kernel SVM, :class:`~sklearn.svm.SVC` with an RBF kernel and
            ``gamma="scale"``). The linear models are fast; ``"rbf_svc"`` captures
            non-linear boundaries -- and finds genuinely harder splits on
            non-linearly-separable embeddings -- but costs O(n^2)+ per fold, so
            reserve it for smaller n. Each is standardized via a per-fold
            :class:`~sklearn.preprocessing.StandardScaler`.
        score: hardness measure. ``"confidence"`` (default) uses the top-1 minus
            top-2 margin and works for both models; ``"entropy"`` uses predictive
            entropy and requires probabilities (``model="logistic"`` only).
        stratify: ``"per_class"`` (default; per-class test quota -> label-balanced
            test) or ``"global"`` (purest-hard; may unbalance classes).
        cv: number of cross-validation folds, clamped to the smallest class
            count. Each fold supplies the out-of-fold margins for its held-out
            samples.
        random_state: seeds the fold shuffling and the classifier; the split is
            deterministic (ties broken by index).

    Returns:
        train_indices: ndarray of confident, class-interior sample indices.
        test_indices: ndarray of near-boundary, model-confusable indices.

    Raises:
        ValueError: on a ``y``/embeddings length mismatch, fewer than two
            distinct classes, an unknown ``model``/``score``/``stratify``, an
            ``entropy`` score with ``linear_svc``, or a class with fewer than two
            samples (too small to hold out).

    References:
        Uncertainty / margin sampling from active learning (Lewis & Gale, 1994;
        Scheffer et al., 2001; Settles, 2009). Like :func:`class_boundary_split`,
        a label-aware variant of the hard-split motivation of Soegaard et al.
        (2021), "We Need to Talk About Random Splits"
        (https://aclanthology.org/2021.eacl-main.156), but using a *learned*
        boundary rather than embedding geometry.
    """
    from sklearn.linear_model import LogisticRegression
    from sklearn.model_selection import StratifiedKFold
    from sklearn.pipeline import make_pipeline
    from sklearn.preprocessing import StandardScaler
    from sklearn.svm import SVC, LinearSVC

    if model not in ("linear_svc", "logistic", "rbf_svc"):
        raise ValueError(
            f"model must be 'linear_svc', 'logistic', or 'rbf_svc', got {model!r}"
        )
    if score not in ("confidence", "entropy"):
        raise ValueError(f"score must be 'confidence' or 'entropy', got {score!r}")
    if stratify not in ("per_class", "global"):
        raise ValueError(f"stratify must be 'per_class' or 'global', got {stratify!r}")
    if score == "entropy" and model != "logistic":
        raise ValueError("score='entropy' needs predict_proba; use model='logistic'.")

    embeddings = validate_split_inputs(embeddings, train_size)
    y = np.asarray(y)
    n_samples = len(embeddings)
    if len(y) != n_samples:
        raise ValueError(f"y has length {len(y)} but embeddings has {n_samples} rows")

    classes, counts = np.unique(y, return_counts=True)
    if len(classes) < 2:
        raise ValueError(
            f"need at least 2 distinct classes for a boundary split, got {len(classes)}"
        )
    if counts.min() < 2:
        raise ValueError(
            "every class needs >= 2 samples for out-of-fold margins; smallest "
            f"class has {int(counts.min())}."
        )

    n_folds = min(cv, int(counts.min()))

    def make_clf():
        if model == "linear_svc":
            clf = LinearSVC(random_state=random_state)
        elif model == "rbf_svc":
            clf = SVC(kernel="rbf", decision_function_shape="ovr",
                      random_state=random_state)
        else:
            clf = LogisticRegression(max_iter=1000, random_state=random_state)
        return make_pipeline(StandardScaler(), clf)

    # Out-of-fold hardness: higher == closer to the boundary == harder.
    hardness = np.empty(n_samples, dtype=float)
    skf = StratifiedKFold(n_splits=n_folds, shuffle=True, random_state=random_state)
    for fit_idx, hold_idx in skf.split(embeddings, y):
        pipe = make_clf()
        pipe.fit(embeddings[fit_idx], y[fit_idx])
        Xh = embeddings[hold_idx]

        if score == "entropy":
            P = pipe.predict_proba(Xh)
            hardness[hold_idx] = -np.sum(P * np.log(P + 1e-12), axis=1)
            continue

        if model == "logistic":
            margins = pipe.predict_proba(Xh)
        else:  # linear_svc or rbf_svc
            D = pipe.decision_function(Xh)
            if D.ndim == 1:  # binary: distance to the single boundary
                hardness[hold_idx] = -np.abs(D)
                continue
            if model == "linear_svc":
                # Normalize each OvR column to a geometric distance by its own
                # ||w_k|| so the cross-class top1 - top2 is comparable. The RBF
                # SVC has no coef_; its OvR decision_function is already
                # vote-aggregated across classes, so it is used as-is.
                D = D / np.linalg.norm(pipe[-1].coef_, axis=1)
            margins = D

        top2 = np.sort(margins, axis=1)[:, -2:]
        hardness[hold_idx] = -(top2[:, 1] - top2[:, 0])

    # Route the hardest samples to test (stable sort -> ties broken by index).
    if stratify == "global":
        n_test = n_samples - resolve_n_train(n_samples, train_size)
        test_idx = np.argsort(-hardness, kind="stable")[:n_test]
        test_set = set(test_idx.tolist())
    else:
        test_set = set()
        for k in classes:
            members = np.flatnonzero(y == k)
            n_train_k = min(resolve_n_train(len(members), train_size), len(members))
            n_test_k = len(members) - n_train_k
            if n_test_k <= 0:
                continue
            order = np.argsort(-hardness[members], kind="stable")
            test_set.update(members[order[:n_test_k]].tolist())

    train = [i for i in range(n_samples) if i not in test_set]
    return as_index_array(train), as_index_array(sorted(test_set))

distance_adversarial_split

distance_adversarial_split(embeddings: ArrayLike, train_size: float | int = 0.7, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split based on individual sample distance from centroid.

Samples closest to centroid go to train, furthest go to test. Unlike cluster-based methods, this operates on individual samples.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
metric	str	distance metric	'euclidean'
random_state	int	accepted for API consistency; this split is deterministic	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def distance_adversarial_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split based on individual sample distance from centroid.

    Samples closest to centroid go to train, furthest go to test.
    Unlike cluster-based methods, this operates on individual samples.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        metric: distance metric
        random_state: accepted for API consistency; this split is deterministic

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    centroid = compute_centroid(embeddings)

    # Compute distance from centroid for each sample
    distances = np.linalg.norm(embeddings - centroid, axis=1)

    # Sort by distance (closest first)
    sorted_indices = np.argsort(distances)

    # Split
    n_train = resolve_n_train(len(embeddings), train_size)
    return as_index_array(sorted_indices[:n_train]), as_index_array(
        sorted_indices[n_train:]
    )

density_adversarial_split

density_adversarial_split(embeddings: ArrayLike, train_size: float | int = 0.7, metric: str = 'euclidean', k: int = 10, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split based on local density.

Samples in dense regions go to train, isolated samples go to test.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
metric	str	distance metric	'euclidean'
k	int	number of neighbors for density estimation	10
random_state	int	accepted for API consistency; this split is deterministic	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def density_adversarial_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    k: int = 10,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split based on local density.

    Samples in dense regions go to train, isolated samples go to test.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        metric: distance metric
        k: number of neighbors for density estimation
        random_state: accepted for API consistency; this split is deterministic

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    # TODO: Replace full pairwise matrix with NearestNeighbors(k) to reduce
    # memory from O(n²) to O(nk). Only k distances per point are needed here.
    distances = cdist(embeddings, embeddings, metric=metric)
    np.fill_diagonal(distances, np.inf)

    # Compute density as inverse of mean distance to k nearest neighbors
    k = min(k, len(embeddings) - 1)
    knn_distances = np.sort(distances, axis=1)[:, :k]
    densities = 1.0 / (knn_distances.mean(axis=1) + 1e-10)

    # Sort by density (highest density first -> train)
    sorted_indices = np.argsort(-densities)

    # Split
    n_train = resolve_n_train(len(embeddings), train_size)
    return as_index_array(sorted_indices[:n_train]), as_index_array(
        sorted_indices[n_train:]
    )

outlier_adversarial_split

outlier_adversarial_split(embeddings: ArrayLike, train_size: float | int = 0.7, contamination: float = 0.1, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split using outlier detection.

Normal samples go to train, outliers go to test.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
contamination	float	expected proportion of outliers	0.1
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def outlier_adversarial_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    contamination: float = 0.1,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split using outlier detection.

    Normal samples go to train, outliers go to test.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        contamination: expected proportion of outliers
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    from sklearn.ensemble import IsolationForest

    embeddings = validate_split_inputs(embeddings, train_size)

    detector = IsolationForest(
        contamination=contamination,
        random_state=random_state
    )
    detector.fit(embeddings)

    # Get outlier scores (more negative = more normal)
    scores = detector.score_samples(embeddings)

    # Sort by score (most normal first -> train)
    sorted_indices = np.argsort(-scores)

    # Split
    n_train = resolve_n_train(len(embeddings), train_size)
    return as_index_array(sorted_indices[:n_train]), as_index_array(
        sorted_indices[n_train:]
    )

min_cut_split

min_cut_split(embeddings: ArrayLike, train_size: float | int = 0.7, similarity_threshold: float | None = None, metric: str = 'euclidean', method: str = 'spectral', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split using graph min-cut.

Builds a similarity graph and finds a partition that minimizes edges (similarity) crossing the train/test boundary.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
similarity_threshold	float | None	only connect samples above this kernel similarity (default: median similarity)	None
metric	str	distance metric	'euclidean'
method	str	'spectral' (fast, approximate) or 'stoer_wagner' (exact, slow)	'spectral'
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def min_cut_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    similarity_threshold: float | None = None,
    metric: str = "euclidean",
    method: str = "spectral",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split using graph min-cut.

    Builds a similarity graph and finds a partition that minimizes
    edges (similarity) crossing the train/test boundary.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        similarity_threshold: only connect samples above this kernel similarity
                              (default: median similarity)
        metric: distance metric
        method: 'spectral' (fast, approximate) or 'stoer_wagner' (exact, slow)
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)

    if n_samples < 3:
        # Too few samples for meaningful split
        n_train = resolve_n_train(n_samples, train_size)
        return as_index_array(range(n_train)), as_index_array(
            range(n_train, n_samples)
        )

    # TODO: Build sparse similarity graph directly via kneighbors_graph instead
    # of materializing the full O(n²) distance matrix.
    distances = cdist(embeddings, embeddings, metric=metric)

    # Convert distance to similarity using Gaussian kernel
    sigma = np.median(distances[distances > 0])
    if sigma == 0:
        sigma = 1.0
    similarities = np.exp(-distances**2 / (2 * sigma**2))
    np.fill_diagonal(similarities, 0)  # No self-loops

    # Threshold to create sparse graph
    if similarity_threshold is None:
        nonzero_sims = similarities[similarities > 0]
        if len(nonzero_sims) > 0:
            similarity_threshold = np.median(nonzero_sims)
        else:
            similarity_threshold = 0.0

    similarities[similarities < similarity_threshold] = 0

    if method == "spectral":
        # Spectral partitioning using Fiedler vector
        # (eigenvector corresponding to 2nd smallest eigenvalue of Laplacian)

        L = laplacian(csr_matrix(similarities), normed=True)

        try:
            # Get 2 smallest eigenvalues/vectors
            eigenvalues, eigenvectors = eigsh(L, k=2, which='SM', tol=1e-6)

            # Fiedler vector (2nd eigenvector)
            fiedler = eigenvectors[:, 1]

            # Partition by Fiedler vector values
            # Sort and split to achieve desired train_size
            sorted_indices = np.argsort(fiedler)

        except Exception:
            # Fallback to random if eigendecomposition fails
            rng = check_random_state(random_state)
            sorted_indices = np.arange(n_samples)
            rng.shuffle(sorted_indices)

        n_train = resolve_n_train(n_samples, train_size)
        train_indices = sorted_indices[:n_train]
        test_indices = sorted_indices[n_train:]

    elif method == "stoer_wagner":
        # Exact min-cut using Stoer-Wagner algorithm (slower)
        try:
            import networkx as nx

            # Build weighted graph
            G = nx.Graph()
            G.add_nodes_from(range(n_samples))

            for i in range(n_samples):
                for j in range(i + 1, n_samples):
                    if similarities[i, j] > 0:
                        G.add_edge(i, j, weight=similarities[i, j])

            # Check if graph is connected
            if not nx.is_connected(G):
                # Find connected components and use largest
                components = list(nx.connected_components(G))
                components.sort(key=len, reverse=True)

                # Assign smaller components to test, sample from largest for train
                train_indices = []
                test_indices = []

                for comp in components[1:]:
                    test_indices.extend(comp)

                main_component = list(components[0])
                rng = check_random_state(random_state)
                rng.shuffle(main_component)

                n_train_needed = resolve_n_train(n_samples, train_size) - len(
                    train_indices
                )
                n_train_needed = max(0, min(n_train_needed, len(main_component)))

                train_indices.extend(main_component[:n_train_needed])
                test_indices.extend(main_component[n_train_needed:])

            else:
                # Run Stoer-Wagner
                cut_value, partition = nx.stoer_wagner(G)
                set1, set2 = list(partition[0]), list(partition[1])

                # Adjust to match train_size
                n_train = resolve_n_train(n_samples, train_size)

                if len(set1) >= n_train:
                    train_indices = set1[:n_train]
                    test_indices = set1[n_train:] + set2
                else:
                    train_indices = set1 + set2[:n_train - len(set1)]
                    test_indices = set2[n_train - len(set1):]

        except ImportError as err:
            raise ImportError(
                "networkx is required for method='stoer_wagner'"
            ) from err

    else:
        raise ValueError(f"Unknown method: {method}. Use 'spectral' or 'stoer_wagner'.")

    return as_index_array(train_indices), as_index_array(test_indices)

normalized_cut_split

normalized_cut_split(embeddings: ArrayLike, train_size: float | int = 0.7, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split using normalized graph cut.

Normalized cut balances the cut value with partition sizes, avoiding trivially small partitions.

NCut(A,B) = cut(A,B)/vol(A) + cut(A,B)/vol(B)

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
metric	str	distance metric	'euclidean'
random_state	int	accepted for API consistency; this split is deterministic	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/adversarial.py

def normalized_cut_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split using normalized graph cut.

    Normalized cut balances the cut value with partition sizes,
    avoiding trivially small partitions.

    NCut(A,B) = cut(A,B)/vol(A) + cut(A,B)/vol(B)

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        metric: distance metric
        random_state: accepted for API consistency; this split is deterministic

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)

    if n_samples < 3:
        n_train = resolve_n_train(n_samples, train_size)
        return as_index_array(range(n_train)), as_index_array(
            range(n_train, n_samples)
        )

    # TODO: Build sparse similarity graph directly via kneighbors_graph instead
    # of materializing the full O(n²) distance matrix.
    distances = cdist(embeddings, embeddings, metric=metric)
    sigma = np.median(distances[distances > 0])
    if sigma == 0:
        sigma = 1.0
    W = np.exp(-distances**2 / (2 * sigma**2))
    np.fill_diagonal(W, 0)

    # Compute symmetric normalized Laplacian (D^{-1/2} W D^{-1/2}).
    D_inv_sqrt = np.diag(1.0 / np.sqrt(W.sum(axis=1) + 1e-10))

    L_norm = np.eye(n_samples) - D_inv_sqrt @ W @ D_inv_sqrt

    # Compute eigenvectors
    eigenvalues, eigenvectors = np.linalg.eigh(L_norm)

    # Use second eigenvector (Fiedler vector)
    fiedler = eigenvectors[:, 1]

    # Sort by Fiedler vector to get partition
    sorted_indices = np.argsort(fiedler)

    n_train = resolve_n_train(n_samples, train_size)
    return as_index_array(sorted_indices[:n_train]), as_index_array(
        sorted_indices[n_train:]
    )

wasserstein_adversarial_split

wasserstein_adversarial_split(embeddings: ArrayLike, train_size: float | int = 0.7, leaf_size: int = 40, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Adversarial split via Wasserstein nearest-neighbors.

Treats each embedding row as a 1-D distribution and selects, as the test set, the n_test points nearest (in 1-D Wasserstein / earth-mover distance) to a random anchor — yielding a tight, hard-to-generalize-to test neighborhood.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
leaf_size	int	BallTree leaf size (higher = slower but less memory)	40
random_state	int	for reproducibility (anchor sampling)	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

References

Adapted from the adversarial split of Søgaard, Ebert, Bastings & Filippova (2021), "We Need to Talk About Random Splits," EACL, which constructs hard splits by (approximately) maximizing the Wasserstein distance between train and test. https://aclanthology.org/2021.eacl-main.156

Source code in splytters/adversarial.py

def wasserstein_adversarial_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    leaf_size: int = 40,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split via Wasserstein nearest-neighbors.

    Treats each embedding row as a 1-D distribution and selects, as the test
    set, the ``n_test`` points nearest (in 1-D Wasserstein / earth-mover
    distance) to a random anchor — yielding a tight, hard-to-generalize-to test
    neighborhood.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        leaf_size: BallTree leaf size (higher = slower but less memory)
        random_state: for reproducibility (anchor sampling)

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set

    References:
        Adapted from the adversarial split of Søgaard, Ebert, Bastings &
        Filippova (2021), "We Need to Talk About Random Splits," EACL, which
        constructs hard splits by (approximately) maximizing the Wasserstein
        distance between train and test. https://aclanthology.org/2021.eacl-main.156
    """
    from scipy.stats import wasserstein_distance
    from sklearn.neighbors import NearestNeighbors

    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    n_test = n_samples - resolve_n_train(n_samples, train_size)
    n_test = max(1, min(n_test, n_samples - 1))
    rng = check_random_state(random_state)

    tree = NearestNeighbors(
        n_neighbors=n_test,
        algorithm="ball_tree",
        leaf_size=leaf_size,
        metric=wasserstein_distance,
    )
    tree.fit(embeddings)

    # Sample a random anchor in the per-dimension bounding box (works for any
    # real-valued embeddings, unlike the original integer-only formulation).
    anchor = rng.uniform(
        low=embeddings.min(axis=0), high=embeddings.max(axis=0)
    ).reshape(1, -1)

    test_indices = tree.kneighbors(anchor, return_distance=False)[0]
    test_set = {int(i) for i in test_indices}
    train_indices = [i for i in range(n_samples) if i not in test_set]
    return as_index_array(train_indices), as_index_array(test_indices)

mmd_maximized_split

mmd_maximized_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_iterations: int = 500, kernel: str = 'rbf', gamma: float | None = None, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Maximize Maximum Mean Discrepancy (MMD) between train and test.

The adversarial dual of :func:splytters.mmd_minimized_split: instead of matching the two distributions, it pushes them as far apart as possible in kernel/MMD terms, yielding a hard, domain-shifted held-out set. Motivated by model selection under domain shift -- a maximally-shifted validation set tends to select more robust hyperparameters.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_iterations	int	number of optimization iterations	500
kernel	str	kernel type ('rbf' or 'linear')	'rbf'
gamma	float | None	RBF kernel parameter (default: 1/n_features)	None
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

References

Implements the objective of Napoli & White (2025), "Clustering-Based Validation Splits for Model Selection under Domain Shift," TMLR (https://openreview.net/forum?id=Q692C0WtiD): the train/validation split should maximize the MMD between the two sets. NOTE: this captures that objective via swap optimization (like mmd_minimized_split); it does not implement their full method -- a constrained kernel k-means (max-MMD is shown equivalent to kernel k-means with k=2) solved by linear programming to preserve class/group balance, with convergence guarantees and Nyström scaling. For class-balanced clustering splits, see cluster_split(strategy="subset_sum") and :func:cluster_kfold.

Source code in splytters/adversarial.py

def mmd_maximized_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_iterations: int = 500,
    kernel: str = "rbf",
    gamma: float | None = None,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Maximize Maximum Mean Discrepancy (MMD) between train and test.

    The adversarial dual of :func:`splytters.mmd_minimized_split`: instead of
    matching the two distributions, it pushes them as far apart as possible in
    kernel/MMD terms, yielding a hard, domain-shifted held-out set. Motivated by
    model selection under domain shift -- a maximally-shifted validation set
    tends to select more robust hyperparameters.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of optimization iterations
        kernel: kernel type ('rbf' or 'linear')
        gamma: RBF kernel parameter (default: 1/n_features)
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set

    References:
        Implements the *objective* of Napoli & White (2025), "Clustering-Based
        Validation Splits for Model Selection under Domain Shift," TMLR
        (https://openreview.net/forum?id=Q692C0WtiD): the train/validation split
        should maximize the MMD between the two sets. NOTE: this captures that
        objective via swap optimization (like ``mmd_minimized_split``); it does
        not implement their full method -- a constrained kernel k-means (max-MMD
        is shown equivalent to kernel k-means with k=2) solved by linear
        programming to preserve class/group balance, with convergence guarantees
        and Nyström scaling. For class-balanced clustering splits, see
        ``cluster_split(strategy="subset_sum")`` and :func:`cluster_kfold`.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_dims = embeddings.shape[1]

    _gamma = gamma if gamma is not None else 1.0 / n_dims

    def _rbf_kernel(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
        return np.exp(-_gamma * cdist(X, Y, metric="sqeuclidean"))

    def _linear_kernel(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
        return X @ Y.T

    kernel_fn = _rbf_kernel if kernel == "rbf" else _linear_kernel

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        train_data, test_data = X[train], X[test]
        K_tt = kernel_fn(train_data, train_data)
        K_ss = kernel_fn(test_data, test_data)
        K_ts = kernel_fn(train_data, test_data)
        m, n = len(train_data), len(test_data)
        return (K_tt.sum() / (m * m) +
                K_ss.sum() / (n * n) -
                2 * K_ts.sum() / (m * n))

    # minimize=False -> push the two distributions apart (adversarial dual).
    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state, minimize=False
    )

get_cluster_info

get_cluster_info(embeddings: ArrayLike, train_indices: ArrayLike, test_indices: ArrayLike, n_clusters: int = 10, random_state: int = 42) -> dict[str, Any]

Utility to analyze cluster distribution across train/test split.

Returns:

Type	Description
dict[str, Any]	dict with cluster statistics including leakage info

Source code in splytters/adversarial.py

def get_cluster_info(
    embeddings: ArrayLike,
    train_indices: ArrayLike,
    test_indices: ArrayLike,
    n_clusters: int = 10,
    random_state: int = 42,
) -> dict[str, Any]:
    """
    Utility to analyze cluster distribution across train/test split.

    Returns:
        dict with cluster statistics including leakage info
    """
    embeddings = np.asarray(embeddings)

    clusterer = KMeans(n_clusters=n_clusters, random_state=random_state, n_init="auto")
    labels = clusterer.fit_predict(embeddings)

    train_set = {int(i) for i in train_indices}
    test_set = {int(i) for i in test_indices}

    cluster_stats = {}
    for cluster_id in range(n_clusters):
        cluster_indices = [i for i, l in enumerate(labels) if l == cluster_id]
        in_train = sum(1 for i in cluster_indices if i in train_set)
        in_test = sum(1 for i in cluster_indices if i in test_set)
        cluster_stats[cluster_id] = {
            "total": len(cluster_indices),
            "in_train": in_train,
            "in_test": in_test,
            "leakage": min(in_train, in_test) > 0
        }

    total_leaking = sum(1 for s in cluster_stats.values() if s["leakage"])

    return {
        "cluster_stats": cluster_stats,
        "clusters_with_leakage": total_leaking,
        "leakage_ratio": total_leaking / n_clusters
    }

Overlap splitters

overlap

High-overlap splitting algorithms that maximize train-test similarity.

These methods create "easy" evaluation sets where test samples are similar to training samples, useful for sanity checks and debugging.

cluster_leak_split

cluster_leak_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_clusters: int = 10, random_state: int = 42, **cluster_kwargs: Any) -> tuple[np.ndarray, np.ndarray]

Split clusters across train/test to maximize similarity.

Instead of assigning entire clusters to one set (adversarial), this splits each cluster proportionally between train and test, ensuring similar samples appear in both sets.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_clusters	int	number of clusters	10
random_state	int	for reproducibility	42
**cluster_kwargs	Any	passed to KMeans	{}

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def cluster_leak_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_clusters: int = 10,
    random_state: int = 42,
    **cluster_kwargs: Any,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Split clusters across train/test to maximize similarity.

    Instead of assigning entire clusters to one set (adversarial),
    this splits each cluster proportionally between train and test,
    ensuring similar samples appear in both sets.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_clusters: number of clusters
        random_state: for reproducibility
        **cluster_kwargs: passed to KMeans

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    rng = check_random_state(random_state)

    labels, cluster_to_indices, _ = cluster_embeddings(
        embeddings, n_clusters, "kmeans", random_state, **cluster_kwargs
    )

    n_samples = len(embeddings)
    n_train_total = resolve_n_train(n_samples, train_size)

    # Apportion the global train target across clusters (largest-remainder) so
    # the realized split matches train_size instead of undershooting from
    # per-cluster truncation.
    clusters = [np.array(idx) for idx in cluster_to_indices.values()]
    per_cluster_train = apportion_train([len(c) for c in clusters], n_train_total)

    train_indices = []
    test_indices = []
    for indices, k in zip(clusters, per_cluster_train, strict=True):
        rng.shuffle(indices)
        train_indices.extend(indices[:k].tolist())
        test_indices.extend(indices[k:].tolist())

    return as_index_array(train_indices), as_index_array(test_indices)

neighbor_coverage_split

neighbor_coverage_split(embeddings: ArrayLike, train_size: float | int = 0.7, k: int = 5, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Ensure each test sample has k similar samples in train.

Iteratively assigns samples to test only if they have enough similar samples already in train, maximizing coverage.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
k	int	minimum number of similar train samples for each test sample	5
metric	str	distance metric	'euclidean'
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def neighbor_coverage_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    k: int = 5,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Ensure each test sample has k similar samples in train.

    Iteratively assigns samples to test only if they have enough
    similar samples already in train, maximizing coverage.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        k: minimum number of similar train samples for each test sample
        metric: distance metric
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    # TODO: Replace full pairwise matrix with chunked or BallTree.query_radius
    # computation to reduce memory from O(n²). Only threshold-based neighbor
    # lookups are needed here.
    distances = cdist(embeddings, embeddings, metric=metric)
    np.fill_diagonal(distances, np.inf)

    # Compute similarity threshold (median distance)
    finite_dists = distances[distances < np.inf]
    threshold = np.median(finite_dists)

    # Start with random train set
    all_indices = np.arange(n_samples)
    rng.shuffle(all_indices)

    n_train = resolve_n_train(n_samples, train_size)
    train_set = set(all_indices[:n_train])
    remaining = list(all_indices[n_train:])

    # Iteratively improve: swap samples to maximize coverage
    test_indices = []

    for idx in remaining:
        # Count similar samples in train
        similar_in_train = sum(
            1 for t_idx in train_set
            if distances[idx, t_idx] <= threshold
        )

        if similar_in_train >= k:
            test_indices.append(idx)
        else:
            # Find a train sample to swap
            # Prefer train samples that have many similar train neighbors
            train_list = list(train_set)
            swap_candidate = None
            max_redundancy = -1

            for t_idx in train_list:
                redundancy = sum(
                    1 for other in train_set
                    if other != t_idx and distances[t_idx, other] <= threshold
                )
                if redundancy > max_redundancy:
                    max_redundancy = redundancy
                    swap_candidate = t_idx

            if swap_candidate is not None and max_redundancy > k:
                train_set.remove(swap_candidate)
                train_set.add(idx)
                test_indices.append(swap_candidate)
            else:
                test_indices.append(idx)

    return as_index_array(train_set), as_index_array(test_indices)

centroid_matched_split

centroid_matched_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_iterations: int = 100, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Minimize distance between train and test centroids.

Uses iterative optimization to find a split where the centroids of train and test sets are as close as possible.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_iterations	int	number of swap iterations	100
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def centroid_matched_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_iterations: int = 100,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Minimize distance between train and test centroids.

    Uses iterative optimization to find a split where the centroids
    of train and test sets are as close as possible.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of swap iterations
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        train_centroid = X[train].mean(axis=0)
        test_centroid = X[test].mean(axis=0)
        return float(np.linalg.norm(train_centroid - test_centroid))

    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state
    )

stratified_similarity_split

stratified_similarity_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_bins: int = 10, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Stratify by distance from centroid, ensuring similar distribution in both sets.

Bins samples by their distance from the centroid and samples proportionally from each bin for both train and test.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_bins	int	number of distance bins for stratification	10
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def stratified_similarity_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_bins: int = 10,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Stratify by distance from centroid, ensuring similar distribution in both sets.

    Bins samples by their distance from the centroid and samples
    proportionally from each bin for both train and test.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_bins: number of distance bins for stratification
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    rng = check_random_state(random_state)

    # Compute distances from centroid
    centroid = compute_centroid(embeddings)
    distances = np.linalg.norm(embeddings - centroid, axis=1)

    # Bin samples by distance
    bin_edges = np.percentile(distances, np.linspace(0, 100, n_bins + 1))
    bin_assignments = np.digitize(distances, bin_edges[1:-1])

    n_train_total = resolve_n_train(len(embeddings), train_size)

    # Collect non-empty bins, then apportion the global train target across them
    # (largest-remainder) so the split hits train_size and never empties the
    # test set, even when bins are singletons (n_bins >= n_samples).
    bins = [b for b in (np.where(bin_assignments == i)[0] for i in range(n_bins))
            if len(b) > 0]
    per_bin_train = apportion_train([len(b) for b in bins], n_train_total)

    train_indices = []
    test_indices = []
    for bin_samples, k in zip(bins, per_bin_train, strict=True):
        rng.shuffle(bin_samples)
        train_indices.extend(bin_samples[:k].tolist())
        test_indices.extend(bin_samples[k:].tolist())

    return as_index_array(train_indices), as_index_array(test_indices)

nearest_neighbor_split

nearest_neighbor_split(embeddings: ArrayLike, train_size: float | int = 0.7, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

For each test sample, ensure its nearest neighbor is in train.

Greedily builds test set by moving points whose nearest neighbor is already confirmed in train. Uses sklearn's NearestNeighbors which auto-selects the best algorithm (kd-tree, ball tree, or brute force) based on dimensionality.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
metric	str	distance metric	'euclidean'
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def nearest_neighbor_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    For each test sample, ensure its nearest neighbor is in train.

    Greedily builds test set by moving points whose nearest neighbor
    is already confirmed in train. Uses sklearn's NearestNeighbors
    which auto-selects the best algorithm (kd-tree, ball tree, or
    brute force) based on dimensionality.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        metric: distance metric
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    from sklearn.neighbors import NearestNeighbors

    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    n_test = n_samples - resolve_n_train(n_samples, train_size)
    rng = check_random_state(random_state)

    # Find each point's nearest neighbor
    nn_model = NearestNeighbors(n_neighbors=2, metric=metric, algorithm="auto")
    nn_model.fit(embeddings)
    # k=2 because the first neighbor is the point itself when querying the
    # same dataset — but fit/kneighbors on the same data excludes self only
    # if we use radius; instead just grab second neighbor
    neighbors = nn_model.kneighbors(embeddings, return_distance=False)[:, 1]

    # Start with all points in train, then greedily move points to test.
    # A point can become test only if its NN is staying in train.
    in_train = np.ones(n_samples, dtype=bool)

    # Process in random order so the result isn't biased by index ordering
    order = np.arange(n_samples)
    rng.shuffle(order)

    test_indices = []
    for idx in order:
        if len(test_indices) >= n_test:
            break
        nn = neighbors[idx]
        if in_train[nn]:
            in_train[idx] = False
            test_indices.append(idx)

    train_indices = np.where(in_train)[0]
    return as_index_array(train_indices), as_index_array(test_indices)

duplicate_spread_split

duplicate_spread_split(embeddings: ArrayLike, train_size: float | int = 0.7, similarity_threshold: float | None = None, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Intentionally put near-duplicates in both train and test.

Identifies clusters of near-duplicates and ensures at least one sample from each cluster appears in both sets.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
similarity_threshold	float | None	distance threshold for near-duplicates (default: 10th percentile of distances)	None
metric	str	distance metric	'euclidean'
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def duplicate_spread_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    similarity_threshold: float | None = None,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Intentionally put near-duplicates in both train and test.

    Identifies clusters of near-duplicates and ensures at least one
    sample from each cluster appears in both sets.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        similarity_threshold: distance threshold for near-duplicates
                              (default: 10th percentile of distances)
        metric: distance metric
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    rng = check_random_state(random_state)

    # TODO: Replace full pairwise matrix with BallTree.query_radius to find
    # near-duplicates without materializing O(n²) distances.
    distances = cdist(embeddings, embeddings, metric=metric)
    np.fill_diagonal(distances, np.inf)

    # Set threshold
    if similarity_threshold is None:
        finite_dists = distances[distances < np.inf]
        similarity_threshold = np.percentile(finite_dists, 10)

    # Find near-duplicate groups using connected components
    from scipy.sparse import csr_matrix
    from scipy.sparse.csgraph import connected_components

    adjacency = (distances <= similarity_threshold).astype(int)
    n_components, labels = connected_components(csr_matrix(adjacency))

    # Group samples by component
    component_to_indices = defaultdict(list)
    for idx, label in enumerate(labels):
        component_to_indices[label].append(idx)

    train_fraction = resolve_n_train(len(embeddings), train_size) / len(embeddings)

    train_indices = []
    test_indices = []

    # For each component, split proportionally (ensuring both sets get samples)
    for indices in component_to_indices.values():
        indices = np.array(indices)
        rng.shuffle(indices)

        if len(indices) == 1:
            # Single sample: assign to train
            train_indices.extend(indices.tolist())
        else:
            # Split ensuring both sets get at least one
            n_train = max(1, int(len(indices) * train_fraction))
            n_train = min(n_train, len(indices) - 1)  # Leave at least 1 for test

            train_indices.extend(indices[:n_train].tolist())
            test_indices.extend(indices[n_train:].tolist())

    return as_index_array(train_indices), as_index_array(test_indices)

max_coverage_split

max_coverage_split(embeddings: ArrayLike, train_size: float | int = 0.7, radius: float | None = None, metric: str = 'euclidean', random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Maximize the coverage of test set by train set.

Coverage = fraction of test samples with at least one train sample within radius.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
radius	float | None	distance threshold for coverage (default: median distance)	None
metric	str	distance metric	'euclidean'
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/overlap.py

def max_coverage_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    radius: float | None = None,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Maximize the coverage of test set by train set.

    Coverage = fraction of test samples with at least one train
    sample within radius.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        radius: distance threshold for coverage (default: median distance)
        metric: distance metric
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    # TODO: Replace full pairwise matrix with BallTree.query_radius to check
    # coverage without materializing O(n²) distances.
    distances = cdist(embeddings, embeddings, metric=metric)
    np.fill_diagonal(distances, np.inf)

    # Set radius
    if radius is None:
        finite_dists = distances[distances < np.inf]
        radius = np.median(finite_dists)

    # Start with random split
    all_indices = np.arange(n_samples)
    rng.shuffle(all_indices)

    n_train = resolve_n_train(n_samples, train_size)
    train_set = set(all_indices[:n_train])
    test_set = set(all_indices[n_train:])

    def compute_coverage():
        covered = 0
        for t_idx in test_set:
            for tr_idx in train_set:
                if distances[t_idx, tr_idx] <= radius:
                    covered += 1
                    break
        return covered / len(test_set) if test_set else 1.0

    current_coverage = compute_coverage()

    # Greedy optimization
    improved = True
    max_iterations = n_samples * 2

    for _ in range(max_iterations):
        if not improved:
            break
        improved = False

        for test_idx in list(test_set):
            # Check if covered
            is_covered = any(
                distances[test_idx, tr_idx] <= radius
                for tr_idx in train_set
            )

            if not is_covered:
                # Find best train sample to swap
                best_swap = None
                best_coverage = current_coverage

                for train_idx in list(train_set):
                    # Simulate swap
                    train_set.remove(train_idx)
                    train_set.add(test_idx)
                    test_set.remove(test_idx)
                    test_set.add(train_idx)

                    new_coverage = compute_coverage()

                    # Revert
                    train_set.remove(test_idx)
                    train_set.add(train_idx)
                    test_set.remove(train_idx)
                    test_set.add(test_idx)

                    if new_coverage > best_coverage:
                        best_coverage = new_coverage
                        best_swap = train_idx

                if best_swap is not None:
                    train_set.remove(best_swap)
                    train_set.add(test_idx)
                    test_set.remove(test_idx)
                    test_set.add(best_swap)
                    current_coverage = best_coverage
                    improved = True
                    break

    return as_index_array(train_set), as_index_array(test_set)

Balanced splitters

balanced

Balanced splitting algorithms that match distributions between train/test.

These methods create splits where train and test have similar statistical properties, useful for fair evaluation without distribution shift.

distribution_matched_split

distribution_matched_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_iterations: int = 1000, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Minimize distribution divergence between train and test.

Uses iterative optimization to match the marginal distributions of each feature dimension.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_iterations	int	number of optimization iterations	1000
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/balanced.py

def distribution_matched_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_iterations: int = 1000,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Minimize distribution divergence between train and test.

    Uses iterative optimization to match the marginal distributions
    of each feature dimension.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of optimization iterations
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        """Mean KS statistic across all dimensions."""
        train_data, test_data = X[train], X[test]
        return float(np.mean([
            ks_2samp(train_data[:, d], test_data[:, d]).statistic
            for d in range(X.shape[1])
        ]))

    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state
    )

moment_matched_split

moment_matched_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_iterations: int = 1000, match_variance: bool = True, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Match mean (and optionally variance) between train and test.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_iterations	int	number of optimization iterations	1000
match_variance	bool	if True, also match variance	True
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/balanced.py

def moment_matched_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_iterations: int = 1000,
    match_variance: bool = True,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Match mean (and optionally variance) between train and test.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of optimization iterations
        match_variance: if True, also match variance
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        train_data, test_data = X[train], X[test]
        mean_diff = np.linalg.norm(train_data.mean(axis=0) - test_data.mean(axis=0))
        if match_variance:
            var_diff = np.linalg.norm(train_data.var(axis=0) - test_data.var(axis=0))
            return mean_diff + var_diff
        return mean_diff

    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state
    )

histogram_matched_split

histogram_matched_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_bins: int = 10, n_iterations: int = 1000, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Match feature histograms between train and test.

Minimizes the sum of histogram differences across all dimensions.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_bins	int	number of histogram bins per dimension	10
n_iterations	int	number of optimization iterations	1000
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/balanced.py

def histogram_matched_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_bins: int = 10,
    n_iterations: int = 1000,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Match feature histograms between train and test.

    Minimizes the sum of histogram differences across all dimensions.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_bins: number of histogram bins per dimension
        n_iterations: number of optimization iterations
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_dims = embeddings.shape[1]

    # Precompute bin edges for each dimension. Skip degenerate dimensions whose
    # percentile edges collapse (constant/near-constant features) — a zero-width
    # bin makes density=True divide by zero and poison the score with NaN, which
    # would silently stall the optimizer (every NaN comparison is False).
    bin_edges = [
        np.percentile(embeddings[:, d], np.linspace(0, 100, n_bins + 1))
        for d in range(n_dims)
    ]
    valid_dims = [d for d in range(n_dims) if np.all(np.diff(bin_edges[d]) > 0)]

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        train_data, test_data = X[train], X[test]
        total_diff = 0.0
        for d in valid_dims:
            train_hist, _ = np.histogram(train_data[:, d], bins=bin_edges[d], density=True)
            test_hist, _ = np.histogram(test_data[:, d], bins=bin_edges[d], density=True)
            total_diff += np.sum(np.abs(train_hist - test_hist))
        return total_diff

    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state
    )

stratified_random_split

stratified_random_split(embeddings: ArrayLike, y: ArrayLike, train_size: float | int = 0.7, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Standard stratified split maintaining label proportions.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
y	ArrayLike	array of labels for stratification (aligned with sklearn's split(X, y))	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/balanced.py

def stratified_random_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    train_size: float | int = 0.7,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Standard stratified split maintaining label proportions.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        y: array of labels for stratification (aligned with sklearn's ``split(X, y)``)
        train_size: fraction in (0, 1) or absolute count for the training set
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    from sklearn.model_selection import train_test_split

    embeddings = validate_split_inputs(embeddings, train_size)
    indices = np.arange(len(embeddings))

    train_indices, test_indices = train_test_split(
        indices,
        train_size=train_size,
        stratify=y,
        random_state=random_state
    )

    return as_index_array(train_indices), as_index_array(test_indices)

density_balanced_split

density_balanced_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_bins: int = 10, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Balance local density distribution between train and test.

Bins samples by local density and samples proportionally from each bin for both sets.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_bins	int	number of density bins	10
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/balanced.py

def density_balanced_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_bins: int = 10,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Balance local density distribution between train and test.

    Bins samples by local density and samples proportionally
    from each bin for both sets.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_bins: number of density bins
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    # TODO: Replace full pairwise matrix with NearestNeighbors(k) to reduce
    # memory from O(n²) to O(nk). Only k distances per point are needed here.
    distances = cdist(embeddings, embeddings, metric="euclidean")
    np.fill_diagonal(distances, np.inf)

    # Compute density (inverse of mean distance to 10 nearest neighbors)
    k = min(10, n_samples - 1)
    knn_distances = np.sort(distances, axis=1)[:, :k]
    densities = 1.0 / (knn_distances.mean(axis=1) + 1e-10)

    # Bin by density
    bin_edges = np.percentile(densities, np.linspace(0, 100, n_bins + 1))
    bin_assignments = np.digitize(densities, bin_edges[1:-1])

    n_train_total = resolve_n_train(n_samples, train_size)

    # Apportion the global train target across non-empty density bins
    # (largest-remainder) so the split hits train_size and never empties the
    # test set, even when bins are singletons (n_bins >= n_samples).
    bins = [b for b in (np.where(bin_assignments == i)[0] for i in range(n_bins))
            if len(b) > 0]
    per_bin_train = apportion_train([len(b) for b in bins], n_train_total)

    train_indices = []
    test_indices = []
    for bin_samples, k in zip(bins, per_bin_train, strict=True):
        rng.shuffle(bin_samples)
        train_indices.extend(bin_samples[:k].tolist())
        test_indices.extend(bin_samples[k:].tolist())

    return as_index_array(train_indices), as_index_array(test_indices)

mmd_minimized_split

mmd_minimized_split(embeddings: ArrayLike, train_size: float | int = 0.7, n_iterations: int = 500, kernel: str = 'rbf', gamma: float | None = None, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Minimize Maximum Mean Discrepancy between train and test.

MMD is a kernel-based measure of distribution difference. Lower MMD indicates more similar distributions.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
n_iterations	int	number of optimization iterations	500
kernel	str	kernel type ('rbf' or 'linear')	'rbf'
gamma	float | None	RBF kernel parameter (default: 1/n_features)	None
random_state	int	for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

References

The adversarial dual -- maximizing train/validation MMD for robust model selection under domain shift -- is studied by Napoli & White (2025), "Clustering-Based Validation Splits for Model Selection under Domain Shift," TMLR (https://openreview.net/forum?id=Q692C0WtiD). See :func:splytters.mmd_maximized_split.

Source code in splytters/balanced.py

def mmd_minimized_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    n_iterations: int = 500,
    kernel: str = "rbf",
    gamma: float | None = None,
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Minimize Maximum Mean Discrepancy between train and test.

    MMD is a kernel-based measure of distribution difference.
    Lower MMD indicates more similar distributions.

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of optimization iterations
        kernel: kernel type ('rbf' or 'linear')
        gamma: RBF kernel parameter (default: 1/n_features)
        random_state: for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set

    References:
        The adversarial dual -- *maximizing* train/validation MMD for robust
        model selection under domain shift -- is studied by Napoli & White
        (2025), "Clustering-Based Validation Splits for Model Selection under
        Domain Shift," TMLR (https://openreview.net/forum?id=Q692C0WtiD). See
        :func:`splytters.mmd_maximized_split`.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_dims = embeddings.shape[1]

    _gamma = gamma if gamma is not None else 1.0 / n_dims

    def _rbf_kernel(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
        return np.exp(-_gamma * cdist(X, Y, metric="sqeuclidean"))

    def _linear_kernel(X: np.ndarray, Y: np.ndarray) -> np.ndarray:
        return X @ Y.T

    kernel_fn = _rbf_kernel if kernel == "rbf" else _linear_kernel

    def score_fn(X: np.ndarray, train: list[int], test: list[int]) -> float:
        train_data, test_data = X[train], X[test]
        K_tt = kernel_fn(train_data, train_data)
        K_ss = kernel_fn(test_data, test_data)
        K_ts = kernel_fn(train_data, test_data)
        m, n = len(train_data), len(test_data)
        return (K_tt.sum() / (m * m) +
                K_ss.sum() / (n * n) -
                2 * K_ts.sum() / (m * n))

    return optimized_split(
        embeddings, train_size, n_iterations, score_fn, random_state
    )

Curriculum splits

curriculum

Curriculum / ordering-driven splitting.

Unlike the embedding-based splitters (adversarial / overlap / balanced), these take an explicit sort order — typically a :mod:splytters.sorters ranking — and partition each class along it. The classic use is a "train on easy, test on hard" curriculum split: sort by a difficulty metric, then take the first train_size fraction of each class as train.

sorted_stratified_split

sorted_stratified_split(order: Sequence[Any], y: ArrayLike, train_size: float | int = 0.7, *, largest_first: bool = False) -> tuple[np.ndarray, np.ndarray]

Per-class curriculum split driven by a sort order.

Within each class, samples are visited in the given order and the first train_size fraction (or absolute count) go to train, the rest to test. Pair it with a :mod:splytters.sorters ranking to split along an interpretable difficulty axis — e.g. "train on easy, test on hard".

Parameters:

Name	Type	Description	Default
order	Sequence[Any]	the sort order over samples — either a flat sequence of sample indices, or a sorter result (sequence of (index, score) pairs, as returned by e.g. readability_score). Sorters rank ascending, so the lowest-scoring samples are the train-preferred "first" ones.	required
y	ArrayLike	class labels aligned to the original samples (length n_samples). The split is performed independently within each class.	required
train_size	float | int	fraction in (0, 1) or absolute count, applied per class.	0.7
largest_first	bool	if True, reverse order so the highest-scoring samples are train-preferred.	False

Returns:

Type	Description
tuple[ndarray, ndarray]	(train_indices, test_indices) integer ndarrays, sorted ascending.

Raises:

Type	Description
ValueError	if order and y differ in length, order is not a permutation of range(len(y)), or train_size is out of range.

Source code in splytters/curriculum.py

def sorted_stratified_split(
    order: Sequence[Any],
    y: ArrayLike,
    train_size: float | int = 0.7,
    *,
    largest_first: bool = False,
) -> tuple[np.ndarray, np.ndarray]:
    """Per-class curriculum split driven by a sort order.

    Within each class, samples are visited in the given ``order`` and the first
    ``train_size`` fraction (or absolute count) go to train, the rest to test.
    Pair it with a :mod:`splytters.sorters` ranking to split along an
    interpretable difficulty axis — e.g. "train on easy, test on hard".

    Args:
        order: the sort order over samples — either a flat sequence of sample
            indices, or a sorter result (sequence of ``(index, score)`` pairs,
            as returned by e.g. ``readability_score``). Sorters rank ascending,
            so the lowest-scoring samples are the train-preferred "first" ones.
        y: class labels aligned to the original samples (length n_samples). The
            split is performed independently within each class.
        train_size: fraction in (0, 1) or absolute count, applied *per class*.
        largest_first: if True, reverse ``order`` so the highest-scoring samples
            are train-preferred.

    Returns:
        ``(train_indices, test_indices)`` integer ndarrays, sorted ascending.

    Raises:
        ValueError: if ``order`` and ``y`` differ in length, ``order`` is not a
            permutation of ``range(len(y))``, or ``train_size`` is out of range.
    """
    ordered_idx = _to_index_order(order)
    y = np.asarray(y)

    if len(ordered_idx) != len(y):
        raise ValueError(f"order has {len(ordered_idx)} entries but y has {len(y)}")
    if len(ordered_idx) and (
        ordered_idx.min() < 0
        or ordered_idx.max() >= len(y)
        or len(np.unique(ordered_idx)) != len(ordered_idx)
    ):
        raise ValueError("order must be a permutation of range(len(y))")
    if isinstance(train_size, float) and not 0.0 < train_size < 1.0:
        raise ValueError(f"train_size fraction must be in (0, 1), got {train_size}")

    if largest_first:
        ordered_idx = ordered_idx[::-1]

    labels = y[ordered_idx]  # class labels in sorted order
    train: list[int] = []
    test: list[int] = []
    for c in np.unique(y):
        class_order = ordered_idx[labels == c]  # this class, in sorted order
        n_train = max(0, min(resolve_n_train(len(class_order), train_size), len(class_order)))
        train.extend(class_order[:n_train].tolist())
        test.extend(class_order[n_train:].tolist())

    return as_index_array(sorted(train)), as_index_array(sorted(test))

scikit-learn compatibility

sklearn_api

scikit-learn compatibility layer for splytters.

Provides:

• :class:SplytterSplit — a single-split cross-validator wrapping any splytter splitting function, usable as cv= in :func:sklearn.model_selection.cross_validate and :class:~sklearn.model_selection.GridSearchCV.
• *_train_test_split convenience functions mirroring :func:sklearn.model_selection.train_test_split for drop-in swaps.

The functional splitters in :mod:splitters remain the source of truth; these wrappers are purely additive.

SplytterSplit

Bases: BaseCrossValidator

A scikit-learn cross-validator backed by a splytter splitting function.

Produces n_splits train/test partitions (one by default, like :class:~sklearn.model_selection.PredefinedSplit). Each partition is the output of splitter(embeddings, train_size=..., random_state=...).

Parameters:

Name	Type	Description	Default
splitter	Splitter	any splytter splitter taking (embeddings, train_size=, random_state=) and returning (train_idx, test_idx) integer ndarrays. Defaults to :func:~splytters.adversarial.cluster_split.	cluster_split
embeddings	Any	array-like of shape (n_samples, n_features), optional. Embeddings to split on. If None, the X passed to :meth:split is used (i.e. the estimator is assumed to consume the embeddings).	None
train_size	float | int	fraction in (0, 1) or absolute count for the training set (default 0.7).	0.7
random_state	int | None	base seed (default 42); the i-th repeat uses random_state + i.	42
n_splits	int	number of (repeated) partitions to yield (default 1).	1
**splitter_kwargs	Any	extra keyword arguments forwarded to splitter.	{}

Examples:

>>> from sklearn.linear_model import LogisticRegression
>>> from sklearn.model_selection import cross_validate
>>> cv = SplytterSplit(embeddings=X)
>>> cross_validate(LogisticRegression(), X, y, cv=cv)

Source code in splytters/sklearn_api.py

class SplytterSplit(BaseCrossValidator):
    """A scikit-learn cross-validator backed by a splytter splitting function.

    Produces ``n_splits`` train/test partitions (one by default, like
    :class:`~sklearn.model_selection.PredefinedSplit`). Each partition is the
    output of ``splitter(embeddings, train_size=..., random_state=...)``.

    Args:
        splitter: any splytter splitter taking
            ``(embeddings, train_size=, random_state=)`` and returning
            ``(train_idx, test_idx)`` integer ndarrays. Defaults to
            :func:`~splytters.adversarial.cluster_split`.
        embeddings: array-like of shape (n_samples, n_features), optional.
            Embeddings to split on. If ``None``, the ``X`` passed to
            :meth:`split` is used (i.e. the estimator is assumed to consume the
            embeddings).
        train_size: fraction in (0, 1) or absolute count for the training set
            (default 0.7).
        random_state: base seed (default 42); the i-th repeat uses
            ``random_state + i``.
        n_splits: number of (repeated) partitions to yield (default 1).
        **splitter_kwargs: extra keyword arguments forwarded to ``splitter``.

    Examples:
        >>> from sklearn.linear_model import LogisticRegression
        >>> from sklearn.model_selection import cross_validate
        >>> cv = SplytterSplit(embeddings=X)            # doctest: +SKIP
        >>> cross_validate(LogisticRegression(), X, y, cv=cv)   # doctest: +SKIP
    """

    def __init__(
        self,
        splitter: Splitter = cluster_split,
        *,
        embeddings: Any = None,
        train_size: float | int = 0.7,
        random_state: int | None = 42,
        n_splits: int = 1,
        **splitter_kwargs: Any,
    ) -> None:
        self.splitter = splitter
        self.embeddings = embeddings
        self.train_size = train_size
        self.random_state = random_state
        self.n_splits = n_splits
        self.splitter_kwargs = splitter_kwargs

    def get_n_splits(self, X=None, y=None, groups=None) -> int:
        return self.n_splits

    def split(self, X, y=None, groups=None) -> Iterator[tuple[np.ndarray, np.ndarray]]:
        emb = self.embeddings if self.embeddings is not None else X
        for i in range(self.n_splits):
            seed = _derive_seed(self.random_state, i)
            train_idx, test_idx = self.splitter(
                emb,
                train_size=self.train_size,
                random_state=seed,
                **self.splitter_kwargs,
            )
            yield np.asarray(train_idx, dtype=np.intp), np.asarray(
                test_idx, dtype=np.intp
            )

    def _iter_test_indices(self, X=None, y=None, groups=None):
        # Splytter splits are always a full partition, so test indices fully
        # determine the fold; this keeps the base-class helpers consistent.
        for _, test_idx in self.split(X, y, groups):
            yield test_idx

splytter_train_test_split

splytter_train_test_split(*arrays: Any, splitter: Splitter = cluster_split, embeddings: Any = None, train_size: float | int = 0.7, random_state: int | None = 42, **splitter_kwargs: Any) -> list[Any] | tuple[np.ndarray, np.ndarray]

Split arrays into train/test subsets using a splytter splitter.

Mirrors :func:sklearn.model_selection.train_test_split: for inputs a, b returns [a_train, a_test, b_train, b_test].

Parameters:

Name	Type	Description	Default
*arrays	Any	array-likes (numpy, list, pandas, etc.) to split, all of the same length (n_samples).	()
splitter	Splitter	splytter splitting function. Defaults to :func:~splytters.adversarial.cluster_split.	cluster_split
embeddings	Any	array-like to compute the split on, optional. Defaults to the first array.	None
train_size	float | int	fraction in (0, 1) or absolute count for the training set.	0.7
random_state	int | None	seed forwarded to splitter.	42
**splitter_kwargs	Any	extra keyword arguments forwarded to splitter.	{}

Returns:

Type	Description
list[Any] | tuple[ndarray, ndarray]	len(arrays) * 2 outputs (train/test pairs in order). If no arrays
list[Any] | tuple[ndarray, ndarray]	are passed, returns (train_idx, test_idx).

Source code in splytters/sklearn_api.py

def splytter_train_test_split(
    *arrays: Any,
    splitter: Splitter = cluster_split,
    embeddings: Any = None,
    train_size: float | int = 0.7,
    random_state: int | None = 42,
    **splitter_kwargs: Any,
) -> list[Any] | tuple[np.ndarray, np.ndarray]:
    """Split arrays into train/test subsets using a splytter splitter.

    Mirrors :func:`sklearn.model_selection.train_test_split`: for inputs
    ``a, b`` returns ``[a_train, a_test, b_train, b_test]``.

    Args:
        *arrays: array-likes (numpy, list, pandas, etc.) to split, all of the
            same length (n_samples).
        splitter: splytter splitting function. Defaults to
            :func:`~splytters.adversarial.cluster_split`.
        embeddings: array-like to compute the split on, optional. Defaults to
            the first array.
        train_size: fraction in (0, 1) or absolute count for the training set.
        random_state: seed forwarded to ``splitter``.
        **splitter_kwargs: extra keyword arguments forwarded to ``splitter``.

    Returns:
        ``len(arrays) * 2`` outputs (train/test pairs in order). If no arrays
        are passed, returns ``(train_idx, test_idx)``.
    """
    if embeddings is None:
        if not arrays:
            raise ValueError("Pass at least one array or `embeddings=`.")
        embeddings = arrays[0]

    # All inputs must align with the indices the splitter returns; mismatched
    # lengths otherwise IndexError (longer) or silently drop rows (shorter).
    n = _num_samples(embeddings)
    for i, a in enumerate(arrays):
        if _num_samples(a) != n:
            raise ValueError(
                f"All arrays must have the same length as embeddings ({n}); "
                f"array {i} has length {_num_samples(a)}."
            )

    train_idx, test_idx = splitter(
        embeddings, train_size=train_size, random_state=random_state, **splitter_kwargs
    )

    if not arrays:
        return train_idx, test_idx

    out: list[Any] = []
    for a in arrays:
        out.append(_safe_indexing(a, train_idx))
        out.append(_safe_indexing(a, test_idx))
    return out

Framework interop

interop

Framework interop helpers for splytters.

Thin adapters that take an embedding-based splitter and return native objects for popular ecosystems. Heavy dependencies (pandas, torch, datasets) are imported lazily inside each function, so importing this module never requires them — only calling the relevant helper does.

split_dataframe

split_dataframe(df: Any, embeddings: Any, splitter: Splitter = cluster_split, *, train_size: float | int = 0.7, random_state: int | None = 42, **splitter_kwargs: Any) -> tuple[Any, Any]

Split a pandas DataFrame into (train_df, test_df) by position.

Parameters:

Name	Type	Description	Default
df	Any	pandas DataFrame of n_samples rows.	required
embeddings	Any	array-like of shape (n_samples, n_features).	required
splitter	Splitter	splytter splitting function.	cluster_split
train_size	float | int	fraction in (0, 1) or absolute count for the training set.	0.7
random_state	int | None	seed forwarded to splitter.	42
**splitter_kwargs	Any	extra keyword arguments forwarded to splitter.	{}

Returns:

Type	Description
Any	(train_df, test_df) selected via positional .iloc (original
Any	index labels are preserved).

Source code in splytters/interop.py

def split_dataframe(
    df: Any,
    embeddings: Any,
    splitter: Splitter = cluster_split,
    *,
    train_size: float | int = 0.7,
    random_state: int | None = 42,
    **splitter_kwargs: Any,
) -> tuple[Any, Any]:
    """Split a pandas DataFrame into ``(train_df, test_df)`` by position.

    Args:
        df: pandas ``DataFrame`` of n_samples rows.
        embeddings: array-like of shape (n_samples, n_features).
        splitter: splytter splitting function.
        train_size: fraction in (0, 1) or absolute count for the training set.
        random_state: seed forwarded to ``splitter``.
        **splitter_kwargs: extra keyword arguments forwarded to ``splitter``.

    Returns:
        ``(train_df, test_df)`` selected via positional ``.iloc`` (original
        index labels are preserved).
    """
    if len(df) != len(embeddings):
        raise ValueError(
            f"df has {len(df)} rows but embeddings has {len(embeddings)}"
        )
    train_idx, test_idx = splitter(
        embeddings, train_size=train_size, random_state=random_state, **splitter_kwargs
    )
    return df.iloc[train_idx], df.iloc[test_idx]

to_torch_subsets

to_torch_subsets(dataset: Any, embeddings: Any, splitter: Splitter = cluster_split, *, train_size: float | int = 0.7, random_state: int | None = 42, **splitter_kwargs: Any) -> tuple[Any, Any]

Split a torch Dataset into (train_subset, test_subset).

Returns two :class:torch.utils.data.Subset views over dataset.

Source code in splytters/interop.py

def to_torch_subsets(
    dataset: Any,
    embeddings: Any,
    splitter: Splitter = cluster_split,
    *,
    train_size: float | int = 0.7,
    random_state: int | None = 42,
    **splitter_kwargs: Any,
) -> tuple[Any, Any]:
    """Split a torch ``Dataset`` into ``(train_subset, test_subset)``.

    Returns two :class:`torch.utils.data.Subset` views over ``dataset``.
    """
    from torch.utils.data import Subset

    if len(dataset) != len(embeddings):
        raise ValueError(
            f"dataset has {len(dataset)} items but embeddings has {len(embeddings)}"
        )
    train_idx, test_idx = splitter(
        embeddings, train_size=train_size, random_state=random_state, **splitter_kwargs
    )
    return (
        Subset(dataset, train_idx.tolist()),
        Subset(dataset, test_idx.tolist()),
    )

split_dataset

split_dataset(ds: Any, embeddings: Any, splitter: Splitter = cluster_split, *, train_size: float | int = 0.7, random_state: int | None = 42, **splitter_kwargs: Any) -> Any

Split a HuggingFace datasets.Dataset into a DatasetDict.

Returns DatasetDict({"train": ds.select(train_idx), "test": ds.select(test_idx)}).

Source code in splytters/interop.py

def split_dataset(
    ds: Any,
    embeddings: Any,
    splitter: Splitter = cluster_split,
    *,
    train_size: float | int = 0.7,
    random_state: int | None = 42,
    **splitter_kwargs: Any,
) -> Any:
    """Split a HuggingFace ``datasets.Dataset`` into a ``DatasetDict``.

    Returns ``DatasetDict({"train": ds.select(train_idx),
    "test": ds.select(test_idx)})``.
    """
    from datasets import DatasetDict

    if len(ds) != len(embeddings):
        raise ValueError(
            f"dataset has {len(ds)} rows but embeddings has {len(embeddings)}"
        )
    train_idx, test_idx = splitter(
        embeddings, train_size=train_size, random_state=random_state, **splitter_kwargs
    )
    return DatasetDict(
        {
            "train": ds.select(train_idx.tolist()),
            "test": ds.select(test_idx.tolist()),
        }
    )

Split-quality reporting

report

Split-quality reporting.

Quantifies how adversarial, overlapping, or balanced a produced train/test split actually is — so users can compare splitters and so papers can report a single, interpretable table. Builds on :func:splytters.utils.compute_split_similarity and :func:splytters.adversarial.get_cluster_info, adding distribution-distance metrics (MMD, energy distance, mean 1-D Wasserstein / KS) and label balance.

split_report

split_report(embeddings: ArrayLike, train_indices: ArrayLike, test_indices: ArrayLike, y: ArrayLike | None = None, *, texts: list[str] | None = None, metric: str = 'euclidean', n_clusters: int = 10, max_samples: int | None = 2000, random_state: int | None = 42) -> dict[str, float]

Summarize how similar/dissimilar a train/test split is.

Larger distribution distances (mmd_rbf, energy_distance, wasserstein_mean, ks_mean) and mean_cross_distance indicate a more adversarial (harder) split; values near a random split indicate a balanced split; near-zero indicates overlap.

Alongside these train↔test distance metrics, train_diversity and test_diversity report the within-split spread (mean distance to each side's own centroid) — useful for spotting when a split concentrates one side (e.g. an adversarial test drawn from a single outlier cluster).

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, n_features).	required
train_indices	ArrayLike	integer index array for the training split.	required
test_indices	ArrayLike	integer index array for the test split.	required
y	ArrayLike | None	array-like of labels, optional. If given, adds label_distribution_shift (total-variation distance between train/test class proportions; 0 = identical balance).	None
texts	list[str] | None	per-sample raw strings, optional and aligned with embeddings. If given, adds train_text_diversity / test_text_diversity (mean pairwise n-gram distance within each side), subsampled to max_samples for the O(n²) pairwise computation.	None
metric	str	distance metric (default "euclidean").	'euclidean'
n_clusters	int	clusters used for the leakage statistic (default 10).	10
max_samples	int | None	cap per side for the O(n²) distribution metrics (subsampled for speed); None uses all samples (default 2000).	2000
random_state	int | None	seed for subsampling (default 42).	42

Returns:

Type	Description
dict[str, float]	A dict of structural, geometric, distributional, diversity, and
dict[str, float]	(optional) label metrics.

Source code in splytters/report.py

def split_report(
    embeddings: ArrayLike,
    train_indices: ArrayLike,
    test_indices: ArrayLike,
    y: ArrayLike | None = None,
    *,
    texts: list[str] | None = None,
    metric: str = "euclidean",
    n_clusters: int = 10,
    max_samples: int | None = 2000,
    random_state: int | None = 42,
) -> dict[str, float]:
    """Summarize how similar/dissimilar a train/test split is.

    Larger distribution distances (``mmd_rbf``, ``energy_distance``,
    ``wasserstein_mean``, ``ks_mean``) and ``mean_cross_distance`` indicate a
    *more adversarial* (harder) split; values near a random split indicate a
    *balanced* split; near-zero indicates *overlap*.

    Alongside these train↔test *distance* metrics, ``train_diversity`` and
    ``test_diversity`` report the within-split *spread* (mean distance to each
    side's own centroid) — useful for spotting when a split concentrates one
    side (e.g. an adversarial test drawn from a single outlier cluster).

    Args:
        embeddings: array-like of shape (n_samples, n_features).
        train_indices: integer index array for the training split.
        test_indices: integer index array for the test split.
        y: array-like of labels, optional. If given, adds
            ``label_distribution_shift`` (total-variation distance between
            train/test class proportions; 0 = identical balance).
        texts: per-sample raw strings, optional and aligned with ``embeddings``.
            If given, adds ``train_text_diversity`` / ``test_text_diversity``
            (mean pairwise n-gram distance within each side), subsampled to
            ``max_samples`` for the O(n²) pairwise computation.
        metric: distance metric (default ``"euclidean"``).
        n_clusters: clusters used for the leakage statistic (default 10).
        max_samples: cap per side for the O(n²) distribution metrics
            (subsampled for speed); ``None`` uses all samples (default 2000).
        random_state: seed for subsampling (default 42).

    Returns:
        A dict of structural, geometric, distributional, diversity, and
        (optional) label metrics.
    """
    X = validate_split_inputs(embeddings, 0.5)  # reuse finite/2-D validation
    train_indices = np.asarray(train_indices, dtype=np.intp)
    test_indices = np.asarray(test_indices, dtype=np.intp)
    rng = check_random_state(random_state)

    n_train, n_test = len(train_indices), len(test_indices)
    # A report compares two non-empty sides; an empty side otherwise divides by
    # zero here or crashes downstream in compute_split_similarity.
    if n_train == 0 or n_test == 0:
        raise ValueError(
            f"split_report requires non-empty train and test splits "
            f"(got n_train={n_train}, n_test={n_test})."
        )
    report: dict[str, float] = {
        "n_train": float(n_train),
        "n_test": float(n_test),
        "train_fraction": float(n_train / (n_train + n_test)),
    }

    # Geometric similarity (centroid distance, nearest-train distance, coverage).
    report.update(compute_split_similarity(X, train_indices, test_indices, metric))

    # Cluster leakage.
    info = get_cluster_info(
        X, train_indices, test_indices, n_clusters=n_clusters,
        random_state=random_state if isinstance(random_state, int) else 42,
    )
    report["cluster_leakage_ratio"] = float(info["leakage_ratio"])

    # Distribution distances on (optionally subsampled) embeddings.
    A = _subsample(X, train_indices, max_samples, rng)
    B = _subsample(X, test_indices, max_samples, rng)
    gamma = 1.0 / X.shape[1]
    report["mmd_rbf"] = _rbf_mmd(A, B, gamma)
    report["energy_distance"] = _energy_distance(A, B)
    report["wasserstein_mean"] = float(
        np.mean([wasserstein_distance(A[:, d], B[:, d]) for d in range(X.shape[1])])
    )
    report["ks_mean"] = float(
        np.mean([ks_2samp(A[:, d], B[:, d]).statistic for d in range(X.shape[1])])
    )

    # Within-split diversity (spread): how varied is each side on its own.
    report["train_diversity"] = mean_dist(X[train_indices])
    report["test_diversity"] = mean_dist(X[test_indices])

    # Optional text diversity: mean pairwise n-gram distance within each side.
    if texts is not None:
        texts = list(texts)
        seed = random_state if isinstance(random_state, int) else 42
        report["train_text_diversity"] = diversity_text(
            [texts[i] for i in train_indices],
            sample_size=max_samples, random_state=seed,
        )
        report["test_text_diversity"] = diversity_text(
            [texts[i] for i in test_indices],
            sample_size=max_samples, random_state=seed,
        )

    # Label balance: total-variation distance between class proportions.
    if y is not None:
        y = np.asarray(y)
        y_train, y_test = y[train_indices], y[test_indices]
        tv = 0.0
        for c in np.unique(y):
            p_train = (y_train == c).mean() if n_train else 0.0
            p_test = (y_test == c).mean() if n_test else 0.0
            tv += abs(p_train - p_test)
        report["label_distribution_shift"] = float(0.5 * tv)

    return report

compare_splitters

compare_splitters(embeddings: ArrayLike, splitters: dict[str, Splitter], y: ArrayLike | None = None, *, train_size: float | int = 0.7, random_state: int | None = 42, **report_kwargs: Any) -> dict[str, dict[str, float]]

Run several splitters and return {name: split_report(...)}.

Handy for generating a comparison table (e.g. random vs adversarial vs balanced) in one call.

Source code in splytters/report.py

def compare_splitters(
    embeddings: ArrayLike,
    splitters: dict[str, Splitter],
    y: ArrayLike | None = None,
    *,
    train_size: float | int = 0.7,
    random_state: int | None = 42,
    **report_kwargs: Any,
) -> dict[str, dict[str, float]]:
    """Run several splitters and return ``{name: split_report(...)}``.

    Handy for generating a comparison table (e.g. random vs adversarial vs
    balanced) in one call.
    """
    out: dict[str, dict[str, float]] = {}
    for name, splitter in splitters.items():
        train_idx, test_idx = splitter(
            embeddings, train_size=train_size, random_state=random_state
        )
        out[name] = split_report(
            embeddings, train_idx, test_idx, y, random_state=random_state,
            **report_kwargs,
        )
    return out

Sorters

Ranking functions that order samples by interpretable difficulty/quality metrics, grouped by modality. Pair a ranking with sorted_stratified_split for a curriculum split. Each modality imports its heavy dependency lazily, so importing a sorter module needs no optional extra.

Embedding sorters

embedding_sorters

Sorting algorithms for adversarial dataset partitioning based on embeddings.

These functions rank samples by embedding-based criteria (distance to centroid, nearest neighbors, density, outlier scores) to enable train-test splits that maximize dissimilarity.

dist_euclidean

dist_euclidean(u: ArrayLike, v: ArrayLike) -> float

Compute Euclidean distance between two vectors.

Source code in splytters/sorters/embedding_sorters.py

def dist_euclidean(u: ArrayLike, v: ArrayLike) -> float:
    """Compute Euclidean distance between two vectors."""
    return _dist_euclidean(u, v)

distance_to_mean

distance_to_mean(embeddings: ndarray, distance: Callable[[ndarray, ndarray], float] = dist_euclidean) -> list[tuple[int, float]]

Sort samples by distance from the dataset centroid.

Samples closest to the centroid (most "typical") appear first. Useful for adversarial splits: assign nearby samples to train, distant samples to test.

Parameters:

Name	Type	Description	Default
embeddings	ndarray	np.array of shape (n_samples, embedding_dim)	required
distance	Callable[[ndarray, ndarray], float]	distance function taking two vectors, default Euclidean	dist_euclidean

Returns:

Type	Description
list[tuple[int, float]]	List of (index, distance) tuples sorted by distance ascending.

Source code in splytters/sorters/embedding_sorters.py

def distance_to_mean(
    embeddings: np.ndarray,
    distance: Callable[[np.ndarray, np.ndarray], float] = dist_euclidean,
) -> list[tuple[int, float]]:
    """
    Sort samples by distance from the dataset centroid.

    Samples closest to the centroid (most "typical") appear first.
    Useful for adversarial splits: assign nearby samples to train,
    distant samples to test.

    Args:
        embeddings: np.array of shape (n_samples, embedding_dim)
        distance: distance function taking two vectors, default Euclidean

    Returns:
        List of (index, distance) tuples sorted by distance ascending.
    """
    (n, d) = embeddings.shape
    centroid = embeddings.mean(0)
    distances = []
    for i in range(n):
        dist = distance(embeddings[i], centroid)
        distances.append((i, dist))
    distances.sort(key=lambda p: p[1])
    return distances

distance_to_nearest_neighbor

distance_to_nearest_neighbor(embeddings: ArrayLike, metric: str = 'euclidean') -> list[tuple[int, float]]

Sort samples by distance to their nearest neighbor.

Samples in dense regions (close to neighbors) appear first. Isolated samples (far from all neighbors) appear last.

Useful for adversarial splits: train on samples in dense clusters, test on isolated/unique samples.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	np.array of shape (n_samples, embedding_dim)	required
metric	str	distance metric for cdist (default 'euclidean')	'euclidean'

Returns:

Type	Description
list[tuple[int, float]]	List of (index, distance) tuples sorted by nearest neighbor distance ascending.

Source code in splytters/sorters/embedding_sorters.py

def distance_to_nearest_neighbor(
    embeddings: ArrayLike, metric: str = "euclidean"
) -> list[tuple[int, float]]:
    """
    Sort samples by distance to their nearest neighbor.

    Samples in dense regions (close to neighbors) appear first.
    Isolated samples (far from all neighbors) appear last.

    Useful for adversarial splits: train on samples in dense clusters,
    test on isolated/unique samples.

    Args:
        embeddings: np.array of shape (n_samples, embedding_dim)
        metric: distance metric for cdist (default 'euclidean')

    Returns:
        List of (index, distance) tuples sorted by nearest neighbor distance ascending.
    """
    embeddings = np.asarray(embeddings)
    n = len(embeddings)
    if n < 2:
        # No neighbor to measure against.
        return [(i, float("inf")) for i in range(n)]

    # O(n·k) memory via NearestNeighbors instead of a full O(n²) matrix.
    # Querying the fit data puts each point's self at column 0 (distance 0),
    # so column 1 is exactly the 1-nearest-neighbor distance.
    from sklearn.neighbors import NearestNeighbors

    nn = NearestNeighbors(n_neighbors=2, metric=metric)
    nn.fit(embeddings)
    dists, _ = nn.kneighbors(embeddings)
    min_distances = dists[:, 1]

    scores = [(i, float(min_distances[i])) for i in range(n)]
    scores.sort(key=lambda p: p[1])

    return scores

local_density

local_density(embeddings: ArrayLike, radius: float | None = None, metric: str = 'euclidean', low_first: bool = True) -> list[tuple[int, int]]

Sort samples by local density (number of neighbors within radius).

Samples in dense regions have many neighbors; isolated samples have few.

Useful for adversarial splits: train on high-density regions, test on sparse/low-density regions.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	np.array of shape (n_samples, embedding_dim)	required
radius	float | None	distance threshold for counting neighbors. If None, uses median pairwise distance.	None
metric	str	distance metric for cdist (default 'euclidean')	'euclidean'
low_first	bool	if True, sparse/isolated samples first; if False, dense samples first	True

Returns:

Type	Description
list[tuple[int, int]]	List of (index, neighbor_count) tuples sorted by density.

Source code in splytters/sorters/embedding_sorters.py

def local_density(
    embeddings: ArrayLike,
    radius: float | None = None,
    metric: str = "euclidean",
    low_first: bool = True,
) -> list[tuple[int, int]]:
    """
    Sort samples by local density (number of neighbors within radius).

    Samples in dense regions have many neighbors; isolated samples have few.

    Useful for adversarial splits: train on high-density regions,
    test on sparse/low-density regions.

    Args:
        embeddings: np.array of shape (n_samples, embedding_dim)
        radius: distance threshold for counting neighbors.
                If None, uses median pairwise distance.
        metric: distance metric for cdist (default 'euclidean')
        low_first: if True, sparse/isolated samples first;
                   if False, dense samples first

    Returns:
        List of (index, neighbor_count) tuples sorted by density.
    """
    embeddings = np.asarray(embeddings)

    # TODO: Replace full pairwise matrix with BallTree.query_radius to count
    # neighbors without materializing O(n²) distances.
    pairwise_dist = cdist(embeddings, embeddings, metric=metric)

    # Set diagonal to infinity so we don't count self
    np.fill_diagonal(pairwise_dist, np.inf)

    # Auto-select radius if not provided
    if radius is None:
        # Use median of all pairwise distances as default radius
        finite_dists = pairwise_dist[pairwise_dist < np.inf]
        radius = np.median(finite_dists)

    # Count neighbors within radius for each sample
    neighbor_counts = (pairwise_dist <= radius).sum(axis=1)

    # Create sorted list of (index, count)
    scores = [(i, int(neighbor_counts[i])) for i in range(len(embeddings))]
    scores.sort(key=lambda p: p[1], reverse=not low_first)

    return scores

outlier_score

outlier_score(embeddings: ArrayLike, method: str = 'isolation_forest', low_first: bool = True, **kwargs: Any) -> list[tuple[int, float]]

Sort samples by anomaly/outlier score.

Higher outlier scores indicate samples that are unusual or don't fit the overall data distribution.

Useful for adversarial splits: train on normal/typical samples, test on outliers/anomalies.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	np.array of shape (n_samples, embedding_dim)	required
method	str	outlier detection algorithm, one of: - 'isolation_forest': Isolation Forest (fast, good for high dimensions) - 'lof': Local Outlier Factor (density-based)	'isolation_forest'
low_first	bool	if True, normal/inlier samples first; if False, outliers/anomalies first	True
**kwargs	Any	additional arguments passed to the outlier detector	{}

Returns:

Type	Description
list[tuple[int, float]]	List of (index, outlier_score) tuples sorted by outlier score.
list[tuple[int, float]]	For isolation_forest: more negative = more normal, more positive = more outlier.
list[tuple[int, float]]	For lof: scores > 1 indicate outliers, < 1 indicate inliers.

Source code in splytters/sorters/embedding_sorters.py

def outlier_score(
    embeddings: ArrayLike,
    method: str = "isolation_forest",
    low_first: bool = True,
    **kwargs: Any,
) -> list[tuple[int, float]]:
    """
    Sort samples by anomaly/outlier score.

    Higher outlier scores indicate samples that are unusual or don't fit
    the overall data distribution.

    Useful for adversarial splits: train on normal/typical samples,
    test on outliers/anomalies.

    Args:
        embeddings: np.array of shape (n_samples, embedding_dim)
        method: outlier detection algorithm, one of:
            - 'isolation_forest': Isolation Forest (fast, good for high dimensions)
            - 'lof': Local Outlier Factor (density-based)
        low_first: if True, normal/inlier samples first;
                   if False, outliers/anomalies first
        **kwargs: additional arguments passed to the outlier detector

    Returns:
        List of (index, outlier_score) tuples sorted by outlier score.
        For isolation_forest: more negative = more normal, more positive = more outlier.
        For lof: scores > 1 indicate outliers, < 1 indicate inliers.
    """
    embeddings = np.asarray(embeddings)

    if method == "isolation_forest":
        detector = IsolationForest(random_state=42, **kwargs)
        detector.fit(embeddings)
        # score_samples returns negative scores; more negative = more normal
        # We negate so higher = more outlier
        raw_scores = -detector.score_samples(embeddings)

    elif method == "lof":
        detector = LocalOutlierFactor(novelty=False, **kwargs)
        detector.fit_predict(embeddings)
        # negative_outlier_factor_ is negative; more negative = more outlier
        # We negate so higher = more outlier
        raw_scores = -detector.negative_outlier_factor_

    else:
        raise ValueError(f"Unknown outlier detection method: {method}")

    # Create sorted list of (index, score)
    scores = [(i, raw_scores[i]) for i in range(len(embeddings))]
    scores.sort(key=lambda p: p[1], reverse=not low_first)

    return scores

Text sorters

text_sorters

Sorting algorithms for adversarial text dataset partitioning.

These functions rank text samples by various criteria (length, complexity, perplexity, readability, vocabulary) to enable train-test splits that maximize dissimilarity.

simple_tokenizer

simple_tokenizer(s: str) -> list[str]

Split text on whitespace into a list of tokens.

Source code in splytters/sorters/text_sorters.py

def simple_tokenizer(s: str) -> list[str]:
    """Split text on whitespace into a list of tokens."""
    return s.split()

character_length

character_length(texts: list[str], low_first: bool = True) -> list[tuple[int, int]]

Sort texts by character count.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings	required
low_first	bool	if True, shortest texts first; if False, longest first	True

Returns:

Type	Description
list[tuple[int, int]]	List of (index, character_count) tuples sorted by length.

References

Splitting by length to test generalization (put long texts in test): - Søgaard, Ebert, Bastings & Filippova (2021), "We Need to Talk About Random Splits," EACL. Heuristic splits via a sentence-length threshold. https://aclanthology.org/2021.eacl-main.156 - Varis & Bojar (2021), "Sequence Length is a Domain: Length-based Overfitting in Transformer Models," EMNLP. https://aclanthology.org/2021.emnlp-main.650 - Lake & Baroni (2018), "Generalization without Systematicity," ICML, introduced the SCAN length split (train short, test long) -- a canonical length-generalization test (on output sequence length, synthetic data). https://arxiv.org/abs/1711.00350 - Søgaard et al. trace the idea of testing generalization to longer sequences back to Siegelmann & Sontag (1992), "On the Computational Power of Neural Nets," COLT, pp. 440-449. https://doi.org/10.1145/130385.130432

Source code in splytters/sorters/text_sorters.py

def character_length(texts: list[str], low_first: bool = True) -> list[tuple[int, int]]:
    """
    Sort texts by character count.

    Args:
        texts: list of strings
        low_first: if True, shortest texts first; if False, longest first

    Returns:
        List of (index, character_count) tuples sorted by length.

    References:
        Splitting by length to test generalization (put long texts in test):
          - Søgaard, Ebert, Bastings & Filippova (2021), "We Need to Talk
            About Random Splits," EACL. Heuristic splits via a sentence-length
            threshold. https://aclanthology.org/2021.eacl-main.156
          - Varis & Bojar (2021), "Sequence Length is a Domain: Length-based
            Overfitting in Transformer Models," EMNLP.
            https://aclanthology.org/2021.emnlp-main.650
          - Lake & Baroni (2018), "Generalization without Systematicity," ICML,
            introduced the SCAN length split (train short, test long) -- a
            canonical length-generalization test (on output sequence length,
            synthetic data). https://arxiv.org/abs/1711.00350
          - Søgaard et al. trace the idea of testing generalization to longer
            sequences back to Siegelmann & Sontag (1992), "On the Computational
            Power of Neural Nets," COLT, pp. 440-449.
            https://doi.org/10.1145/130385.130432
    """
    scores = [(i, len(text)) for i, text in enumerate(texts)]
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

tokens_length

tokens_length(texts: list[str], low_first: bool = True, tokenizer: Callable[[str], list[str]] = simple_tokenizer) -> list[tuple[int, int]]

Sort texts by token count.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings	required
low_first	bool	if True, fewest tokens first; if False, most tokens first	True
tokenizer	Callable[[str], list[str]]	function that splits text into tokens, default whitespace split	simple_tokenizer

Returns:

Type	Description
list[tuple[int, int]]	List of (index, token_count) tuples sorted by token count.

References

A token-count variant of length-based splitting; see character_length for the motivation and references (Søgaard et al. 2021, EACL; Varis & Bojar 2021, EMNLP; Lake & Baroni 2018, ICML; Siegelmann & Sontag 1992, COLT).

Source code in splytters/sorters/text_sorters.py

def tokens_length(
    texts: list[str],
    low_first: bool = True,
    tokenizer: Callable[[str], list[str]] = simple_tokenizer,
) -> list[tuple[int, int]]:
    """
    Sort texts by token count.

    Args:
        texts: list of strings
        low_first: if True, fewest tokens first; if False, most tokens first
        tokenizer: function that splits text into tokens, default whitespace split

    Returns:
        List of (index, token_count) tuples sorted by token count.

    References:
        A token-count variant of length-based splitting; see ``character_length``
        for the motivation and references (Søgaard et al. 2021, EACL; Varis &
        Bojar 2021, EMNLP; Lake & Baroni 2018, ICML; Siegelmann & Sontag 1992,
        COLT).
    """
    scores = [(i, len(tokenizer(text))) for i, text in enumerate(texts)]
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

sentence_count

sentence_count(texts: list[str], language: str = 'en', low_first: bool = True) -> list[tuple[int, int]]

Sort texts by number of sentences.

Uses pysbd for robust sentence boundary detection across languages.

Useful for adversarial splits: train on single-sentence texts, test on multi-sentence/complex texts.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings	required
language	str	language code for sentence segmentation (default 'en')	'en'
low_first	bool	if True, fewer sentences first; if False, more sentences first	True

Returns:

Type	Description
list[tuple[int, int]]	List of (index, sentence_count) tuples sorted by sentence count.

Source code in splytters/sorters/text_sorters.py

def sentence_count(
    texts: list[str], language: str = "en", low_first: bool = True
) -> list[tuple[int, int]]:
    """
    Sort texts by number of sentences.

    Uses pysbd for robust sentence boundary detection across languages.

    Useful for adversarial splits: train on single-sentence texts,
    test on multi-sentence/complex texts.

    Args:
        texts: list of strings
        language: language code for sentence segmentation (default 'en')
        low_first: if True, fewer sentences first; if False, more sentences first

    Returns:
        List of (index, sentence_count) tuples sorted by sentence count.
    """
    import pysbd

    segmenter = pysbd.Segmenter(language=language, clean=False)

    scores = []
    for i, text in enumerate(texts):
        sentences = segmenter.segment(text)
        scores.append((i, len(sentences)))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

lexical_diversity

lexical_diversity(texts: list[str], tokenizer: Callable[[str], list[str]] = simple_tokenizer, low_first: bool = True) -> list[tuple[int, float]]

Sort texts by lexical diversity (type-token ratio).

Type-token ratio = unique tokens / total tokens. Higher values indicate more diverse vocabulary; lower values indicate more repetitive text.

Useful for adversarial splits: train on repetitive/simple vocabulary, test on diverse/rich vocabulary.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings	required
tokenizer	Callable[[str], list[str]]	function that splits text into tokens	simple_tokenizer
low_first	bool	if True, repetitive texts first; if False, diverse texts first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, ttr) tuples sorted by type-token ratio.
list[tuple[int, float]]	Texts with no tokens receive a score of 0.

Source code in splytters/sorters/text_sorters.py

def lexical_diversity(
    texts: list[str],
    tokenizer: Callable[[str], list[str]] = simple_tokenizer,
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort texts by lexical diversity (type-token ratio).

    Type-token ratio = unique tokens / total tokens.
    Higher values indicate more diverse vocabulary; lower values indicate
    more repetitive text.

    Useful for adversarial splits: train on repetitive/simple vocabulary,
    test on diverse/rich vocabulary.

    Args:
        texts: list of strings
        tokenizer: function that splits text into tokens
        low_first: if True, repetitive texts first; if False, diverse texts first

    Returns:
        List of (index, ttr) tuples sorted by type-token ratio.
        Texts with no tokens receive a score of 0.
    """
    scores = []
    for i, text in enumerate(texts):
        tokens = tokenizer(text)
        if len(tokens) == 0:
            scores.append((i, 0.0))
        else:
            ttr = len(set(tokens)) / len(tokens)
            scores.append((i, ttr))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

vocabulary_rarity

vocabulary_rarity(texts: list[str], language: str = 'en', tokenizer: Callable[[str], list[str]] = simple_tokenizer, low_first: bool = True) -> list[tuple[int, float]]

Sort texts by average word rarity.

Uses word frequency data to score each word's rarity. Rarer words have lower frequency, so we use (1 - frequency) as the rarity score.

Useful for adversarial splits: train on common vocabulary, test on rare/specialized vocabulary.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings	required
language	str	language code for frequency lookup (default 'en')	'en'
tokenizer	Callable[[str], list[str]]	function that splits text into tokens	simple_tokenizer
low_first	bool	if True, common vocabulary first; if False, rare vocabulary first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, avg_rarity) tuples sorted by average word rarity.
list[tuple[int, float]]	Texts with no tokens receive a score of 0.

Source code in splytters/sorters/text_sorters.py

def vocabulary_rarity(
    texts: list[str],
    language: str = "en",
    tokenizer: Callable[[str], list[str]] = simple_tokenizer,
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort texts by average word rarity.

    Uses word frequency data to score each word's rarity. Rarer words have
    lower frequency, so we use (1 - frequency) as the rarity score.

    Useful for adversarial splits: train on common vocabulary,
    test on rare/specialized vocabulary.

    Args:
        texts: list of strings
        language: language code for frequency lookup (default 'en')
        tokenizer: function that splits text into tokens
        low_first: if True, common vocabulary first; if False, rare vocabulary first

    Returns:
        List of (index, avg_rarity) tuples sorted by average word rarity.
        Texts with no tokens receive a score of 0.
    """
    from wordfreq import word_frequency

    scores = []
    for i, text in enumerate(texts):
        tokens = tokenizer(text.lower())
        if len(tokens) == 0:
            scores.append((i, 0.0))
            continue

        # Calculate average rarity (1 - frequency) for all tokens
        # word_frequency returns 0 for unknown words, so rarity = 1 for unknown
        rarities = [1.0 - word_frequency(token, language) for token in tokens]
        avg_rarity = sum(rarities) / len(rarities)
        scores.append((i, avg_rarity))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

perplexity_score

perplexity_score(texts: list[str], model: GPT2LMHeadModel | None = None, tokenizer: GPT2TokenizerFast | None = None, low_first: bool = True, scoring: str = 'perplexity') -> list[tuple[int, float]]

Sort texts by how typical a language model finds them.

A causal LM is run over each text to measure how "surprised" it is. Two scoring modes are offered (see scoring); both rank the same way — typical text is ordered before unusual text — but they put different things in the difficult tail, so the returned score values differ.

Useful for adversarial / curriculum splits: train on typical text, test on unusual text. Pair with :func:splytters.sorted_stratified_split to build "Likelihood Splits" (train on the high-likelihood head, test on the tail).

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings to score	required
model	GPT2LMHeadModel | None	HuggingFace causal LM (defaults to GPT-2 if None)	None
tokenizer	GPT2TokenizerFast | None	HuggingFace tokenizer (defaults to GPT-2 if None)	None
low_first	bool	if True, typical/predictable texts first; if False, unusual/surprising texts first. The ordering is by typicality, so this holds for both scoring modes despite their opposite raw-value directions.	True
scoring	str	how to score each text. - "perplexity" (default): exp of the mean per-token negative log-likelihood. Length-normalized; lower is more typical. - "log_likelihood": total (un-normalized) log-likelihood summed over the predicted tokens; higher is more typical. This is the score used by Likelihood Splits (Godbole & Jia, 2023).	'perplexity'

Returns:

Type	Description
list[tuple[int, float]]	List of (index, score) tuples in the chosen metric's natural units,
list[tuple[int, float]]	ordered by typicality per low_first. Texts too short to score receive
list[tuple[int, float]]	the most-difficult sentinel (+inf perplexity / -inf
list[tuple[int, float]]	log-likelihood), so they always land in the tail.

Raises:

Type	Description
ValueError	if scoring is not "perplexity" or "log_likelihood".

References

Godbole & Jia (2023), "Benchmarking Long-tail Generalization with Likelihood Splits," Findings of the ACL: EACL, pp. 963-983 — https://aclanthology.org/2023.findings-eacl.71 . They split datasets by LM likelihood (head -> train, tail -> test). Note they deliberately use total log-likelihood rather than perplexity: perplexity normalizes for length and over-corrects toward short examples in the tail (their fn. 3). Use scoring="log_likelihood" to reproduce their choice.

Source code in splytters/sorters/text_sorters.py

def perplexity_score(
    texts: list[str],
    model: GPT2LMHeadModel | None = None,
    tokenizer: GPT2TokenizerFast | None = None,
    low_first: bool = True,
    scoring: str = "perplexity",
) -> list[tuple[int, float]]:
    """
    Sort texts by how typical a language model finds them.

    A causal LM is run over each text to measure how "surprised" it is. Two
    scoring modes are offered (see ``scoring``); both rank the same way — typical
    text is ordered before unusual text — but they put different things in the
    difficult tail, so the returned score values differ.

    Useful for adversarial / curriculum splits: train on typical text, test on
    unusual text. Pair with :func:`splytters.sorted_stratified_split` to build
    "Likelihood Splits" (train on the high-likelihood head, test on the tail).

    Args:
        texts: list of strings to score
        model: HuggingFace causal LM (defaults to GPT-2 if None)
        tokenizer: HuggingFace tokenizer (defaults to GPT-2 if None)
        low_first: if True, typical/predictable texts first; if False,
            unusual/surprising texts first. The ordering is by *typicality*, so
            this holds for both scoring modes despite their opposite raw-value
            directions.
        scoring: how to score each text.
            - ``"perplexity"`` (default): ``exp`` of the mean per-token negative
              log-likelihood. Length-normalized; *lower* is more typical.
            - ``"log_likelihood"``: total (un-normalized) log-likelihood summed
              over the predicted tokens; *higher* is more typical. This is the
              score used by Likelihood Splits (Godbole & Jia, 2023).

    Returns:
        List of ``(index, score)`` tuples in the chosen metric's natural units,
        ordered by typicality per ``low_first``. Texts too short to score receive
        the most-difficult sentinel (``+inf`` perplexity / ``-inf``
        log-likelihood), so they always land in the tail.

    Raises:
        ValueError: if ``scoring`` is not ``"perplexity"`` or ``"log_likelihood"``.

    References:
        Godbole & Jia (2023), "Benchmarking Long-tail Generalization with
        Likelihood Splits," Findings of the ACL: EACL, pp. 963-983 —
        https://aclanthology.org/2023.findings-eacl.71 . They split datasets by
        LM likelihood (head -> train, tail -> test). Note they deliberately use
        *total log-likelihood* rather than perplexity: perplexity normalizes for
        length and over-corrects toward short examples in the tail (their fn. 3).
        Use ``scoring="log_likelihood"`` to reproduce their choice.
    """
    if scoring not in ("perplexity", "log_likelihood"):
        raise ValueError(
            f"scoring must be 'perplexity' or 'log_likelihood', got {scoring!r}"
        )

    import torch
    from transformers import GPT2LMHeadModel, GPT2TokenizerFast

    if model is None or tokenizer is None:
        tokenizer = GPT2TokenizerFast.from_pretrained("gpt2")
        model = GPT2LMHeadModel.from_pretrained("gpt2")
        model.eval()

    device = next(model.parameters()).device

    # Texts too short to score sort into the difficult tail regardless of mode:
    # maximal perplexity (+inf) or minimal log-likelihood (-inf).
    unscorable = float("inf") if scoring == "perplexity" else float("-inf")

    scores = []
    with torch.no_grad():
        for i, text in enumerate(texts):
            encodings = tokenizer(text, return_tensors="pt").to(device)
            input_ids = encodings.input_ids

            if input_ids.size(1) < 2:
                # Too short to compute a per-token loss
                scores.append((i, unscorable))
                continue

            outputs = model(input_ids, labels=input_ids)
            loss = outputs.loss.item()  # mean per-token negative log-likelihood
            if scoring == "perplexity":
                score = torch.exp(torch.tensor(loss)).item()
            else:
                n_tokens = input_ids.size(1) - 1  # number of predicted tokens
                score = -loss * n_tokens  # total log-likelihood
            scores.append((i, score))

    # Sort by difficulty (higher == more atypical/tail) so that low_first means
    # typical-first in both modes: high perplexity == low log-likelihood == hard.
    def difficulty(pair: tuple[int, float]) -> float:
        score = pair[1]
        return score if scoring == "perplexity" else -score

    scores.sort(key=difficulty, reverse=not low_first)
    return scores

readability_score

readability_score(texts: list[str], metric: str = 'flesch_kincaid', low_first: bool = True) -> list[tuple[int, float]]

Sort texts by readability score.

Readability scores estimate the education level required to understand text. Higher scores generally indicate more complex/difficult text.

Useful for adversarial splits: train on simple text, test on complex text.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings to score	required
metric	str	readability formula to use, one of: - 'flesch_kincaid': Flesch-Kincaid Grade Level - 'flesch': Flesch Reading Ease (higher = easier, inverted scale) - 'gunning_fog': Gunning Fog Index - 'coleman_liau': Coleman-Liau Index - 'dale_chall': Dale-Chall Readability Score - 'ari': Automated Readability Index - 'linsear_write': Linsear Write Formula - 'smog': SMOG Index	'flesch_kincaid'
low_first	bool	if True, easier/simpler texts first; if False, harder/complex texts first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, score) tuples sorted by readability.
list[tuple[int, float]]	Texts too short to score receive a score of infinity.

Source code in splytters/sorters/text_sorters.py

def readability_score(
    texts: list[str], metric: str = "flesch_kincaid", low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort texts by readability score.

    Readability scores estimate the education level required to understand text.
    Higher scores generally indicate more complex/difficult text.

    Useful for adversarial splits: train on simple text, test on complex text.

    Args:
        texts: list of strings to score
        metric: readability formula to use, one of:
            - 'flesch_kincaid': Flesch-Kincaid Grade Level
            - 'flesch': Flesch Reading Ease (higher = easier, inverted scale)
            - 'gunning_fog': Gunning Fog Index
            - 'coleman_liau': Coleman-Liau Index
            - 'dale_chall': Dale-Chall Readability Score
            - 'ari': Automated Readability Index
            - 'linsear_write': Linsear Write Formula
            - 'smog': SMOG Index
        low_first: if True, easier/simpler texts first;
                   if False, harder/complex texts first

    Returns:
        List of (index, score) tuples sorted by readability.
        Texts too short to score receive a score of infinity.
    """
    valid_metrics = {
        "flesch_kincaid", "flesch", "gunning_fog", "coleman_liau",
        "dale_chall", "ari", "linsear_write", "smog",
    }
    # Validate up front so an unknown metric raises instead of being swallowed
    # by the per-text "too short -> inf" exception handler below.
    if metric not in valid_metrics:
        raise ValueError(f"Unknown readability metric: {metric}")

    from readability import Readability

    scores = []

    for i, text in enumerate(texts):
        try:
            r = Readability(text)
            if metric == "flesch_kincaid":
                score = r.flesch_kincaid().score
            elif metric == "flesch":
                score = r.flesch().score
            elif metric == "gunning_fog":
                score = r.gunning_fog().score
            elif metric == "coleman_liau":
                score = r.coleman_liau().score
            elif metric == "dale_chall":
                score = r.dale_chall().score
            elif metric == "ari":
                score = r.ari().score
            elif metric == "linsear_write":
                score = r.linsear_write().score
            elif metric == "smog":
                score = r.smog().score
            else:  # pragma: no cover - unreachable; metric validated above
                raise ValueError(f"Unknown readability metric: {metric}")
            scores.append((i, score))
        except Exception:
            # Readability requires minimum text length
            scores.append((i, float("inf")))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

Image sorters

image_sorters

Sorting algorithms for adversarial image dataset partitioning.

These functions rank images by visual properties (brightness, contrast, color, complexity) to enable train-test splits that maximize dissimilarity.

All functions accept either file paths or PIL Image objects.

mean_brightness

mean_brightness(images: list[ImageInput], low_first: bool = True) -> list[tuple[int, float]]

Sort images by mean brightness (average pixel intensity).

Converts images to grayscale and computes mean intensity (0-255 scale).

Useful for adversarial splits: train on mid-tone images, test on very dark or very bright images.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
low_first	bool	if True, darker images first; if False, brighter images first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, brightness) tuples sorted by mean brightness.

Source code in splytters/sorters/image_sorters.py

def mean_brightness(
    images: list[ImageInput], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by mean brightness (average pixel intensity).

    Converts images to grayscale and computes mean intensity (0-255 scale).

    Useful for adversarial splits: train on mid-tone images,
    test on very dark or very bright images.

    Args:
        images: list of file paths or PIL Image objects
        low_first: if True, darker images first; if False, brighter images first

    Returns:
        List of (index, brightness) tuples sorted by mean brightness.
    """
    scores = []
    for i, img in enumerate(images):
        gray = _to_array(img, mode="L")
        brightness = gray.mean()
        scores.append((i, brightness))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

contrast

contrast(images: list[ImageInput], low_first: bool = True) -> list[tuple[int, float]]

Sort images by contrast (standard deviation of pixel intensities).

Higher contrast means more variation between light and dark regions. Lower contrast means flatter, more uniform images.

Useful for adversarial splits: train on normal contrast images, test on very flat or very high contrast images.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
low_first	bool	if True, low contrast (flat) images first; if False, high contrast images first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, contrast) tuples sorted by contrast.

Source code in splytters/sorters/image_sorters.py

def contrast(
    images: list[ImageInput], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by contrast (standard deviation of pixel intensities).

    Higher contrast means more variation between light and dark regions.
    Lower contrast means flatter, more uniform images.

    Useful for adversarial splits: train on normal contrast images,
    test on very flat or very high contrast images.

    Args:
        images: list of file paths or PIL Image objects
        low_first: if True, low contrast (flat) images first;
                   if False, high contrast images first

    Returns:
        List of (index, contrast) tuples sorted by contrast.
    """
    scores = []
    for i, img in enumerate(images):
        gray = _to_array(img, mode="L")
        std = gray.std()
        scores.append((i, std))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

color_variance

color_variance(images: list[ImageInput], low_first: bool = True) -> list[tuple[int, float]]

Sort images by color variance (spread across RGB channels).

Measures how much the RGB channels differ from each other. Higher variance indicates more colorful/saturated images. Lower variance indicates grayscale-like or desaturated images.

Useful for adversarial splits: train on neutral/desaturated images, test on highly colorful/saturated images.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
low_first	bool	if True, desaturated images first; if False, colorful images first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, color_variance) tuples sorted by color variance.

Source code in splytters/sorters/image_sorters.py

def color_variance(
    images: list[ImageInput], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by color variance (spread across RGB channels).

    Measures how much the RGB channels differ from each other.
    Higher variance indicates more colorful/saturated images.
    Lower variance indicates grayscale-like or desaturated images.

    Useful for adversarial splits: train on neutral/desaturated images,
    test on highly colorful/saturated images.

    Args:
        images: list of file paths or PIL Image objects
        low_first: if True, desaturated images first; if False, colorful images first

    Returns:
        List of (index, color_variance) tuples sorted by color variance.
    """
    scores = []
    for i, img in enumerate(images):
        rgb = _to_array(img, mode="RGB")

        # Compute variance across color channels for each pixel, then average
        # This measures how much R, G, B differ from each other
        channel_variance = rgb.var(axis=2).mean()
        scores.append((i, channel_variance))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

dominant_color

dominant_color(images: list[ImageInput], n_clusters: int = 5, target_color: list[int] | None = None, low_first: bool = True, random_state: int = 42) -> list[tuple[int, float]]

Sort images by their dominant color.

Uses k-means clustering to find dominant colors, then sorts by distance from a target color (default: neutral gray).

Useful for adversarial splits: train on images with one color palette, test on images with different dominant colors.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
n_clusters	int	number of color clusters to find (default 5)	5
target_color	list[int] | None	RGB tuple to measure distance from (default [128, 128, 128] gray)	None
low_first	bool	if True, images closest to target color first; if False, images furthest from target first	True
random_state	int	seed for the pixel subsample and k-means (default 42)	42

Returns:

Type	Description
list[tuple[int, float]]	List of (index, distance) tuples sorted by distance from target color.

Source code in splytters/sorters/image_sorters.py

def dominant_color(
    images: list[ImageInput],
    n_clusters: int = 5,
    target_color: list[int] | None = None,
    low_first: bool = True,
    random_state: int = 42,
) -> list[tuple[int, float]]:
    """
    Sort images by their dominant color.

    Uses k-means clustering to find dominant colors, then sorts by
    distance from a target color (default: neutral gray).

    Useful for adversarial splits: train on images with one color palette,
    test on images with different dominant colors.

    Args:
        images: list of file paths or PIL Image objects
        n_clusters: number of color clusters to find (default 5)
        target_color: RGB tuple to measure distance from (default [128, 128, 128] gray)
        low_first: if True, images closest to target color first;
                   if False, images furthest from target first
        random_state: seed for the pixel subsample and k-means (default 42)

    Returns:
        List of (index, distance) tuples sorted by distance from target color.
    """
    if target_color is None:
        target_color = np.array([128, 128, 128])
    else:
        target_color = np.array(target_color)

    scores = []
    for i, img in enumerate(images):
        rgb = _to_array(img, mode="RGB")

        # Reshape to list of pixels
        pixels = rgb.reshape(-1, 3)

        # Sample pixels if image is large (for speed). Use a seeded generator
        # so the subsample — and therefore the dominant color and the resulting
        # ranking — is reproducible across calls (the global RNG was not).
        if len(pixels) > 10000:
            rng = np.random.default_rng(random_state)
            indices = rng.choice(len(pixels), 10000, replace=False)
            pixels = pixels[indices]

        # Find dominant colors via k-means
        kmeans = KMeans(n_clusters=n_clusters, n_init="auto", random_state=random_state)
        kmeans.fit(pixels)

        # Get the most common cluster (dominant color)
        labels, counts = np.unique(kmeans.labels_, return_counts=True)
        dominant_idx = labels[counts.argmax()]
        dominant_rgb = kmeans.cluster_centers_[dominant_idx]

        # Distance from target color
        distance = np.linalg.norm(dominant_rgb - target_color)
        scores.append((i, float(distance)))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

compression_ratio

compression_ratio(images: list[ImageInput], quality: int = 85, low_first: bool = True) -> list[tuple[int, float]]

Sort images by JPEG compression ratio.

Simple/uniform images compress well (high ratio). Complex/detailed images compress poorly (low ratio).

Useful for adversarial splits: train on simple/compressible images, test on complex/incompressible images.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
quality	int	JPEG quality setting for compression test (default 85)	85
low_first	bool	if True, complex images first (low ratio); if False, simple images first (high ratio)	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, ratio) tuples sorted by compression ratio.
list[tuple[int, float]]	Ratio = uncompressed size / compressed size.

Source code in splytters/sorters/image_sorters.py

def compression_ratio(
    images: list[ImageInput], quality: int = 85, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by JPEG compression ratio.

    Simple/uniform images compress well (high ratio).
    Complex/detailed images compress poorly (low ratio).

    Useful for adversarial splits: train on simple/compressible images,
    test on complex/incompressible images.

    Args:
        images: list of file paths or PIL Image objects
        quality: JPEG quality setting for compression test (default 85)
        low_first: if True, complex images first (low ratio);
                   if False, simple images first (high ratio)

    Returns:
        List of (index, ratio) tuples sorted by compression ratio.
        Ratio = uncompressed size / compressed size.
    """
    scores = []
    for i, img in enumerate(images):
        pil_img = _load_image(img).convert("RGB")

        # Uncompressed size (width * height * 3 channels)
        uncompressed_size = pil_img.width * pil_img.height * 3

        # Compress to JPEG in memory
        buffer = io.BytesIO()
        pil_img.save(buffer, format="JPEG", quality=quality)
        compressed_size = buffer.tell()

        ratio = uncompressed_size / compressed_size
        scores.append((i, ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

frequency_content

frequency_content(images: list[ImageInput], low_first: bool = True) -> list[tuple[int, float]]

Sort images by high-frequency content (texture/detail level).

Uses FFT to analyze frequency distribution. Higher values indicate more high-frequency content (edges, textures, fine details). Lower values indicate smoother, low-frequency images.

Useful for adversarial splits: train on smooth/simple images, test on textured/detailed images.

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects	required
low_first	bool	if True, smooth/low-frequency images first; if False, textured/high-frequency images first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, high_freq_ratio) tuples sorted by frequency content.

Source code in splytters/sorters/image_sorters.py

def frequency_content(
    images: list[ImageInput], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by high-frequency content (texture/detail level).

    Uses FFT to analyze frequency distribution. Higher values indicate
    more high-frequency content (edges, textures, fine details).
    Lower values indicate smoother, low-frequency images.

    Useful for adversarial splits: train on smooth/simple images,
    test on textured/detailed images.

    Args:
        images: list of file paths or PIL Image objects
        low_first: if True, smooth/low-frequency images first;
                   if False, textured/high-frequency images first

    Returns:
        List of (index, high_freq_ratio) tuples sorted by frequency content.
    """
    scores = []
    for i, img in enumerate(images):
        gray = _to_array(img, mode="L")

        # Compute 2D FFT and shift zero frequency to center
        f_transform = fft.fft2(gray)
        f_shifted = fft.fftshift(f_transform)
        magnitude = np.abs(f_shifted)

        # Create a mask for high frequencies (outer region)
        rows, cols = gray.shape
        center_row, center_col = rows // 2, cols // 2

        # Define "low frequency" as inner 25% of frequency space
        low_freq_radius = min(rows, cols) // 8

        y, x = np.ogrid[:rows, :cols]
        distance_from_center = np.sqrt((x - center_col) ** 2 + (y - center_row) ** 2)

        low_freq_mask = distance_from_center <= low_freq_radius
        high_freq_mask = ~low_freq_mask

        # Ratio of high frequency energy to total energy
        total_energy = magnitude.sum()
        if total_energy == 0:
            scores.append((i, 0.0))
            continue

        high_freq_energy = magnitude[high_freq_mask].sum()
        high_freq_ratio = high_freq_energy / total_energy
        scores.append((i, high_freq_ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

Audio sorters

audio_sorters

Sorting algorithms for adversarial audio dataset partitioning.

These functions rank audio samples by various criteria (loudness, spectral features, rhythm, timbre) to enable train-test splits that maximize dissimilarity.

All functions accept either file paths or pre-loaded audio arrays. Requires librosa for audio processing.

duration

duration(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by total duration in seconds.

Useful for adversarial / curriculum splits: train on short utterances, test on long ones (or vice versa). Pair with :func:splytters.sorted_stratified_split for a length-based curriculum.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, shortest audio first; if False, longest first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, seconds) tuples sorted by duration.

References

Liu, Spence & Prud'hommeaux (2023), "Investigating data partitioning strategies for crosslinguistic low-resource ASR evaluation," EACL, pp. 123-131. https://aclanthology.org/2023.eacl-main.10 . They build heuristic test sets from per-utterance features -- utterance duration, average intensity, average pitch, transcript token count and perplexity -- and find utterance duration and intensity to be the most predictive factors of word-error-rate variability across data splits. Caveat: in their very low-resource setting, heuristic and adversarial splits behaved much like random splits, which were the more reliable choice under data sparsity -- so prefer random/averaged splits when data is scarce.

Source code in splytters/sorters/audio_sorters.py

def duration(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by total duration in seconds.

    Useful for adversarial / curriculum splits: train on short utterances,
    test on long ones (or vice versa). Pair with
    :func:`splytters.sorted_stratified_split` for a length-based curriculum.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, shortest audio first; if False, longest first

    Returns:
        List of (index, seconds) tuples sorted by duration.

    References:
        Liu, Spence & Prud'hommeaux (2023), "Investigating data partitioning
        strategies for crosslinguistic low-resource ASR evaluation," EACL,
        pp. 123-131. https://aclanthology.org/2023.eacl-main.10 . They build
        heuristic test sets from per-utterance features -- utterance duration,
        average intensity, average pitch, transcript token count and perplexity
        -- and find utterance duration and intensity to be the most predictive
        factors of word-error-rate variability across data splits. Caveat: in
        their very low-resource setting, heuristic and adversarial splits behaved
        much like random splits, which were the more reliable choice under data
        sparsity -- so prefer random/averaged splits when data is scarce.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        seconds = len(y) / sample_rate if sample_rate else 0.0
        scores.append((i, seconds))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

mean_amplitude

mean_amplitude(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by mean absolute amplitude.

Useful for adversarial splits: train on mid-volume audio, test on very quiet or very loud audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, quieter audio first; if False, louder first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, amplitude) tuples sorted by mean amplitude.

References

A loudness/intensity feature; see :func:duration for the Liu, Spence & Prud'hommeaux (2023) ASR data-partitioning study, which uses average intensity as a heuristic split feature.

Source code in splytters/sorters/audio_sorters.py

def mean_amplitude(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by mean absolute amplitude.

    Useful for adversarial splits: train on mid-volume audio,
    test on very quiet or very loud audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, quieter audio first; if False, louder first

    Returns:
        List of (index, amplitude) tuples sorted by mean amplitude.

    References:
        A loudness/intensity feature; see :func:`duration` for the
        Liu, Spence & Prud'hommeaux (2023) ASR data-partitioning study, which
        uses average intensity as a heuristic split feature.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, _ = _load_audio(audio, sr)
        amp = np.abs(y).mean()
        scores.append((i, amp))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

rms_energy

rms_energy(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by root mean square energy.

RMS energy is a common measure of audio loudness that better correlates with perceived volume than simple amplitude.

Useful for adversarial splits: train on consistent energy levels, test on very quiet or very loud audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, lower energy first; if False, higher energy first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, rms) tuples sorted by RMS energy.

References

RMS energy is an "average intensity" measure; see :func:duration for the Liu, Spence & Prud'hommeaux (2023) ASR data-partitioning study, which found intensity among the most predictive features of WER variability.

Source code in splytters/sorters/audio_sorters.py

def rms_energy(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by root mean square energy.

    RMS energy is a common measure of audio loudness that better
    correlates with perceived volume than simple amplitude.

    Useful for adversarial splits: train on consistent energy levels,
    test on very quiet or very loud audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, lower energy first; if False, higher energy first

    Returns:
        List of (index, rms) tuples sorted by RMS energy.

    References:
        RMS energy is an "average intensity" measure; see :func:`duration` for
        the Liu, Spence & Prud'hommeaux (2023) ASR data-partitioning study, which
        found intensity among the most predictive features of WER variability.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, _ = _load_audio(audio, sr)
        rms = np.sqrt(np.mean(y ** 2))
        scores.append((i, rms))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

dynamic_range

dynamic_range(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by dynamic range (difference between max and min amplitude).

Higher dynamic range indicates more variation between loud and quiet parts. Lower dynamic range indicates more compressed/consistent audio.

Useful for adversarial splits: train on compressed audio, test on highly dynamic audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, compressed audio first; if False, dynamic audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, range) tuples sorted by dynamic range.

Source code in splytters/sorters/audio_sorters.py

def dynamic_range(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by dynamic range (difference between max and min amplitude).

    Higher dynamic range indicates more variation between loud and quiet parts.
    Lower dynamic range indicates more compressed/consistent audio.

    Useful for adversarial splits: train on compressed audio,
    test on highly dynamic audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, compressed audio first; if False, dynamic audio first

    Returns:
        List of (index, range) tuples sorted by dynamic range.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, _ = _load_audio(audio, sr)
        y = np.asarray(y, dtype=float)
        # Dynamic range = spread between the loudest and quietest parts *over
        # time*. Measure it as the range of frame-wise RMS energy. Taking
        # |amplitude|.min() instead is ~0 for any zero-crossing waveform, so it
        # collapses to peak loudness and can't tell compressed from dynamic
        # audio.
        frame, hop = 2048, 512
        if len(y) >= frame:
            frames = np.lib.stride_tricks.sliding_window_view(y, frame)[::hop]
            rms = np.sqrt(np.mean(frames ** 2, axis=1))
        elif len(y) > 0:
            rms = np.array([np.sqrt(np.mean(y ** 2))])
        else:
            rms = np.zeros(1)
        dyn_range = float(rms.max() - rms.min())
        scores.append((i, dyn_range))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

peak_to_average_ratio

peak_to_average_ratio(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by peak-to-average ratio (crest factor).

Higher values indicate peaky/transient audio (e.g., percussion). Lower values indicate more sustained/compressed audio.

Useful for adversarial splits: train on sustained sounds, test on transient/percussive sounds.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, sustained audio first; if False, peaky audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, ratio) tuples sorted by peak-to-average ratio.

Source code in splytters/sorters/audio_sorters.py

def peak_to_average_ratio(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by peak-to-average ratio (crest factor).

    Higher values indicate peaky/transient audio (e.g., percussion).
    Lower values indicate more sustained/compressed audio.

    Useful for adversarial splits: train on sustained sounds,
    test on transient/percussive sounds.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, sustained audio first; if False, peaky audio first

    Returns:
        List of (index, ratio) tuples sorted by peak-to-average ratio.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, _ = _load_audio(audio, sr)
        peak = np.abs(y).max()
        avg = np.abs(y).mean()
        ratio = peak / avg if avg > 0 else 0
        scores.append((i, ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

spectral_centroid

spectral_centroid(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by spectral centroid ("center of mass" of the spectrum).

Higher values indicate brighter, treble-heavy audio. Lower values indicate darker, bass-heavy audio.

Useful for adversarial splits: train on mid-frequency audio, test on very bright or very dark audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, darker audio first; if False, brighter audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, centroid_hz) tuples sorted by spectral centroid.

Source code in splytters/sorters/audio_sorters.py

def spectral_centroid(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by spectral centroid ("center of mass" of the spectrum).

    Higher values indicate brighter, treble-heavy audio.
    Lower values indicate darker, bass-heavy audio.

    Useful for adversarial splits: train on mid-frequency audio,
    test on very bright or very dark audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, darker audio first; if False, brighter audio first

    Returns:
        List of (index, centroid_hz) tuples sorted by spectral centroid.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        centroid = librosa.feature.spectral_centroid(y=y, sr=sample_rate)
        mean_centroid = centroid.mean()
        scores.append((i, mean_centroid))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

spectral_bandwidth

spectral_bandwidth(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by spectral bandwidth (spread around the centroid).

Higher values indicate wider frequency spread (broadband). Lower values indicate narrower frequency content (tonal).

Useful for adversarial splits: train on narrow-band audio, test on broadband audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, narrow-band first; if False, broadband first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, bandwidth_hz) tuples sorted by spectral bandwidth.

Source code in splytters/sorters/audio_sorters.py

def spectral_bandwidth(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by spectral bandwidth (spread around the centroid).

    Higher values indicate wider frequency spread (broadband).
    Lower values indicate narrower frequency content (tonal).

    Useful for adversarial splits: train on narrow-band audio,
    test on broadband audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, narrow-band first; if False, broadband first

    Returns:
        List of (index, bandwidth_hz) tuples sorted by spectral bandwidth.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        bandwidth = librosa.feature.spectral_bandwidth(y=y, sr=sample_rate)
        mean_bandwidth = bandwidth.mean()
        scores.append((i, mean_bandwidth))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

spectral_rolloff

spectral_rolloff(audios: list[AudioInput], sr: int = 22050, roll_percent: float = 0.85, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by spectral rolloff frequency.

The rolloff frequency is the frequency below which a specified percentage of total spectral energy lies.

Higher values indicate more high-frequency content. Lower values indicate more low-frequency content.

Useful for adversarial splits: train on one frequency profile, test on different frequency profiles.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
roll_percent	float	energy percentage threshold (default 0.85)	0.85
low_first	bool	if True, low rolloff first; if False, high rolloff first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, rolloff_hz) tuples sorted by spectral rolloff.

Source code in splytters/sorters/audio_sorters.py

def spectral_rolloff(
    audios: list[AudioInput],
    sr: int = 22050,
    roll_percent: float = 0.85,
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort audio by spectral rolloff frequency.

    The rolloff frequency is the frequency below which a specified
    percentage of total spectral energy lies.

    Higher values indicate more high-frequency content.
    Lower values indicate more low-frequency content.

    Useful for adversarial splits: train on one frequency profile,
    test on different frequency profiles.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        roll_percent: energy percentage threshold (default 0.85)
        low_first: if True, low rolloff first; if False, high rolloff first

    Returns:
        List of (index, rolloff_hz) tuples sorted by spectral rolloff.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        rolloff = librosa.feature.spectral_rolloff(
            y=y, sr=sample_rate, roll_percent=roll_percent
        )
        mean_rolloff = rolloff.mean()
        scores.append((i, mean_rolloff))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

spectral_flatness

spectral_flatness(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by spectral flatness (Wiener entropy).

Values close to 1 indicate noise-like audio (flat spectrum). Values close to 0 indicate tonal audio (peaked spectrum).

Useful for adversarial splits: train on tonal audio, test on noisy audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, tonal audio first; if False, noisy audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, flatness) tuples sorted by spectral flatness.

Source code in splytters/sorters/audio_sorters.py

def spectral_flatness(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by spectral flatness (Wiener entropy).

    Values close to 1 indicate noise-like audio (flat spectrum).
    Values close to 0 indicate tonal audio (peaked spectrum).

    Useful for adversarial splits: train on tonal audio,
    test on noisy audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, tonal audio first; if False, noisy audio first

    Returns:
        List of (index, flatness) tuples sorted by spectral flatness.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        flatness = librosa.feature.spectral_flatness(y=y)
        mean_flatness = flatness.mean()
        scores.append((i, mean_flatness))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

zero_crossing_rate

zero_crossing_rate(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by zero crossing rate.

Higher values indicate noisier or more percussive audio. Lower values indicate smoother, more tonal audio.

Useful for adversarial splits: train on smooth audio, test on noisy/percussive audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, smooth audio first; if False, noisy audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, zcr) tuples sorted by zero crossing rate.

Source code in splytters/sorters/audio_sorters.py

def zero_crossing_rate(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by zero crossing rate.

    Higher values indicate noisier or more percussive audio.
    Lower values indicate smoother, more tonal audio.

    Useful for adversarial splits: train on smooth audio,
    test on noisy/percussive audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, smooth audio first; if False, noisy audio first

    Returns:
        List of (index, zcr) tuples sorted by zero crossing rate.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, _ = _load_audio(audio, sr)
        zcr = librosa.feature.zero_crossing_rate(y)
        mean_zcr = zcr.mean()
        scores.append((i, mean_zcr))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

fundamental_frequency

fundamental_frequency(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by estimated fundamental frequency (pitch).

Uses pYIN algorithm for robust pitch estimation.

Useful for adversarial splits: train on mid-pitch audio, test on very high or very low pitch audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, lower pitch first; if False, higher pitch first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, f0_hz) tuples sorted by fundamental frequency.
list[tuple[int, float]]	Audio with no detectable pitch receives f0 of 0.

References

Pitch is one of the per-utterance heuristic split features in Liu, Spence & Prud'hommeaux (2023); see :func:duration for the full citation.

Source code in splytters/sorters/audio_sorters.py

def fundamental_frequency(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by estimated fundamental frequency (pitch).

    Uses pYIN algorithm for robust pitch estimation.

    Useful for adversarial splits: train on mid-pitch audio,
    test on very high or very low pitch audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, lower pitch first; if False, higher pitch first

    Returns:
        List of (index, f0_hz) tuples sorted by fundamental frequency.
        Audio with no detectable pitch receives f0 of 0.

    References:
        Pitch is one of the per-utterance heuristic split features in
        Liu, Spence & Prud'hommeaux (2023); see :func:`duration` for the full
        citation.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        f0, voiced_flag, voiced_probs = librosa.pyin(
            y, fmin=librosa.note_to_hz('C2'), fmax=librosa.note_to_hz('C7'), sr=sample_rate
        )
        # Take mean of voiced frames only
        voiced_f0 = f0[~np.isnan(f0)]
        mean_f0 = voiced_f0.mean() if len(voiced_f0) > 0 else 0
        scores.append((i, mean_f0))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

mfcc_mean

mfcc_mean(audios: list[AudioInput], sr: int = 22050, n_mfcc: int = 13, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by mean of first MFCC coefficient.

The first MFCC coefficient relates to overall energy/loudness. Higher MFCCs capture timbral characteristics.

Useful for adversarial splits: separate by timbral characteristics.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
n_mfcc	int	number of MFCC coefficients (default 13)	13
low_first	bool	if True, lower mean MFCC first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, mfcc_mean) tuples sorted by mean of first MFCC.

Source code in splytters/sorters/audio_sorters.py

def mfcc_mean(
    audios: list[AudioInput],
    sr: int = 22050,
    n_mfcc: int = 13,
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort audio by mean of first MFCC coefficient.

    The first MFCC coefficient relates to overall energy/loudness.
    Higher MFCCs capture timbral characteristics.

    Useful for adversarial splits: separate by timbral characteristics.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        n_mfcc: number of MFCC coefficients (default 13)
        low_first: if True, lower mean MFCC first

    Returns:
        List of (index, mfcc_mean) tuples sorted by mean of first MFCC.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        mfccs = librosa.feature.mfcc(y=y, sr=sample_rate, n_mfcc=n_mfcc)
        # Use mean of first coefficient (energy-related)
        mean_mfcc = mfccs[0].mean()
        scores.append((i, mean_mfcc))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

mfcc_variance

mfcc_variance(audios: list[AudioInput], sr: int = 22050, n_mfcc: int = 13, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by variance of MFCC coefficients over time.

Higher variance indicates more timbral variation over time. Lower variance indicates more consistent timbre.

Useful for adversarial splits: train on stable timbre, test on varying timbre.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
n_mfcc	int	number of MFCC coefficients (default 13)	13
low_first	bool	if True, stable timbre first; if False, varying timbre first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, mfcc_var) tuples sorted by MFCC variance.

Source code in splytters/sorters/audio_sorters.py

def mfcc_variance(
    audios: list[AudioInput],
    sr: int = 22050,
    n_mfcc: int = 13,
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort audio by variance of MFCC coefficients over time.

    Higher variance indicates more timbral variation over time.
    Lower variance indicates more consistent timbre.

    Useful for adversarial splits: train on stable timbre,
    test on varying timbre.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        n_mfcc: number of MFCC coefficients (default 13)
        low_first: if True, stable timbre first; if False, varying timbre first

    Returns:
        List of (index, mfcc_var) tuples sorted by MFCC variance.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        mfccs = librosa.feature.mfcc(y=y, sr=sample_rate, n_mfcc=n_mfcc)
        # Compute mean variance across all coefficients
        var = mfccs.var(axis=1).mean()
        scores.append((i, var))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

tempo

tempo(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by estimated tempo (beats per minute).

Useful for adversarial splits: train on mid-tempo audio, test on very slow or very fast audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, slower tempo first; if False, faster tempo first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, bpm) tuples sorted by tempo.

Source code in splytters/sorters/audio_sorters.py

def tempo(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by estimated tempo (beats per minute).

    Useful for adversarial splits: train on mid-tempo audio,
    test on very slow or very fast audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, slower tempo first; if False, faster tempo first

    Returns:
        List of (index, bpm) tuples sorted by tempo.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        # librosa >=0.10 returns tempo as an array; take the scalar.
        bpm, _ = librosa.beat.beat_track(y=y, sr=sample_rate)
        scores.append((i, float(np.atleast_1d(bpm)[0])))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

beat_strength

beat_strength(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by beat/onset strength.

Higher values indicate more percussive, rhythmically strong audio. Lower values indicate more ambient, less rhythmic audio.

Useful for adversarial splits: train on strong beat audio, test on ambient/weak beat audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, weak beats first; if False, strong beats first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, strength) tuples sorted by beat strength.

Source code in splytters/sorters/audio_sorters.py

def beat_strength(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by beat/onset strength.

    Higher values indicate more percussive, rhythmically strong audio.
    Lower values indicate more ambient, less rhythmic audio.

    Useful for adversarial splits: train on strong beat audio,
    test on ambient/weak beat audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, weak beats first; if False, strong beats first

    Returns:
        List of (index, strength) tuples sorted by beat strength.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        onset_env = librosa.onset.onset_strength(y=y, sr=sample_rate)
        mean_strength = onset_env.mean()
        scores.append((i, mean_strength))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

harmonic_ratio

harmonic_ratio(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by harmonic-to-percussive ratio.

Higher values indicate more melodic/harmonic content. Lower values indicate more percussive/rhythmic content.

Useful for adversarial splits: train on melodic audio, test on percussive audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, percussive audio first; if False, harmonic audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, ratio) tuples sorted by harmonic ratio.

Source code in splytters/sorters/audio_sorters.py

def harmonic_ratio(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by harmonic-to-percussive ratio.

    Higher values indicate more melodic/harmonic content.
    Lower values indicate more percussive/rhythmic content.

    Useful for adversarial splits: train on melodic audio,
    test on percussive audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, percussive audio first; if False, harmonic audio first

    Returns:
        List of (index, ratio) tuples sorted by harmonic ratio.
    """
    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)
        harmonic, percussive = librosa.effects.hpss(y)
        h_energy = np.sum(harmonic ** 2)
        p_energy = np.sum(percussive ** 2)
        ratio = h_energy / (p_energy + 1e-10)  # Avoid division by zero
        scores.append((i, ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

compression_ratio

compression_ratio(audios: list[AudioInput], sr: int = 22050, low_first: bool = True) -> list[tuple[int, float]]

Sort audio by compression ratio (uncompressed / compressed size).

Simple/repetitive audio compresses well (high ratio). Complex/varied audio compresses poorly (low ratio).

Useful for adversarial splits: train on simple audio, test on complex audio.

Parameters:

Name	Type	Description	Default
audios	list[AudioInput]	list of file paths or (samples, sr) tuples	required
sr	int	sample rate for loading files (default 22050)	22050
low_first	bool	if True, complex audio first; if False, simple audio first	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, ratio) tuples sorted by compression ratio.

Source code in splytters/sorters/audio_sorters.py

def compression_ratio(
    audios: list[AudioInput], sr: int = 22050, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort audio by compression ratio (uncompressed / compressed size).

    Simple/repetitive audio compresses well (high ratio).
    Complex/varied audio compresses poorly (low ratio).

    Useful for adversarial splits: train on simple audio,
    test on complex audio.

    Args:
        audios: list of file paths or (samples, sr) tuples
        sr: sample rate for loading files (default 22050)
        low_first: if True, complex audio first; if False, simple audio first

    Returns:
        List of (index, ratio) tuples sorted by compression ratio.
    """
    import soundfile as sf

    scores = []
    for i, audio in enumerate(audios):
        y, sample_rate = _load_audio(audio, sr)

        # Uncompressed size (samples * bytes per sample)
        uncompressed_size = len(y) * 4  # float32 = 4 bytes

        # Compress to OGG in memory
        buffer = io.BytesIO()
        sf.write(buffer, y, sample_rate, format='OGG')
        compressed_size = buffer.tell()

        ratio = uncompressed_size / compressed_size if compressed_size > 0 else 0
        scores.append((i, ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return [(idx, float(value)) for idx, value in scores]

Tabular sorters

tabular_sorters

Sorting algorithms for adversarial tabular dataset partitioning.

These functions rank rows in tabular data (pandas DataFrames) by various criteria (missing values, outliers, sparsity, column values) to enable train-test splits that maximize dissimilarity.

All functions accept pandas DataFrames and return sorted index lists.

column_value

column_value(df: DataFrame, column: str, low_first: bool = True) -> list[tuple[Hashable, Any]]

Sort rows by values in a specified column.

Generic sorting function that allows users to sort by any column.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
column	str	column name to sort by	required
low_first	bool	if True, lowest values first; if False, highest values first	True

Returns:

Type	Description
list[tuple[Hashable, Any]]	List of (index, value) tuples sorted by column value.
list[tuple[Hashable, Any]]	NaN values are placed at the end.

Source code in splytters/sorters/tabular_sorters.py

def column_value(
    df: pd.DataFrame, column: str, low_first: bool = True
) -> list[tuple[Hashable, Any]]:
    """
    Sort rows by values in a specified column.

    Generic sorting function that allows users to sort by any column.

    Args:
        df: pandas DataFrame
        column: column name to sort by
        low_first: if True, lowest values first; if False, highest values first

    Returns:
        List of (index, value) tuples sorted by column value.
        NaN values are placed at the end.
    """
    scores = []
    for idx in df.index:
        value = df.loc[idx, column]
        # Handle NaN by placing at end
        if pd.isna(value):
            sort_value = float('inf') if low_first else float('-inf')
        else:
            sort_value = value
        scores.append((idx, value, sort_value))

    scores.sort(key=lambda p: p[2], reverse=not low_first)
    return [(idx, val) for idx, val, _ in scores]

column_rank

column_rank(df: DataFrame, column: str, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by percentile rank within a column.

Useful when you want to split by relative position rather than absolute values.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
column	str	column name to rank by	required
low_first	bool	if True, lowest percentile first; if False, highest first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, percentile) tuples sorted by percentile rank (0-100).

Source code in splytters/sorters/tabular_sorters.py

def column_rank(
    df: pd.DataFrame, column: str, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by percentile rank within a column.

    Useful when you want to split by relative position rather than
    absolute values.

    Args:
        df: pandas DataFrame
        column: column name to rank by
        low_first: if True, lowest percentile first; if False, highest first

    Returns:
        List of (index, percentile) tuples sorted by percentile rank (0-100).
    """
    series = df[column]
    ranks = series.rank(pct=True, na_option='bottom') * 100

    scores = [(idx, ranks.loc[idx]) for idx in df.index]
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

column_zscore

column_zscore(df: DataFrame, column: str, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by z-score (standard deviations from mean) in a column.

Useful for finding outliers in a specific column.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
column	str	column name to compute z-scores for	required
low_first	bool	if True, lowest z-scores first (below mean); if False, highest z-scores first (above mean)	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, zscore) tuples sorted by z-score.
list[tuple[Hashable, float]]	NaN values receive z-score of infinity.

Source code in splytters/sorters/tabular_sorters.py

def column_zscore(
    df: pd.DataFrame, column: str, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by z-score (standard deviations from mean) in a column.

    Useful for finding outliers in a specific column.

    Args:
        df: pandas DataFrame
        column: column name to compute z-scores for
        low_first: if True, lowest z-scores first (below mean);
                   if False, highest z-scores first (above mean)

    Returns:
        List of (index, zscore) tuples sorted by z-score.
        NaN values receive z-score of infinity.
    """
    series = df[column]
    mean = series.mean()
    std = series.std()

    if std == 0:
        # All values are the same
        return [(idx, 0.0) for idx in df.index]

    scores = []
    for idx in df.index:
        value = series.loc[idx]
        if pd.isna(value):
            zscore = float('inf')
        else:
            zscore = (value - mean) / std
        scores.append((idx, zscore))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

column_absolute_zscore

column_absolute_zscore(df: DataFrame, column: str, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by absolute z-score (distance from mean) in a column.

Useful for adversarial splits: train on typical values (low |z|), test on extreme values (high |z|).

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
column	str	column name to compute z-scores for	required
low_first	bool	if True, typical values first; if False, extreme values first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, abs_zscore) tuples sorted by absolute z-score.

Source code in splytters/sorters/tabular_sorters.py

def column_absolute_zscore(
    df: pd.DataFrame, column: str, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by absolute z-score (distance from mean) in a column.

    Useful for adversarial splits: train on typical values (low |z|),
    test on extreme values (high |z|).

    Args:
        df: pandas DataFrame
        column: column name to compute z-scores for
        low_first: if True, typical values first; if False, extreme values first

    Returns:
        List of (index, abs_zscore) tuples sorted by absolute z-score.
    """
    series = df[column]
    mean = series.mean()
    std = series.std()

    if std == 0:
        return [(idx, 0.0) for idx in df.index]

    scores = []
    for idx in df.index:
        value = series.loc[idx]
        if pd.isna(value):
            abs_zscore = float('inf')
        else:
            abs_zscore = abs((value - mean) / std)
        scores.append((idx, abs_zscore))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

missing_value_ratio

missing_value_ratio(df: DataFrame, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by proportion of missing (NaN/None) values.

Useful for adversarial splits: train on complete rows, test on rows with missing data.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
low_first	bool	if True, complete rows first; if False, sparse rows first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, missing_ratio) tuples sorted by missing ratio (0-1).

Source code in splytters/sorters/tabular_sorters.py

def missing_value_ratio(
    df: pd.DataFrame, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by proportion of missing (NaN/None) values.

    Useful for adversarial splits: train on complete rows,
    test on rows with missing data.

    Args:
        df: pandas DataFrame
        low_first: if True, complete rows first; if False, sparse rows first

    Returns:
        List of (index, missing_ratio) tuples sorted by missing ratio (0-1).
    """
    scores = []
    n_cols = len(df.columns)

    for idx in df.index:
        row = df.loc[idx]
        n_missing = row.isna().sum()
        ratio = n_missing / n_cols if n_cols > 0 else 0
        scores.append((idx, ratio))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

row_sparsity

row_sparsity(df: DataFrame, zero_threshold: float = 1e-10, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by proportion of zero (or near-zero) values.

Useful for adversarial splits: train on dense rows, test on sparse rows.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame (numeric columns only)	required
zero_threshold	float	values with abs < threshold are considered zero	1e-10
low_first	bool	if True, dense rows first; if False, sparse rows first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, sparsity) tuples sorted by sparsity ratio (0-1).

Source code in splytters/sorters/tabular_sorters.py

def row_sparsity(
    df: pd.DataFrame, zero_threshold: float = 1e-10, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by proportion of zero (or near-zero) values.

    Useful for adversarial splits: train on dense rows,
    test on sparse rows.

    Args:
        df: pandas DataFrame (numeric columns only)
        zero_threshold: values with abs < threshold are considered zero
        low_first: if True, dense rows first; if False, sparse rows first

    Returns:
        List of (index, sparsity) tuples sorted by sparsity ratio (0-1).
    """
    # Select only numeric columns
    numeric_df = df.select_dtypes(include=[np.number])
    n_cols = len(numeric_df.columns)

    if n_cols == 0:
        return [(idx, 0.0) for idx in df.index]

    scores = []
    for idx in df.index:
        row = numeric_df.loc[idx]
        n_zeros = (row.abs() < zero_threshold).sum()
        sparsity = n_zeros / n_cols
        scores.append((idx, sparsity))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

outlier_score

outlier_score(df: DataFrame, method: str = 'isolation_forest', low_first: bool = True, **kwargs: Any) -> list[tuple[Hashable, float]]

Sort rows by anomaly/outlier score.

Uses outlier detection algorithms to score how unusual each row is.

Useful for adversarial splits: train on normal/typical rows, test on outliers/anomalies.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame (numeric columns only, NaN will be filled with median)	required
method	str	outlier detection algorithm, one of: - 'isolation_forest': Isolation Forest (fast, good for high dimensions) - 'lof': Local Outlier Factor (density-based) - 'zscore': Mean absolute z-score across columns	'isolation_forest'
low_first	bool	if True, normal rows first; if False, outliers first	True
**kwargs	Any	additional arguments passed to the outlier detector	{}

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, outlier_score) tuples sorted by outlier score.

Source code in splytters/sorters/tabular_sorters.py

def outlier_score(
    df: pd.DataFrame,
    method: str = "isolation_forest",
    low_first: bool = True,
    **kwargs: Any,
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by anomaly/outlier score.

    Uses outlier detection algorithms to score how unusual each row is.

    Useful for adversarial splits: train on normal/typical rows,
    test on outliers/anomalies.

    Args:
        df: pandas DataFrame (numeric columns only, NaN will be filled with median)
        method: outlier detection algorithm, one of:
            - 'isolation_forest': Isolation Forest (fast, good for high dimensions)
            - 'lof': Local Outlier Factor (density-based)
            - 'zscore': Mean absolute z-score across columns
        low_first: if True, normal rows first; if False, outliers first
        **kwargs: additional arguments passed to the outlier detector

    Returns:
        List of (index, outlier_score) tuples sorted by outlier score.
    """
    # Select only numeric columns and fill NaN
    numeric_df = df.select_dtypes(include=[np.number]).copy()
    numeric_df = numeric_df.fillna(numeric_df.median())

    if len(numeric_df.columns) == 0:
        return [(idx, 0.0) for idx in df.index]

    X = numeric_df.values

    if method == "isolation_forest":
        detector = IsolationForest(random_state=42, **kwargs)
        detector.fit(X)
        # Negate so higher = more outlier
        raw_scores = -detector.score_samples(X)

    elif method == "lof":
        detector = LocalOutlierFactor(novelty=False, **kwargs)
        detector.fit_predict(X)
        # Negate so higher = more outlier
        raw_scores = -detector.negative_outlier_factor_

    elif method == "zscore":
        # Compute mean absolute z-score across all columns
        zscores = np.abs(stats.zscore(X, nan_policy='omit'))
        raw_scores = np.nanmean(zscores, axis=1)

    else:
        raise ValueError(f"Unknown outlier detection method: {method}")

    scores = [(df.index[i], raw_scores[i]) for i in range(len(df))]
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

numerical_range_score

numerical_range_score(df: DataFrame, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by how extreme their numerical values are (distance from median).

Computes the mean percentile distance from 50th percentile across all numeric columns.

Useful for adversarial splits: train on typical values, test on extreme values.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
low_first	bool	if True, typical rows first; if False, extreme rows first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, extremity) tuples sorted by extremity score (0-50).

Source code in splytters/sorters/tabular_sorters.py

def numerical_range_score(
    df: pd.DataFrame, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by how extreme their numerical values are (distance from median).

    Computes the mean percentile distance from 50th percentile across all
    numeric columns.

    Useful for adversarial splits: train on typical values,
    test on extreme values.

    Args:
        df: pandas DataFrame
        low_first: if True, typical rows first; if False, extreme rows first

    Returns:
        List of (index, extremity) tuples sorted by extremity score (0-50).
    """
    numeric_df = df.select_dtypes(include=[np.number])

    if len(numeric_df.columns) == 0:
        return [(idx, 0.0) for idx in df.index]

    # Compute percentile ranks for each column
    ranks = numeric_df.rank(pct=True) * 100

    # Distance from 50th percentile (0 = median, 50 = extreme)
    extremity = (ranks - 50).abs().mean(axis=1)

    scores = [(idx, extremity.loc[idx]) for idx in df.index]
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

categorical_rarity

categorical_rarity(df: DataFrame, column: str, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by rarity of categorical value in a column.

Useful for adversarial splits: train on common categories, test on rare categories.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
column	str	categorical column name	required
low_first	bool	if True, common categories first; if False, rare categories first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, frequency) tuples sorted by category frequency.
list[tuple[Hashable, float]]	Higher frequency = more common.

Note

Unlike most sorters (whose score is a difficulty/rarity measure where low_first orders ascending), the score here is frequency (high = common). low_first=True therefore orders by commonness — common categories first — which is the natural "train on common, test on rare" direction for an adversarial split.

Source code in splytters/sorters/tabular_sorters.py

def categorical_rarity(
    df: pd.DataFrame, column: str, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by rarity of categorical value in a column.

    Useful for adversarial splits: train on common categories,
    test on rare categories.

    Args:
        df: pandas DataFrame
        column: categorical column name
        low_first: if True, common categories first; if False, rare categories first

    Returns:
        List of (index, frequency) tuples sorted by category frequency.
        Higher frequency = more common.

    Note:
        Unlike most sorters (whose ``score`` is a difficulty/rarity measure
        where ``low_first`` orders ascending), the score here is *frequency*
        (high = common). ``low_first=True`` therefore orders by commonness —
        common categories first — which is the natural "train on common,
        test on rare" direction for an adversarial split.
    """
    # Compute value counts as frequencies
    value_counts = df[column].value_counts(normalize=True)

    scores = []
    for idx in df.index:
        value = df.loc[idx, column]
        if pd.isna(value):
            freq = 0.0  # Treat NaN as rarest
        else:
            freq = value_counts.get(value, 0.0)
        scores.append((idx, freq))

    # Note: low_first=True means common first (high frequency first)
    # So we reverse the logic
    scores.sort(key=lambda p: p[1], reverse=low_first)
    return scores

feature_entropy

feature_entropy(df: DataFrame, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by entropy of categorical features.

Rows with rare combinations of categorical values have higher "entropy" in the sense that they're less predictable.

Approximated by summing -log(freq) for each categorical value.

Useful for adversarial splits: train on common patterns, test on rare patterns.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
low_first	bool	if True, common patterns first; if False, rare patterns first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, neg_log_freq) tuples sorted by rarity score.

Source code in splytters/sorters/tabular_sorters.py

def feature_entropy(
    df: pd.DataFrame, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by entropy of categorical features.

    Rows with rare combinations of categorical values have higher "entropy"
    in the sense that they're less predictable.

    Approximated by summing -log(freq) for each categorical value.

    Useful for adversarial splits: train on common patterns,
    test on rare patterns.

    Args:
        df: pandas DataFrame
        low_first: if True, common patterns first; if False, rare patterns first

    Returns:
        List of (index, neg_log_freq) tuples sorted by rarity score.
    """
    # Select categorical columns
    cat_cols = df.select_dtypes(include=['object', 'category']).columns

    if len(cat_cols) == 0:
        return [(idx, 0.0) for idx in df.index]

    # Compute frequencies for each categorical column
    freq_maps = {}
    for col in cat_cols:
        freq_maps[col] = df[col].value_counts(normalize=True).to_dict()

    scores = []
    for idx in df.index:
        neg_log_freq = 0.0
        for col in cat_cols:
            value = df.loc[idx, col]
            if pd.isna(value):
                freq = 1e-10  # Very rare
            else:
                freq = freq_maps[col].get(value, 1e-10)
            neg_log_freq += -np.log(freq + 1e-10)
        scores.append((idx, neg_log_freq))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

row_distance_to_mean

row_distance_to_mean(df: DataFrame, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by Euclidean distance from the column-wise mean.

Similar to embedding distance_to_mean but operates directly on tabular features.

Useful for adversarial splits: train on typical rows, test on atypical rows.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame (numeric columns only)	required
low_first	bool	if True, typical rows first; if False, atypical rows first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, distance) tuples sorted by distance from mean.

Source code in splytters/sorters/tabular_sorters.py

def row_distance_to_mean(
    df: pd.DataFrame, low_first: bool = True
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by Euclidean distance from the column-wise mean.

    Similar to embedding distance_to_mean but operates directly on
    tabular features.

    Useful for adversarial splits: train on typical rows,
    test on atypical rows.

    Args:
        df: pandas DataFrame (numeric columns only)
        low_first: if True, typical rows first; if False, atypical rows first

    Returns:
        List of (index, distance) tuples sorted by distance from mean.
    """
    numeric_df = df.select_dtypes(include=[np.number]).copy()
    numeric_df = numeric_df.fillna(numeric_df.median())

    if len(numeric_df.columns) == 0:
        return [(idx, 0.0) for idx in df.index]

    # Compute column-wise mean
    mean_vector = numeric_df.mean().values

    scores = []
    for idx in df.index:
        row_vector = numeric_df.loc[idx].values
        distance = np.linalg.norm(row_vector - mean_vector)
        scores.append((idx, distance))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

multi_column_sort

multi_column_sort(df: DataFrame, columns: list[str], weights: list[float] | None = None, low_first: bool = True) -> list[tuple[Hashable, float]]

Sort rows by weighted combination of multiple columns.

Useful when you want to combine multiple criteria into a single sort.

Parameters:

Name	Type	Description	Default
df	DataFrame	pandas DataFrame	required
columns	list[str]	list of column names to sort by	required
weights	list[float] | None	list of weights for each column (default: equal weights) Positive weight = higher value increases score Negative weight = higher value decreases score	None
low_first	bool	if True, lowest combined score first	True

Returns:

Type	Description
list[tuple[Hashable, float]]	List of (index, combined_score) tuples sorted by combined score.

Source code in splytters/sorters/tabular_sorters.py

def multi_column_sort(
    df: pd.DataFrame,
    columns: list[str],
    weights: list[float] | None = None,
    low_first: bool = True,
) -> list[tuple[Hashable, float]]:
    """
    Sort rows by weighted combination of multiple columns.

    Useful when you want to combine multiple criteria into a single sort.

    Args:
        df: pandas DataFrame
        columns: list of column names to sort by
        weights: list of weights for each column (default: equal weights)
                 Positive weight = higher value increases score
                 Negative weight = higher value decreases score
        low_first: if True, lowest combined score first

    Returns:
        List of (index, combined_score) tuples sorted by combined score.
    """
    if weights is None:
        weights = [1.0] * len(columns)

    if len(columns) != len(weights):
        raise ValueError("Number of columns must match number of weights")

    # Normalize each column to 0-1 range
    normalized = pd.DataFrame(index=df.index)
    for col in columns:
        series = df[col]
        min_val = series.min()
        max_val = series.max()
        if max_val > min_val:
            # NaN cells normalize to NaN; map them to the neutral midpoint so
            # the weighted sum stays finite instead of poisoning the score.
            normalized[col] = ((series - min_val) / (max_val - min_val)).fillna(0.5)
        else:
            normalized[col] = 0.5

    # Compute weighted sum
    scores = []
    for idx in df.index:
        combined = sum(
            weights[i] * normalized.loc[idx, col]
            for i, col in enumerate(columns)
        )
        scores.append((idx, combined))

    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return scores

Utilities

utils

Shared utilities for splitting algorithms.

to_numpy

to_numpy(X: ArrayLike) -> Any

Best-effort conversion of framework tensors to numpy.

Handles PyTorch tensors (including those on GPU / requiring grad) by detaching, moving to CPU, and converting. Any other input is returned unchanged so the caller's downstream check_array/np.asarray can handle lists, numpy arrays, pandas DataFrames, etc.

Source code in splytters/utils.py

def to_numpy(X: ArrayLike) -> Any:
    """Best-effort conversion of framework tensors to numpy.

    Handles PyTorch tensors (including those on GPU / requiring grad) by
    detaching, moving to CPU, and converting. Any other input is returned
    unchanged so the caller's downstream ``check_array``/``np.asarray`` can
    handle lists, numpy arrays, pandas DataFrames, etc.
    """
    # Duck-type torch.Tensor without importing torch.
    if hasattr(X, "detach") and hasattr(X, "cpu") and hasattr(X, "numpy"):
        try:
            return X.detach().cpu().numpy()
        except Exception:  # pragma: no cover - fall through to generic handling
            pass
    return X

validate_split_inputs

validate_split_inputs(embeddings: ArrayLike, train_size: float | int, min_samples: int = 2) -> np.ndarray

Validate and coerce inputs shared by every splitting function.

Accepts any array-like (numpy, list, pandas, torch tensor), converts it to a finite 2-D float ndarray, and validates train_size.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, n_features).	required
train_size	float | int	fraction in the open interval (0, 1), or an absolute count in [1, n_samples). Mirrors sklearn.model_selection.train_test_split.	required
min_samples	int	minimum number of samples required to form a split (default 2).	2

Returns:

Type	Description
ndarray	The validated, finite, float embeddings as an ndarray of shape
ndarray	(n_samples, n_features).

Raises:

Type	Description
ValueError	if train_size is out of range, there are too few samples, or the embeddings contain NaN/inf or are not 2-D.

Source code in splytters/utils.py

def validate_split_inputs(
    embeddings: ArrayLike, train_size: float | int, min_samples: int = 2
) -> np.ndarray:
    """Validate and coerce inputs shared by every splitting function.

    Accepts any array-like (numpy, list, pandas, torch tensor), converts it to
    a finite 2-D float ``ndarray``, and validates ``train_size``.

    Args:
        embeddings: array-like of shape (n_samples, n_features).
        train_size: fraction in the open interval (0, 1), or an absolute count
            in ``[1, n_samples)``. Mirrors
            ``sklearn.model_selection.train_test_split``.
        min_samples: minimum number of samples required to form a split
            (default 2).

    Returns:
        The validated, finite, float embeddings as an ndarray of shape
        (n_samples, n_features).

    Raises:
        ValueError: if ``train_size`` is out of range, there are too few
            samples, or the embeddings contain NaN/inf or are not 2-D.
    """
    # Validate train_size first so its message takes priority (matches the
    # historical "between 0 and 1" contract for fractional sizes).
    is_int = isinstance(train_size, (int, np.integer)) and not isinstance(
        train_size, bool
    )
    if not is_int:
        if not (isinstance(train_size, (float, np.floating)) and 0 < train_size < 1):
            raise ValueError(
                "train_size as a fraction must be between 0 and 1 exclusive, "
                f"got {train_size!r}"
            )

    X = to_numpy(embeddings)
    n = len(X)
    if n < min_samples:
        raise ValueError(f"Need at least {min_samples} samples to split, got {n}")

    if is_int and not (1 <= train_size < n):
        raise ValueError(
            f"train_size as an absolute count must be in [1, {n}), got {train_size}"
        )

    # Coerce to a finite 2-D numeric array (rejects NaN/inf, 1-D, ragged, sparse).
    X = check_array(X, ensure_2d=True, allow_nd=False, dtype="numeric")
    return X

resolve_n_train

resolve_n_train(n_samples: int, train_size: float | int) -> int

Resolve train_size (fraction or absolute count) to an int count.

A fractional train_size is clamped to [1, n_samples - 1] so the resolved count never collapses one side of the split to empty (e.g. n_samples=2, train_size=0.3 would otherwise truncate to 0).

Source code in splytters/utils.py

def resolve_n_train(n_samples: int, train_size: float | int) -> int:
    """Resolve ``train_size`` (fraction or absolute count) to an int count.

    A fractional ``train_size`` is clamped to ``[1, n_samples - 1]`` so the
    resolved count never collapses one side of the split to empty (e.g.
    ``n_samples=2, train_size=0.3`` would otherwise truncate to 0).
    """
    if isinstance(train_size, (int, np.integer)) and not isinstance(train_size, bool):
        return int(train_size)
    n_train = int(n_samples * train_size)
    return min(max(n_train, 1), n_samples - 1)

apportion_train

apportion_train(bin_sizes: list[int], n_train_total: int) -> np.ndarray

Split n_train_total train slots across bins via largest-remainder.

Each bin gets floor(size * n_train_total / total) slots, then the leftover slots go to the bins with the largest fractional remainders. The per-bin counts sum exactly to n_train_total (clamped to [0, total]), avoiding both the per-bin int() truncation bias that undershoots the requested train fraction and the all-to-train rounding that can empty the test set.

Parameters:

Name	Type	Description	Default
bin_sizes	list[int]	number of samples in each (non-empty) bin.	required
n_train_total	int	total samples to assign to train across all bins.	required

Returns:

Type	Description
ndarray	Integer ndarray of per-bin train counts, aligned with bin_sizes.

Source code in splytters/utils.py

def apportion_train(bin_sizes: list[int], n_train_total: int) -> np.ndarray:
    """Split ``n_train_total`` train slots across bins via largest-remainder.

    Each bin gets ``floor(size * n_train_total / total)`` slots, then the
    leftover slots go to the bins with the largest fractional remainders. The
    per-bin counts sum exactly to ``n_train_total`` (clamped to ``[0, total]``),
    avoiding both the per-bin ``int()`` truncation bias that undershoots the
    requested train fraction and the all-to-train rounding that can empty the
    test set.

    Args:
        bin_sizes: number of samples in each (non-empty) bin.
        n_train_total: total samples to assign to train across all bins.

    Returns:
        Integer ndarray of per-bin train counts, aligned with ``bin_sizes``.
    """
    sizes = np.asarray(bin_sizes, dtype=np.intp)
    total = int(sizes.sum())
    if total == 0:
        return np.zeros(len(sizes), dtype=np.intp)
    n_train_total = int(min(max(n_train_total, 0), total))

    ideal = sizes * (n_train_total / total)
    counts = np.floor(ideal).astype(np.intp)
    remainder = n_train_total - int(counts.sum())
    if remainder > 0:
        # Hand leftover slots to the largest fractional remainders that still
        # have spare capacity (ideal <= size, so capacity always exists here).
        order = np.argsort(-(ideal - np.floor(ideal)))
        for i in order:
            if remainder == 0:
                break
            if counts[i] < sizes[i]:
                counts[i] += 1
                remainder -= 1
    return counts

as_index_array

as_index_array(indices: ArrayLike) -> np.ndarray

Return indices as a 1-D integer ndarray (the canonical split output).

Source code in splytters/utils.py

def as_index_array(indices: ArrayLike) -> np.ndarray:
    """Return ``indices`` as a 1-D integer ndarray (the canonical split output)."""
    return np.asarray(list(indices), dtype=np.intp)

compute_pairwise_distances

compute_pairwise_distances(X: ArrayLike, metric: str = 'euclidean') -> np.ndarray

Compute pairwise distance matrix.

.. warning:: Materializes a full O(n²) matrix. For large datasets, prefer NearestNeighbors or chunked computation.

Source code in splytters/utils.py

def compute_pairwise_distances(X: ArrayLike, metric: str = "euclidean") -> np.ndarray:
    """Compute pairwise distance matrix.

    .. warning::
        Materializes a full O(n²) matrix. For large datasets, prefer
        NearestNeighbors or chunked computation.
    """
    # TODO: Add an optional max_samples guard or return a sparse representation
    # for large inputs to avoid OOM.
    X = np.asarray(X)
    return cdist(X, X, metric=metric)

compute_centroid

compute_centroid(X: ArrayLike) -> np.ndarray

Compute centroid of embeddings.

Source code in splytters/utils.py

def compute_centroid(X: ArrayLike) -> np.ndarray:
    """Compute centroid of embeddings."""
    X = np.asarray(X)
    return X.mean(axis=0)

compute_split_centroids

compute_split_centroids(X: ArrayLike, train_indices: ArrayLike, test_indices: ArrayLike) -> tuple[np.ndarray | None, np.ndarray | None]

Compute centroids of train and test sets.

Source code in splytters/utils.py

def compute_split_centroids(
    X: ArrayLike, train_indices: ArrayLike, test_indices: ArrayLike
) -> tuple[np.ndarray | None, np.ndarray | None]:
    """Compute centroids of train and test sets."""
    X = np.asarray(X)
    train_indices = np.asarray(train_indices, dtype=np.intp)
    test_indices = np.asarray(test_indices, dtype=np.intp)
    train_centroid = X[train_indices].mean(axis=0) if len(train_indices) else None
    test_centroid = X[test_indices].mean(axis=0) if len(test_indices) else None
    return train_centroid, test_centroid

cluster_embeddings

cluster_embeddings(X: ArrayLike, n_clusters: int = 10, method: str = 'kmeans', random_state: int = 42, **kwargs: Any) -> tuple[np.ndarray, dict[int, list[int]], np.ndarray]

Cluster embeddings and return labels and cluster info.

Returns:

Name	Type	Description
labels	ndarray	cluster label for each sample
cluster_to_indices	dict[int, list[int]]	dict mapping cluster_id to list of indices
cluster_centers	ndarray	cluster centroids (for kmeans)

Source code in splytters/utils.py

def cluster_embeddings(
    X: ArrayLike,
    n_clusters: int = 10,
    method: str = "kmeans",
    random_state: int = 42,
    **kwargs: Any,
) -> tuple[np.ndarray, dict[int, list[int]], np.ndarray]:
    """
    Cluster embeddings and return labels and cluster info.

    Returns:
        labels: cluster label for each sample
        cluster_to_indices: dict mapping cluster_id to list of indices
        cluster_centers: cluster centroids (for kmeans)
    """
    X = np.asarray(X)

    if method == "kmeans":
        clusterer = KMeans(
            n_clusters=n_clusters,
            random_state=random_state,
            n_init="auto",
            **kwargs
        )
        labels = clusterer.fit_predict(X)
        cluster_centers = clusterer.cluster_centers_
    else:
        raise ValueError(f"Unknown clustering method: {method}")

    cluster_to_indices = defaultdict(list)
    for idx, label in enumerate(labels):
        cluster_to_indices[label].append(idx)

    return labels, cluster_to_indices, cluster_centers

random_split

random_split(embeddings: ArrayLike, train_size: float | int = 0.7, random_state: int = 42) -> tuple[np.ndarray, np.ndarray]

Simple random train/test split (baseline).

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array-like of shape (n_samples, embedding_dim)	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	0.7
random_state	int	int, RandomState, or None for reproducibility	42

Returns:

Name	Type	Description
train_indices	ndarray	ndarray of indices for training set
test_indices	ndarray	ndarray of indices for test set

Source code in splytters/utils.py

def random_split(
    embeddings: ArrayLike, train_size: float | int = 0.7, random_state: int = 42
) -> tuple[np.ndarray, np.ndarray]:
    """
    Simple random train/test split (baseline).

    Args:
        embeddings: array-like of shape (n_samples, embedding_dim)
        train_size: fraction in (0, 1) or absolute count for the training set
        random_state: int, RandomState, or None for reproducibility

    Returns:
        train_indices: ndarray of indices for training set
        test_indices: ndarray of indices for test set
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    n_train = resolve_n_train(n_samples, train_size)
    rng = check_random_state(random_state)
    indices = np.arange(n_samples)
    rng.shuffle(indices)
    return indices[:n_train], indices[n_train:]

compute_split_similarity

compute_split_similarity(X: ArrayLike, train_indices: ArrayLike, test_indices: ArrayLike, metric: str = 'euclidean') -> dict[str, float]

Compute similarity metrics between train and test splits.

Returns dict with

• centroid_distance: distance between train/test centroids
• mean_cross_distance: mean distance from test to nearest train
• coverage: fraction of test samples with train neighbor within median distance

Source code in splytters/utils.py

def compute_split_similarity(
    X: ArrayLike,
    train_indices: ArrayLike,
    test_indices: ArrayLike,
    metric: str = "euclidean",
) -> dict[str, float]:
    """
    Compute similarity metrics between train and test splits.

    Returns dict with:
        - centroid_distance: distance between train/test centroids
        - mean_cross_distance: mean distance from test to nearest train
        - coverage: fraction of test samples with train neighbor within median distance
    """
    X = np.asarray(X)
    train_indices = np.asarray(train_indices, dtype=np.intp)
    test_indices = np.asarray(test_indices, dtype=np.intp)

    if len(train_indices) == 0 or len(test_indices) == 0:
        raise ValueError(
            "compute_split_similarity requires non-empty train and test sets"
        )

    train_X = X[train_indices]
    test_X = X[test_indices]

    # Centroid distance
    train_centroid = train_X.mean(axis=0)
    test_centroid = test_X.mean(axis=0)
    centroid_distance = np.linalg.norm(train_centroid - test_centroid)

    # Cross-set distances
    cross_distances = cdist(test_X, train_X, metric=metric)
    min_distances = cross_distances.min(axis=1)
    mean_cross_distance = min_distances.mean()

    # Coverage (fraction of test with nearby train sample)
    # TODO: Replace full pairwise matrix with sampled median estimation and
    # NearestNeighbors for coverage check to reduce O(n²) memory.
    all_distances = cdist(X, X, metric=metric)
    np.fill_diagonal(all_distances, np.inf)
    median_dist = np.median(all_distances[all_distances < np.inf])
    coverage = (min_distances <= median_dist).mean()

    return {
        "centroid_distance": float(centroid_distance),
        "mean_cross_distance": float(mean_cross_distance),
        "coverage": float(coverage),
    }

optimized_split

optimized_split(embeddings: ndarray, train_size: float | int, n_iterations: int, score_fn: Callable[[ndarray, list[int], list[int]], float], random_state: int = 42, minimize: bool = True) -> tuple[np.ndarray, np.ndarray]

Iterative swap-optimization to find a split that optimizes a score.

Starts from a random split and repeatedly swaps one train/test pair, keeping the swap only if the score improves.

Parameters:

Name	Type	Description	Default
embeddings	ndarray	array of shape (n_samples, embedding_dim), already np.ndarray	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set	required
n_iterations	int	number of swap attempts	required
score_fn	Callable[[ndarray, list[int], list[int]], float]	callable(embeddings, train_indices, test_indices) -> float	required
random_state	int	int, RandomState, or None for reproducibility	42
minimize	bool	if True, accept swaps that lower the score; if False, accept swaps that raise it	True

Source code in splytters/utils.py

def optimized_split(
    embeddings: np.ndarray,
    train_size: float | int,
    n_iterations: int,
    score_fn: Callable[[np.ndarray, list[int], list[int]], float],
    random_state: int = 42,
    minimize: bool = True,
) -> tuple[np.ndarray, np.ndarray]:
    """Iterative swap-optimization to find a split that optimizes a score.

    Starts from a random split and repeatedly swaps one train/test pair,
    keeping the swap only if the score improves.

    Args:
        embeddings: array of shape (n_samples, embedding_dim), already np.ndarray
        train_size: fraction in (0, 1) or absolute count for the training set
        n_iterations: number of swap attempts
        score_fn: callable(embeddings, train_indices, test_indices) -> float
        random_state: int, RandomState, or None for reproducibility
        minimize: if True, accept swaps that lower the score;
                  if False, accept swaps that raise it
    """
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    all_indices = np.arange(n_samples)
    rng.shuffle(all_indices)

    n_train = resolve_n_train(n_samples, train_size)
    train_indices = set(all_indices[:n_train].tolist())
    test_indices = set(all_indices[n_train:].tolist())

    current_score = score_fn(embeddings, list(train_indices), list(test_indices))

    for _ in range(n_iterations):
        train_sample = rng.choice(list(train_indices))
        test_sample = rng.choice(list(test_indices))

        # swap
        train_indices.remove(train_sample)
        train_indices.add(test_sample)
        test_indices.remove(test_sample)
        test_indices.add(train_sample)

        new_score = score_fn(embeddings, list(train_indices), list(test_indices))

        improved = new_score < current_score if minimize else new_score > current_score
        if improved:
            current_score = new_score
        else:
            # revert
            train_indices.remove(test_sample)
            train_indices.add(train_sample)
            test_indices.remove(train_sample)
            test_indices.add(test_sample)

    return as_index_array(sorted(train_indices)), as_index_array(sorted(test_indices))

greedy_assign_to_target

greedy_assign_to_target(items_with_sizes: list[tuple[int, int]], target_size: int) -> tuple[list[int], list[int]]

Greedily assign items to reach target size.

Parameters:

Name	Type	Description	Default
items_with_sizes	list[tuple[int, int]]	list of (item_id, size) tuples	required
target_size	int	target total size	required

Returns:

Name	Type	Description
selected	list[int]	list of item_ids assigned
remaining	list[int]	list of item_ids not assigned

Source code in splytters/utils.py

def greedy_assign_to_target(
    items_with_sizes: list[tuple[int, int]], target_size: int
) -> tuple[list[int], list[int]]:
    """
    Greedily assign items to reach target size.

    Args:
        items_with_sizes: list of (item_id, size) tuples
        target_size: target total size

    Returns:
        selected: list of item_ids assigned
        remaining: list of item_ids not assigned
    """
    selected = []
    remaining = []
    current_size = 0

    for item_id, size in items_with_sizes:
        if current_size + size <= target_size:
            selected.append(item_id)
            current_size += size
        else:
            remaining.append(item_id)

    return selected, remaining

Embedders

embedders

Embedder

Bases: ABC

Base class for all embedders.

Source code in splytters/embedders.py

class Embedder(ABC):
    """Base class for all embedders."""

    @abstractmethod
    def embed(self, inputs: Sequence[Any]) -> np.ndarray:
        """Embed a list of inputs. Returns np.ndarray of shape (n, dim)."""

embed abstractmethod

embed(inputs: Sequence[Any]) -> np.ndarray

Embed a list of inputs. Returns np.ndarray of shape (n, dim).

Source code in splytters/embedders.py

@abstractmethod
def embed(self, inputs: Sequence[Any]) -> np.ndarray:
    """Embed a list of inputs. Returns np.ndarray of shape (n, dim)."""

TextEmbedder

Bases: Embedder

Embed text using a SentenceTransformer model.

Source code in splytters/embedders.py

class TextEmbedder(Embedder):
    """Embed text using a SentenceTransformer model."""

    def __init__(self, model_name: str = "all-MiniLM-L6-v2") -> None:
        from sentence_transformers import SentenceTransformer

        self.model = SentenceTransformer(model_name)

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        return self.model.encode(texts, convert_to_numpy=True)

CLIPTextEmbedder

Bases: Embedder

Embed text using a CLIP model.

Source code in splytters/embedders.py

class CLIPTextEmbedder(Embedder):
    """Embed text using a CLIP model."""

    def __init__(self, model_name: str = "openai/clip-vit-base-patch32") -> None:
        from transformers import CLIPModel, CLIPTokenizerFast

        self.model = CLIPModel.from_pretrained(model_name)
        self.tokenizer = CLIPTokenizerFast.from_pretrained(model_name)

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        inputs = self.tokenizer(texts, padding=True, truncation=True, return_tensors="pt")
        outputs = self.model.get_text_features(**inputs)
        return _features_to_numpy(outputs)

CLIPImageEmbedder

Bases: Embedder

Embed images using a CLIP model.

Source code in splytters/embedders.py

class CLIPImageEmbedder(Embedder):
    """Embed images using a CLIP model."""

    def __init__(self, model_name: str = "openai/clip-vit-base-patch32") -> None:
        from transformers import CLIPModel, CLIPProcessor

        self.model = CLIPModel.from_pretrained(model_name)
        self.processor = CLIPProcessor.from_pretrained(model_name)

    def embed(self, images: Sequence[Any]) -> np.ndarray:
        inputs = self.processor(images=images, return_tensors="pt")
        outputs = self.model.get_image_features(**inputs)
        return _features_to_numpy(outputs)

OpenAIEmbedder

Bases: Embedder

Embed text using the OpenAI embeddings API.

Source code in splytters/embedders.py

class OpenAIEmbedder(Embedder):
    """Embed text using the OpenAI embeddings API."""

    def __init__(self, model_name: str = "text-embedding-3-small") -> None:
        from openai import OpenAI

        self.client = OpenAI()
        self.model_name = model_name

    def embed(self, texts: Sequence[str]) -> np.ndarray:
        response = self.client.embeddings.create(input=texts, model=self.model_name)
        return np.array([item.embedding for item in response.data])

list_embedders

list_embedders() -> list[str]

Return the names of the available concrete embedder classes.

These live in :mod:splytters.embedders (the [embedders] extra). The underlying model libraries (sentence-transformers, transformers, openai) are imported lazily only when an embedder is constructed, so listing pulls in no optional dependency.

Returns:

Type	Description
list[str]	Embedder class names, e.g. ``["TextEmbedder", "CLIPTextEmbedder",
list[str]	"CLIPImageEmbedder", "OpenAIEmbedder"]``.

Source code in splytters/embedders.py

def list_embedders() -> list[str]:
    """Return the names of the available concrete embedder classes.

    These live in :mod:`splytters.embedders` (the ``[embedders]`` extra). The
    underlying model libraries (sentence-transformers, transformers, openai)
    are imported lazily only when an embedder is constructed, so listing pulls
    in no optional dependency.

    Returns:
        Embedder class names, e.g. ``["TextEmbedder", "CLIPTextEmbedder",
        "CLIPImageEmbedder", "OpenAIEmbedder"]``.
    """
    return [cls.__name__ for cls in Embedder.__subclasses__()]

Introspection

list_splitters

list_splitters(by_family: bool = False) -> list[str] | dict[str, list[str]]

Return the names of all available splitter functions.

Parameters:

Name	Type	Description	Default
by_family	bool	if True, return a dict mapping each family ("adversarial", "overlap", "balanced", "baseline") to its list of splitter names. If False (default), return a flat list of every splitter name, in family order.	False

Returns:

Type	Description
list[str] | dict[str, list[str]]	A flat list of names, or a dict of family -> names.

Source code in splytters/__init__.py

def list_splitters(by_family: bool = False) -> list[str] | dict[str, list[str]]:
    """Return the names of all available splitter functions.

    Args:
        by_family: if True, return a dict mapping each family
            ("adversarial", "overlap", "balanced", "baseline") to its list of
            splitter names. If False (default), return a flat list of every
            splitter name, in family order.

    Returns:
        A flat list of names, or a dict of family -> names.
    """
    if by_family:
        return {family: list(names) for family, names in _SPLITTER_FAMILIES.items()}
    return [name for names in _SPLITTER_FAMILIES.values() for name in names]

list_sorters

list_sorters(by_modality: bool = False) -> list[str] | dict[str, list[str]]

Return the names of all available sorters.

Listing pulls in no optional dependency — only the names are returned; the sorters themselves stay lazily imported until first accessed.

Parameters:

Name	Type	Description	Default
by_modality	bool	if True, return a dict mapping each modality ("embedding", "text", "image", "audio", "tabular") to its sorted list of sorter names. If False (default), return a flat sorted list of every sorter name.	False

Returns:

Type	Description
list[str] | dict[str, list[str]]	A sorted list of names, or a dict of modality -> sorted names.

Source code in splytters/sorters/__init__.py

def list_sorters(by_modality: bool = False) -> list[str] | dict[str, list[str]]:
    """Return the names of all available sorters.

    Listing pulls in no optional dependency — only the names are returned;
    the sorters themselves stay lazily imported until first accessed.

    Args:
        by_modality: if True, return a dict mapping each modality
            ("embedding", "text", "image", "audio", "tabular") to its sorted
            list of sorter names. If False (default), return a flat sorted
            list of every sorter name.

    Returns:
        A sorted list of names, or a dict of modality -> sorted names.
    """
    if by_modality:
        groups: dict[str, list[str]] = {}
        for name, (module_suffix, _attr) in _LAZY.items():
            modality = module_suffix.removesuffix("_sorters")
            groups.setdefault(modality, []).append(name)
        return {modality: sorted(names) for modality, names in groups.items()}
    return sorted(_LAZY)

list_embedders

list_embedders() -> list[str]

Return the names of the available concrete embedder classes.

These live in :mod:splytters.embedders (the [embedders] extra). The underlying model libraries (sentence-transformers, transformers, openai) are imported lazily only when an embedder is constructed, so listing pulls in no optional dependency.

Returns:

Type	Description
list[str]	Embedder class names, e.g. ``["TextEmbedder", "CLIPTextEmbedder",
list[str]	"CLIPImageEmbedder", "OpenAIEmbedder"]``.

Source code in splytters/embedders.py

def list_embedders() -> list[str]:
    """Return the names of the available concrete embedder classes.

    These live in :mod:`splytters.embedders` (the ``[embedders]`` extra). The
    underlying model libraries (sentence-transformers, transformers, openai)
    are imported lazily only when an embedder is constructed, so listing pulls
    in no optional dependency.

    Returns:
        Embedder class names, e.g. ``["TextEmbedder", "CLIPTextEmbedder",
        "CLIPImageEmbedder", "OpenAIEmbedder"]``.
    """
    return [cls.__name__ for cls in Embedder.__subclasses__()]

Types

Splitter module-attribute

Splitter = Callable[..., tuple[np.ndarray, np.ndarray]]