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_range: tuple[int, int] | None = None, fill_individual: bool = False, fill_anchor: str = 'cluster_centroids', **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 the subset of clusters whose pooled per-class counts stay within, and best cover, a class-balanced target test set (exact multidimensional subset-sum, greedy fallback for large problems), then complete the per-class shortfall with randomly chosen individual examples so the test set hits the exact class-balanced target. Requires y. - "closest": seed the cluster farthest from the mean of all cluster centroids and grow it by single-linkage nearest-neighbor into a single coherent test "pocket" (adversarial); with fill_individual it is then topped up to the exact target size.	'size'
y	ArrayLike | None	class labels of shape (n_samples,). Required for strategy="subset_sum"; for the other strategies it is optional and, when given, drives the class-distribution selection criterion of the cluster_range search and the class-aware individual fill.	None
cluster_range	tuple[int, int] | None	None (default) uses the fixed n_clusters path. When set to (low, high) (the paper uses (3, 50)), the split search is repeated for every k in that inclusive range (KMeans only) and the best split is kept. For the paper strategies the criterion is the paper's: keep the k requiring the fewest individually-added examples to reach the target (subset_sum: the smallest whole-cluster coverage gap; closest: the largest whole-cluster pocket). For size/centroid it is the L1 distance between the whole-cluster test set's per-class counts and target_test * class_proportions (absolute test-size deviation without y). Ties are broken toward the smaller k.	None
fill_individual	bool	False (default) keeps the whole-cluster "closest" pocket. When True and strategy="closest", the pocket is topped up to exactly target_test with individual examples nearest (cosine) to the nearest selected test-cluster centroid, class-aware toward target_test * class_proportions when y is given -- the paper's final individual-example fill step. Ignored by the other strategies (subset_sum always completes to the exact target).	False
fill_anchor	str	what the fill measures proximity to. The paper's prose and its released code disagree here, so both are offered: - "cluster_centroids" (default): distance to the nearest selected test-cluster centroid, per the paper's wording ("individual examples that are closest to one of the test set centroids"). - "test_mean": distance to the single mean of all embeddings currently in the test pocket, which is what the authors' released code (closest_clusters.py) actually computes. Only used with strategy="closest" and fill_individual=True.	'cluster_centroids'
**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, strategy, or fill_anchor; if strategy="subset_sum" is used without a valid y; or if cluster_range is combined with a non-KMeans method or is not a valid (low, high) pair.

Warns:

Type	Description
UserWarning	if whole-cluster assignment would leave train or test empty. The function returns a seeded exact-size sample split because no cluster-coherent non-empty split exists.

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.

Faithful reproduction (matched to the authors' reference code): SUBSET-SUM-SPLIT solves the multidimensional subset-sum over cluster class-count vectors exactly by dynamic programming (a greedy fallback engages only past an internal state cap for large / high-class-count problems) and completes the remaining per-class shortfall with random individual examples, so the test set always hits the exact class-balanced target. CLOSEST-SPLIT seeds the cluster farthest from the centroid of all centroids, grows the pocket by single-linkage nearest neighbour, then (with fill_individual=True) fills to the exact target with the individual examples closest to the fill_anchor. On the fill anchor the paper and its released code diverge: the paper's prose says examples "closest to one of the test set centroids" (the default, fill_anchor="cluster_centroids"), while the released code measures distance to the mean of all current test embeddings (fill_anchor="test_mean"). To reproduce the paper's cluster-count search, pass cluster_range=(3, 50): it keeps the k requiring the fewest individually-added examples (Section "Splitting methods"). subset_sum enforces equal train/test class distributions as a hard constraint (the test set exactly matches target_test * class_proportions); closest tracks that target class-aware during the fill but pocket growth is capped by total size only, so a class can overshoot its per-class target by a small margin.

Precondition for full faithfulness: the paper extracts the [CLS] representations of a language model finetuned on the task (its RQ1 shows this matters); pass such embeddings to reproduce the paper. On arbitrary or pretrained-only embeddings the same procedure runs but is an adaptation. The fixed-n_clusters path (cluster_range=None) runs the same exact assignment at a single caller-chosen k instead of the paper's search.

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.

Seed stability: varies with the seed -- the KMeans clustering, and which whole clusters are held out, shift with the seed, so the test set can differ substantially between runs (from near-identical to nearly disjoint) even though the clusters themselves are similar.

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_range: tuple[int, int] | None = None,
    fill_individual: bool = False,
    fill_anchor: str = "cluster_centroids",
    **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 the subset of clusters whose pooled
              per-class counts stay within, and best cover, a class-balanced
              target test set (exact multidimensional subset-sum, greedy fallback
              for large problems), then complete the per-class shortfall with
              randomly chosen individual examples so the test set hits the exact
              class-balanced target. Requires ``y``.
            - ``"closest"``: seed the cluster farthest from the mean of all
              cluster centroids and grow it by single-linkage nearest-neighbor
              into a single coherent test "pocket" (adversarial); with
              ``fill_individual`` it is then topped up to the exact target size.
        y: class labels of shape (n_samples,). Required for ``strategy="subset_sum"``;
            for the other strategies it is optional and, when given, drives the
            class-distribution selection criterion of the ``cluster_range`` search
            and the class-aware individual fill.
        cluster_range: ``None`` (default) uses the fixed ``n_clusters`` path.
            When set to ``(low, high)`` (the paper uses ``(3, 50)``), the split
            search is repeated for every ``k`` in that inclusive range (KMeans
            only) and the best split is kept. For the paper strategies the
            criterion is the paper's: keep the ``k`` requiring the *fewest*
            individually-added examples to reach the target (``subset_sum``: the
            smallest whole-cluster coverage gap; ``closest``: the largest
            whole-cluster pocket). For ``size``/``centroid`` it is the L1 distance
            between the whole-cluster test set's per-class counts and
            ``target_test * class_proportions`` (absolute test-size deviation
            without ``y``). Ties are broken toward the smaller ``k``.
        fill_individual: ``False`` (default) keeps the whole-cluster ``"closest"``
            pocket. When ``True`` and ``strategy="closest"``, the pocket is topped
            up to exactly ``target_test`` with individual examples nearest (cosine)
            to the nearest selected test-cluster centroid, class-aware toward
            ``target_test * class_proportions`` when ``y`` is given -- the paper's
            final individual-example fill step. Ignored by the other strategies
            (``subset_sum`` always completes to the exact target).
        fill_anchor: what the fill measures proximity to. The paper's prose and
            its released code disagree here, so both are offered:
            - ``"cluster_centroids"`` (default): distance to the *nearest selected
              test-cluster centroid*, per the paper's wording ("individual
              examples that are closest to one of the test set centroids").
            - ``"test_mean"``: distance to the single mean of all embeddings
              currently in the test pocket, which is what the authors' released
              code (``closest_clusters.py``) actually computes.
            Only used with ``strategy="closest"`` and ``fill_individual=True``.
        **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``, ``strategy``, or ``fill_anchor``; if
            ``strategy="subset_sum"`` is used without a valid ``y``; or if
            ``cluster_range`` is combined with a non-KMeans ``method`` or is not a
            valid ``(low, high)`` pair.

    Warns:
        UserWarning: if whole-cluster assignment would leave train or test empty.
            The function returns a seeded exact-size sample split because no
            cluster-coherent non-empty split exists.

    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`.

        Faithful reproduction (matched to the authors' reference code):
        SUBSET-SUM-SPLIT solves the multidimensional subset-sum over cluster
        class-count vectors *exactly* by dynamic programming (a greedy fallback
        engages only past an internal state cap for large / high-class-count
        problems) and completes the remaining per-class shortfall with random
        individual examples, so the test set always hits the exact class-balanced
        target. CLOSEST-SPLIT seeds the cluster farthest from the centroid of all
        centroids, grows the pocket by single-linkage nearest neighbour, then
        (with ``fill_individual=True``) fills to the exact target with the
        individual examples closest to the ``fill_anchor``. On the fill anchor
        the paper and its released code diverge: the paper's prose says examples
        "closest to one of the test set centroids" (the default,
        ``fill_anchor="cluster_centroids"``), while the released code measures
        distance to the mean of all current test embeddings
        (``fill_anchor="test_mean"``). To reproduce the paper's cluster-count
        search, pass ``cluster_range=(3, 50)``: it keeps the ``k`` requiring the
        fewest individually-added examples (Section "Splitting methods").
        ``subset_sum`` enforces equal train/test class distributions as a hard
        constraint (the test set exactly matches
        ``target_test * class_proportions``); ``closest`` tracks that target
        class-aware during the fill but pocket growth is capped by total size
        only, so a class can overshoot its per-class target by a small margin.

        Precondition for full faithfulness: the paper extracts the ``[CLS]``
        representations of a language model *finetuned on the task* (its RQ1 shows
        this matters); pass such embeddings to reproduce the paper. On arbitrary
        or pretrained-only embeddings the same procedure runs but is an adaptation.
        The fixed-``n_clusters`` path (``cluster_range=None``) runs the same exact
        assignment at a single caller-chosen ``k`` instead of the paper's search.

        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.

    Seed stability: varies with the seed -- the KMeans clustering, and which
    whole clusters are held out, shift with the seed, so the test set can differ
    substantially between runs (from near-identical to nearly disjoint) even
    though the clusters themselves are similar.
    """
    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" and y is None:
        raise ValueError("strategy='subset_sum' requires class labels `y`")
    if y is not None:
        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 not in {"kmeans", "dbscan"}:
        raise ValueError(f"Unknown clustering method: {method}")

    if fill_anchor not in {"cluster_centroids", "test_mean"}:
        raise ValueError(
            f"Unknown fill_anchor: {fill_anchor!r}. "
            "Choose 'cluster_centroids' (paper prose) or 'test_mean' (reference code)"
        )

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

    classes = np.unique(y) if y is not None else None
    subset_target = (
        _class_target_vector(y, classes, target_test)
        if strategy == "subset_sum"
        else None
    )
    completion_rng = check_random_state(random_state)

    def cluster_at(k: int) -> dict[int, list[int]]:
        if method == "kmeans":
            clusterer: Any = KMeans(
                n_clusters=k,
                random_state=random_state,
                n_init="auto",
                **cluster_kwargs,
            )
        else:  # dbscan (ignores k)
            clusterer = DBSCAN(**cluster_kwargs)
        labels = clusterer.fit_predict(embeddings)
        mapping: dict[int, list[int]] = defaultdict(list)
        for idx, label in enumerate(labels):
            mapping[int(label)].append(idx)
        return mapping

    def core(
        mapping: dict[int, list[int]],
    ) -> tuple[list[int], list[int], float, np.ndarray | None]:
        """Pre-completion assignment, its k-search score, and any completion aid.

        Score is what the paper minimises when searching k: for the paper methods
        it is the number of individual examples still needed to reach the target
        (fewest completions / smallest coverage gap); for the size/centroid
        strategies it is the L1 distance of the whole-cluster test set to the
        class-balanced target. Lower is better; ties keep the smaller k.
        """
        if strategy == "size":
            tr, te = _assign_by_size(mapping, target_train)
            return tr, te, _test_target_distance(te, y, target_test, n_samples), None
        if strategy == "centroid":
            tr, te = _assign_by_centroid(embeddings, mapping, target_train)
            return tr, te, _test_target_distance(te, y, target_test, n_samples), None
        if strategy == "closest":
            tr, te, test_centroids = _closest_core(embeddings, mapping, target_test)
            gap = target_test - len(te)
            # Individual examples still to add; an overshooting pocket (every
            # cluster larger than the target -- only at very small k) can't be
            # filled down, so rank it below every pocket that fits.
            score = float(gap) if gap >= 0 else float(target_test - gap)
            return tr, te, score, test_centroids
        # subset_sum
        tr, te = _subset_sum_core(mapping, y, classes, subset_target)
        return tr, te, float(target_test - len(te)), None

    def complete(
        tr: list[int], te: list[int], aid: np.ndarray | None
    ) -> tuple[list[int], list[int]]:
        """Apply the paper's per-example completion to the chosen split."""
        if strategy == "closest" and fill_individual:
            # "test_mean" passes None so the fill falls back to the mean of the
            # current test pocket -- the released reference-code behavior.
            return _fill_individual_examples(
                embeddings,
                tr,
                te,
                target_test,
                y=y,
                test_centroids=aid if fill_anchor == "cluster_centroids" else None,
            )
        if strategy == "subset_sum":
            return _complete_subset_sum(
                tr, te, y, classes, subset_target, completion_rng
            )
        return tr, te

    if cluster_range is None:
        train, test, _, aid = core(cluster_at(n_clusters))
    else:
        if method != "kmeans":
            raise ValueError("cluster_range requires method='kmeans'")
        low, high = cluster_range
        if not (isinstance(low, int) and isinstance(high, int)) or low < 2 or low > high:
            raise ValueError(
                f"cluster_range must be a (low, high) pair with 2 <= low <= high; "
                f"got {cluster_range!r}"
            )
        high = min(high, n_samples)  # cannot ask KMeans for more clusters than points
        low = min(low, high)  # keep the range non-empty after clamping
        best: tuple[list[int], list[int], float, np.ndarray | None] | None = None
        best_score = float("inf")
        for k in range(low, high + 1):
            cand = core(cluster_at(k))
            if cand[2] < best_score:  # strict: ties keep the smaller k
                best_score = cand[2]
                best = cand
        assert best is not None  # range is non-empty (low <= high after clamp)
        train, test, _, aid = best

    train, test = complete(train, test, aid)

    if not train or not test:
        warnings.warn(
            "cluster_split's whole-cluster assignment left train or test empty; "
            "falling back to a seeded exact-size sample split because preserving "
            "the effective clusters cannot satisfy the non-empty split contract.",
            UserWarning,
            stacklevel=2,
        )
        train, test = _seeded_exact_fallback_split(
            n_samples, target_train, random_state
        )

    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

Seed stability: varies with the seed -- built on KMeans clustering, so which whole clusters land in test shifts between seeds and the held-out set can change substantially.

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

    Seed stability: varies with the seed -- built on KMeans clustering, so which
    whole clusters land in test shifts between seeds and the held-out set can
    change substantially.
    """
    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' (default) or 'dbscan' use the cluster-coherent greedy fold-packing heuristic described below. 'sds_kmeans' instead runs the paper's SDS K-means, where the n_folds clusters ARE the folds (see the SDS note in References).	'kmeans'
n_clusters	int | None	number of clusters for the 'kmeans' heuristic. Defaults to max(10, n_folds * 3) -- comfortably above n_folds so every fold can receive whole clusters. Ignored by 'sds_kmeans', which always uses exactly n_folds clusters.	None
random_state	int	for reproducibility	42
**cluster_kwargs	Any	passed to the clustering algorithm. For method='sds_kmeans' the accepted keys are n_init (number of k-means++ initializations, default 8) and max_iter (max swap rounds, default 100).	{}

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

Wecker, Friedrich & Adel (2020), "ClusterDataSplit: Exploring Challenging Clustering-Based Data Splits for Model Performance Evaluation," Eval4NLP @ COLING. https://aclanthology.org/2020.eval4nlp-1.15 Reference code (archived): https://github.com/boschresearch/clusterdatasplit_eval4nlp-2020

The default 'kmeans'/'dbscan' modes are a cluster-coherent, label-aware fold-assignment heuristic inspired by ClusterDataSplit: they cluster with off-the-shelf KMeans/DBSCAN and greedily pack whole clusters into folds that are lexically distinct from train while preserving label balance. They do NOT implement the paper's SDS K-means.

method='sds_kmeans' implements the paper's SDS K-means (Size and Distribution Sensitive K-means, an extension of Same-Size K-means). The n_folds clusters are the folds directly. Each label's examples are spread over the clusters as evenly as possible (label-specific capacities), so every fold is ~n/n_folds in size with a label distribution close to global. Points are placed by an order-measure assignment (nearest-minus-farthest squared distance) that fills clusters up to their per-label capacity, then refined by 1-on-1 swaps: two points in different clusters that each want the other's cluster are exchanged, which lowers overall cluster-internal variance while keeping sizes fixed. Iteration stops when a full pass yields no swaps; the best of n_init k-means++ initializations (by inertia) is returned.

Fidelity notes vs. the reference code: the reference exposes max_ids_labels for the caller to specify per-cluster/per-label capacities; here they are derived automatically from the global label distribution (the paper's default SDS case, which also gives equal-size clusters as in Same-Size K-means). The reference uses sklearn's private _init_centroids; this uses the public kmeans_plusplus. Assignment, capacity handling, the distance-sorted 1-on-1 swap update, the inertia objective (sum of per-cluster mean squared distance), and the "stop when no swaps" rule follow the reference implementation.

Seed stability: structure-stable. For the default modes, fold assignment varies with the KMeans clustering, though the cluster-coherent, label-balanced structure is preserved. For 'sds_kmeans', fold SIZES and per-label counts are fixed by the capacities (seed-invariant); only which samples land in which fold varies with the seed.

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'`` (default) or ``'dbscan'`` use
            the cluster-coherent greedy fold-packing heuristic described below.
            ``'sds_kmeans'`` instead runs the paper's SDS K-means, where the
            ``n_folds`` clusters ARE the folds (see the SDS note in References).
        n_clusters: number of clusters for the ``'kmeans'`` heuristic. Defaults to
            ``max(10, n_folds * 3)`` -- comfortably above ``n_folds`` so every
            fold can receive whole clusters. Ignored by ``'sds_kmeans'``, which
            always uses exactly ``n_folds`` clusters.
        random_state: for reproducibility
        **cluster_kwargs: passed to the clustering algorithm. For
            ``method='sds_kmeans'`` the accepted keys are ``n_init`` (number of
            k-means++ initializations, default 8) and ``max_iter`` (max swap
            rounds, default 100).

    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:
        Wecker, Friedrich & Adel (2020), "ClusterDataSplit: Exploring Challenging
        Clustering-Based Data Splits for Model Performance Evaluation,"
        Eval4NLP @ COLING. https://aclanthology.org/2020.eval4nlp-1.15
        Reference code (archived):
        https://github.com/boschresearch/clusterdatasplit_eval4nlp-2020

        The default ``'kmeans'``/``'dbscan'`` modes are a cluster-coherent,
        label-aware fold-assignment heuristic inspired by ClusterDataSplit: they
        cluster with off-the-shelf KMeans/DBSCAN and greedily pack whole clusters
        into folds that are lexically distinct from train while preserving label
        balance. They do NOT implement the paper's SDS K-means.

        ``method='sds_kmeans'`` implements the paper's SDS K-means (Size and
        Distribution Sensitive K-means, an extension of Same-Size K-means). The
        ``n_folds`` clusters are the folds directly. Each label's examples are
        spread over the clusters as evenly as possible (label-specific
        capacities), so every fold is ~``n/n_folds`` in size with a label
        distribution close to global. Points are placed by an order-measure
        assignment (nearest-minus-farthest squared distance) that fills clusters
        up to their per-label capacity, then refined by 1-on-1 swaps: two points
        in different clusters that each want the other's cluster are exchanged,
        which lowers overall cluster-internal variance while keeping sizes fixed.
        Iteration stops when a full pass yields no swaps; the best of ``n_init``
        k-means++ initializations (by inertia) is returned.

        Fidelity notes vs. the reference code: the reference exposes
        ``max_ids_labels`` for the caller to specify per-cluster/per-label
        capacities; here they are derived automatically from the global label
        distribution (the paper's default SDS case, which also gives equal-size
        clusters as in Same-Size K-means). The reference uses
        sklearn's private ``_init_centroids``; this uses the public
        ``kmeans_plusplus``. Assignment, capacity handling, the distance-sorted
        1-on-1 swap update, the inertia objective (sum of per-cluster mean
        squared distance), and the "stop when no swaps" rule follow the reference
        implementation.

    Seed stability: structure-stable. For the default modes, fold assignment
    varies with the KMeans clustering, though the cluster-coherent,
    label-balanced structure is preserved. For ``'sds_kmeans'``, fold SIZES and
    per-label counts are fixed by the capacities (seed-invariant); only which
    samples land in which fold varies with the seed.
    """
    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 method == "sds_kmeans":
        # SDS K-means: the n_folds clusters ARE the folds (no separate
        # clustering + greedy packing). Handled up front so the existing
        # kmeans/dbscan paths below are untouched.
        return _sds_kmeans_folds(
            embeddings, y, n_folds, random_state, **cluster_kwargs
        )

    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_route

minority_route(cluster_labels: ArrayLike, y: ArrayLike, minority_labels: str = 'all_but_majority') -> tuple[np.ndarray, np.ndarray]

Route a precomputed clustering into a bias-amplified minority split.

The label-only half of :func:minority_split: given an existing cluster assignment (from any source -- kmeans, Ward, or an external clusterer such as a faithful DEEP CLUSTER pass), send each cluster's majority-label instances to train and its minority-label instances to test. Pure and gradient-free -- it inspects only per-cluster label counts. Exposed so heavier, out-of-library clusterers can reuse the exact same faithful routing (incl. footnote 10).

Parameters:

Name	Type	Description	Default
cluster_labels	ArrayLike	integer cluster id per sample, shape (n_samples,).	required
y	ArrayLike	class labels, shape (n_samples,).	required
minority_labels	str	'all_but_majority' (default) or 'least_only' -- see :func:minority_split.	'all_but_majority'

Returns:

Type	Description
tuple[ndarray, ndarray]	(train_indices, test_indices), each a sorted index array.

Raises:

Type	Description
ValueError	on an unknown minority_labels, a length mismatch, or if no minority examples exist (every cluster is label-pure).

Source code in splytters/adversarial.py

def minority_route(
    cluster_labels: ArrayLike,
    y: ArrayLike,
    minority_labels: str = "all_but_majority",
) -> tuple[np.ndarray, np.ndarray]:
    """Route a precomputed clustering into a bias-amplified minority split.

    The label-only half of :func:`minority_split`: given an existing cluster
    assignment (from any source -- kmeans, Ward, or an external clusterer such as a
    faithful DEEP CLUSTER pass), send each cluster's *majority*-label instances to
    train and its *minority*-label instances to test. Pure and gradient-free -- it
    inspects only per-cluster label counts. Exposed so heavier, out-of-library
    clusterers can reuse the exact same faithful routing (incl. footnote 10).

    Args:
        cluster_labels: integer cluster id per sample, shape (n_samples,).
        y: class labels, shape (n_samples,).
        minority_labels: ``'all_but_majority'`` (default) or ``'least_only'`` --
            see :func:`minority_split`.

    Returns:
        (train_indices, test_indices), each a sorted index array.

    Raises:
        ValueError: on an unknown ``minority_labels``, a length mismatch, or if no
            minority examples exist (every cluster is label-pure).
    """
    if minority_labels not in ("all_but_majority", "least_only"):
        raise ValueError(f"Unknown minority_labels: {minority_labels}")
    cluster_labels = np.asarray(cluster_labels)
    y = np.asarray(y)
    if len(cluster_labels) != len(y):
        raise ValueError(
            f"cluster_labels has length {len(cluster_labels)} but y has {len(y)}"
        )

    cluster_to_indices: dict[int, list[int]] = defaultdict(list)
    for idx, label in enumerate(cluster_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)
        if len(vals) < 2:
            train.extend(idx_arr.tolist())  # label-pure cluster: no minority here
            continue
        if minority_labels == "least_only":
            # Footnote 10 (Reif & Schwartz, 2023): only the single least-frequent
            # label per cluster is anti-biased. Keeps the test set small on
            # many-class data, where "all but majority" would flood it.
            minority = vals[np.argmin(counts)]  # ties -> lowest label (deterministic)
            is_test = cls_labels == minority
        else:  # "all_but_majority": every non-majority label is anti-biased
            majority = vals[np.argmax(counts)]  # ties -> lowest label (deterministic)
            is_test = cls_labels != majority
        train.extend(idx_arr[~is_test].tolist())
        test.extend(idx_arr[is_test].tolist())

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

minority_split

minority_split(embeddings: ArrayLike, y: ArrayLike, n_clusters: int = 10, method: str = 'kmeans', random_state: int = 42, minority_labels: str = 'all_but_majority', **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 (ignored by 'dbscan')	10
method	str	clustering algorithm. 'kmeans' (default, fast) or 'dbscan' are standard clusterers that -- as the paper notes -- tend to produce label-homogeneous clusters, so few minority examples are found and the test set can be degenerately small. 'ward' is the paper's deterministic base clusterer (agglomerative, Ward linkage; O(n^2) memory). 'deepcluster-lite' is a light surrogate for the paper's DEEP CLUSTER step that seeks label-diverse clusters (hence a larger, non-degenerate minority set); see the References note below on faithfulness.	'kmeans'
random_state	int	for reproducibility (used by 'kmeans' and 'deepcluster-lite')	42
minority_labels	str	which of a cluster's labels are anti-biased (test). 'all_but_majority' (default) sends every non-majority label to test; 'least_only' (the paper's footnote 10) sends only the single least-frequent label. Prefer 'least_only' on many-class data: with few clusters relative to the class count, 'all_but_majority' floods test (e.g. ~93% on 150-class CLINC with 10 clusters), which the paper avoids once minority examples would exceed ~40% of the data.	'all_but_majority'
**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 or minority_labels, 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.

Faithfulness caveat: the paper deliberately does NOT cluster with a standard algorithm. It reports that k-means/Ward produce label-homogeneous clusters (too few minority examples) and instead uses DEEP CLUSTER (Caron et al., 2018) -- a fine-tuned encoder trained on cluster pseudo-labels -- to obtain label-diverse clusters. Our default (method='kmeans') is the standard clusterer the paper rejects, so it can under-produce minority examples and yield a small test set (hence the degenerate-size warning above). method='ward' matches the paper's base clusterer (still homogeneous); method='deepcluster-lite' is a static-feature stand-in for the full DEEP CLUSTER step (no encoder fine-tuning -- it under-fits a small MLP on pseudo-labels and re-clusters its hidden representation) that recovers some label diversity at CPU cost.

Seed stability: with 'kmeans'/'deepcluster-lite', nearly deterministic given a fixed random_state (the per-cluster minority labels are stable; only the clustering wobbles slightly). 'ward' is fully deterministic (no seed).

Source code in splytters/adversarial.py

def minority_split(
    embeddings: ArrayLike,
    y: ArrayLike,
    n_clusters: int = 10,
    method: str = "kmeans",
    random_state: int = 42,
    minority_labels: str = "all_but_majority",
    **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 (ignored by 'dbscan')
        method: clustering algorithm. ``'kmeans'`` (default, fast) or ``'dbscan'``
            are standard clusterers that -- as the paper notes -- tend to produce
            *label-homogeneous* clusters, so few minority examples are found and the
            test set can be degenerately small. ``'ward'`` is the paper's
            deterministic base clusterer (agglomerative, Ward linkage; O(n^2) memory).
            ``'deepcluster-lite'`` is a light surrogate for the paper's DEEP CLUSTER
            step that seeks *label-diverse* clusters (hence a larger, non-degenerate
            minority set); see the References note below on faithfulness.
        random_state: for reproducibility (used by 'kmeans' and 'deepcluster-lite')
        minority_labels: which of a cluster's labels are anti-biased (test).
            ``'all_but_majority'`` (default) sends every non-majority label to test;
            ``'least_only'`` (the paper's footnote 10) sends only the single
            least-frequent label. Prefer ``'least_only'`` on many-class data: with
            few clusters relative to the class count, ``'all_but_majority'`` floods
            test (e.g. ~93% on 150-class CLINC with 10 clusters), which the paper
            avoids once minority examples would exceed ~40% of the data.
        **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`` or ``minority_labels``, 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.

        Faithfulness caveat: the paper deliberately does NOT cluster with a standard
        algorithm. It reports that k-means/Ward produce label-homogeneous clusters
        (too few minority examples) and instead uses DEEP CLUSTER (Caron et al.,
        2018) -- a fine-tuned encoder trained on cluster pseudo-labels -- to obtain
        label-diverse clusters. Our default (``method='kmeans'``) is the standard
        clusterer the paper rejects, so it can under-produce minority examples and
        yield a small test set (hence the degenerate-size warning above).
        ``method='ward'`` matches the paper's *base* clusterer (still homogeneous);
        ``method='deepcluster-lite'`` is a static-feature stand-in for the full DEEP
        CLUSTER step (no encoder fine-tuning -- it under-fits a small MLP on
        pseudo-labels and re-clusters its hidden representation) that recovers some
        label diversity at CPU cost.

    Seed stability: with ``'kmeans'``/``'deepcluster-lite'``, nearly deterministic given a
    fixed ``random_state`` (the per-cluster minority labels are stable; only the
    clustering wobbles slightly). ``'ward'`` is fully deterministic (no seed).
    """
    if minority_labels not in ("all_but_majority", "least_only"):
        # Fail before the (possibly expensive, e.g. deepcluster-lite) clustering
        # pass below rather than after it -- minority_route validates this too,
        # but only once clustering has already run.
        raise ValueError(f"Unknown minority_labels: {minority_labels}")
    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")

    labels = _minority_cluster_labels(
        embeddings, method, n_clusters, random_state, **cluster_kwargs
    )
    train_idx, test_idx = minority_route(labels, y, minority_labels)

    test_fraction = len(test_idx) / n_samples
    if test_fraction < _MINORITY_DEGENERATE_FRACTION:
        warnings.warn(
            f"minority_split produced a degenerate test set: {len(test_idx)} 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 train_idx, test_idx

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', minority_labels: str = 'all_but_majority', **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 (ignored by 'dbscan')	10
method	str	clustering algorithm for the seed -- see :func:minority_split ('kmeans', 'dbscan', 'ward', or 'deepcluster-lite')	'kmeans'
metric	str	distance metric (scipy.spatial.distance.cdist) for growing	'euclidean'
random_state	int	for reproducibility of the clustering seed (and of any oversized-seed subsampling)	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'
minority_labels	str	which cluster labels seed the test set, 'all_but_majority' (default) or 'least_only' (footnote 10) -- see :func:minority_split. Prefer 'least_only' on many-class data; note that either way an oversized seed is subsampled back to the target size (below).	'all_but_majority'
**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). A seed larger than the target is subsampled down to it (round-robin across classes to preserve coverage); a seed equal to the target is returned as-is. Under a stratified mode the test set may still fall short of the target if every class hits its quota first.

Raises:

Type	Description
ValueError	on an unknown method, minority_labels 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").

Seed stability: nearly deterministic -- the minority seeds and proximity growth are largely fixed; only the underlying clustering wobbles between seeds. This holds in every stratify mode (none / global / per-class).

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",
    minority_labels: str = "all_but_majority",
    **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 (ignored by 'dbscan')
        method: clustering algorithm for the seed -- see :func:`minority_split`
            ('kmeans', 'dbscan', 'ward', or 'deepcluster-lite')
        metric: distance metric (``scipy.spatial.distance.cdist``) for growing
        random_state: for reproducibility of the clustering seed (and of any
            oversized-seed subsampling)
        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.
        minority_labels: which cluster labels seed the test set, ``'all_but_majority'``
            (default) or ``'least_only'`` (footnote 10) -- see :func:`minority_split`.
            Prefer ``'least_only'`` on many-class data; note that either way an
            oversized seed is subsampled back to the target size (below).
        **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). A seed larger than the target is subsampled
            down to it (round-robin across classes to preserve coverage); a seed
            equal to the target is returned as-is. Under a stratified mode the test
            set may still fall short of the target if every class hits its quota
            first.

    Raises:
        ValueError: on an unknown ``method``, ``minority_labels`` 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")``.

    Seed stability: nearly deterministic -- the minority seeds and proximity
    growth are largely fixed; only the underlying clustering wobbles between
    seeds. This holds in every ``stratify`` mode (none / global / per-class).
    """
    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,
            minority_labels=minority_labels,
            random_state=random_state,
            **cluster_kwargs,
        )

    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 len(seed) > target_n_test:
        # Grow can only add, never remove, so an oversized seed (e.g. the
        # 'all_but_majority' flood on many-class data) would otherwise blow past
        # train_size and starve train of classes. Shrink it back to target.
        seed = _subsample_seed(seed, y, target_n_test, random_state)

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

    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).

Seed stability: deterministic -- closeness to other classes is a fixed geometric quantity, so the seed has no effect (it is accepted only for API consistency).

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).

    Seed stability: deterministic -- closeness to other classes is a fixed
    geometric quantity, so the seed has no effect (it is accepted only for API
    consistency).
    """
    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.

Seed stability: structure-stable -- the cross-validation folds depend on the seed, so which samples sit nearest the learned boundary (and go to test) shifts somewhat between seeds. This holds in both stratify modes (per-class and global wobble about equally).

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.

    Seed stability: structure-stable -- the cross-validation folds depend on the
    seed, so which samples sit nearest the learned boundary (and go to test)
    shifts somewhat between seeds. This holds in both ``stratify`` modes
    (per-class and global wobble about equally).
    """
    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

Seed stability: deterministic -- samples are ranked by distance from the centroid, so the seed has no effect (it is accepted only for API consistency).

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

    Seed stability: deterministic -- samples are ranked by distance from the
    centroid, so the seed has no effect (it is accepted only for API
    consistency).
    """
    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

Seed stability: deterministic -- samples are ranked by local density, so the seed has no effect.

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

    Seed stability: deterministic -- samples are ranked by local density, so the
    seed has no effect.
    """
    embeddings = validate_split_inputs(embeddings, train_size)

    if not isinstance(k, (int, np.integer)) or isinstance(k, bool) or k < 1:
        raise ValueError(f"k must be a positive integer, got {k!r}")

    # Compute density as inverse of mean distance to k nearest neighbors
    k = min(k, len(embeddings) - 1)
    knn_distances, _ = kneighbors_excluding_self(embeddings, k, metric)
    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

Seed stability: nearly deterministic -- the outlier ranking is fixed; only samples right at the train/test cutoff can change between seeds.

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

    Seed stability: nearly deterministic -- the outlier ranking is fixed; only
    samples right at the train/test cutoff can change between seeds.
    """
    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 approximation) or 'stoer_wagner' (exact unconstrained cut used as a seed, then greedily resized to the requested train size; the resized fixed-size cut is approximate)	'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

Seed stability: deterministic for successful spectral and Stoer-Wagner paths. The spectral split follows a Fiedler vector from a dense eigendecomposition (no random start vector), whose otherwise-arbitrary sign is oriented deterministically. The only seeded branch is the warned random fallback after a spectral solver failure.

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 approximation) or ``'stoer_wagner'``
            (exact unconstrained cut used as a seed, then greedily resized to the
            requested train size; the resized fixed-size cut is approximate)
        random_state: for reproducibility

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

    Seed stability: deterministic for successful spectral and Stoer-Wagner paths.
    The spectral split follows a Fiedler vector from a dense eigendecomposition
    (no random start vector), whose otherwise-arbitrary sign is oriented
    deterministically. The only seeded branch is the warned random fallback after
    a spectral solver failure.
    """
    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
    positive_distances = distances[np.isfinite(distances) & (distances > 0)]
    if len(positive_distances) == 0:
        warnings.warn(
            "min_cut_split found no positive pairwise distances; returning a "
            "deterministic index split because the similarity graph has no "
            "geometric separation.",
            UserWarning,
            stacklevel=2,
        )
        n_train = resolve_n_train(n_samples, train_size)
        return as_index_array(range(n_train)), as_index_array(
            range(n_train, n_samples)
        )
    sigma = float(np.median(positive_distances))
    # Divide before squaring: ``distances**2`` and ``sigma**2`` can overflow
    # independently even when their ratio is well-defined.
    with np.errstate(over="ignore", invalid="ignore", under="ignore"):
        similarities = np.exp(-0.5 * (distances / sigma) ** 2)
    similarities[~np.isfinite(similarities)] = 0.0
    np.fill_diagonal(similarities, 0)  # No self-loops

    # Threshold to create sparse graph
    if similarity_threshold is None:
        nonzero_sims = similarities[similarities > 0]
        # A finite positive distance always contributes exp(-0.5) at the
        # median scale, so the earlier ``positive_distances`` guard makes this
        # non-empty.
        similarity_threshold = np.median(nonzero_sims)

    similarities[similarities < similarity_threshold] = 0

    if method == "spectral":
        # Spectral partitioning using Fiedler vector
        # (eigenvector corresponding to 2nd smallest eigenvalue of Laplacian)
        #
        # Solved with a DENSE partial eigendecomposition (scipy.linalg.eigh,
        # subset_by_index=[0, 1]) rather than sparse ARPACK (eigsh, which='SM').
        # ARPACK aimed at the *small* end of the spectrum converges very slowly
        # and, on larger/messier graphs, effectively hangs — tens of thousands of
        # iterations pegging a core for hours before it finally raises. The
        # distance matrix above is already dense and O(n²), so a dense direct
        # solve costs no extra memory class, always terminates, computes only the
        # two eigenpairs we need, and is fully deterministic (no random ARPACK
        # start vector). (A future sparse kNN-graph build — see the cdist TODO
        # above — should pair with a shift-invert solver to scale past O(n²).)
        L = laplacian(similarities, normed=True)

        try:
            # Two smallest eigenpairs only (subset_by_index is 0-based, inclusive).
            eigenvalues, eigenvectors = dense_eigh(L, subset_by_index=[0, 1])

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

            # The Fiedler sign is mathematically arbitrary and, with eigsh's
            # random start vector, flips between runs — which would flip which
            # half of the cut becomes test (giving a completely different, seed-
            # dependent held-out set). Orient it deterministically (sign of the
            # largest-magnitude component, as in sklearn's svd_flip) so the split
            # is reproducible.
            if fiedler[np.argmax(np.abs(fiedler))] < 0:
                fiedler = -fiedler

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

        except Exception as err:
            # Eigendecomposition failed; fall back to a random split. Warn
            # loudly — silently returning a *random* split from an adversarial
            # splitter would misreport the result's difficulty.
            warnings.warn(
                f"min_cut_split: spectral eigendecomposition failed ({err!r}); "
                "falling back to a random split, which is NOT adversarial. "
                "Try method='stoer_wagner' or adjust similarity_threshold.",
                stacklevel=2,
            )
            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":
        # Stoer-Wagner gives the exact unconstrained global min-cut. Enforcing an
        # arbitrary train_size is a different, cardinality-constrained problem,
        # so use the exact cut as a seed and resize both orientations greedily.
        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])

            n_train = resolve_n_train(n_samples, train_size)

            # Stoer-Wagner requires connectivity. On a disconnected graph, any
            # connected component is already one side of a zero-weight cut.
            if not nx.is_connected(G):
                components = list(nx.connected_components(G))
                components.sort(key=lambda comp: (-len(comp), min(comp)))
                seed = components[0]
                other = set(range(n_samples)) - seed
                train_indices, test_indices = _best_resized_cut(
                    similarities, (set(seed), other), n_train
                )

            else:
                _, partition = nx.stoer_wagner(G)
                train_indices, test_indices = _best_resized_cut(
                    similarities, (set(partition[0]), set(partition[1])), n_train
                )

        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

Seed stability: deterministic -- the normalized-cut partition uses a dense eigendecomposition (eigh), which is fixed, so the seed has no effect.

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

    Seed stability: deterministic -- the normalized-cut partition uses a dense
    eigendecomposition (eigh), which is fixed, so the seed has no effect.
    """
    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)
    positive_distances = distances[np.isfinite(distances) & (distances > 0)]
    if len(positive_distances) == 0:
        warnings.warn(
            "normalized_cut_split found no positive pairwise distances; "
            "returning a deterministic index split because the similarity "
            "graph has no geometric separation.",
            UserWarning,
            stacklevel=2,
        )
        n_train = resolve_n_train(n_samples, train_size)
        return as_index_array(range(n_train)), as_index_array(
            range(n_train, n_samples)
        )
    sigma = float(np.median(positive_distances))
    with np.errstate(over="ignore", invalid="ignore", under="ignore"):
        W = np.exp(-0.5 * (distances / sigma) ** 2)
    W[~np.isfinite(W)] = 0.0
    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]

    # The Fiedler sign is mathematically arbitrary; eigh's choice is LAPACK-
    # dependent and can differ across platforms/builds, which would flip which
    # half becomes test. Orient it deterministically (sign of the largest-
    # magnitude component, as in sklearn's svd_flip) so the "deterministic"
    # claim above holds portably, not just within one machine.
    if fiedler[np.argmax(np.abs(fiedler))] < 0:
        fiedler = -fiedler

    # 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

Seed stability: varies with the seed like a random split -- the test pocket is grown around a randomly chosen anchor point.

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

    Seed stability: varies with the seed like a random split -- the test pocket
    is grown around a randomly chosen anchor point.
    """
    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, method: str = 'swap', y: ArrayLike | None = None, groups: ArrayLike | None = None) -> 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.

Two methods are available:

• method="swap" (default): the historical swap-optimized approximation of the max-MMD objective -- a random split refined by accepting train/test swaps that raise the MMD.
• method="kernel_kmeans": a full-kernel implementation derived from Napoli & White (see References) -- constrained kernel k-means with k = 2 whose assignment step is a linear program (LP). Pass optional y and/or groups to enforce the paper's per-label / per-group distribution constraints in both sides of the split; the size constraint comes from train_size.

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	swap attempts (method="swap") or maximum Lloyd-style iterations (method="kernel_kmeans")	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
method	str	"swap" (default) or "kernel_kmeans"	'swap'
y	ArrayLike | None	optional labels; only valid with method="kernel_kmeans", where it constrains per-label proportions in train and test	None
groups	ArrayLike | None	optional group ids; only valid with method="kernel_kmeans", where it constrains per-group proportions in train and test	None

Returns:

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

References

Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation Splits for Model Selection under Domain Shift" (https://openreview.net/forum?id=Q692C0WtiD): the train/validation split should maximize the MMD between the two sets. Under method="kernel_kmeans" implements their kernel k-means objective and solves both disjunctive LP orientations from Eq. 15 before keeping the lower-cost assignment. Maximizing MMD is equivalent to minimizing the k = 2 kernel k-means scatter (their Theorem 1). This full-kernel variant uses exact-size largest-remainder group targets and omits their Nyström approximation for very large n; those departures are documented in :func:splytters.utils.constrained_kernel_kmeans_split. The default method="swap" remains a swap-optimized approximation of the same objective.

Seed stability: with method="swap", varies with the seed like a random split -- the swap optimization starts from a random split and many assignments reach a similar MMD. With method="kernel_kmeans" the result is deterministic given random_state (which seeds only the initial partition).

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,
    method: str = "swap",
    y: ArrayLike | None = None,
    groups: ArrayLike | None = None,
) -> 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.

    Two methods are available:

    - ``method="swap"`` (default): the historical swap-optimized approximation of
      the max-MMD objective -- a random split refined by accepting train/test
      swaps that raise the MMD.
    - ``method="kernel_kmeans"``: a full-kernel implementation derived from
      Napoli & White (see References) -- constrained kernel k-means with ``k = 2``
      whose assignment step is a linear program (LP). Pass optional ``y`` and/or
      ``groups`` to enforce the paper's per-label / per-group distribution
      constraints in both sides of the split; the size constraint comes from
      ``train_size``.

    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: swap attempts (``method="swap"``) or maximum Lloyd-style
            iterations (``method="kernel_kmeans"``)
        kernel: kernel type ('rbf' or 'linear')
        gamma: RBF kernel parameter (default: 1/n_features)
        random_state: for reproducibility
        method: ``"swap"`` (default) or ``"kernel_kmeans"``
        y: optional labels; only valid with ``method="kernel_kmeans"``, where it
            constrains per-label proportions in train and test
        groups: optional group ids; only valid with ``method="kernel_kmeans"``,
            where it constrains per-group proportions in train and test

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

    References:
        Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation
        Splits for Model Selection under Domain Shift"
        (https://openreview.net/forum?id=Q692C0WtiD): the train/validation split
        should maximize the MMD between the two sets. Under
        ``method="kernel_kmeans"`` implements their kernel k-means objective and
        solves both disjunctive LP orientations from Eq. 15 before keeping the
        lower-cost assignment. Maximizing MMD is equivalent to minimizing the
        ``k = 2`` kernel k-means scatter (their Theorem 1). This full-kernel variant
        uses exact-size largest-remainder group targets and omits their Nyström
        approximation for very large ``n``; those departures are documented in
        :func:`splytters.utils.constrained_kernel_kmeans_split`. The default
        ``method="swap"`` remains a swap-optimized approximation of the same
        objective.

    Seed stability: with ``method="swap"``, varies with the seed like a random
    split -- the swap optimization starts from a random split and many
    assignments reach a similar MMD. With ``method="kernel_kmeans"`` the result
    is deterministic given ``random_state`` (which seeds only the initial
    partition).
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    if kernel not in {"rbf", "linear"}:
        raise ValueError(f"kernel must be 'rbf' or 'linear', got {kernel!r}")
    if method not in {"swap", "kernel_kmeans"}:
        raise ValueError(
            f"method must be 'swap' or 'kernel_kmeans', got {method!r}"
        )
    if method == "swap" and (y is not None or groups is not None):
        raise ValueError(
            "y and groups are only supported with method='kernel_kmeans'"
        )
    if method == "kernel_kmeans":
        return constrained_kernel_kmeans_split(
            embeddings,
            train_size,
            kernel=kernel,
            gamma=gamma,
            y=y,
            groups=groups,
            random_state=random_state,
            n_iterations=n_iterations,
            maximize_mmd=True,
        )

    return optimized_mmd_split(
        embeddings,
        train_size,
        n_iterations,
        kernel=kernel,
        gamma=gamma,
        random_state=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)

    # KMeans requires n_clusters <= n_samples; clamp so this helper (and its
    # caller split_report) works on small inputs instead of raising.
    n_clusters = max(1, min(n_clusters, len(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
    }

maximin_split

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

Adversarial split via farthest-point (k-center) test selection.

Builds the test set by greedy farthest-first traversal: seed with one point, then repeatedly add the point farthest (metric) from everything already selected. The result is a maximally spread-out test set that covers the corners and extremes of the embedding space — a diverse, hard evaluation that, unlike a random sample, never under-represents sparse / outlying regions. Train is the (denser) remainder.

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 passed to scipy.spatial.distance.cdist.	'euclidean'
random_state	int	seeds the starting point (the traversal is otherwise deterministic).	42

Returns:

Name	Type	Description
train_indices	ndarray	the denser interior remainder.
test_indices	ndarray	a farthest-point-sampled, space-covering test set.

Seed stability: structure-stable -- only the starting point of the farthest-first traversal is random, so different seeds give overlapping but not identical space-covering test sets.

Source code in splytters/adversarial.py

def maximin_split(
    embeddings: ArrayLike,
    train_size: float | int = 0.7,
    metric: str = "euclidean",
    random_state: int = 42,
) -> tuple[np.ndarray, np.ndarray]:
    """
    Adversarial split via farthest-point (k-center) test selection.

    Builds the test set by greedy farthest-first traversal: seed with one point,
    then repeatedly add the point farthest (``metric``) from everything already
    selected. The result is a maximally *spread-out* test set that covers the
    corners and extremes of the embedding space — a diverse, hard evaluation
    that, unlike a random sample, never under-represents sparse / outlying
    regions. Train is the (denser) remainder.

    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 passed to ``scipy.spatial.distance.cdist``.
        random_state: seeds the starting point (the traversal is otherwise
            deterministic).

    Returns:
        train_indices: the denser interior remainder.
        test_indices: a farthest-point-sampled, space-covering test set.

    Seed stability: structure-stable -- only the starting point of the
    farthest-first traversal is random, so different seeds give overlapping but
    not identical space-covering test sets.
    """
    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)

    # TODO: Replace the full pairwise matrix with incremental NearestNeighbors
    # queries to avoid materializing O(n²) distances.
    distances = cdist(embeddings, embeddings, metric=metric)

    start = int(rng.randint(n_samples))
    selected = [start]
    selected_mask = np.zeros(n_samples, dtype=bool)
    selected_mask[start] = True
    # Min distance from each point to the current test set; -1 marks selected.
    min_dist = distances[start].copy()
    min_dist[selected_mask] = -1.0
    while len(selected) < n_test:
        nxt = int(np.argmax(min_dist))
        selected.append(nxt)
        selected_mask[nxt] = True
        min_dist = np.minimum(min_dist, distances[nxt])
        min_dist[selected_mask] = -1.0

    test_set = set(selected)
    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))

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

Seed stability: varies with the seed like a random split -- each cluster is shuffled before being split, so which members go to test differs run to run.

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

    Seed stability: varies with the seed like a random split -- each cluster is
    shuffled before being split, so which members go to test differs run to run.
    """
    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.

Solves a binary feasibility problem whose constraints require every test sample to have at least k train neighbors within the median pairwise distance. Raises when no split of the requested size can meet that promise.

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

Seed stability: varies with the seed. The constraints are fixed, while a seeded random linear objective selects among multiple feasible solutions.

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.

    Solves a binary feasibility problem whose constraints require every test
    sample to have at least ``k`` train neighbors within the median pairwise
    distance. Raises when no split of the requested size can meet that promise.

    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

    Seed stability: varies with the seed. The constraints are fixed, while a
    seeded random linear objective selects among multiple feasible solutions.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    if not isinstance(k, (int, np.integer)) or isinstance(k, bool) or k < 1:
        raise ValueError(f"k must be a positive integer, got {k!r}")

    # 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)

    n_train = resolve_n_train(n_samples, train_size)
    if k > n_train:
        raise ValueError(
            "no feasible neighbor-coverage split: "
            f"k={k} exceeds the {n_train} available train samples"
        )

    # Binary variable x_i is 1 when sample i is in train. For every sample i,
    #     sum_{j in N(i)} x_j + k*x_i >= k
    # activates the coverage requirement only when i is in test (x_i == 0).
    # Together with sum_i x_i == n_train this is the exact documented contract.
    from scipy import sparse
    from scipy.optimize import Bounds, LinearConstraint, milp

    neighbor_matrix = (distances <= threshold).astype(float)
    coverage_matrix = neighbor_matrix + k * np.eye(n_samples)
    constraint_matrix = sparse.vstack(
        [np.ones((1, n_samples)), sparse.csr_matrix(coverage_matrix)],
        format="csr",
    )
    lower = np.concatenate(([n_train], np.full(n_samples, float(k))))
    upper = np.concatenate(([n_train], np.full(n_samples, np.inf)))

    # The objective only breaks ties between feasible solutions. Seeded random
    # costs retain this splitter's documented seed variation.
    result = milp(
        c=rng.uniform(size=n_samples),
        integrality=np.ones(n_samples),
        bounds=Bounds(0.0, 1.0),
        constraints=LinearConstraint(constraint_matrix, lower, upper),
    )
    if not result.success or result.x is None:
        raise ValueError(
            "no feasible neighbor-coverage split for the requested train_size, "
            f"k={k}, and median-distance neighborhood graph"
        )

    train_mask = result.x >= 0.5
    train_indices = np.flatnonzero(train_mask)
    test_indices = np.flatnonzero(~train_mask)
    coverage = neighbor_matrix[np.ix_(test_indices, train_indices)].sum(axis=1)
    if len(train_indices) != n_train or np.any(coverage < k):
        raise RuntimeError("neighbor-coverage MILP returned an invalid solution")
    return as_index_array(train_indices), 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

Seed stability: varies with the seed like a random split -- the swap optimization has many assignments with equally matched centroids.

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

    Seed stability: varies with the seed like a random split -- the swap
    optimization has many assignments with equally matched centroids.
    """
    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

Seed stability: varies with the seed like a random split -- samples are shuffled within each distance bin before the proportional split.

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

    Seed stability: varies with the seed like a random split -- samples are
    shuffled within each distance bin before the proportional split.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    rng = check_random_state(random_state)

    if (
        not isinstance(n_bins, (int, np.integer))
        or isinstance(n_bins, bool)
        or n_bins < 1
    ):
        raise ValueError(f"n_bins must be a positive integer, got {n_bins!r}")

    # 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

Seed stability: varies with the seed like a random split -- points are processed in a random order when building the 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

    Seed stability: varies with the seed like a random split -- points are
    processed in a random order when building the 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 (and stays) in train.
    in_train = np.ones(n_samples, dtype=bool)
    # `pinned` marks train points that some test point relies on as its nearest
    # neighbor; they must never move to test, or that test point's invariant
    # would break. Without this, a later iteration could pull a point's NN into
    # test, silently violating the documented guarantee.
    pinned = np.zeros(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
        if pinned[idx]:
            continue  # must stay in train to satisfy a dependent test point
        nn = neighbors[idx]
        if in_train[nn]:
            in_train[idx] = False
            pinned[nn] = True  # lock idx's NN into train
            test_indices.append(idx)

    if len(test_indices) < n_test:
        warnings.warn(
            f"nearest_neighbor_split could only place {len(test_indices)} of the "
            f"requested {n_test} test samples while keeping every test point's "
            "nearest neighbor in train; the test set is smaller than train_size "
            "implies. Use a larger dataset or a smaller test fraction.",
            stacklevel=2,
        )

    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

Seed stability: varies with the seed like a random split -- each near-duplicate group is shuffled before being split across both sides.

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

    Seed stability: varies with the seed like a random split -- each
    near-duplicate group is shuffled before being split across both sides.
    """
    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)

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

    train_indices = []
    test_indices = []
    singletons: list[int] = []

    # For each near-duplicate group of size >= 2, split it across both sides so
    # the duplicates appear in train *and* test (the point of this splitter).
    for indices in component_to_indices.values():
        indices = np.array(indices)
        rng.shuffle(indices)

        if len(indices) == 1:
            singletons.append(int(indices[0]))
            continue

        # 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())

    # Singletons have no near-duplicate to spread, so where they land doesn't
    # affect the duplicates-on-both-sides property. Distribute them to bring the
    # realized split close to train_size, instead of dumping them all in train
    # (which skewed the split and could starve the test set).
    rng.shuffle(singletons)
    n_needed = max(0, n_train_total - len(train_indices))
    train_indices.extend(singletons[:n_needed])
    test_indices.extend(singletons[n_needed:])

    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

Seed stability: varies with the seed like a random split -- it starts from a random split and greedily swaps, reaching different high-coverage solutions.

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

    Seed stability: varies with the seed like a random split -- it starts from a
    random split and greedily swaps, reaching different high-coverage solutions.
    """
    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)  # a point never covers itself

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

    # Boolean "within radius" graph, plus a running count of how many *train*
    # points cover each sample. Maintaining cover_count incrementally makes each
    # candidate swap an O(n) update rather than an O(n²) full recompute of
    # coverage (the previous nested-loop version was O(n⁴) worst case).
    within = distances <= radius

    # Start with a random split.
    all_indices = np.arange(n_samples)
    rng.shuffle(all_indices)
    n_train = resolve_n_train(n_samples, train_size)
    train_mask = np.zeros(n_samples, dtype=bool)
    train_mask[all_indices[:n_train]] = True

    cover_count = within[:, train_mask].sum(axis=1)  # train points within radius

    def coverage(test_mask: np.ndarray, cc: np.ndarray) -> float:
        return float((cc[test_mask] > 0).mean()) if test_mask.any() else 1.0

    current_coverage = coverage(~train_mask, cover_count)

    # Greedy optimization: repeatedly take an uncovered test point and, if some
    # train point can swap in to raise overall coverage, apply the best such swap.
    max_iterations = n_samples * 2
    for _ in range(max_iterations):
        test_mask = ~train_mask
        uncovered = np.flatnonzero(test_mask & (cover_count == 0))
        if uncovered.size == 0:
            break  # every test point is covered; nothing left to improve

        improved = False
        train_idx = np.flatnonzero(train_mask)
        for u in uncovered:
            base = cover_count + within[:, u]  # effect of moving u into train
            best_v, best_cov = None, current_coverage
            for v in train_idx:
                new_cc = base - within[:, v]  # and moving v out of train
                new_test = test_mask.copy()
                new_test[u] = False
                new_test[v] = True
                cov = coverage(new_test, new_cc)
                if cov > best_cov:
                    best_cov, best_v = cov, v
            if best_v is not None:
                cover_count = cover_count + within[:, u] - within[:, best_v]
                train_mask[u] = True
                train_mask[best_v] = False
                current_coverage = best_cov
                improved = True
                break
        if not improved:
            break

    return (
        as_index_array(np.flatnonzero(train_mask)),
        as_index_array(np.flatnonzero(~train_mask)),
    )

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

Seed stability: varies with the seed like a random split -- the swap optimization starts from a random split and many assignments match the distribution equally well.

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

    Seed stability: varies with the seed like a random split -- the swap
    optimization starts from a random split and many assignments match the
    distribution equally well.
    """
    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

Seed stability: varies with the seed like a random split -- the swap optimization reaches different assignments with equally matched moments.

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

    Seed stability: varies with the seed like a random split -- the swap
    optimization reaches different assignments with equally matched moments.
    """
    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

Seed stability: varies with the seed like a random split -- the swap optimization reaches different assignments with equally matched histograms.

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

    Seed stability: varies with the seed like a random split -- the swap
    optimization reaches different assignments with equally matched histograms.
    """
    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

Seed stability: varies with the seed like a random split -- it is a seeded stratified random draw.

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

    Seed stability: varies with the seed like a random split -- it is a seeded
    stratified random draw.
    """
    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

Seed stability: varies with the seed like a random split -- samples are shuffled within each density bin before the proportional split.

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

    Seed stability: varies with the seed like a random split -- samples are
    shuffled within each density bin before the proportional split.
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    n_samples = len(embeddings)
    rng = check_random_state(random_state)

    if (
        not isinstance(n_bins, (int, np.integer))
        or isinstance(n_bins, bool)
        or n_bins < 1
    ):
        raise ValueError(f"n_bins must be a positive integer, got {n_bins!r}")

    # Compute density (inverse of mean distance to 10 nearest neighbors)
    k = min(10, n_samples - 1)
    knn_distances, _ = kneighbors_excluding_self(embeddings, 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, method: str = 'swap', y: ArrayLike | None = None, groups: ArrayLike | None = None) -> 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.

Two methods are available:

• method="swap" (default): the historical swap-optimized approximation that lowers the MMD by accepting train/test swaps that reduce it.
• method="kernel_kmeans": the min-MMD dual of the constrained kernel k-means (k = 2) method of Napoli & White (see References). The assignment step is a linear program (LP); it maximizes the kernel scatter (anti-clustering) so the two sides resemble each other. Pass optional y and/or groups to enforce per-label / per-group distribution constraints; the size constraint comes from train_size. Note: while the resulting MMD is far below a random split's, in practice it is typically around 2x the MMD reached by the default swap optimizer, so prefer method="swap" when the absolute lowest MMD matters and method="kernel_kmeans" when the label/group constraints do.

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	swap attempts (method="swap") or maximum Lloyd-style iterations (method="kernel_kmeans")	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
method	str	"swap" (default) or "kernel_kmeans"	'swap'
y	ArrayLike | None	optional labels; only valid with method="kernel_kmeans", where it constrains per-label proportions in train and test	None
groups	ArrayLike | None	optional group ids; only valid with method="kernel_kmeans", where it constrains per-group proportions in train and test	None

Returns:

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

References

Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation Splits for Model Selection under Domain Shift" (https://openreview.net/forum?id=Q692C0WtiD) study the adversarial dual -- maximizing train/validation MMD -- via constrained kernel k-means with k = 2 solved by linear programming (their Algorithm 1). This splitter minimizes MMD instead. Under method="kernel_kmeans" it uses the same constrained-LP kernel k-means machinery with the objective sign flipped (maximizing the k = 2 scatter, the min-MMD dual). The default method="swap" remains a swap-optimized approximation of the same objective, matching :func:splytters.mmd_maximized_split. Their Nyström scaling for very large n is not implemented (out of scope).

Seed stability: with method="swap", varies with the seed like a random split -- the swap optimization reaches different assignments with similarly low MMD. With method="kernel_kmeans" the result is deterministic given random_state (which seeds only the initial partition).

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,
    method: str = "swap",
    y: ArrayLike | None = None,
    groups: ArrayLike | None = None,
) -> 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.

    Two methods are available:

    - ``method="swap"`` (default): the historical swap-optimized approximation
      that lowers the MMD by accepting train/test swaps that reduce it.
    - ``method="kernel_kmeans"``: the min-MMD dual of the constrained kernel
      k-means (``k = 2``) method of Napoli & White (see References). The
      assignment step is a linear program (LP); it maximizes the kernel scatter
      (anti-clustering) so the two sides resemble each other. Pass optional ``y``
      and/or ``groups`` to enforce per-label / per-group distribution
      constraints; the size constraint comes from ``train_size``. Note: while
      the resulting MMD is far below a random split's, in practice it is
      typically around 2x the MMD reached by the default swap optimizer, so
      prefer ``method="swap"`` when the absolute lowest MMD matters and
      ``method="kernel_kmeans"`` when the label/group constraints do.

    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: swap attempts (``method="swap"``) or maximum Lloyd-style
            iterations (``method="kernel_kmeans"``)
        kernel: kernel type ('rbf' or 'linear')
        gamma: RBF kernel parameter (default: 1/n_features)
        random_state: for reproducibility
        method: ``"swap"`` (default) or ``"kernel_kmeans"``
        y: optional labels; only valid with ``method="kernel_kmeans"``, where it
            constrains per-label proportions in train and test
        groups: optional group ids; only valid with ``method="kernel_kmeans"``,
            where it constrains per-group proportions in train and test

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

    References:
        Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation
        Splits for Model Selection under Domain Shift"
        (https://openreview.net/forum?id=Q692C0WtiD) study the adversarial dual
        -- *maximizing* train/validation MMD -- via constrained kernel k-means
        with ``k = 2`` solved by linear programming (their Algorithm 1). This
        splitter minimizes MMD instead. Under ``method="kernel_kmeans"`` it uses
        the same constrained-LP kernel k-means machinery with the objective sign
        flipped (maximizing the ``k = 2`` scatter, the min-MMD dual). The default
        ``method="swap"`` remains a swap-optimized approximation of the same
        objective, matching :func:`splytters.mmd_maximized_split`. Their Nyström
        scaling for very large ``n`` is not implemented (out of scope).

    Seed stability: with ``method="swap"``, varies with the seed like a random
    split -- the swap optimization reaches different assignments with similarly
    low MMD. With ``method="kernel_kmeans"`` the result is deterministic given
    ``random_state`` (which seeds only the initial partition).
    """
    embeddings = validate_split_inputs(embeddings, train_size)
    if kernel not in {"rbf", "linear"}:
        raise ValueError(f"kernel must be 'rbf' or 'linear', got {kernel!r}")
    if method not in {"swap", "kernel_kmeans"}:
        raise ValueError(
            f"method must be 'swap' or 'kernel_kmeans', got {method!r}"
        )
    if method == "swap" and (y is not None or groups is not None):
        raise ValueError(
            "y and groups are only supported with method='kernel_kmeans'"
        )
    if method == "kernel_kmeans":
        return constrained_kernel_kmeans_split(
            embeddings,
            train_size,
            kernel=kernel,
            gamma=gamma,
            y=y,
            groups=groups,
            random_state=random_state,
            n_iterations=n_iterations,
            maximize_mmd=False,
        )

    return optimized_mmd_split(
        embeddings,
        train_size,
        n_iterations,
        kernel=kernel,
        gamma=gamma,
        random_state=random_state,
        minimize=True,
    )

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
        pass_seed = accepts_random_state(self.splitter)
        for i in range(self.n_splits):
            seed = _derive_seed(self.random_state, i)
            extra = {"random_state": seed} if pass_seed else {}
            train_idx, test_idx = self.splitter(
                emb,
                train_size=self.train_size,
                **extra,
                **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)}."
            )

    extra = {"random_state": random_state} if accepts_random_state(splitter) else {}
    train_idx, test_idx = splitter(
        embeddings, train_size=train_size, **extra, **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, mean 1-D and sliced Wasserstein, KS, Fréchet), a classifier two-sample test (c2st_auc), k-NN manifold precision/recall, 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, sliced_wasserstein, ks_mean, frechet_distance) and mean_cross_distance indicate a more adversarial (harder) split; values near a random split indicate a balanced split; near-zero indicates overlap.

c2st_auc is a classifier two-sample test — the out-of-fold AUC of a random-forest classifier trained to tell train from test points: ~0.5 means the two sides are indistinguishable (overlap), →1.0 means trivially separable (adversarial). manifold_precision / manifold_recall (k-NN support coverage) report how much of test lies inside the train manifold and vice versa.

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``, ``sliced_wasserstein``, ``ks_mean``,
    ``frechet_distance``) and ``mean_cross_distance`` indicate a *more
    adversarial* (harder) split; values near a random split indicate a
    *balanced* split; near-zero indicates *overlap*.

    ``c2st_auc`` is a classifier two-sample test — the out-of-fold AUC of a
    random-forest classifier trained to tell train from test points: ~0.5 means
    the two sides are indistinguishable (overlap), →1.0 means trivially separable
    (adversarial). ``manifold_precision`` / ``manifold_recall`` (k-NN support
    coverage) report how much of test lies inside the train manifold and vice
    versa.

    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])])
    )
    report["sliced_wasserstein"] = _sliced_wasserstein(
        A, B, _SLICED_WASSERSTEIN_PROJECTIONS, rng
    )
    report["frechet_distance"] = _frechet_distance(A, B)

    # Learned separability (can a classifier tell train from test?) and k-NN
    # support coverage between the two manifolds.
    report["c2st_auc"] = _c2st_auc(
        A, B, random_state if isinstance(random_state, int) else 42
    )
    precision, recall = _manifold_precision_recall(A, B, _MANIFOLD_K)
    report["manifold_precision"] = precision
    report["manifold_recall"] = recall

    # 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 NearestNeighbors (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 ``NearestNeighbors`` (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)]

    dists, _ = kneighbors_excluding_self(embeddings, 1, metric)
    min_distances = dists[:, 0]

    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, random_state: int = 42, **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
random_state	int	seed for isolation_forest (ignored by lof, which is deterministic). Exposed as its own parameter so it can't collide with a random_state passed through **kwargs.	42
**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,
    random_state: int = 42,
    **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
        random_state: seed for isolation_forest (ignored by lof, which is
            deterministic). Exposed as its own parameter so it can't collide
            with a ``random_state`` passed through ``**kwargs``.
        **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=random_state, **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

mahalanobis_distance_to_mean

mahalanobis_distance_to_mean(embeddings: ArrayLike, low_first: bool = True) -> list[tuple[int, float]]

Sort samples by Mahalanobis distance from the dataset centroid.

Like :func:distance_to_mean, but covariance-aware: atypicality is measured relative to the data's own spread and feature correlations rather than in raw Euclidean units. Most typical samples first.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array of shape (n_samples, embedding_dim).	required
low_first	bool	if True, the most typical (smallest distance) samples first.	True

Returns:

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

Source code in splytters/sorters/embedding_sorters.py

def mahalanobis_distance_to_mean(
    embeddings: ArrayLike, low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort samples by Mahalanobis distance from the dataset centroid.

    Like :func:`distance_to_mean`, but covariance-aware: atypicality is measured
    relative to the data's own spread and feature correlations rather than in
    raw Euclidean units. Most typical samples first.

    Args:
        embeddings: array of shape (n_samples, embedding_dim).
        low_first: if True, the most typical (smallest distance) samples first.

    Returns:
        List of (index, mahalanobis_distance) tuples.
    """
    X = np.asarray(embeddings, dtype=float)
    centroid = X.mean(axis=0)
    cov = np.atleast_2d(np.cov(X, rowvar=False)) + 1e-6 * np.eye(X.shape[1])
    inv_cov = np.linalg.pinv(cov)
    diff = X - centroid
    d2 = np.einsum("ij,jk,ik->i", diff, inv_cov, diff)
    dists = np.sqrt(np.maximum(d2, 0.0))
    scores = sorted(
        enumerate(dists.tolist()), key=lambda p: p[1], reverse=not low_first
    )
    return [(i, float(v)) for i, v in scores]

knn_label_disagreement

knn_label_disagreement(embeddings: ArrayLike, y: ArrayLike, k: int = 5, metric: str = 'euclidean', low_first: bool = True) -> list[tuple[int, float]]

Sort samples by the fraction of their k nearest neighbors with a different label — a class-boundary / ambiguity score.

A point surrounded by same-label neighbors is class-typical (interior); one whose neighbors disagree sits on a boundary and is harder. This is the label-aware (supervised) sorter — pair it with :func:splytters.sorted_stratified_split for a difficulty curriculum.

Parameters:

Name	Type	Description	Default
embeddings	ArrayLike	array of shape (n_samples, embedding_dim).	required
y	ArrayLike	class labels of shape (n_samples,).	required
k	int	number of neighbors to inspect (the point itself is excluded).	5
metric	str	distance metric for the neighbor search.	'euclidean'
low_first	bool	if True, the most class-typical (low-disagreement) first.	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, disagreement_fraction) tuples, each in [0, 1].

Source code in splytters/sorters/embedding_sorters.py

def knn_label_disagreement(
    embeddings: ArrayLike,
    y: ArrayLike,
    k: int = 5,
    metric: str = "euclidean",
    low_first: bool = True,
) -> list[tuple[int, float]]:
    """
    Sort samples by the fraction of their k nearest neighbors with a *different*
    label — a class-boundary / ambiguity score.

    A point surrounded by same-label neighbors is class-typical (interior); one
    whose neighbors disagree sits on a boundary and is harder. This is the
    label-aware (supervised) sorter — pair it with
    :func:`splytters.sorted_stratified_split` for a difficulty curriculum.

    Args:
        embeddings: array of shape (n_samples, embedding_dim).
        y: class labels of shape (n_samples,).
        k: number of neighbors to inspect (the point itself is excluded).
        metric: distance metric for the neighbor search.
        low_first: if True, the most class-typical (low-disagreement) first.

    Returns:
        List of (index, disagreement_fraction) tuples, each in [0, 1].
    """
    X = np.asarray(embeddings)
    y = np.asarray(y)
    n = len(X)
    _, neighbors = kneighbors_excluding_self(X, k, metric)
    scores = [(i, float(np.mean(y[neighbors[i]] != y[i]))) for i in range(n)]
    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}"
        )

    # Require model and tokenizer together: defaulting only the missing one to
    # GPT-2 would pair a user's model with a mismatched tokenizer (or vice
    # versa) and silently produce wrong scores. Checked before the heavy imports
    # so the caller gets this error even without torch/transformers installed.
    if (model is None) != (tokenizer is None):
        raise ValueError(
            "pass both `model` and `tokenizer`, or neither (to default to GPT-2)."
        )

    if not texts:
        return []

    import torch
    from transformers import GPT2LMHeadModel, GPT2TokenizerFast

    if model is None:  # tokenizer is None here too
        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

gzip_complexity

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

Sort texts by gzip compression ratio — a model-free complexity proxy.

Redundant / repetitive / simple text compresses well (low ratio); varied, information-dense text compresses poorly (high ratio). Needs no model, so unlike :func:perplexity_score it runs without any heavy dependency.

Parameters:

Name	Type	Description	Default
texts	list[str]	list of strings.	required
low_first	bool	if True, the simplest (most compressible) texts first.	True

Returns:

Type	Description
list[tuple[int, float]]	List of (index, compression_ratio) tuples (compressed / raw bytes).

Note

gzip carries ~18 bytes of header overhead, so the ratio is noisy for very short strings; it is most meaningful across texts of comparable length.

Source code in splytters/sorters/text_sorters.py

def gzip_complexity(
    texts: list[str], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort texts by gzip compression ratio — a model-free complexity proxy.

    Redundant / repetitive / simple text compresses well (low ratio); varied,
    information-dense text compresses poorly (high ratio). Needs no model, so
    unlike :func:`perplexity_score` it runs without any heavy dependency.

    Args:
        texts: list of strings.
        low_first: if True, the simplest (most compressible) texts first.

    Returns:
        List of (index, compression_ratio) tuples (compressed / raw bytes).

    Note:
        gzip carries ~18 bytes of header overhead, so the ratio is noisy for
        very short strings; it is most meaningful across texts of comparable
        length.
    """
    import gzip

    scores = []
    for i, text in enumerate(texts):
        raw = text.encode("utf-8")
        ratio = len(gzip.compress(raw)) / len(raw) if raw else 0.0
        scores.append((i, float(ratio)))
    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, method: str = 'quantize') -> list[tuple[int, float]]

Sort images by their dominant color.

Uses deterministic color quantization by default to find dominant colors, then sorts by distance from a target color (default: neutral gray). Pass method="kmeans" for the legacy k-means estimator.

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
method	str	dominant-color estimator, either "quantize" (deterministic, default) or "kmeans" (legacy k-means behavior)	'quantize'

Returns:

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

Raises:

Type	Description
ValueError	if method is not "quantize" or "kmeans".

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,
    method: str = "quantize",
) -> list[tuple[int, float]]:
    """
    Sort images by their dominant color.

    Uses deterministic color quantization by default to find dominant colors,
    then sorts by distance from a target color (default: neutral gray). Pass
    ``method="kmeans"`` for the legacy k-means estimator.

    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)
        method: dominant-color estimator, either "quantize" (deterministic,
            default) or "kmeans" (legacy k-means behavior)

    Returns:
        List of (index, distance) tuples sorted by distance from target color.

    Raises:
        ValueError: if ``method`` is not "quantize" or "kmeans".
    """
    if method not in {"quantize", "kmeans"}:
        raise ValueError(f"method must be 'quantize' or 'kmeans', got {method!r}")

    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):
        if method == "quantize":
            pil_img = _load_image(img).convert("RGB")
            palette_img = pil_img.quantize(colors=n_clusters)
            color_counts = palette_img.getcolors(maxcolors=pil_img.width * pil_img.height)
            if not color_counts:  # pragma: no cover - quantized images should fit maxcolors
                dominant_rgb = np.array([0.0, 0.0, 0.0])
            else:
                _, palette_index = max(color_counts, key=lambda p: (p[0], -p[1]))
                palette = palette_img.getpalette()
                offset = palette_index * 3
                dominant_rgb = np.array(palette[offset:offset + 3], dtype=np.float32)
        else:
            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 resulting ranking are reproducible across calls.
            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]

sharpness

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

Sort images by sharpness — the variance of the Laplacian (a blur detector).

A blurry image has little high-frequency detail and a low Laplacian variance; a sharp, detailed image has a high one. Distinct from contrast / frequency_content in isolating edge focus.

Useful for adversarial / curriculum splits: train on sharp images, test on blurry ones (or vice versa).

Parameters:

Name	Type	Description	Default
images	list[ImageInput]	list of file paths or PIL Image objects.	required
low_first	bool	if True, the blurriest images first.	True

Returns:

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

Source code in splytters/sorters/image_sorters.py

def sharpness(
    images: list[ImageInput], low_first: bool = True
) -> list[tuple[int, float]]:
    """
    Sort images by sharpness — the variance of the Laplacian (a blur detector).

    A blurry image has little high-frequency detail and a low Laplacian
    variance; a sharp, detailed image has a high one. Distinct from
    ``contrast`` / ``frequency_content`` in isolating edge focus.

    Useful for adversarial / curriculum splits: train on sharp images, test on
    blurry ones (or vice versa).

    Args:
        images: list of file paths or PIL Image objects.
        low_first: if True, the blurriest images first.

    Returns:
        List of (index, laplacian_variance) tuples.
    """
    from scipy.ndimage import laplace

    scores = []
    for i, img in enumerate(images):
        gray = _to_array(img, mode="L")
        scores.append((i, float(laplace(gray).var())))
    scores.sort(key=lambda p: p[1], reverse=not low_first)
    return 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):
            # Push NaN to the end regardless of sort direction (a fixed +inf
            # sentinel would sort NaN rows to the *front* when low_first=False).
            zscore = float('inf') if low_first else 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):
            # Keep NaN at the end for both sort directions (see column_zscore).
            abs_zscore = float('inf') if low_first else 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).
    """
    n_cols = len(df.columns)
    if n_cols == 0:
        return [(idx, 0.0) for idx in df.index]

    # Vectorized: one pass over the whole frame instead of a per-row .loc.
    ratios = df.isna().mean(axis=1)
    scores = list(zip(df.index, ratios.tolist(), strict=True))
    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]

    # Vectorized: fraction of near-zero cells per row in one pass.
    sparsity = (numeric_df.abs() < zero_threshold).mean(axis=1)
    scores = list(zip(df.index, sparsity.tolist(), strict=True))
    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, random_state: int = 42, **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
random_state	int	seed for isolation_forest (ignored by lof/zscore). Exposed separately so it can't collide with a random_state passed through **kwargs.	42
**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,
    random_state: int = 42,
    **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
        random_state: seed for isolation_forest (ignored by lof/zscore).
            Exposed separately so it can't collide with a ``random_state``
            passed through ``**kwargs``.
        **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=random_state, **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]

    # Vectorized Euclidean distance of every row to the column-wise mean.
    mean_vector = numeric_df.mean().values
    diffs = numeric_df.values - mean_vector
    distances = np.linalg.norm(diffs, axis=1)
    scores = list(zip(df.index, distances.tolist(), strict=True))
    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.

accepts_random_state

accepts_random_state(fn: Any) -> bool

Whether fn takes a random_state keyword (directly or via **kwargs).

Used to decide whether to forward random_state to a user-supplied splitter: passing it to one that doesn't accept it raises TypeError.

Source code in splytters/utils.py

def accepts_random_state(fn: Any) -> bool:
    """Whether ``fn`` takes a ``random_state`` keyword (directly or via **kwargs).

    Used to decide whether to forward ``random_state`` to a user-supplied
    splitter: passing it to one that doesn't accept it raises ``TypeError``.
    """
    try:
        params = inspect.signature(fn).parameters
    except (TypeError, ValueError):
        return True  # can't introspect (e.g. a C callable) — assume it does
    if any(p.kind == p.VAR_KEYWORD for p in params.values()):
        return True
    return "random_state" in params

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)

kneighbors_excluding_self

kneighbors_excluding_self(X: ArrayLike, k: int, metric: str = 'euclidean') -> tuple[np.ndarray, np.ndarray]

Return each sample's k nearest other samples.

NearestNeighbors.kneighbors(X) does not guarantee that the query sample appears in column zero when several points tie at distance zero. Query one extra neighbor, remove the query by its index, and only then take the first k results. This keeps duplicate embeddings correct without materializing an n x n distance matrix.

Parameters:

Name	Type	Description	Default
X	ArrayLike	array-like of shape (n_samples, n_features).	required
k	int	positive number of non-self neighbors, clamped to n_samples - 1.	required
metric	str	metric accepted by :class:sklearn.neighbors.NearestNeighbors.	'euclidean'

Returns:

Type	Description
tuple[ndarray, ndarray]	(distances, indices) arrays of shape (n_samples, effective_k).

Source code in splytters/utils.py

def kneighbors_excluding_self(
    X: ArrayLike, k: int, metric: str = "euclidean"
) -> tuple[np.ndarray, np.ndarray]:
    """Return each sample's ``k`` nearest *other* samples.

    ``NearestNeighbors.kneighbors(X)`` does not guarantee that the query sample
    appears in column zero when several points tie at distance zero. Query one
    extra neighbor, remove the query by its index, and only then take the first
    ``k`` results. This keeps duplicate embeddings correct without materializing
    an ``n x n`` distance matrix.

    Args:
        X: array-like of shape ``(n_samples, n_features)``.
        k: positive number of non-self neighbors, clamped to ``n_samples - 1``.
        metric: metric accepted by :class:`sklearn.neighbors.NearestNeighbors`.

    Returns:
        ``(distances, indices)`` arrays of shape ``(n_samples, effective_k)``.
    """
    from sklearn.neighbors import NearestNeighbors

    points = np.asarray(X)
    n_samples = len(points)
    if not isinstance(k, (int, np.integer)) or isinstance(k, bool) or k < 1:
        raise ValueError(f"k must be a positive integer, got {k!r}")
    if n_samples < 2:
        raise ValueError("at least 2 samples are required for a neighbor query")

    effective_k = min(int(k), n_samples - 1)
    query_k = min(effective_k + 1, n_samples)
    distances, indices = NearestNeighbors(
        n_neighbors=query_k, metric=metric
    ).fit(points).kneighbors(points)

    out_distances = np.empty((n_samples, effective_k), dtype=float)
    out_indices = np.empty((n_samples, effective_k), dtype=np.intp)
    for i in range(n_samples):
        keep = indices[i] != i
        row_distances = distances[i, keep][:effective_k]
        row_indices = indices[i, keep][:effective_k]
        if len(row_indices) != effective_k:  # defensive: sklearn returns query_k
            raise RuntimeError(
                f"neighbor query returned only {len(row_indices)} non-self "
                f"neighbors for sample {i}; expected {effective_k}"
            )
        out_distances[i] = row_distances
        out_indices[i] = row_indices
    return out_distances, out_indices

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

Seed stability: fully random -- every seed gives a different train/test split; this is the baseline the other splitters' seed stability is measured against.

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

    Seed stability: fully random -- every seed gives a different train/test
    split; this is the baseline the other splitters' seed stability is measured
    against.
    """
    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))

optimized_mmd_split

optimized_mmd_split(embeddings: ndarray, train_size: float | int, n_iterations: int, *, kernel: str = 'rbf', gamma: float | None = None, random_state: int = 42, minimize: bool = True) -> tuple[np.ndarray, np.ndarray]

Optimize biased empirical MMD with exact incremental swap updates.

The kernel matrix is invariant across swap attempts, so compute it once and maintain the train/train, test/test, and train/test kernel sums. A proposed swap is then scored from row sums instead of rebuilding three kernel matrices. Accepted swaps update those row sums in O(n); rejected swaps preserve the historical set-mutation order so seeded results stay backward compatible.

Source code in splytters/utils.py

def optimized_mmd_split(
    embeddings: np.ndarray,
    train_size: float | int,
    n_iterations: int,
    *,
    kernel: str = "rbf",
    gamma: float | None = None,
    random_state: int = 42,
    minimize: bool = True,
) -> tuple[np.ndarray, np.ndarray]:
    """Optimize biased empirical MMD with exact incremental swap updates.

    The kernel matrix is invariant across swap attempts, so compute it once and
    maintain the train/train, test/test, and train/test kernel sums. A proposed
    swap is then scored from row sums instead of rebuilding three kernel matrices.
    Accepted swaps update those row sums in ``O(n)``; rejected swaps preserve the
    historical set-mutation order so seeded results stay backward compatible.
    """
    n_samples, n_dims = embeddings.shape
    if kernel == "rbf":
        kernel_gamma = gamma if gamma is not None else 1.0 / n_dims
        K = np.exp(
            -kernel_gamma * cdist(
                embeddings, embeddings, metric="sqeuclidean"
            )
        )
    elif kernel == "linear":
        K = embeddings @ embeddings.T
    else:
        raise ValueError(f"kernel must be 'rbf' or 'linear', got {kernel!r}")

    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)
    n_test = n_samples - n_train
    train_indices = set(all_indices[:n_train].tolist())
    test_indices = set(all_indices[n_train:].tolist())

    train_array = np.fromiter(train_indices, dtype=np.intp, count=n_train)
    test_array = np.fromiter(test_indices, dtype=np.intp, count=n_test)
    row_to_train = K[:, train_array].sum(axis=1)
    row_to_test = K[:, test_array].sum(axis=1)
    sum_train_train = float(row_to_train[train_array].sum())
    sum_test_test = float(row_to_test[test_array].sum())
    sum_train_test = float(row_to_test[train_array].sum())

    def score(within_train: float, within_test: float, cross: float) -> float:
        return (
            within_train / (n_train * n_train)
            + within_test / (n_test * n_test)
            - 2.0 * cross / (n_train * n_test)
        )

    current_score = score(sum_train_train, sum_test_test, sum_train_test)
    diagonal = np.diag(K)

    for _ in range(n_iterations):
        train_sample = rng.choice(list(train_indices))
        test_sample = rng.choice(list(test_indices))
        cross_pair = K[train_sample, test_sample]

        new_train_train = (
            sum_train_train
            - 2.0 * row_to_train[train_sample]
            + diagonal[train_sample]
            + 2.0 * (row_to_train[test_sample] - cross_pair)
            + diagonal[test_sample]
        )
        new_test_test = (
            sum_test_test
            - 2.0 * row_to_test[test_sample]
            + diagonal[test_sample]
            + 2.0 * (row_to_test[train_sample] - cross_pair)
            + diagonal[train_sample]
        )
        new_train_test = (
            sum_train_test
            - row_to_test[train_sample]
            - row_to_train[test_sample]
            + cross_pair
            + row_to_test[test_sample]
            - diagonal[test_sample]
            + row_to_train[train_sample]
            - diagonal[train_sample]
            + cross_pair
        )
        new_score = score(new_train_train, new_test_test, new_train_test)

        # Match optimized_split's mutation/reversion order so a fixed seed draws
        # the same later candidates even when a proposed swap is rejected.
        train_indices.remove(train_sample)
        train_indices.add(test_sample)
        test_indices.remove(test_sample)
        test_indices.add(train_sample)

        improved = new_score < current_score if minimize else new_score > current_score
        if improved:
            row_to_train += K[:, test_sample] - K[:, train_sample]
            row_to_test += K[:, train_sample] - K[:, test_sample]
            sum_train_train = new_train_train
            sum_test_test = new_test_test
            sum_train_test = new_train_test
            current_score = new_score
        else:
            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))

constrained_kernel_kmeans_split

constrained_kernel_kmeans_split(embeddings: ndarray, train_size: float | int, kernel: str = 'rbf', gamma: float | None = None, y: ArrayLike | None = None, groups: ArrayLike | None = None, random_state: int = 42, n_iterations: int = 100, maximize_mmd: bool = True) -> tuple[np.ndarray, np.ndarray]

Constrained kernel k-means (k=2) train/validation split via an LP.

Full-kernel, paper-derived implementation of the partitioning method of Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation Splits for Model Selection under Domain Shift" (https://openreview.net/forum?id=Q692C0WtiD).

The paper (their Theorem 1, Eq. 8-11) proves that maximizing the MMD between the two sets is equivalent to minimizing the kernel k-means objective Psi(T, V) = SSq(T) + SSq(V) with k = 2 (the within-cluster kernel scatter), because MMD^2 = c * (SSq(S) - Psi(T, V)) with SSq(S) and the cluster sizes held constant. It solves the constrained clustering with a Lloyd-style alternation (their Algorithm 1):

1. Distance update -- kernel distance from every point i to each cluster centroid j via the kernel trick, D_ij = K_ii - (2 / s_j) * sum_l U_lj K_il + (1 / s_j^2) * sum_lm U_lj U_mj K_lm where s_j = sum_l U_lj is the cluster mass.
2. Constrained assignment -- solve the assignment LP (their Eq. 12-16) argmin_U sum_ij U_ij D_ij subject to sum_j U_ij = 1 (Eq. 14), the group-balance equalities sum_{i in g} U_ij = round(h |g|) for every group g and either cluster j (the disjunction in Eq. 15), and the relaxed box constraint 0 <= U_ij <= 1 (Eq. 16). As prescribed by the paper, both cluster orientations are solved and the lower-cost LP is kept. The paper drops the integrality constraint (Eq. 13) because the constraint matrix meets Hoffman's total-unimodularity conditions, so the LP has integral optima.

Groups g are the Cartesian product of the label y and groups inputs (their Y x D); with neither, there is a single group and Eq. 15 reduces to the size constraint alone. The holdout fraction h is derived from train_size; per-group targets are apportioned with largest-remainder so they sum exactly to the requested validation size.

Rounding: even though the LP optima are (near-)integral, fractional values can appear under solver degeneracy, so each iteration rounds by selecting, within every group, the round(h |g|) points with the largest validation weight U[:, 1]. This guarantees an exact, feasible hard partition.

Convergence (their Proposition 2): Psi is bounded below and monotone under the alternation and there are finitely many partitions, so the max-MMD iteration (full-step assignment) is run until the hard assignment stops changing (capped at n_iterations), keeping the best partition seen.

The min-MMD case (maximize_mmd=False) is the library's dual, not part of the paper. MMD^2 is convex in the relaxed assignment, so minimizing it is a convex program: a full assignment step oscillates between the two segregated vertices, so this uses a damped Frank-Wolfe update v <- (1 - gamma) v + gamma * u* with gamma = 2 / (k + 2) (u* the LP vertex), then rounds the fractional membership at the end. The rounded result lands far below a random split's MMD but, in practice, typically around 2x the MMD reached by the swap optimizer -- prefer the swap method when the absolute lowest MMD matters and this path when the label/group constraints do.

Fidelity notes: the paper's Nyström scaling (Algorithm 1, Step 1; an O(qn) random submatrix approximation for very large n) is not implemented here, so this helper materializes the full n x n kernel and is intended for moderate n. It also uses largest-remainder apportionment rather than independently rounding every group target, guaranteeing that the group constraints sum to the caller's exact requested validation size.

Parameters:

Name	Type	Description	Default
embeddings	ndarray	validated float ndarray of shape (n_samples, n_features).	required
train_size	float | int	fraction in (0, 1) or absolute count for the training set.	required
kernel	str	'rbf' or 'linear'.	'rbf'
gamma	float | None	RBF kernel parameter (default: 1 / n_features).	None
y	ArrayLike | None	optional labels; enforces per-label proportions in both sides.	None
groups	ArrayLike | None	optional group ids; enforces per-group proportions in both sides.	None
random_state	int	seeds the initial random (feasible) partition.	42
n_iterations	int	maximum Lloyd-style iterations.	100
maximize_mmd	bool	if True, minimize the kernel k-means scatter (max-MMD, the paper's objective); if False, maximize the scatter (the min-MMD dual, an anti-clustering that makes the two sets resemble each other).	True

Returns:

Type	Description
tuple[ndarray, ndarray]	(train_indices, test_indices) as sorted 1-D integer ndarrays.

Source code in splytters/utils.py

def constrained_kernel_kmeans_split(
    embeddings: np.ndarray,
    train_size: float | int,
    kernel: str = "rbf",
    gamma: float | None = None,
    y: ArrayLike | None = None,
    groups: ArrayLike | None = None,
    random_state: int = 42,
    n_iterations: int = 100,
    maximize_mmd: bool = True,
) -> tuple[np.ndarray, np.ndarray]:
    """Constrained kernel k-means (k=2) train/validation split via an LP.

    Full-kernel, paper-derived implementation of the partitioning method of
    Napoli & White (TMLR 2025; arXiv 2024), "Clustering-Based Validation Splits
    for Model Selection under Domain Shift"
    (https://openreview.net/forum?id=Q692C0WtiD).

    The paper (their Theorem 1, Eq. 8-11) proves that maximizing the MMD between
    the two sets is equivalent to *minimizing* the kernel k-means objective
    ``Psi(T, V) = SSq(T) + SSq(V)`` with ``k = 2`` (the within-cluster kernel
    scatter), because ``MMD^2 = c * (SSq(S) - Psi(T, V))`` with ``SSq(S)`` and the
    cluster sizes held constant. It solves the constrained clustering with a
    Lloyd-style alternation (their Algorithm 1):

    1. **Distance update** -- kernel distance from every point ``i`` to each
       cluster centroid ``j`` via the kernel trick,
       ``D_ij = K_ii - (2 / s_j) * sum_l U_lj K_il + (1 / s_j^2) * sum_lm U_lj U_mj K_lm``
       where ``s_j = sum_l U_lj`` is the cluster mass.
    2. **Constrained assignment** -- solve the assignment LP (their Eq. 12-16)
       ``argmin_U sum_ij U_ij D_ij`` subject to ``sum_j U_ij = 1`` (Eq. 14),
       the group-balance equalities ``sum_{i in g} U_ij = round(h |g|)`` for every
       group ``g`` and either cluster ``j`` (the disjunction in Eq. 15), and the
       relaxed box constraint ``0 <= U_ij <= 1`` (Eq. 16). As prescribed by the
       paper, both cluster orientations are solved and the lower-cost LP is kept.
       The paper drops the integrality constraint
       (Eq. 13) because the constraint matrix meets Hoffman's total-unimodularity
       conditions, so the LP has integral optima.

    Groups ``g`` are the Cartesian product of the label ``y`` and ``groups``
    inputs (their ``Y x D``); with neither, there is a single group and Eq. 15
    reduces to the size constraint alone. The holdout fraction ``h`` is derived
    from ``train_size``; per-group targets are apportioned with largest-remainder
    so they sum exactly to the requested validation size.

    Rounding: even though the LP optima are (near-)integral, fractional values can
    appear under solver degeneracy, so each iteration rounds by selecting, within
    every group, the ``round(h |g|)`` points with the largest validation weight
    ``U[:, 1]``. This guarantees an exact, feasible hard partition.

    Convergence (their Proposition 2): ``Psi`` is bounded below and monotone under
    the alternation and there are finitely many partitions, so the max-MMD
    iteration (full-step assignment) is run until the hard assignment stops
    changing (capped at ``n_iterations``), keeping the best partition seen.

    The min-MMD case (``maximize_mmd=False``) is the library's dual, not part of
    the paper. ``MMD^2`` is convex in the relaxed assignment, so *minimizing* it
    is a convex program: a full assignment step oscillates between the two
    segregated vertices, so this uses a damped Frank-Wolfe update
    ``v <- (1 - gamma) v + gamma * u*`` with ``gamma = 2 / (k + 2)`` (``u*`` the
    LP vertex), then rounds the fractional membership at the end. The rounded
    result lands far below a random split's MMD but, in practice, typically
    around 2x the MMD reached by the swap optimizer -- prefer the swap method
    when the absolute lowest MMD matters and this path when the label/group
    constraints do.

    Fidelity notes: the paper's Nyström scaling (Algorithm 1, Step 1; an
    ``O(qn)`` random submatrix approximation for very large ``n``) is not
    implemented here, so this helper materializes the full ``n x n`` kernel and
    is intended for moderate ``n``. It also uses largest-remainder apportionment
    rather than independently rounding every group target, guaranteeing that the
    group constraints sum to the caller's exact requested validation size.

    Args:
        embeddings: validated float ndarray of shape (n_samples, n_features).
        train_size: fraction in (0, 1) or absolute count for the training set.
        kernel: 'rbf' or 'linear'.
        gamma: RBF kernel parameter (default: 1 / n_features).
        y: optional labels; enforces per-label proportions in both sides.
        groups: optional group ids; enforces per-group proportions in both sides.
        random_state: seeds the initial random (feasible) partition.
        n_iterations: maximum Lloyd-style iterations.
        maximize_mmd: if True, minimize the kernel k-means scatter (max-MMD, the
            paper's objective); if False, maximize the scatter (the min-MMD dual,
            an anti-clustering that makes the two sets resemble each other).

    Returns:
        (train_indices, test_indices) as sorted 1-D integer ndarrays.
    """
    from scipy import sparse
    from scipy.optimize import linprog

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

    if kernel == "rbf":
        _gamma = gamma if gamma is not None else 1.0 / embeddings.shape[1]
        K = np.exp(-_gamma * cdist(embeddings, embeddings, metric="sqeuclidean"))
    elif kernel == "linear":
        K = embeddings @ embeddings.T
    else:
        raise ValueError(f"kernel must be 'rbf' or 'linear', got {kernel!r}")
    K_diag = np.diag(K).copy()

    # Build group ids as the Cartesian product of y and groups (paper: Y x D).
    keys: list[Any] = []
    for i in range(n_samples):
        key = []
        if y is not None:
            key.append(y[i])
        if groups is not None:
            key.append(groups[i])
        keys.append(tuple(key))
    unique_keys = sorted(set(keys), key=lambda k: repr(k))
    group_members = [
        np.array([i for i in range(n_samples) if keys[i] == k], dtype=np.intp)
        for k in unique_keys
    ]
    group_sizes = [len(m) for m in group_members]

    # Per-group train targets (largest-remainder) -> validation targets. This
    # preserves each group's label/group proportion at the global train fraction
    # and makes the totals sum exactly to n_train / n_val.
    per_group_train = apportion_train(group_sizes, n_train)
    val_targets = [
        sz - tr for sz, tr in zip(group_sizes, per_group_train, strict=True)
    ]

    # Initial feasible partition: pick val_targets[g] members of each group at
    # random for the validation cluster. ``v`` is the (fractional) validation
    # membership weight per point; train membership is ``1 - v``.
    rng = check_random_state(random_state)
    v = np.zeros(n_samples)
    for members, vt in zip(group_members, val_targets, strict=True):
        if vt > 0:
            chosen = rng.choice(members, size=vt, replace=False)
            v[chosen] = 1.0

    # Precompute both Eq. 15 orientations. Variables are ordered u[2*i + j] for
    # point i, cluster j. The paper treats the constrained validation cluster as
    # disjunctive because cluster labels are arbitrary, so every assignment step
    # solves one LP with column 0 constrained and one with column 1 constrained.
    n_vars = 2 * n_samples

    def assignment_constraints(
        validation_column: int,
    ) -> tuple[sparse.csr_matrix, np.ndarray]:
        rows: list[int] = []
        cols: list[int] = []
        data: list[float] = []
        b_eq: list[float] = []
        # Eq. 14: each point assigned once.
        for i in range(n_samples):
            rows += [i, i]
            cols += [2 * i, 2 * i + 1]
            data += [1.0, 1.0]
            b_eq.append(1.0)
        # Eq. 15: per-group validation mass fixed in this orientation. The other
        # cluster's mass follows from Eq. 14.
        row = n_samples
        for members, vt in zip(group_members, val_targets, strict=True):
            for i in members:
                rows.append(row)
                cols.append(2 * i + validation_column)
                data.append(1.0)
            b_eq.append(float(vt))
            row += 1
        matrix = sparse.csr_matrix(
            (data, (rows, cols)),
            shape=(n_samples + len(unique_keys), n_vars),
        )
        return matrix, np.asarray(b_eq)

    constraint_systems = [
        (column, *assignment_constraints(column)) for column in (0, 1)
    ]

    def round_membership(weights: np.ndarray) -> set[int]:
        """Round fractional weights to an exact, feasible validation set: within
        each group, take the ``val_targets[g]`` points with the largest weight."""
        chosen_all: set[int] = set()
        for members, vt in zip(group_members, val_targets, strict=True):
            if vt <= 0:
                continue
            order = np.argsort(-weights[members], kind="stable")
            chosen_all.update(int(i) for i in members[order[:vt]])
        return chosen_all

    def mmd_sq(val_set: set[int]) -> float:
        """True (biased) MMD^2 between the current train/validation partition."""
        val = np.zeros(n_samples)
        if val_set:
            val[np.fromiter(val_set, dtype=np.intp, count=len(val_set))] = 1.0
        trn = 1.0 - val
        s_v, s_t = val.sum(), trn.sum()
        Kv, Kt = K @ val, K @ trn
        return float(
            (trn @ Kt) / (s_t * s_t)
            + (val @ Kv) / (s_v * s_v)
            - 2.0 * (trn @ Kv) / (s_t * s_v)
        )

    # Lloyd-style alternation. For max-MMD this is exactly constrained kernel
    # k-means (full-step assignment). MMD^2 is convex in the (relaxed) assignment,
    # so a full step maximizes it toward a vertex, but *minimizing* it (the
    # min-MMD dual, an anti-clustering) needs a damped Frank-Wolfe step
    # ``gamma = 2 / (k + 2)`` to avoid the two-vertex oscillation a full step
    # produces. Track the best hard partition seen and return it.
    best_val: set[int] = round_membership(v)
    best_obj = mmd_sq(best_val)
    prev_val: set[int] | None = None
    for k in range(max(1, n_iterations)):
        # Step 1: kernel distances to the current soft centroids.
        D = np.empty((n_samples, 2))
        for j, w in enumerate((1.0 - v, v)):
            s = w.sum()
            if s <= 0:
                D[:, j] = K_diag  # empty cluster: distance is just K_ii
                continue
            Kw = K @ w
            D[:, j] = K_diag - 2.0 * Kw / s + float(w @ Kw) / (s * s)

        # Step 2: constrained assignment LP. Minimize scatter for max-MMD;
        # maximize it (negate objective) for the min-MMD dual.
        c = D.reshape(-1).copy()
        if not maximize_mmd:
            c = -c
        lp_solutions = []
        # Both orientations are part of the paper's max-MMD algorithm. The
        # min-MMD Frank-Wolfe dual is library-specific and keeps a fixed
        # validation column so its damped iterates do not relabel every step.
        systems_to_solve = (
            constraint_systems if maximize_mmd else constraint_systems[1:]
        )
        for validation_column, A_eq, b_eq_arr in systems_to_solve:
            res = linprog(
                c, A_eq=A_eq, b_eq=b_eq_arr,
                bounds=(0.0, 1.0), method="highs",
            )
            if not res.success:  # pragma: no cover - LP is always feasible here
                raise RuntimeError(
                    f"assignment LP failed for validation column "
                    f"{validation_column}: {res.message}"
                )
            lp_solutions.append((float(res.fun), validation_column, res.x))

        _, validation_column, solution = min(
            lp_solutions, key=lambda item: (item[0], item[1])
        )
        vertex = solution.reshape(n_samples, 2)[:, validation_column]

        gamma = 1.0 if maximize_mmd else 2.0 / (k + 2)
        v = (1.0 - gamma) * v + gamma * vertex

        val_set = round_membership(v)
        obj = mmd_sq(val_set)
        improved = obj > best_obj if maximize_mmd else obj < best_obj
        if improved:
            best_obj, best_val = obj, val_set
        # Full-step kernel k-means converges to a stable hard assignment; stop
        # early there. Damped Frank-Wolfe keeps moving, so run the full budget.
        if maximize_mmd and val_set == prev_val:
            break
        prev_val = val_set

    test_indices = sorted(best_val)
    train_indices = [i for i in range(n_samples) if i not in best_val]
    return as_index_array(train_indices), as_index_array(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", "grouped") 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", "grouped") 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]]