mwatershed

A rusty mutex watershed

Documentation: https://pattonw.github.io/mwatershed

Install

pip install mwatershed    # or: uv add mwatershed

Wheels are published for Linux, macOS and Windows; building from source needs a Rust toolchain.

Features

• A mutex watershed implementation for affinities
• Should handle up to 6 dimensional arrays of fragments (7 dimensions if you include offsets in affinities) but only tested on 2 and 3 dimensions

Usage

components = mwatershed.agglom(
    affinities: NDArray[np.float64],
    offsets: list[list[int]],
    seeds: NDArray[np.uint64],
    edges: Optional[list[tuple[usize, usize, f64]]] = None,
)

where:

• affinities is a k+1 dimensional array of non nan affinities with leading dimension having size n
• offsets is a list of length n of offset tuples of k integers
• seeds is a k dimensional array of fragment ids. Note seeds.shape must be equal to affinities.shape[1:]. Any entry not equal to 0 is guaranteed to stay that way, any entry equal to zero has no priors.
• edges is a list of (u, v, aff) tuples to insert arbitrary extra affinities between fragment ids

Graphs as arrays

For a graph that already lives in numpy arrays, cluster_arrays runs the mutex watershed without building a Python object per edge:

nodes, labels = mwatershed.cluster_arrays(
    u,
    v,
    weights,  # one edge per position; positive attracts, else repels
    nodes=None,  # sorted unique node ids; default: those in the edges
    sorted=False,  # set if the edges are already in processing order
    prefix=None,  # (u, v, positive) hard decisions processed first
    linkage="abs_max",  # or any GASP linkage below
    cannot_link=True,
)

Node ids are arbitrary uint64. Edges are processed by descending |weight| with ties in input order, after the prefix edges, which are processed in their own order. The result is one label per entry of nodes: the smallest node id in that node's cluster, so the partition equals cluster on the same edges in the same order.

Linkage criteria (GASP)

agglom, cluster_edges and cluster_arrays accept a linkage argument that selects how the interaction between two clusters is computed from the affinities connecting them, following the GASP framework of Bailoni et al. (2022):

linkage	Interaction between two clusters	Known as
"abs_max"	the connecting edge with the largest magnitude (default)	Mutex watershed
"sum"	sum of the connecting affinities	GAEC / HCC-Sum
"mean"	mean of the connecting affinities	average linkage (HC-Avg / HCC-Avg)
"max"	largest connecting affinity	single linkage
"min"	smallest connecting affinity	complete linkage

components = mwatershed.agglom(affinities, offsets, linkage="mean")

Pairs of clusters are visited in order of decreasing interaction magnitude. An attractive pair is merged; a repulsive pair receives a cannot-link constraint when cannot_link=True (the default). Once every remaining pair is repulsive or constrained, the constraints are lifted and any pairs that are still attractive are merged. With "abs_max" the constraints never change the result, which is exactly the mutex watershed. Seeds and the edges argument always act as hard constraints regardless of the linkage.

The generic linkages use a contracted-graph priority queue and are several times slower than the dedicated mutex watershed path; see LINKAGE_REPORT.md for timings (cargo run --release --example benchmark_linkages) or run cargo bench for the Criterion comparison.

Contributing

The commands below are the ones CI runs, so a green local run means a green PR.

git clone https://github.com/pattonw/mwatershed
cd mwatershed
uv sync                                             # builds the extension (needs Rust) + dev tools
uv run pytest                                       # Python tests with coverage
uv run pytest --no-cov tests/<file>.py              # one file, without the coverage floor
cargo test                                          # Rust tests
uv run --only-group dev ruff check . && uv run --only-group dev ruff format --check .
uv run ty check --error-on-warning                  # types
uv run --group docs sphinx-build docs/source docs/build/html -b html

uv.lock pins the development environment; the lock-upgrade workflow refreshes it monthly and opens a pull request.