spacr.ops_accel

The three array operations OPS uses, run on a GPU when one is available.

The GPU-suitable steps of OPS are the per-window affine warp, peak finding, the unmixing matrix multiply and Cellpose segmentation (which has its own GPU path), and, for the merge, ORB detection with Hamming descriptor matching and mutual-nearest-neighbour cell matching. These reduce to three primitives: a windowed maximum, a matrix multiply and a nearest-neighbour search. They are collected here so device handling lives in one place.

Backend order: PyTorch first, CuPy second, NumPy always. Torch is already a spaCR dependency through Cellpose, so nothing new is installed, and one code path reaches CUDA, ROCm and Apple MPS.

Every function takes and returns NumPy arrays, so callers do not depend on which backend ran. Each has a test that runs both paths on the same input and asserts they agree.

The GPU may be shared, so ops_gpu=False refuses a device that exists, and a backend that raises at runtime (an out-of-memory, a driver that went away) falls through to the next rather than failing the plate.

The CPU path is a supported path, not an error path: it runs on laptops, in CI and when the GPU is busy.

Functions

accelerated_backends(→ Tuple[str, ...])	Which backends can run here, best first, always ending in "numpy".
matmul(→ numpy.ndarray)	first @ second, on the first available backend.
maximum_filter(→ numpy.ndarray)	The maximum of each size x size window, centred on every pixel.
nearest_neighbours(→ Tuple[numpy.ndarray, numpy.ndarray])	For each row of source, the closest row of target.

Module Contents

spacr.ops_accel.accelerated_backends(gpu: bool = True) → Tuple[str, ...]

Which backends can run here, best first, always ending in “numpy”.

Parameters:

gpu – False refuses the accelerated ones outright.

Returns:

the usable names.

spacr.ops_accel.matmul(first: numpy.ndarray, second: numpy.ndarray, *, gpu: bool = True, backend: str | None = None) → numpy.ndarray

first @ second, on the first available backend.

Used for unmixing: correcting dye bleed-through applies one small matrix to every spot of every cycle, a large multiply by a tiny operand that suits a GPU.

Parameters:

• first – left operand.

• second – right operand.

• gpu – False keeps it on the CPU.

• backend – force one.

Returns:

the product, as float32 NumPy.

spacr.ops_accel.maximum_filter(field: numpy.ndarray, size: int, *, gpu: bool = True, backend: str | None = None) → numpy.ndarray

The maximum of each size x size window, centred on every pixel.

Used for peak finding: a read is a local maximum of the spot score, and “local” means this. On a well-sized field it is also the most expensive thing in the decode chain, and it is a max-pool, which GPUs run efficiently.

Edges repeat the border pixel, matching scipy.ndimage.maximum_filter with mode="nearest", because the alternative is a rim of false maxima all the way round the field.

Parameters:

• field – a 2-D array.

• size – the window edge, in pixels. Even sizes are raised to the next odd one: a window with no centre cannot be centred.

• gpu – False keeps it on the CPU.

• backend – force one.

Returns:

the filtered field, same shape and dtype-compatible.

spacr.ops_accel.nearest_neighbours(source: numpy.ndarray, target: numpy.ndarray, *, gpu: bool = True, backend: str | None = None, chunk: int = CHUNK) → Tuple[numpy.ndarray, numpy.ndarray]

For each row of source, the closest row of target.

Always chunked. A well can carry a hundred thousand nuclei, and the full distance matrix between two such sets is 40 GB – more than the card and more than the host. Chunking the first set bounds the working set at chunk x M however large the well is.

Ties go to the lower index, on every backend, so a caller that pairs mutual nearest neighbours gets the same pairs whatever ran.

Parameters:

• source – (N, D) points.

• target – (M, D) points.

• gpu – False keeps it on the CPU.

• backend – force one.

• chunk – rows of source per pass.

Returns:

(indices, distances), both length N.

Nested helpers

matmul.run(name)

The matrix product on one named backend.

Parameters:

name – “torch”, “cupy” or “numpy”.

Returns:

the product as float32 NumPy, wherever it was computed.

Raises:

RuntimeError – when that backend is not installed.

spacr/ops_accel.py:210

maximum_filter.run(name)

The windowed maximum on one named backend.

Parameters:

name – “torch”, “cupy” or “numpy”.

Returns:

the filtered field as a 2-D array.

Raises:

RuntimeError – when that backend is not installed, which is what _try reads to fall through to the next one.

spacr/ops_accel.py:156

nearest_neighbours.run(name)

Nearest neighbours on one named backend.

Parameters:

name – “torch”, “cupy” or “numpy”.

Returns:

(indices, distances), one row per source point.

Raises:

RuntimeError – when that backend is not installed.

Chunked at CHUNK on every backend, because the pairwise distance matrix is the memory cost here and it is quadratic.

spacr/ops_accel.py:268