diff --git a/CMakeLists.txt b/CMakeLists.txt index e56cde5..23391cd 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -11,6 +11,7 @@ find_package(nanobind CONFIG REQUIRED) nanobind_add_module(_core NB_STATIC + src/bindings/affinities.cxx src/bindings/blocking.cxx src/bindings/module.cxx src/bindings/graph.cxx diff --git a/development/affinities/check_compute_affinities.py b/development/affinities/check_compute_affinities.py new file mode 100644 index 0000000..1056067 --- /dev/null +++ b/development/affinities/check_compute_affinities.py @@ -0,0 +1,133 @@ +"""Cross-check bioimage-cpp's compute_affinities against affogato. + +Benchmarks on the registered ISBI ground-truth segmentation volume +(30 × 512 × 512 = 7.86 M voxels, ~660 distinct labels) using the same +17-channel offset configuration that elf uses for mutex-watershed on this +data. Small enough to fit in memory, big enough that initialization and +allocation effects don't dominate. + +Not part of the pytest suite (per CLAUDE.md: external-library comparisons +live under ``development/``, not under ``tests/``). +""" + +from __future__ import annotations + +import argparse +import sys +from statistics import mean, median +from time import perf_counter + +import numpy as np + +import bioimage_cpp as bic +from bioimage_cpp._data import ISBI_AFFINITY_OFFSETS, load_isbi_gt_segmentation + +try: + import affogato.affinities as affo +except ImportError as error: # pragma: no cover - dev script + sys.stderr.write(f"affogato not installed: {error}\n") + sys.exit(1) + + +# Subsets of ISBI_AFFINITY_OFFSETS. The "nearest neighbours" subset is the +# typical multicut input (3 channels); the "full 17" subset matches +# elf's mutex-watershed proposal generator and exercises long-range offsets. +OFFSET_SUBSETS = { + "nearest": [(-1, 0, 0), (0, -1, 0), (0, 0, -1)], + "full17": list(ISBI_AFFINITY_OFFSETS), +} + + +def time_call(fn, repeats): + timings = [] + result = None + for _ in range(repeats): + start = perf_counter() + result = fn() + timings.append(perf_counter() - start) + return timings, result + + +def run_case(labels, offsets, *, repeats, ignore_label=None): + offsets_list = [list(offset) for offset in offsets] + + bic_timings, (bic_affs, bic_mask) = time_call( + lambda: bic.affinities.compute_affinities( + labels, + offsets_list, + ignore_label=ignore_label, + return_mask=True, + number_of_threads=1, + ), + repeats, + ) + affo_timings, (affo_affs, affo_mask) = time_call( + lambda: affo.compute_affinities( + labels, + offsets_list, + have_ignore_label=ignore_label is not None, + ignore_label=ignore_label if ignore_label is not None else 0, + ), + repeats, + ) + + return { + "n_offsets": len(offsets_list), + "ignore": ignore_label, + "ok_affs": np.array_equal(bic_affs, affo_affs), + "ok_mask": np.array_equal(bic_mask, affo_mask), + "bic_median_s": median(bic_timings), + "affo_median_s": median(affo_timings), + "bic_mean_s": mean(bic_timings), + "affo_mean_s": mean(affo_timings), + "max_abs_diff": float(np.max(np.abs(bic_affs - affo_affs))) if bic_affs.shape == affo_affs.shape else float("nan"), + } + + +def main(): + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--repeats", type=int, default=3) + parser.add_argument("--timeout", type=float, default=60.0) + args = parser.parse_args() + + labels = load_isbi_gt_segmentation(timeout=args.timeout) + n_labels = int(labels.max()) + 1 + n_voxels = int(np.prod(labels.shape)) + print( + f"labels: shape={labels.shape}, dtype={labels.dtype}, " + f"n_voxels={n_voxels:,}, n_labels={n_labels}", + flush=True, + ) + + rows = [] + for name, offsets in OFFSET_SUBSETS.items(): + for ig in (None, 0): + row = run_case(labels, offsets, repeats=args.repeats, ignore_label=ig) + row["offsets_name"] = name + rows.append(row) + + print() + print( + f"{'offsets':>10} {'n':>3} {'ignore':>6} {'affs':>5} {'mask':>5}" + f" {'bic_s':>9} {'affo_s':>9} {'speedup':>8}" + ) + print("-" * 68) + all_ok = True + for r in rows: + speedup = r["affo_median_s"] / r["bic_median_s"] if r["bic_median_s"] > 0 else float("inf") + print( + f"{r['offsets_name']:>10} {r['n_offsets']:>3d} {str(r['ignore']):>6}" + f" {'OK' if r['ok_affs'] else 'FAIL':>5}" + f" {'OK' if r['ok_mask'] else 'FAIL':>5}" + f" {r['bic_median_s']:>9.4f} {r['affo_median_s']:>9.4f}" + f" {speedup:>7.2f}x" + ) + all_ok = all_ok and r["ok_affs"] and r["ok_mask"] + + if not all_ok: + print("\nFAIL: output mismatch", file=sys.stderr) + sys.exit(1) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/affinities/compute_affinities.hxx b/include/bioimage_cpp/affinities/compute_affinities.hxx new file mode 100644 index 0000000..92950d2 --- /dev/null +++ b/include/bioimage_cpp/affinities/compute_affinities.hxx @@ -0,0 +1,191 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/threading.hxx" + +#include +#include +#include +#include +#include +#include + +// Pairwise affinity computation from a label volume. Mirrors affogato's +// `affinities/affinities.hxx::compute_affinities` but without the xtensor +// dependency, with proper input validation at the binding boundary, and with +// optional output mask + optional per-offset parallelism. +// +// For each spatial coordinate ``c`` and offset index ``oi``, +// ``affs[oi, c] = 1`` iff ``labels[c] == labels[c + offsets[oi]]``, else 0. +// ``mask[oi, c] = 1`` iff the offset stays in bounds and neither endpoint +// equals ``ignore_label``; out-of-bounds and ignore-label positions produce +// ``affs = 0`` and ``mask = 0``. +namespace bioimage_cpp::affinities { + +// Boolean affinities on a 2D label volume. Preconditions (validated in the +// binding layer): +// * labels.ndim() == 2 +// * affs.shape == {n_offsets, labels.shape[0], labels.shape[1]} +// * if mask is non-null: mask->shape == affs.shape +// * each entry of offsets has length 2 +template +void compute_affinities_2d( + const ConstArrayView &labels, + const std::vector> &offsets, + const ArrayView &affs, + const ArrayView *mask, + const std::optional ignore_label, + const std::size_t number_of_threads = 1 +) { + const auto height = labels.shape[0]; + const auto width = labels.shape[1]; + const auto plane = height * width; + const auto number_of_offsets = offsets.size(); + + const auto n_threads = ::bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_offsets + ); + + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_offsets, + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t oi = begin; oi < end; ++oi) { + const auto dy = offsets[oi][0]; + const auto dx = offsets[oi][1]; + + AffT * const affs_channel = + affs.data + static_cast(oi) * plane; + std::uint8_t * const mask_channel = (mask != nullptr) + ? mask->data + static_cast(oi) * plane + : nullptr; + + std::fill_n(affs_channel, plane, AffT{0}); + if (mask_channel != nullptr) { + std::fill_n(mask_channel, plane, std::uint8_t{0}); + } + + // Sub-rectangle where (y, x) AND (y+dy, x+dx) are both in bounds. + const auto y_begin = std::max(0, -dy); + const auto y_end = height - std::max(0, dy); + const auto x_begin = std::max(0, -dx); + const auto x_end = width - std::max(0, dx); + if (y_begin >= y_end || x_begin >= x_end) { + continue; + } + + for (std::ptrdiff_t y = y_begin; y < y_end; ++y) { + const auto ny = y + dy; + const LabelT * const row = labels.data + y * width; + const LabelT * const neighbor_row = labels.data + ny * width; + AffT * const out_row = affs_channel + y * width; + std::uint8_t * const mask_row = (mask_channel != nullptr) + ? mask_channel + y * width : nullptr; + for (std::ptrdiff_t x = x_begin; x < x_end; ++x) { + const LabelT a = row[x]; + const LabelT b = neighbor_row[x + dx]; + if (ignore_label.has_value()) { + const LabelT ig = *ignore_label; + if (a == ig || b == ig) { + continue; // affs/mask already 0 + } + } + out_row[x] = (a == b) ? AffT{1} : AffT{0}; + if (mask_row != nullptr) { + mask_row[x] = 1; + } + } + } + } + } + ); +} + +// Boolean affinities on a 3D label volume. Preconditions identical to the 2D +// case with one extra axis. +template +void compute_affinities_3d( + const ConstArrayView &labels, + const std::vector> &offsets, + const ArrayView &affs, + const ArrayView *mask, + const std::optional ignore_label, + const std::size_t number_of_threads = 1 +) { + const auto depth = labels.shape[0]; + const auto height = labels.shape[1]; + const auto width = labels.shape[2]; + const auto volume = depth * height * width; + const auto plane = height * width; + const auto number_of_offsets = offsets.size(); + + const auto n_threads = ::bioimage_cpp::detail::normalize_thread_count( + number_of_threads, number_of_offsets + ); + + ::bioimage_cpp::detail::parallel_for_chunks( + n_threads, + number_of_offsets, + [&](const std::size_t, const std::size_t begin, const std::size_t end) { + for (std::size_t oi = begin; oi < end; ++oi) { + const auto dz = offsets[oi][0]; + const auto dy = offsets[oi][1]; + const auto dx = offsets[oi][2]; + + AffT * const affs_channel = + affs.data + static_cast(oi) * volume; + std::uint8_t * const mask_channel = (mask != nullptr) + ? mask->data + static_cast(oi) * volume + : nullptr; + + std::fill_n(affs_channel, volume, AffT{0}); + if (mask_channel != nullptr) { + std::fill_n(mask_channel, volume, std::uint8_t{0}); + } + + const auto z_begin = std::max(0, -dz); + const auto z_end = depth - std::max(0, dz); + const auto y_begin = std::max(0, -dy); + const auto y_end = height - std::max(0, dy); + const auto x_begin = std::max(0, -dx); + const auto x_end = width - std::max(0, dx); + if (z_begin >= z_end || y_begin >= y_end || x_begin >= x_end) { + continue; + } + + for (std::ptrdiff_t z = z_begin; z < z_end; ++z) { + const auto nz = z + dz; + const LabelT * const slab = labels.data + z * plane; + const LabelT * const neighbor_slab = labels.data + nz * plane; + AffT * const out_slab = affs_channel + z * plane; + std::uint8_t * const mask_slab = (mask_channel != nullptr) + ? mask_channel + z * plane : nullptr; + for (std::ptrdiff_t y = y_begin; y < y_end; ++y) { + const auto ny = y + dy; + const LabelT * const row = slab + y * width; + const LabelT * const neighbor_row = neighbor_slab + ny * width; + AffT * const out_row = out_slab + y * width; + std::uint8_t * const mask_row = (mask_slab != nullptr) + ? mask_slab + y * width : nullptr; + for (std::ptrdiff_t x = x_begin; x < x_end; ++x) { + const LabelT a = row[x]; + const LabelT b = neighbor_row[x + dx]; + if (ignore_label.has_value()) { + const LabelT ig = *ignore_label; + if (a == ig || b == ig) { + continue; + } + } + out_row[x] = (a == b) ? AffT{1} : AffT{0}; + if (mask_row != nullptr) { + mask_row[x] = 1; + } + } + } + } + } + } + ); +} + +} // namespace bioimage_cpp::affinities diff --git a/src/bindings/affinities.cxx b/src/bindings/affinities.cxx new file mode 100644 index 0000000..7b3465c --- /dev/null +++ b/src/bindings/affinities.cxx @@ -0,0 +1,232 @@ +#include "affinities.hxx" + +#include "bioimage_cpp/affinities/compute_affinities.hxx" +#include "bioimage_cpp/array_view.hxx" + +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +template +using LabelArray = nb::ndarray; +using FloatArray = nb::ndarray; +using UInt8Array = nb::ndarray; + +template +std::vector ndarray_shape(const T &array) { + std::vector shape(array.ndim()); + for (std::size_t axis = 0; axis < array.ndim(); ++axis) { + shape[axis] = static_cast(array.shape(axis)); + } + return shape; +} + +FloatArray make_float_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new float[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return FloatArray(data, shape.size(), shape.data(), owner); +} + +UInt8Array make_uint8_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new std::uint8_t[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return UInt8Array(data, shape.size(), shape.data(), owner); +} + +template +std::vector> validate_offsets( + const std::vector> &offsets +) { + std::vector> result; + result.reserve(offsets.size()); + for (std::size_t i = 0; i < offsets.size(); ++i) { + if (offsets[i].size() != D) { + throw std::invalid_argument( + "each offset must have length matching the spatial ndim, got " + "spatial ndim=" + std::to_string(D) + + ", offset[" + std::to_string(i) + "].length=" + + std::to_string(offsets[i].size()) + ); + } + std::array entry{}; + for (std::size_t axis = 0; axis < D; ++axis) { + entry[axis] = offsets[i][axis]; + } + result.push_back(entry); + } + return result; +} + +template +std::pair compute_affinities_nd_t( + LabelArray labels, + const std::vector> &offsets, + const std::optional ignore_label, + const bool return_mask, + const std::size_t number_of_threads +) { + if (labels.ndim() != D) { + throw std::invalid_argument( + "labels must have ndim=" + std::to_string(D) + + ", got ndim=" + std::to_string(labels.ndim()) + ); + } + if (offsets.empty()) { + throw std::invalid_argument("offsets must not be empty"); + } + auto offsets_typed = validate_offsets(offsets); + + const auto labels_shape = ndarray_shape(labels); + + std::vector out_shape; + out_shape.reserve(D + 1); + out_shape.push_back(offsets.size()); + for (std::size_t axis = 0; axis < D; ++axis) { + out_shape.push_back(static_cast(labels_shape[axis])); + } + + std::vector out_view_shape(out_shape.size()); + for (std::size_t axis = 0; axis < out_shape.size(); ++axis) { + out_view_shape[axis] = static_cast(out_shape[axis]); + } + + auto affs = make_float_array(out_shape); + UInt8Array mask = return_mask + ? make_uint8_array(out_shape) + : UInt8Array(nullptr, 0, nullptr); + + ConstArrayView labels_view{ + labels.data(), + labels_shape, + {}, + }; + ArrayView affs_view{ + affs.data(), + out_view_shape, + {}, + }; + ArrayView mask_view{ + return_mask ? mask.data() : nullptr, + out_view_shape, + {}, + }; + const ArrayView *mask_ptr = return_mask ? &mask_view : nullptr; + + { + nb::gil_scoped_release release; + if constexpr (D == 2) { + affinities::compute_affinities_2d( + labels_view, offsets_typed, affs_view, mask_ptr, + ignore_label, number_of_threads + ); + } else if constexpr (D == 3) { + affinities::compute_affinities_3d( + labels_view, offsets_typed, affs_view, mask_ptr, + ignore_label, number_of_threads + ); + } + } + + return {std::move(affs), std::move(mask)}; +} + +} // namespace + +void bind_affinities(nb::module_ &m) { + m.def( + "_compute_affinities_2d_uint32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_uint64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_int32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_2d_int64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_uint32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_uint64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_int32", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); + m.def( + "_compute_affinities_3d_int64", + &compute_affinities_nd_t, + nb::arg("labels"), + nb::arg("offsets"), + nb::arg("ignore_label"), + nb::arg("return_mask"), + nb::arg("number_of_threads") + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/affinities.hxx b/src/bindings/affinities.hxx new file mode 100644 index 0000000..e261a4b --- /dev/null +++ b/src/bindings/affinities.hxx @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace bioimage_cpp::bindings { + +void bind_affinities(nanobind::module_ &m); + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx index cee23a9..bab0a1e 100644 --- a/src/bindings/module.cxx +++ b/src/bindings/module.cxx @@ -1,3 +1,4 @@ +#include "affinities.hxx" #include "blocking.hxx" #include "graph.hxx" #include "ground_truth.hxx" @@ -10,6 +11,7 @@ namespace nb = nanobind; NB_MODULE(_core, m) { m.doc() = "C++ extension module for bioimage_cpp."; + bioimage_cpp::bindings::bind_affinities(m); bioimage_cpp::bindings::bind_blocking(m); bioimage_cpp::bindings::bind_graph(m); bioimage_cpp::bindings::bind_ground_truth(m); diff --git a/src/bioimage_cpp/__init__.py b/src/bioimage_cpp/__init__.py index f8a5e6f..0fe3a83 100644 --- a/src/bioimage_cpp/__init__.py +++ b/src/bioimage_cpp/__init__.py @@ -2,6 +2,7 @@ from ._version import __version__ from ._core import Block, Blocking, BlockWithHalo +from . import affinities from . import graph from . import ground_truth from . import segmentation @@ -12,6 +13,7 @@ "Block", "Blocking", "BlockWithHalo", + "affinities", "graph", "ground_truth", "segmentation", diff --git a/src/bioimage_cpp/_data.py b/src/bioimage_cpp/_data.py index 7993f68..b234406 100644 --- a/src/bioimage_cpp/_data.py +++ b/src/bioimage_cpp/_data.py @@ -214,3 +214,26 @@ def load_isbi_raw( with h5py.File(affinity_path(timeout=timeout), "r") as f: raw = f["raw"][:] return np.ascontiguousarray(raw) + + +def load_isbi_gt_segmentation( + *, + timeout: Optional[float] = None, +) -> np.ndarray: + """Load the registered ISBI ground-truth segmentation volume. + + The labels are stored in the same HDF5 file as the affinities under key + ``labels/gt_segmentation``. Returned as a C-contiguous ``uint64`` array + with shape ``(30, 512, 512)`` (~7.9 M voxels). + """ + try: + import h5py + except ModuleNotFoundError as error: + raise ModuleNotFoundError( + "h5py is required to load the registered ISBI segmentation file. " + "Install it with `pip install h5py`." + ) from error + + with h5py.File(affinity_path(timeout=timeout), "r") as f: + labels = f["labels/gt_segmentation"][:] + return np.ascontiguousarray(labels) diff --git a/src/bioimage_cpp/affinities/__init__.py b/src/bioimage_cpp/affinities/__init__.py new file mode 100644 index 0000000..0406fc7 --- /dev/null +++ b/src/bioimage_cpp/affinities/__init__.py @@ -0,0 +1,5 @@ +"""Pairwise affinities from label volumes.""" + +from .compute_affinities import compute_affinities + +__all__ = ["compute_affinities"] diff --git a/src/bioimage_cpp/affinities/compute_affinities.py b/src/bioimage_cpp/affinities/compute_affinities.py new file mode 100644 index 0000000..9d2c870 --- /dev/null +++ b/src/bioimage_cpp/affinities/compute_affinities.py @@ -0,0 +1,127 @@ +"""Pairwise boolean affinities from a label volume.""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import overload + +import numpy as np + +from .. import _core + + +_COMPUTE_AFFINITIES_2D_BY_DTYPE = { + np.dtype("uint32"): _core._compute_affinities_2d_uint32, + np.dtype("uint64"): _core._compute_affinities_2d_uint64, + np.dtype("int32"): _core._compute_affinities_2d_int32, + np.dtype("int64"): _core._compute_affinities_2d_int64, +} + +_COMPUTE_AFFINITIES_3D_BY_DTYPE = { + np.dtype("uint32"): _core._compute_affinities_3d_uint32, + np.dtype("uint64"): _core._compute_affinities_3d_uint64, + np.dtype("int32"): _core._compute_affinities_3d_int32, + np.dtype("int64"): _core._compute_affinities_3d_int64, +} + + +@overload +def compute_affinities( + labels: np.ndarray, + offsets: Sequence[Sequence[int]] | np.ndarray, + *, + ignore_label: int | None = None, + return_mask: bool = True, + number_of_threads: int = 1, +) -> tuple[np.ndarray, np.ndarray]: ... + + +def compute_affinities( + labels: np.ndarray, + offsets: Sequence[Sequence[int]] | np.ndarray, + *, + ignore_label: int | None = None, + return_mask: bool = True, + number_of_threads: int = 1, +): + """Compute boolean pairwise affinities from a label volume. + + For each spatial coordinate ``c`` and offset index ``oi``, + ``affinities[oi, c]`` is ``1.0`` if ``labels[c] == labels[c + offsets[oi]]`` + (the two voxels are in the same cluster) and ``0.0`` otherwise. + + Parameters + ---------- + labels: + 2D or 3D integer label volume. Supported dtypes are ``uint32``, + ``uint64``, ``int32``, ``int64``. Non-contiguous arrays are copied + to a C-contiguous buffer first. + offsets: + Shape ``(n_offsets, ndim)``. Each offset is a per-axis displacement, + in NumPy axis order, applied at each voxel to find the neighbor. + ignore_label: + If given, any pair where either endpoint has this label produces + ``affinity = 0`` and ``mask = 0`` (treated as out-of-volume). + return_mask: + When ``True`` (default), also return a ``uint8`` validity mask of + the same shape as the affinities: ``1`` for in-bounds non-ignored + pairs, ``0`` otherwise. Set to ``False`` to skip the allocation + when only the affinities are needed. + number_of_threads: + Number of threads to parallelize over the offset channels. + + Returns + ------- + affinities : np.ndarray + ``float32`` array of shape ``(n_offsets, *labels.shape)``. + mask : np.ndarray, only if ``return_mask`` is ``True`` + ``uint8`` array of shape ``(n_offsets, *labels.shape)``. + """ + array = np.ascontiguousarray(labels) + if array.ndim not in (2, 3): + raise ValueError( + "labels must be 2D or 3D, got ndim=" + str(array.ndim) + ) + + table = ( + _COMPUTE_AFFINITIES_2D_BY_DTYPE if array.ndim == 2 + else _COMPUTE_AFFINITIES_3D_BY_DTYPE + ) + try: + run = table[array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in table) + raise TypeError( + f"labels must have one of dtypes ({supported}), got dtype={array.dtype}" + ) from error + + normalized_offsets = [ + [int(value) for value in offset] for offset in np.asarray(offsets).tolist() + ] + if len(normalized_offsets) == 0: + raise ValueError("offsets must not be empty") + if any(len(offset) != array.ndim for offset in normalized_offsets): + raise ValueError( + "each offset must have length matching the spatial ndim, got " + f"spatial ndim={array.ndim}" + ) + + n_threads = int(number_of_threads) + if n_threads < 1: + raise ValueError("number_of_threads must be >= 1") + + if ignore_label is None: + typed_ignore: int | None = None + else: + typed_ignore = int(ignore_label) + + affs, mask = run( + array, + normalized_offsets, + typed_ignore, + bool(return_mask), + n_threads, + ) + if return_mask: + return affs, mask + return affs diff --git a/tests/affinities/test_compute_affinities.py b/tests/affinities/test_compute_affinities.py new file mode 100644 index 0000000..6e25db5 --- /dev/null +++ b/tests/affinities/test_compute_affinities.py @@ -0,0 +1,154 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def _numpy_reference(labels, offsets, ignore_label=None): + """Slow but obvious reference: nested Python loops over voxels.""" + labels = np.asarray(labels) + n = len(offsets) + affs = np.zeros((n,) + labels.shape, dtype=np.float32) + mask = np.zeros((n,) + labels.shape, dtype=np.uint8) + for oi, offset in enumerate(offsets): + it = np.ndindex(*labels.shape) + for coord in it: + neighbor = tuple(c + d for c, d in zip(coord, offset)) + if any(n < 0 or n >= s for n, s in zip(neighbor, labels.shape)): + continue + a = labels[coord] + b = labels[neighbor] + if ignore_label is not None and (a == ignore_label or b == ignore_label): + continue + affs[(oi, *coord)] = 1.0 if a == b else 0.0 + mask[(oi, *coord)] = 1 + return affs, mask + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_2d_matches_numpy_reference(dtype): + rng = np.random.default_rng(0) + labels = rng.integers(0, 5, size=(7, 11)).astype(dtype) + offsets = [[0, 1], [1, 0], [1, 1], [2, -3], [-1, 2]] + + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + + assert affs.dtype == np.float32 + assert mask.dtype == np.uint8 + assert affs.shape == (len(offsets), *labels.shape) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +@pytest.mark.parametrize("dtype", [np.uint32, np.uint64, np.int32, np.int64]) +def test_3d_matches_numpy_reference(dtype): + rng = np.random.default_rng(1) + labels = rng.integers(0, 4, size=(4, 5, 6)).astype(dtype) + offsets = [[0, 0, 1], [0, 1, 0], [1, 0, 0], [1, 1, 1], [-2, 0, 3]] + + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_ignore_label_masks_pairs_at_ignored_voxels(): + labels = np.array( + [ + [0, 1, 1, 2], + [0, 1, 2, 2], + ], + dtype=np.uint32, + ) + offsets = [[0, 1]] + + affs, mask = bic.affinities.compute_affinities( + labels, offsets, ignore_label=0 + ) + ref_affs, ref_mask = _numpy_reference(labels, offsets, ignore_label=0) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + assert mask[0, 0, 0] == 0 # voxel has ignore label + assert affs[0, 0, 0] == 0.0 + + +def test_offset_completely_out_of_bounds_yields_zero_mask(): + labels = np.ones((3, 3), dtype=np.uint32) + offsets = [[10, 0]] # never in bounds + affs, mask = bic.affinities.compute_affinities(labels, offsets) + assert affs.sum() == 0 + assert mask.sum() == 0 + + +def test_return_mask_false_returns_only_affinities(): + labels = np.array([[0, 0], [1, 1]], dtype=np.uint32) + offsets = [[0, 1], [1, 0]] + + affs = bic.affinities.compute_affinities(labels, offsets, return_mask=False) + assert isinstance(affs, np.ndarray) + assert affs.shape == (2, 2, 2) + assert affs.dtype == np.float32 + + +def test_negative_offsets_work(): + labels = np.array([[1, 2], [1, 2]], dtype=np.uint32) + offsets = [[0, -1], [-1, 0]] + affs, mask = bic.affinities.compute_affinities(labels, offsets) + ref_affs, ref_mask = _numpy_reference(labels, offsets) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_threading_does_not_change_output(): + rng = np.random.default_rng(42) + labels = rng.integers(0, 7, size=(8, 12)).astype(np.uint32) + offsets = [[0, 1], [1, 0], [1, 1], [2, 3], [-1, 2]] + + affs_single, mask_single = bic.affinities.compute_affinities( + labels, offsets, number_of_threads=1 + ) + affs_multi, mask_multi = bic.affinities.compute_affinities( + labels, offsets, number_of_threads=4 + ) + np.testing.assert_array_equal(affs_single, affs_multi) + np.testing.assert_array_equal(mask_single, mask_multi) + + +def test_non_contiguous_input_is_handled(): + labels = np.array([[0, 0, 1], [0, 1, 1]], dtype=np.uint32) + # Transpose to produce a non-contiguous view. + labels_T = labels.T + assert not labels_T.flags["C_CONTIGUOUS"] + affs, mask = bic.affinities.compute_affinities(labels_T, [[0, 1]]) + # Should give same result as feeding a contiguous copy. + ref_affs, ref_mask = bic.affinities.compute_affinities( + np.ascontiguousarray(labels_T), [[0, 1]] + ) + np.testing.assert_array_equal(affs, ref_affs) + np.testing.assert_array_equal(mask, ref_mask) + + +def test_rejects_unsupported_dtype(): + labels = np.zeros((4, 4), dtype=np.float32) + with pytest.raises(TypeError, match="dtypes"): + bic.affinities.compute_affinities(labels, [[0, 1]]) + + +def test_rejects_1d_input(): + labels = np.zeros(8, dtype=np.uint32) + with pytest.raises(ValueError, match="2D or 3D"): + bic.affinities.compute_affinities(labels, [[1]]) + + +def test_rejects_empty_offsets(): + labels = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="offsets"): + bic.affinities.compute_affinities(labels, []) + + +def test_rejects_offset_with_wrong_length(): + labels = np.zeros((4, 4), dtype=np.uint32) + with pytest.raises(ValueError, match="spatial ndim"): + bic.affinities.compute_affinities(labels, [[0, 1, 2]])