Affinities (#19)

* Optimize lifted affinity features

* More optimization

* Add affinity computation

Author: constantinpape

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

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()

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 <algorithm>
+#include <array>
+#include <cstddef>
+#include <cstdint>
+#include <optional>
+#include <vector>
+
+// 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 <class LabelT, class AffT>
+void compute_affinities_2d(
+    const ConstArrayView<LabelT> &labels,
+    const std::vector<std::array<std::ptrdiff_t, 2>> &offsets,
+    const ArrayView<AffT> &affs,
+    const ArrayView<std::uint8_t> *mask,
+    const std::optional<LabelT> 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<std::ptrdiff_t>(oi) * plane;
+                std::uint8_t * const mask_channel = (mask != nullptr)
+                    ? mask->data + static_cast<std::ptrdiff_t>(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<std::ptrdiff_t>(0, -dy);
+                const auto y_end = height - std::max<std::ptrdiff_t>(0, dy);
+                const auto x_begin = std::max<std::ptrdiff_t>(0, -dx);
+                const auto x_end = width - std::max<std::ptrdiff_t>(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 <class LabelT, class AffT>
+void compute_affinities_3d(
+    const ConstArrayView<LabelT> &labels,
+    const std::vector<std::array<std::ptrdiff_t, 3>> &offsets,
+    const ArrayView<AffT> &affs,
+    const ArrayView<std::uint8_t> *mask,
+    const std::optional<LabelT> 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<std::ptrdiff_t>(oi) * volume;
+                std::uint8_t * const mask_channel = (mask != nullptr)
+                    ? mask->data + static_cast<std::ptrdiff_t>(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<std::ptrdiff_t>(0, -dz);
+                const auto z_end = depth - std::max<std::ptrdiff_t>(0, dz);
+                const auto y_begin = std::max<std::ptrdiff_t>(0, -dy);
+                const auto y_end = height - std::max<std::ptrdiff_t>(0, dy);
+                const auto x_begin = std::max<std::ptrdiff_t>(0, -dx);
+                const auto x_end = width - std::max<std::ptrdiff_t>(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