API

Assignment designs: each returns one group label (0 to k - 1) per unit.

whiten(x)

Center the covariates and rotate them to identity sample covariance.

Euclidean distance on the result is Mahalanobis distance on x.

Parameters:

x (ndarray) – Covariates, one row per unit.

Returns:

The whitened covariates.

Raises:

ValueError – If covariates are nonfinite, have invalid dimensions, or their sample covariance is singular.

Return type:

ndarray

crd(n, k, rng)

Complete randomization into k equal groups.

Parameters:

• n (int) – Number of units, a multiple of k.

• k (int) – Number of groups.

• rng (Generator) – Random number generator.

Returns:

Group labels.

Raises:

ValueError – If sizes are not positive integers, n is not a multiple of k, or k < 2.

Return type:

ndarray

randomized_aba(x, k, rng, *, jitter=1.0)

Randomized assignment-based anticlustering into k equal groups.

Units are whitened and sorted by squared distance from the center, after multiplying each distance by exp(jitter * N(0, 1)). The sorted list is cut into shells of k consecutive units, and each shell is spread one unit per group by the assignment that maximizes the total squared distance from units to their group’s current centroid, as in assignment-based anticlustering (Baumann et al. 2026). The group labels are then shuffled, so every unit has probability 1 / k of each group.

Parameters:

• x (ndarray) – Covariates, one row per unit.

• k (int) – Number of groups.

• rng (Generator) – Random number generator.

• jitter (float) – Log-scale noise on the sort key; 0 gives a single partition, randomized only by the label shuffle.

Returns:

Group labels.

Raises:

ValueError – If jitter is negative or nonfinite.

Return type:

ndarray

sequential_car(x, k, rng, *, k0=3.0)

Covariate-adaptive randomization over k arms (Zhang 2023, section 4).

Units arrive in a random order. Each goes to arm t with probability proportional to 1 - Phi(clip(d_t, -k0, k0)), where d_t is the imbalance after assigning the unit to t minus the average of that imbalance over arms (Zhang’s allocation function 4.3). The imbalance is sum_t |S_t - S / k|^2, with S_t the sum of the features over arm t and S over all units so far; the features are a constant and the whitened x, so arm sizes are balanced too, though not forced equal. The centered difference d_t reduces to 2 <S_t - S / k, phi> for the arriving unit’s features phi.

Parameters:

• x (ndarray) – Covariates, one row per unit.

• k (int) – Number of arms.

• rng (Generator) – Random number generator.

• k0 (float) – Clipping constant of the allocation function.

Returns:

Arm labels; sizes may differ slightly.

Raises:

ValueError – If k0 is negative or nonfinite.

Return type:

ndarray

exchange_anticlust(x, k, rng, *, init=None, local_max=False)

The exchange method on the k-means anticlustering objective.

Starting from a random equal split (or init), each unit in turn is swapped with the unit in another group whose swap most reduces the spread of the group centroids, if any swap reduces it. One pass is the procedure Papenberg and Angelike (2026) use, after R anticlust’s exchange method with objective = "variance". With local_max the passes repeat until one makes no swap, as in anticlust’s local-maximum method; for two groups this is greedy pair-switching (Krieger, Azriel and Kapelner 2019) with the unweighted Euclidean objective.

Parameters:

• x (ndarray) – Covariates, one row per unit, used as given (not standardized).

• k (int) – Number of groups.

• rng (Generator) – Random number generator, used for the initial split.

• init (ndarray | None) – Optional integer labels, with equal counts of 0 through k - 1.

• local_max (bool) – Repeat passes until a local optimum.

Returns:

Group labels.

Raises:

ValueError – If initial labels do not form k equal groups.

Return type:

ndarray

greedy_switch(x, rng, *, features=None)

Two equal arms by greedy pair-switching from a random start.

A random equal split is improved by swapping units between the arms until no swap reduces the distance between the arms’ means (Krieger, Azriel and Kapelner 2019). Distances are Mahalanobis on features (default x); pass the covariates with their squares to balance spread as well. Each call starts from a fresh random split, so repeated calls draw from the design, which is what a randomization test needs.

Parameters:

• x (ndarray) – Covariates, one row per unit.

• rng (Generator) – Random number generator.

• features (ndarray | None) – Optional matrix to balance instead of x.

Returns:

Treatment indicators (0 or 1), half of each, with labels shuffled.

Raises:

ValueError – If features does not have one row per unit in x.

Return type:

ndarray

blocked(x, k, rng, *, index=None)

Blocked randomization: sort on an index, randomize within blocks of k.

Units are sorted on index (default the first principal component of the standardized covariates), cut into consecutive blocks of k, and each block’s units are randomly permuted across the k groups: matched k-tuples on a one-dimensional index.

Parameters:

• x (ndarray) – Covariates, one row per unit.

• k (int) – Number of groups.

• rng (Generator) – Random number generator.

• index (ndarray | None) – Optional sorting index, one value per unit.

Returns:

Group labels.

Raises:

ValueError – If the index is not a finite vector with one value per unit, or a covariate is constant when computing the default index.

Return type:

ndarray

rerandomize(x, k, rng, *, accept=0.01, batch=64, max_draws=100000)

Rerandomization for k equal groups on a multi-group Mahalanobis criterion.

Complete randomizations are drawn until B = sum_g m |wbar_g - wbar|^2 on whitened covariates, which is approximately chi-squared with (k - 1) p degrees of freedom under complete randomization, falls below that distribution’s accept quantile.

Parameters:

• x (ndarray) – Covariates, one row per unit.

• k (int) – Number of groups.

• rng (Generator) – Random number generator.

• accept (float) – Target acceptance probability.

• batch (int) – Candidate splits drawn at a time.

• max_draws (int) – Maximum candidate splits before raising RuntimeError.

Returns:

Group labels.

Raises:

• ValueError – If accept is outside (0, 1), or batch or max_draws is not a positive integer.

• RuntimeError – If no acceptable assignment is found within max_draws.

Return type:

ndarray

Randomization inference for contrasts between arms.

contrast(y, z, a, b)

Difference in means between arms a and b, for one or many draws.

Parameters:

• y (ndarray) – Outcomes.

• z (ndarray) – Arm labels, one draw (n,) or many (d, n).

• a (int) – First arm.

• b (int) – Second arm.

Returns:

Mean outcome in arm a minus mean in arm b, per draw.

Raises:

ValueError – If dimensions do not match, outcomes are nonfinite, or either requested arm is absent from any draw.

Return type:

ndarray

ri_pvalue(stat_obs, stat_ref)

Two-sided randomization p-value under the sharp null of no effect.

stat_ref holds the statistic recomputed on the observed outcomes for fresh draws of the same design. The observed draw counts as one draw.

Parameters:

• stat_obs (float) – Statistic for the observed draw.

• stat_ref (ndarray) – Statistics for the reference draws.

Returns:

(1 + #{|ref| >= |obs|}) / (1 + d), counting ties within a relative tolerance.

Raises:

ValueError – If the observed statistic is not a finite scalar, or the reference statistics are not a nonempty finite vector.

Return type:

float