From 25fa6e8e30c52e48276bd295e30aa00d2358bad7 Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Sat, 16 May 2026 21:48:42 +0200 Subject: [PATCH 1/2] Optimize lifted affinity features --- development/graph/_rag_compatibility.py | 35 ++- .../graph/feature_accumulation.hxx | 200 +++++++++++++++--- tests/graph/test_rag_features.py | 83 ++++++++ 3 files changed, 278 insertions(+), 40 deletions(-) diff --git a/development/graph/_rag_compatibility.py b/development/graph/_rag_compatibility.py index aedc231..d8c6561 100644 --- a/development/graph/_rag_compatibility.py +++ b/development/graph/_rag_compatibility.py @@ -85,6 +85,10 @@ def make_watershed_labels( def time_call(function, repeats: int): + # One untimed warm-up call before the measured loop so the first sample + # doesn't carry nanobind tuple-shape caching, numpy ufunc init, code-page + # faults, etc. Mirrors the grid-affinity helper for consistency. + function() timings = [] result = None for _ in range(repeats): @@ -136,7 +140,9 @@ def compare_boundary_features( import bioimage_cpp as bic import nifty.graph.rag as nrag - block_shape = list(labels.shape) + # Don't force blockShape on the nifty side — its default block layout is + # what its parallelism is tuned for, and forcing a single block (== labels + # shape) starves nifty's worker pool. bic_timings, bic_features = time_call( lambda: bic.graph.edge_map_features( bic_rag, labels, boundary_map, number_of_threads=threads @@ -147,7 +153,6 @@ def compare_boundary_features( lambda: nrag.accumulateEdgeMeanAndLength( nifty_rag, np.ascontiguousarray(boundary_map, dtype=np.float32), - blockShape=block_shape, numberOfThreads=threads, ), repeats, @@ -193,9 +198,12 @@ def compare_affinity_features( ), repeats, ) - # The installed Nifty wrapper can crash or fail when an explicit - # numberOfThreads is passed here. The default path is still the reference - # implementation and keeps this compatibility script robust. + # nifty's accumulateAffinityStandartFeatures crashes inside vigra's + # UserRangeHistogram when numberOfThreads=1 (the accumulators never get + # setMinMax). For threads >= 2 we can pass the value through; for the + # single-thread case we fall back to nifty's default (-1, all cores) and + # surface the unfairness in the printed report. + nifty_threads_effective = threads if threads != 1 else -1 nifty_timings, nifty_features_full = time_call( lambda: nrag.accumulateAffinityStandartFeatures( nifty_rag, @@ -203,6 +211,7 @@ def compare_affinity_features( offsets_for_nifty, min_val, max_val, + numberOfThreads=nifty_threads_effective, ), repeats, ) @@ -213,6 +222,8 @@ def compare_affinity_features( np.testing.assert_allclose(aligned_bic, nifty_features, rtol=1.0e-5, atol=1.0e-6) return bic_timings, nifty_timings, { "max_abs_diff": float(np.max(np.abs(aligned_bic - nifty_features))), + "nifty_threads_effective": nifty_threads_effective, + "nifty_threads_requested": threads, } @@ -290,7 +301,15 @@ def run_compatibility_check( print(f"boundary-map size convention: {boundary_summary['size_convention']}") _print_timing("affinity features", affinity_bic_timings, affinity_nifty_timings) print(f"affinity feature max abs diff: {affinity_summary['max_abs_diff']:.6g}") - print("nifty affinity timing uses the wrapper default thread handling") + requested = affinity_summary["nifty_threads_requested"] + effective = affinity_summary["nifty_threads_effective"] + if requested != effective: + print( + f"WARNING: affinity features — requested {requested} thread(s) but " + f"nifty was called with numberOfThreads={effective} " + "(its single-thread path crashes inside vigra's UserRangeHistogram). " + f"bioimage-cpp used {requested} thread(s); the timing comparison is NOT apples-to-apples." + ) def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float]): @@ -303,7 +322,9 @@ def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float def add_common_arguments(parser: argparse.ArgumentParser) -> None: - parser.add_argument("--repeats", type=int, default=3) + # 5 matches the grid-affinity helper; median of 3 is noisy because one + # GC stall in the middle sample becomes the median. + parser.add_argument("--repeats", type=int, default=5) parser.add_argument("--threads", type=int, default=1) parser.add_argument("--watershed-min-distance", type=int, default=5) parser.add_argument("--watershed-grid-spacing", type=int, default=12) diff --git a/include/bioimage_cpp/graph/feature_accumulation.hxx b/include/bioimage_cpp/graph/feature_accumulation.hxx index edf4297..e0403c6 100644 --- a/include/bioimage_cpp/graph/feature_accumulation.hxx +++ b/include/bioimage_cpp/graph/feature_accumulation.hxx @@ -2,6 +2,7 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/profile.hxx" #include "bioimage_cpp/detail/threading.hxx" #include "bioimage_cpp/graph/region_adjacency_graph.hxx" @@ -204,29 +205,123 @@ void scan_edge_map_3d_chunk( } } +// Given an offset along one axis and the axis length, return the half-open +// range of axis coordinates `[lo, hi)` for which `coord + delta` stays in +// `[0, length)`. Returns `lo >= hi` if the offset is larger than the axis. +inline void valid_axis_range( + const std::ptrdiff_t delta, + const std::size_t length, + std::size_t &lo, + std::size_t &hi +) { + if (delta >= 0) { + lo = 0; + const auto d = static_cast(delta); + hi = (d >= length) ? 0 : (length - d); + } else { + const auto d = static_cast(-delta); + lo = (d >= length) ? length : d; + hi = length; + } +} + template -void scan_affinity_chunk( +void scan_affinity_2d_chunk( const RegionAdjacencyGraph &rag, const LabelT *labels, const ValueT *affinities, const std::vector> &offsets, - const std::vector &shape, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::vector &stats ) { - const auto spatial_strides = bioimage_cpp::detail::c_order_strides(shape); - const auto number_of_nodes = static_cast(number_of_pixels(shape)); + const auto number_of_nodes = static_cast(height * width); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { - const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offsets[channel], shape, spatial_strides, target)) { - continue; + const auto &off = offsets[channel]; + std::size_t y_lo_full, y_hi_full, x_lo, x_hi; + valid_axis_range(off[0], height, y_lo_full, y_hi_full); + valid_axis_range(off[1], width, x_lo, x_hi); + const auto y_lo = std::max(y_lo_full, y_begin); + const auto y_hi = std::min(y_hi_full, y_end); + if (y_lo >= y_hi || x_lo >= x_hi) { + continue; + } + const auto offset_stride = + off[0] * static_cast(width) + off[1]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] + ); + } } - const auto edge = edge_for_labels(rag, label_at(labels, node), label_at(labels, target)); - if (edge >= 0) { - stats[static_cast(edge)].add(affinities[channel_offset + node]); + } + } +} + +template +void scan_affinity_3d_chunk( + const RegionAdjacencyGraph &rag, + const LabelT *labels, + const ValueT *affinities, + const std::vector> &offsets, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::vector &stats +) { + const auto slice_size = height * width; + const auto number_of_nodes = static_cast(depth * slice_size); + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + const auto &off = offsets[channel]; + std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; + valid_axis_range(off[0], depth, z_lo_full, z_hi_full); + valid_axis_range(off[1], height, y_lo, y_hi); + valid_axis_range(off[2], width, x_lo, x_hi); + const auto z_lo = std::max(z_lo_full, z_begin); + const auto z_hi = std::min(z_hi_full, z_end); + if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { + continue; + } + const auto offset_stride = + off[0] * static_cast(slice_size) + + off[1] * static_cast(width) + + off[2]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + for (std::size_t z = z_lo; z < z_hi; ++z) { + const auto slice_offset = z * slice_size; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = slice_offset + y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] + ); + } + } } } } @@ -408,40 +503,79 @@ void accumulate_affinity_features( throw std::invalid_argument("out shape must be (number_of_edges, number_of_features)"); } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); - const auto n_threads = detail::normalize_thread_count(number_of_threads, number_of_nodes); + const auto work_items = static_cast(labels.shape[0]); + const auto n_threads = detail::normalize_thread_count(number_of_threads, work_items); const auto number_of_edges = static_cast(rag.number_of_edges()); + BIOIMAGE_PROFILE_INIT(aff_profiler); + const auto run_scan = [&](auto &per_thread_stats) { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:scan"); bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_features::scan_affinity_chunk( - rag, labels.data, affinities.data, offsets, labels.shape, - begin, end, per_thread_stats[thread_id] - ); + if (labels.ndim() == 2) { + detail_features::scan_affinity_2d_chunk( + rag, labels.data, affinities.data, offsets, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread_stats[thread_id] + ); + } else { + detail_features::scan_affinity_3d_chunk( + rag, labels.data, affinities.data, offsets, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread_stats[thread_id] + ); + } } ); }; if (compute_complex_features) { - std::vector> per_thread_stats( - n_threads, - std::vector(number_of_edges) - ); + std::vector> per_thread_stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc"); + per_thread_stats.assign( + n_threads, + std::vector(number_of_edges) + ); + } run_scan(per_thread_stats); - auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges); - detail_features::write_complex_features(stats, out); + std::vector stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge"); + stats = detail_features::merge_stats(per_thread_stats, number_of_edges); + } + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write"); + detail_features::write_complex_features(stats, out); + } } else { - std::vector> per_thread_stats( - n_threads, - std::vector(number_of_edges) - ); + std::vector> per_thread_stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:alloc"); + per_thread_stats.assign( + n_threads, + std::vector(number_of_edges) + ); + } run_scan(per_thread_stats); - auto stats = detail_features::merge_stats(per_thread_stats, number_of_edges); - detail_features::write_simple_features(stats, out); + std::vector stats; + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:merge"); + stats = detail_features::merge_stats(per_thread_stats, number_of_edges); + } + { + BIOIMAGE_PROFILE_SCOPE(aff_profiler, "aff:write"); + detail_features::write_simple_features(stats, out); + } } + + BIOIMAGE_PROFILE_REPORT(aff_profiler); } } // namespace bioimage_cpp::graph diff --git a/tests/graph/test_rag_features.py b/tests/graph/test_rag_features.py index 263d32d..8f8c42c 100644 --- a/tests/graph/test_rag_features.py +++ b/tests/graph/test_rag_features.py @@ -89,6 +89,89 @@ def test_affinity_features_simple(): ) +def _reference_affinity_features(labels, rag, affinities, offsets): + """Trivial Python reference: per-channel sweep with explicit bounds checks.""" + uv_ids = np.asarray(rag.uv_ids(), dtype=np.uint64) + edge_lookup = {(int(u), int(v)): i for i, (u, v) in enumerate(uv_ids)} + out = np.zeros((len(edge_lookup), 2), dtype=np.float64) + shape = labels.shape + for channel, offset in enumerate(offsets): + for index in np.ndindex(*shape): + target = tuple(int(c) + int(d) for c, d in zip(index, offset)) + if any(t < 0 or t >= s for t, s in zip(target, shape)): + continue + u, v = int(labels[index]), int(labels[target]) + if u == v: + continue + key = (min(u, v), max(u, v)) + edge = edge_lookup.get(key) + if edge is None: + # Pair (u,v) doesn't exist as a RAG edge (e.g. when the offset + # reaches further than direct neighbors). The kernel skips + # these via find_edge -> -1; the reference must match. + continue + out[edge, 0] += float(affinities[(channel,) + index]) + out[edge, 1] += 1.0 + for edge in range(len(edge_lookup)): + if out[edge, 1] > 0: + out[edge, 0] /= out[edge, 1] + return out + + +@pytest.mark.parametrize( + "offsets", + [ + # axis-aligned positive offsets (the only case the previous tests covered) + [[0, 1], [1, 0]], + # negative offsets — must still hit the valid box correctly + [[0, -1], [-1, 0]], + # large magnitude that crops most of the image + [[0, 3], [3, 0]], + # diagonal / multi-axis nonzero — exercises the new valid-box math + [[1, 1], [1, -1], [-1, 1]], + # offset larger than an axis — must produce zero contributions + [[10, 0], [0, 10]], + ], +) +def test_affinity_features_2d_offsets_match_reference(offsets): + rng = np.random.default_rng(42) + labels = rng.integers(0, 5, size=(5, 7)).astype(np.uint32) + rag = bic.graph.region_adjacency_graph(labels) + affinities = rng.standard_normal( + (len(offsets),) + labels.shape + ).astype(np.float32) + + expected = _reference_affinity_features(labels, rag, affinities, offsets) + got = bic.graph.affinity_features(rag, labels, affinities, offsets=offsets) + + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + +@pytest.mark.parametrize( + "offsets", + [ + [[1, 0, 0], [0, 1, 0], [0, 0, 1]], + [[-1, 0, 0], [0, -1, 0], [0, 0, -1]], + # 3D diagonal + [[1, 1, 1], [1, -1, 1], [-1, 1, -1]], + # mixed magnitudes + [[2, 0, 0], [0, 0, -3]], + ], +) +def test_affinity_features_3d_offsets_match_reference(offsets): + rng = np.random.default_rng(7) + labels = rng.integers(0, 4, size=(4, 5, 6)).astype(np.uint32) + rag = bic.graph.region_adjacency_graph(labels) + affinities = rng.standard_normal( + (len(offsets),) + labels.shape + ).astype(np.float32) + + expected = _reference_affinity_features(labels, rag, affinities, offsets) + got = bic.graph.affinity_features(rag, labels, affinities, offsets=offsets) + + np.testing.assert_allclose(got, expected, rtol=1e-5, atol=1e-6) + + def test_affinity_features_complex_parallel_matches_single_thread(): labels = np.array([[0, 1, 1], [2, 2, 3], [2, 4, 3]], dtype=np.uint64) rag = bic.graph.region_adjacency_graph(labels) From 78eb1d59b276b6b5999df4fb4ab3bca1a1caebdf Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Sat, 16 May 2026 21:59:28 +0200 Subject: [PATCH 2/2] More optimization --- .../graph/feature_accumulation.hxx | 139 ++++++++---- .../graph/lifted_from_affinities.hxx | 212 +++++++++++++----- 2 files changed, 246 insertions(+), 105 deletions(-) diff --git a/include/bioimage_cpp/graph/feature_accumulation.hxx b/include/bioimage_cpp/graph/feature_accumulation.hxx index e0403c6..4b2db1f 100644 --- a/include/bioimage_cpp/graph/feature_accumulation.hxx +++ b/include/bioimage_cpp/graph/feature_accumulation.hxx @@ -225,6 +225,83 @@ inline void valid_axis_range( } } +// Sweep every (node, target) pair on a 2D grid for which `node + offset` stays +// in bounds, restricted to the half-open y-slab [y_begin, y_end). The body +// receives flat C-order indices for both endpoints and is expected to inline +// at -O2 since this is a header-only template with a fully-known callable +// type at instantiation. +template +void sweep_offset_box_2d( + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, + const Body &body +) { + std::size_t y_lo_full, y_hi_full, x_lo, x_hi; + valid_axis_range(dy, height, y_lo_full, y_hi_full); + valid_axis_range(dx, width, x_lo, x_hi); + const auto y_lo = std::max(y_lo_full, y_begin); + const auto y_hi = std::min(y_hi_full, y_end); + if (y_lo >= y_hi || x_lo >= x_hi) { + return; + } + const auto offset_stride = dy * static_cast(width) + dx; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + body(static_cast(node), target); + } + } +} + +// 3D variant of `sweep_offset_box_2d`. Restricts the sweep to a z-slab. +template +void sweep_offset_box_3d( + const std::ptrdiff_t dz, + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + const Body &body +) { + std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; + valid_axis_range(dz, depth, z_lo_full, z_hi_full); + valid_axis_range(dy, height, y_lo, y_hi); + valid_axis_range(dx, width, x_lo, x_hi); + const auto z_lo = std::max(z_lo_full, z_begin); + const auto z_hi = std::min(z_hi_full, z_end); + if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { + return; + } + const auto slice_size = height * width; + const auto offset_stride = + dz * static_cast(slice_size) + + dy * static_cast(width) + dx; + for (std::size_t z = z_lo; z < z_hi; ++z) { + const auto slice_offset = z * slice_size; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = slice_offset + y * width; + for (std::size_t x = x_lo; x < x_hi; ++x) { + const auto node = row_offset + x; + const auto target = static_cast( + static_cast(node) + offset_stride + ); + body(static_cast(node), target); + } + } + } +} + template void scan_affinity_2d_chunk( const RegionAdjacencyGraph &rag, @@ -240,25 +317,11 @@ void scan_affinity_2d_chunk( const auto number_of_nodes = static_cast(height * width); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { const auto &off = offsets[channel]; - std::size_t y_lo_full, y_hi_full, x_lo, x_hi; - valid_axis_range(off[0], height, y_lo_full, y_hi_full); - valid_axis_range(off[1], width, x_lo, x_hi); - const auto y_lo = std::max(y_lo_full, y_begin); - const auto y_hi = std::min(y_hi_full, y_end); - if (y_lo >= y_hi || x_lo >= x_hi) { - continue; - } - const auto offset_stride = - off[0] * static_cast(width) + off[1]; const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::size_t y = y_lo; y < y_hi; ++y) { - const auto row_offset = y * width; - for (std::size_t x = x_lo; x < x_hi; ++x) { - const auto node = row_offset + x; - const auto target = static_cast( - static_cast(node) + offset_stride - ); + sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { const auto u = label_at(labels, node); const auto v = label_at(labels, target); const auto edge = edge_for_labels(rag, u, v); @@ -268,7 +331,7 @@ void scan_affinity_2d_chunk( ); } } - } + ); } } @@ -289,41 +352,21 @@ void scan_affinity_3d_chunk( const auto number_of_nodes = static_cast(depth * slice_size); for (std::size_t channel = 0; channel < offsets.size(); ++channel) { const auto &off = offsets[channel]; - std::size_t z_lo_full, z_hi_full, y_lo, y_hi, x_lo, x_hi; - valid_axis_range(off[0], depth, z_lo_full, z_hi_full); - valid_axis_range(off[1], height, y_lo, y_hi); - valid_axis_range(off[2], width, x_lo, x_hi); - const auto z_lo = std::max(z_lo_full, z_begin); - const auto z_hi = std::min(z_hi_full, z_end); - if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { - continue; - } - const auto offset_stride = - off[0] * static_cast(slice_size) + - off[1] * static_cast(width) + - off[2]; const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::size_t z = z_lo; z < z_hi; ++z) { - const auto slice_offset = z * slice_size; - for (std::size_t y = y_lo; y < y_hi; ++y) { - const auto row_offset = slice_offset + y * width; - for (std::size_t x = x_lo; x < x_hi; ++x) { - const auto node = row_offset + x; - const auto target = static_cast( - static_cast(node) + offset_stride + sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = label_at(labels, node); + const auto v = label_at(labels, target); + const auto edge = edge_for_labels(rag, u, v); + if (edge >= 0) { + stats[static_cast(edge)].add( + affinities[channel_offset + node] ); - const auto u = label_at(labels, node); - const auto v = label_at(labels, target); - const auto edge = edge_for_labels(rag, u, v); - if (edge >= 0) { - stats[static_cast(edge)].add( - affinities[channel_offset + node] - ); - } } } - } + ); } } diff --git a/include/bioimage_cpp/graph/lifted_from_affinities.hxx b/include/bioimage_cpp/graph/lifted_from_affinities.hxx index 4fed26c..8f968c4 100644 --- a/include/bioimage_cpp/graph/lifted_from_affinities.hxx +++ b/include/bioimage_cpp/graph/lifted_from_affinities.hxx @@ -65,34 +65,65 @@ inline void validate_affinity_inputs( } template -void discover_lifted_chunk( +void discover_lifted_2d_chunk( const RegionAdjacencyGraph &rag, const LabelT *labels, const std::vector> &offsets, const std::vector &long_range_channels, - const std::vector &shape, - const std::vector &strides, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::unordered_set &out ) { for (const auto channel : long_range_channels) { - const auto &offset = offsets[channel]; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { - continue; + const auto &off = offsets[channel]; + detail_features::sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + if (rag.find_edge(u, v) >= 0) { + return; + } + out.insert(bioimage_cpp::detail::edge_key(u, v)); } - const auto u = detail_features::label_at(labels, node); - const auto v = detail_features::label_at(labels, target); - if (u == v) { - continue; - } - if (rag.find_edge(u, v) >= 0) { - continue; + ); + } +} + +template +void discover_lifted_3d_chunk( + const RegionAdjacencyGraph &rag, + const LabelT *labels, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::unordered_set &out +) { + for (const auto channel : long_range_channels) { + const auto &off = offsets[channel]; + detail_features::sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + if (rag.find_edge(u, v) >= 0) { + return; + } + out.insert(bioimage_cpp::detail::edge_key(u, v)); } - out.insert(bioimage_cpp::detail::edge_key(u, v)); - } + ); } } @@ -133,11 +164,10 @@ std::vector lifted_edges_from_offsets( return {}; } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto work_items = static_cast(labels.shape[0]); const auto n_threads = bioimage_cpp::detail::normalize_thread_count( - number_of_threads, number_of_nodes + number_of_threads, work_items ); - const auto strides = bioimage_cpp::detail::c_order_strides(labels.shape); using EdgeSet = std::unordered_set< bioimage_cpp::detail::Edge, bioimage_cpp::detail::EdgeHash @@ -146,12 +176,24 @@ std::vector lifted_edges_from_offsets( bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_lifted::discover_lifted_chunk( - rag, labels.data, offsets, long_range_channels, - labels.shape, strides, begin, end, per_thread[thread_id] - ); + if (labels.ndim() == 2) { + detail_lifted::discover_lifted_2d_chunk( + rag, labels.data, offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread[thread_id] + ); + } else { + detail_lifted::discover_lifted_3d_chunk( + rag, labels.data, offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread[thread_id] + ); + } } ); @@ -173,39 +215,82 @@ std::vector lifted_edges_from_offsets( namespace detail_lifted { template -void scan_lifted_affinity_chunk( +void scan_lifted_affinity_2d_chunk( const std::unordered_map &lifted_index, const LabelT *labels, const ValueT *affinities, const std::vector> &offsets, const std::vector &long_range_channels, - const std::vector &shape, - const std::size_t node_begin, - const std::size_t node_end, + const std::size_t height, + const std::size_t width, + const std::size_t y_begin, + const std::size_t y_end, std::vector &stats ) { - const auto strides = bioimage_cpp::detail::c_order_strides(shape); - const auto number_of_nodes = static_cast(detail_features::number_of_pixels(shape)); + const auto number_of_nodes = static_cast(height * width); for (const auto channel : long_range_channels) { - const auto &offset = offsets[channel]; - const auto channel_offset = static_cast(channel) * number_of_nodes; - for (std::uint64_t node = node_begin; node < node_end; ++node) { - std::uint64_t target = 0; - if (!bioimage_cpp::detail::valid_offset_target(node, offset, shape, strides, target)) { - continue; - } - const auto u = detail_features::label_at(labels, node); - const auto v = detail_features::label_at(labels, target); - if (u == v) { - continue; + const auto &off = offsets[channel]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + detail_features::sweep_offset_box_2d( + off[0], off[1], height, width, y_begin, y_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + const auto found = lifted_index.find( + bioimage_cpp::detail::edge_key(u, v) + ); + if (found == lifted_index.end()) { + return; + } + stats[found->second].add(affinities[channel_offset + node]); } - const auto found = lifted_index.find(bioimage_cpp::detail::edge_key(u, v)); - if (found == lifted_index.end()) { - continue; + ); + } +} + +template +void scan_lifted_affinity_3d_chunk( + const std::unordered_map + &lifted_index, + const LabelT *labels, + const ValueT *affinities, + const std::vector> &offsets, + const std::vector &long_range_channels, + const std::size_t depth, + const std::size_t height, + const std::size_t width, + const std::size_t z_begin, + const std::size_t z_end, + std::vector &stats +) { + const auto slice_size = height * width; + const auto number_of_nodes = static_cast(depth * slice_size); + for (const auto channel : long_range_channels) { + const auto &off = offsets[channel]; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + detail_features::sweep_offset_box_3d( + off[0], off[1], off[2], depth, height, width, z_begin, z_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = detail_features::label_at(labels, node); + const auto v = detail_features::label_at(labels, target); + if (u == v) { + return; + } + const auto found = lifted_index.find( + bioimage_cpp::detail::edge_key(u, v) + ); + if (found == lifted_index.end()) { + return; + } + stats[found->second].add(affinities[channel_offset + node]); } - stats[found->second].add(affinities[channel_offset + node]); - } + ); } } @@ -297,21 +382,34 @@ void accumulate_lifted_affinity_features( return; } - const auto number_of_nodes = detail_features::number_of_pixels(labels.shape); + const auto work_items = static_cast(labels.shape[0]); const auto n_threads = bioimage_cpp::detail::normalize_thread_count( - number_of_threads, number_of_nodes + number_of_threads, work_items ); const auto run_scan = [&](auto &per_thread_stats) { bioimage_cpp::detail::parallel_for_chunks( n_threads, - number_of_nodes, + work_items, [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { - detail_lifted::scan_lifted_affinity_chunk( - lifted_index, labels.data, affinities.data, - offsets, long_range_channels, labels.shape, - begin, end, per_thread_stats[thread_id] - ); + if (labels.ndim() == 2) { + detail_lifted::scan_lifted_affinity_2d_chunk( + lifted_index, labels.data, affinities.data, + offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + begin, end, per_thread_stats[thread_id] + ); + } else { + detail_lifted::scan_lifted_affinity_3d_chunk( + lifted_index, labels.data, affinities.data, + offsets, long_range_channels, + static_cast(labels.shape[0]), + static_cast(labels.shape[1]), + static_cast(labels.shape[2]), + begin, end, per_thread_stats[thread_id] + ); + } } ); };