Module batch_mixing 

Spread batches across cell states.

A batch (a donor, a sample, a run) shifts every one of its cells, and it also has its own cell-state composition. Batch correction must remove the first and keep the second. Subtracting each batch’s plain mean removes both, so a batch rich in one state has that state pulled to the origin, and cells of one state stop lining up across batches. Here the projection is fit as

  proj_j = state(bin_j) + shift(batch_j)

by alternating the bins (binary codes of the projection) with the batch shifts, so the shift is what is left after each cell’s state is removed.

The same codes form a binary tree over cell states (bit k is the sign of the k-th component, so dropping the top bit gives the parent). merge_poorly_mixed_bins cuts that tree so every bin holds enough batches.

Constants

CENTRING_ROUNDS

Rounds of the state/shift alternation; it stops early once no column changes bin.

DEFAULT_MERGE_LEVELS

Levels a poorly mixed bin may merge up by default.

DEFAULT_MIN_BATCHES_PER_GROUP

Batches a bin needs by default (capped by the number of batches).

MAX_STATE_BITS

Most state bits used while centring.

Functions

batch_indices

Map arbitrary batch keys to 0..n_batches in order of first appearance.

centre_batches_within_state

Centre each batch within cell state (see the module docs). Starts from the plain per-batch mean; the result has mean zero over all columns.

default_state_bits

State bits for centring: about two columns per batch per bin, at least one, at most the projection dimension and MAX_STATE_BITS.

merge_poorly_mixed_bins

Cut the binary code tree so that bins mix batches. Bottom-up, a node holding fewer than min_batches batches collapses its whole parent subtree into the parent, for at most levels levels or up to the root. A bin still short after that is left as is. Returns one label per column.