diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index f4e7403..9d63ffb 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -490,6 +490,54 @@ Notes: endpoint ids. - Non-contiguous labels are copied to contiguous memory before entering C++. +### Distributed Region Adjacency Graphs and Features + +For volumes too large to hold in memory, `nifty.distributed` builds a RAG and +its edge features block by block and merges the results. `bioimage-cpp` provides +the equivalent **low-level primitives** under `bic.graph.distributed`; the +orchestration nifty bundles (iterating blocks, sizing halos, and serializing the +per-block subgraphs/features to zarr/N5/HDF5) is intentionally left to the +caller, since I/O and block scheduling belong in Python. + +The primitives assume **globally consistent labels** (a segment has the same id +in every block — as after a stitched distributed watershed). A block owns the +pixel-pairs whose reference pixel lies in its inner (non-halo) box, so the +caller reads each block with a halo (≥1 on the forward faces for the region +graph / an edge map; ≥ `max |offset|` per side for affinities) and passes the +owned box as `own_begin` / `own_shape` (e.g. from `bic.utils.Blocking`'s +`get_block_with_halo(...).inner_block_local`). + +```python +d = bic.graph.distributed + +# per block (labels read with a halo): +edges = d.block_region_adjacency_edges(labels_block, own_begin, own_shape) +block_edges, block_stats = d.block_edge_map_stats(labels_block, edge_map_block, own_begin, own_shape) +block_edges, block_stats = d.block_affinity_stats(labels_block, aff_block, offsets, own_begin, own_shape) + +# merge the graph, then build the global graph: +global_edges = d.merge_edges([edges_block_0, edges_block_1, ...]) +graph = bic.graph.UndirectedGraph.from_unique_edges(number_of_nodes, global_edges) + +# fold per-block features onto the global edges, then finalize: +acc = d.empty_edge_stats(graph.number_of_edges) +for be, bs in per_block_stats: + acc = d.merge_block_edge_stats(graph, acc, be, bs) +features = d.finalize_edge_features(acc, compute_complex_features=True) +``` + +Notes: + +- Blocked results reproduce the whole-volume `region_adjacency_graph` / + `features.*_features` exactly for `size`, `min` and `max`, and to + floating-point tolerance for `mean` / `std` (the running sums depend on thread + count and merge order). +- **Median and percentiles are not distributable.** The distributed complex + output is the moment subset `[mean, std, min, max, size]` — the corresponding + columns of the in-core 12-column complex features. +- Making per-block-local ids globally consistent (label stitching) is a separate + step and is not part of these primitives. + ### Breadth-First Search Nifty has an internal `BreadthFirstSearch` template used during lifted-edge diff --git a/include/bioimage_cpp/detail/grid.hxx b/include/bioimage_cpp/detail/grid.hxx index d6afc76..dc677a6 100644 --- a/include/bioimage_cpp/detail/grid.hxx +++ b/include/bioimage_cpp/detail/grid.hxx @@ -66,4 +66,27 @@ inline bool is_valid_grid_edge( return valid_offset_target(node, offset, shape, strides, unused); } +// Given a signed offset `delta` along one axis and the axis `length`, return the +// half-open range of reference coordinates `[lo, hi)` for which `coord + delta` +// stays in `[0, length)`. Returns `lo >= hi` when the offset is larger than the +// axis (no valid reference coordinate). This is the per-axis primitive behind +// the offset-box sweeps in feature accumulation and the distributed block +// extraction — it depends only on grid geometry, not on any array data. +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; + } +} + } // namespace bioimage_cpp::detail diff --git a/include/bioimage_cpp/detail/label_cast.hxx b/include/bioimage_cpp/detail/label_cast.hxx new file mode 100644 index 0000000..aabe5ea --- /dev/null +++ b/include/bioimage_cpp/detail/label_cast.hxx @@ -0,0 +1,23 @@ +#pragma once + +#include +#include +#include + +namespace bioimage_cpp::detail { + +// Convert a label value to a graph node id (`std::uint64_t`), rejecting negative +// values for signed label dtypes. Shared by the region-adjacency-graph scan, +// edge-feature accumulation, and the distributed block-extraction primitives so +// they all treat labels identically. +template +std::uint64_t checked_label_to_node(const T value) { + if constexpr (std::is_signed_v) { + if (value < 0) { + throw std::invalid_argument("labels must not contain negative values"); + } + } + return static_cast(value); +} + +} // namespace bioimage_cpp::detail diff --git a/include/bioimage_cpp/graph/distributed/block_extraction.hxx b/include/bioimage_cpp/graph/distributed/block_extraction.hxx new file mode 100644 index 0000000..4270d99 --- /dev/null +++ b/include/bioimage_cpp/graph/distributed/block_extraction.hxx @@ -0,0 +1,534 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/detail/grid.hxx" +#include "bioimage_cpp/detail/label_cast.hxx" +#include "bioimage_cpp/detail/threading.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include +#include + +// Distributed region-adjacency-graph + edge-feature primitives. +// +// These functions operate on a single (haloed) label block and extract the +// edges / partial edge statistics that the block *owns*, so that per-block +// results can later be merged into a whole-volume result by `merge.hxx`. +// Orchestration (block iteration, halo sizing, I/O, hierarchical merging) lives +// in Python and is intentionally not implemented here. +// +// Ownership rule: a block owns the pixel-pairs whose *reference pixel* lies in +// the block's inner (non-halo) box `[own_begin, own_begin + own_shape)`; the +// neighbor pixel is read from the passed (outer, haloed) array. Because inner +// boxes tile the volume, every contributing pair is counted exactly once, so +// per-block partial statistics (count/sum/sum_of_squares add, min/max reduce) +// reconstruct the whole-volume statistics exactly. The reference pixel is the +// lower pixel of a forward nearest-neighbor edge (region graph / edge map) or +// the pixel an affinity value is stored at (affinities). +// +// Halo requirement (caller's responsibility — not detectable here): the outer +// array must extend at least one pixel past the owned box on the forward faces +// for nearest-neighbor edges, and at least `max |offset component|` past the +// owned box on the relevant faces for affinities. If the halo is too small, +// owned pairs whose neighbor falls outside the block are silently dropped. +namespace bioimage_cpp::graph::distributed { + +// Per-edge partial statistics accumulated over the pixels a block owns. Layout +// mirrors the serialized `(n_edges, 5)` array: count, sum, sum_of_squares, min, +// max. `count/min/max` are exact under any thread/block ordering; the +// floating-point `sum/sum_of_squares` (and hence mean/std after finalization) +// are reproducible only for a fixed thread count and merge order. +struct PartialStats { + double count = 0.0; + double sum = 0.0; + double sum_of_squares = 0.0; + double minimum = std::numeric_limits::infinity(); + double maximum = -std::numeric_limits::infinity(); + + void add(const double value) { + sum += value; + sum_of_squares += value * value; + minimum = std::min(minimum, value); + maximum = std::max(maximum, value); + count += 1.0; + } + + void merge(const PartialStats &other) { + if (other.count == 0.0) { + return; + } + sum += other.sum; + sum_of_squares += other.sum_of_squares; + minimum = std::min(minimum, other.minimum); + maximum = std::max(maximum, other.maximum); + count += other.count; + } +}; + +// A block's owned edges together with their aligned partial statistics. `edges` +// is sorted-unique with `u < v`; `stats` is row-major `(edges.size(), 5)` with +// columns [count, sum, sum_of_squares, min, max] and row `i` describing +// `edges[i]`. +struct BlockEdgeStats { + std::vector edges; + std::vector stats; +}; + +namespace detail_block { + +using bioimage_cpp::detail::checked_label_to_node; +using bioimage_cpp::detail::Edge; +using bioimage_cpp::detail::EdgeHash; +using bioimage_cpp::detail::edge_key; +using bioimage_cpp::detail::valid_axis_range; + +using EdgeSet = std::unordered_set; +using EdgeStatsMap = std::unordered_map; + +inline std::vector to_size_dims(const std::vector &shape) { + std::vector dims(shape.size()); + for (std::size_t axis = 0; axis < shape.size(); ++axis) { + dims[axis] = static_cast(shape[axis]); + } + return dims; +} + +// Forward nearest-neighbor unit offsets, one per axis (2D: (1,0),(0,1); +// 3D: (1,0,0),(0,1,0),(0,0,1)). These are the edges of the region graph. +inline std::vector> forward_nn_offsets(const std::size_t ndim) { + std::vector> offsets; + offsets.reserve(ndim); + for (std::size_t axis = 0; axis < ndim; ++axis) { + std::vector offset(ndim, 0); + offset[axis] = 1; + offsets.push_back(std::move(offset)); + } + return offsets; +} + +inline void validate_owned_box( + const std::vector &shape, + const std::vector &own_begin, + const std::vector &own_shape +) { + const auto ndim = shape.size(); + if (ndim != 2 && ndim != 3) { + throw std::invalid_argument( + "labels must be a 2D or 3D array, got ndim=" + std::to_string(ndim) + ); + } + if (own_begin.size() != ndim) { + throw std::invalid_argument("own_begin length must match labels ndim"); + } + if (own_shape.size() != ndim) { + throw std::invalid_argument("own_shape length must match labels ndim"); + } + for (std::size_t axis = 0; axis < ndim; ++axis) { + if (own_begin[axis] < 0) { + throw std::invalid_argument("own_begin values must be non-negative"); + } + if (own_shape[axis] <= 0) { + throw std::invalid_argument("own_shape values must be positive"); + } + if (own_begin[axis] + own_shape[axis] > shape[axis]) { + throw std::invalid_argument( + "owned box must lie within the block (own_begin + own_shape <= block shape)" + ); + } + } +} + +// Sweep reference nodes over the owned box on a 2D grid, calling +// `body(node, target)` with flat C-order indices into the outer array for each +// reference node whose `+offset` neighbor stays inside the outer array. Axis 0 +// is additionally restricted to the absolute slab `[slab_begin, slab_end)` +// (the caller's thread chunk, already inside the owned box). +template +void sweep_owned_box_2d( + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t outer_h, + const std::size_t outer_w, + const std::int64_t own_begin_y, + const std::int64_t own_begin_x, + const std::int64_t own_shape_y, + const std::int64_t own_shape_x, + const std::size_t slab_begin, + const std::size_t slab_end, + const Body &body +) { + std::size_t y_lo_v, y_hi_v, x_lo_v, x_hi_v; + valid_axis_range(dy, outer_h, y_lo_v, y_hi_v); + valid_axis_range(dx, outer_w, x_lo_v, x_hi_v); + + const auto y_lo = std::max({y_lo_v, static_cast(own_begin_y), slab_begin}); + const auto y_hi = std::min({y_hi_v, static_cast(own_begin_y + own_shape_y), slab_end}); + const auto x_lo = std::max(x_lo_v, static_cast(own_begin_x)); + const auto x_hi = std::min(x_hi_v, static_cast(own_begin_x + own_shape_x)); + if (y_lo >= y_hi || x_lo >= x_hi) { + return; + } + + const auto offset_stride = dy * static_cast(outer_w) + dx; + for (std::size_t y = y_lo; y < y_hi; ++y) { + const auto row_offset = y * outer_w; + 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_owned_box_2d`. +template +void sweep_owned_box_3d( + const std::ptrdiff_t dz, + const std::ptrdiff_t dy, + const std::ptrdiff_t dx, + const std::size_t outer_d, + const std::size_t outer_h, + const std::size_t outer_w, + const std::int64_t own_begin_z, + const std::int64_t own_begin_y, + const std::int64_t own_begin_x, + const std::int64_t own_shape_z, + const std::int64_t own_shape_y, + const std::int64_t own_shape_x, + const std::size_t slab_begin, + const std::size_t slab_end, + const Body &body +) { + std::size_t z_lo_v, z_hi_v, y_lo_v, y_hi_v, x_lo_v, x_hi_v; + valid_axis_range(dz, outer_d, z_lo_v, z_hi_v); + valid_axis_range(dy, outer_h, y_lo_v, y_hi_v); + valid_axis_range(dx, outer_w, x_lo_v, x_hi_v); + + const auto z_lo = std::max({z_lo_v, static_cast(own_begin_z), slab_begin}); + const auto z_hi = std::min({z_hi_v, static_cast(own_begin_z + own_shape_z), slab_end}); + const auto y_lo = std::max(y_lo_v, static_cast(own_begin_y)); + const auto y_hi = std::min(y_hi_v, static_cast(own_begin_y + own_shape_y)); + const auto x_lo = std::max(x_lo_v, static_cast(own_begin_x)); + const auto x_hi = std::min(x_hi_v, static_cast(own_begin_x + own_shape_x)); + if (z_lo >= z_hi || y_lo >= y_hi || x_lo >= x_hi) { + return; + } + + const auto slice_size = outer_h * outer_w; + const auto offset_stride = + dz * static_cast(slice_size) + + dy * static_cast(outer_w) + 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 * outer_w; + 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); + } + } + } +} + +// Dispatch the owned-box sweep for one offset over the correct grid rank. +template +void sweep_owned_box( + const std::vector &offset, + const std::vector &outer_dims, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t slab_begin, + const std::size_t slab_end, + const Body &body +) { + if (outer_dims.size() == 2) { + sweep_owned_box_2d( + offset[0], offset[1], outer_dims[0], outer_dims[1], + own_begin[0], own_begin[1], own_shape[0], own_shape[1], + slab_begin, slab_end, body + ); + } else { + sweep_owned_box_3d( + offset[0], offset[1], offset[2], outer_dims[0], outer_dims[1], outer_dims[2], + own_begin[0], own_begin[1], own_begin[2], + own_shape[0], own_shape[1], own_shape[2], + slab_begin, slab_end, body + ); + } +} + +// Collect the owned region-graph edges of one thread's slab into `out`. +template +void collect_edges_chunk( + const LabelT *labels, + const std::vector &outer_dims, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t slab_begin, + const std::size_t slab_end, + const std::vector> &offsets, + EdgeSet &out +) { + for (const auto &offset : offsets) { + sweep_owned_box( + offset, outer_dims, own_begin, own_shape, slab_begin, slab_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = checked_label_to_node(labels[node]); + const auto v = checked_label_to_node(labels[target]); + if (u != v) { + out.insert(edge_key(u, v)); + } + } + ); + } +} + +// Accumulate one thread's owned partial statistics into `out`. `value_fn` +// receives `(offset_index, node, target)` and returns the value to accumulate +// on the edge between the labels at `node` and `target`. +template +void accumulate_stats_chunk( + const LabelT *labels, + const std::vector &outer_dims, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t slab_begin, + const std::size_t slab_end, + const std::vector> &offsets, + const ValueFn &value_fn, + EdgeStatsMap &out +) { + for (std::size_t offset_index = 0; offset_index < offsets.size(); ++offset_index) { + sweep_owned_box( + offsets[offset_index], outer_dims, own_begin, own_shape, + slab_begin, slab_end, + [&](const std::uint64_t node, const std::uint64_t target) { + const auto u = checked_label_to_node(labels[node]); + const auto v = checked_label_to_node(labels[target]); + if (u != v) { + out[edge_key(u, v)].add(value_fn(offset_index, node, target)); + } + } + ); + } +} + +inline std::vector merge_edge_sets(std::vector &per_thread) { + EdgeSet merged; + for (const auto &edges : per_thread) { + merged.insert(edges.begin(), edges.end()); + } + std::vector sorted_edges(merged.begin(), merged.end()); + std::sort(sorted_edges.begin(), sorted_edges.end()); + return sorted_edges; +} + +inline BlockEdgeStats merge_stats_maps(std::vector &per_thread) { + EdgeStatsMap merged; + for (auto &thread_map : per_thread) { + for (const auto &[edge, stats] : thread_map) { + merged[edge].merge(stats); + } + } + + std::vector sorted_edges; + sorted_edges.reserve(merged.size()); + for (const auto &[edge, stats] : merged) { + sorted_edges.push_back(edge); + } + std::sort(sorted_edges.begin(), sorted_edges.end()); + + BlockEdgeStats result; + result.edges = std::move(sorted_edges); + result.stats.reserve(result.edges.size() * 5); + for (const auto &edge : result.edges) { + const auto &stats = merged[edge]; + result.stats.push_back(stats.count); + result.stats.push_back(stats.sum); + result.stats.push_back(stats.sum_of_squares); + result.stats.push_back(stats.minimum); + result.stats.push_back(stats.maximum); + } + return result; +} + +// Common driver: run `chunk_fn(thread_id, slab_begin, slab_end, container)` +// over the owned box's axis-0 extent split across threads. +template +std::vector run_owned_scan( + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t number_of_threads, + const ChunkFn &chunk_fn +) { + const auto work_items = static_cast(own_shape[0]); + const auto n_threads = bioimage_cpp::detail::normalize_thread_count( + number_of_threads, work_items + ); + const auto axis0_begin = static_cast(own_begin[0]); + + std::vector per_thread(n_threads); + bioimage_cpp::detail::parallel_for_chunks( + n_threads, + work_items, + [&](const std::size_t thread_id, const std::size_t begin, const std::size_t end) { + chunk_fn(axis0_begin + begin, axis0_begin + end, per_thread[thread_id]); + } + ); + return per_thread; +} + +inline void validate_affinities( + const std::vector &labels_shape, + const std::vector &affinities_shape, + const std::vector> &offsets +) { + const auto ndim = labels_shape.size(); + if (affinities_shape.size() != ndim + 1) { + throw std::invalid_argument("affinities must have shape (channels, *labels.shape)"); + } + if (static_cast(affinities_shape[0]) != offsets.size()) { + throw std::invalid_argument("offsets length must match affinities channel count"); + } + for (std::size_t axis = 0; axis < ndim; ++axis) { + if (affinities_shape[axis + 1] != labels_shape[axis]) { + throw std::invalid_argument("affinities spatial shape must match labels shape"); + } + } + for (const auto &offset : offsets) { + if (offset.size() != ndim) { + throw std::invalid_argument("each offset must have length matching labels ndim"); + } + } +} + +} // namespace detail_block + +// Extract the region-adjacency edges a block owns. Returns a sorted-unique edge +// list (`u < v`) using global label ids; concatenating these across blocks and +// passing them through `merge_edges` yields the whole-volume edge set. +template +std::vector block_region_adjacency_edges( + const ConstArrayView &labels_outer, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t number_of_threads +) { + detail_block::validate_owned_box(labels_outer.shape, own_begin, own_shape); + const auto outer_dims = detail_block::to_size_dims(labels_outer.shape); + const auto offsets = detail_block::forward_nn_offsets(outer_dims.size()); + const auto *labels = labels_outer.data; + + auto per_thread = detail_block::run_owned_scan( + own_begin, own_shape, number_of_threads, + [&](const std::size_t slab_begin, const std::size_t slab_end, + detail_block::EdgeSet &out) { + detail_block::collect_edges_chunk( + labels, outer_dims, own_begin, own_shape, slab_begin, slab_end, + offsets, out + ); + } + ); + return detail_block::merge_edge_sets(per_thread); +} + +// Accumulate the owned partial statistics of a scalar edge map. The value on an +// edge is the average of the two endpoint pixel values, `0.5 * (map[node] + +// map[neighbor])`, matching the in-core `accumulate_edge_map_features`. The +// returned edges match `block_region_adjacency_edges` for the same block. +template +BlockEdgeStats block_edge_map_stats( + const ConstArrayView &labels_outer, + const ConstArrayView &edge_map_outer, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t number_of_threads +) { + detail_block::validate_owned_box(labels_outer.shape, own_begin, own_shape); + if (edge_map_outer.shape != labels_outer.shape) { + throw std::invalid_argument("edge_map shape must match labels shape"); + } + const auto outer_dims = detail_block::to_size_dims(labels_outer.shape); + const auto offsets = detail_block::forward_nn_offsets(outer_dims.size()); + const auto *labels = labels_outer.data; + const auto *edge_map = edge_map_outer.data; + + const auto value_fn = [&](const std::size_t, const std::uint64_t node, + const std::uint64_t target) { + return 0.5 * (edge_map[node] + edge_map[target]); + }; + + auto per_thread = detail_block::run_owned_scan( + own_begin, own_shape, number_of_threads, + [&](const std::size_t slab_begin, const std::size_t slab_end, + detail_block::EdgeStatsMap &out) { + detail_block::accumulate_stats_chunk( + labels, outer_dims, own_begin, own_shape, slab_begin, slab_end, + offsets, value_fn, out + ); + } + ); + return detail_block::merge_stats_maps(per_thread); +} + +// Accumulate the owned partial statistics of affinity channels. `affinities` +// has shape `(len(offsets), *labels.shape)`; the value on an edge is the +// affinity stored at the reference node, `affinities[channel * n_nodes + node]`, +// matching the in-core `accumulate_affinity_features`. Values from all offsets +// are aggregated per `(u, v)`, so a nearest-neighbor edge also hit by a +// long-range offset receives that offset's value too. Long-range-only pairs are +// included here and dropped at merge time if absent from the global graph. +template +BlockEdgeStats block_affinity_stats( + const ConstArrayView &labels_outer, + const ConstArrayView &affinities_outer, + const std::vector> &offsets, + const std::vector &own_begin, + const std::vector &own_shape, + const std::size_t number_of_threads +) { + detail_block::validate_owned_box(labels_outer.shape, own_begin, own_shape); + detail_block::validate_affinities(labels_outer.shape, affinities_outer.shape, offsets); + const auto outer_dims = detail_block::to_size_dims(labels_outer.shape); + const auto *labels = labels_outer.data; + const auto *affinities = affinities_outer.data; + + std::uint64_t number_of_nodes = 1; + for (const auto dim : outer_dims) { + number_of_nodes *= static_cast(dim); + } + + const auto value_fn = [&](const std::size_t offset_index, const std::uint64_t node, + const std::uint64_t) { + const auto channel_offset = + static_cast(offset_index) * number_of_nodes; + return affinities[channel_offset + node]; + }; + + auto per_thread = detail_block::run_owned_scan( + own_begin, own_shape, number_of_threads, + [&](const std::size_t slab_begin, const std::size_t slab_end, + detail_block::EdgeStatsMap &out) { + detail_block::accumulate_stats_chunk( + labels, outer_dims, own_begin, own_shape, slab_begin, slab_end, + offsets, value_fn, out + ); + } + ); + return detail_block::merge_stats_maps(per_thread); +} + +} // namespace bioimage_cpp::graph::distributed diff --git a/include/bioimage_cpp/graph/distributed/merge.hxx b/include/bioimage_cpp/graph/distributed/merge.hxx new file mode 100644 index 0000000..040d430 --- /dev/null +++ b/include/bioimage_cpp/graph/distributed/merge.hxx @@ -0,0 +1,170 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include + +// Whole-volume merge primitives for the distributed region-adjacency-graph and +// edge-feature pipeline. These take the per-block artifacts produced by +// `block_extraction.hxx` (which the Python orchestration layer serializes and +// reads back) and combine them. Orchestration — deciding which blocks to merge, +// hierarchical scheduling, and I/O — lives in Python and is not implemented +// here. +namespace bioimage_cpp::graph::distributed { + +// Merge many blocks' edge arrays into the whole-volume edge set. `concatenated` +// is an `(N, 2)` array of `(u, v)` pairs (typically all blocks' edges stacked). +// Edges are canonicalized to `u < v`, self-edges (`u == v`) are dropped, and the +// result is deduplicated and sorted ascending by `(u, v)` — exactly the +// precondition of `UndirectedGraph::from_sorted_unique_edges`, so the merged +// output can be turned into the global graph directly. +inline std::vector merge_edges( + const ConstArrayView &concatenated +) { + if (concatenated.ndim() != 2 || concatenated.shape[1] != 2) { + throw std::invalid_argument("edges must have shape (n_edges, 2)"); + } + + const auto number_of_input_edges = static_cast(concatenated.shape[0]); + std::unordered_set merged; + merged.reserve(number_of_input_edges); + for (std::size_t index = 0; index < number_of_input_edges; ++index) { + const auto u = concatenated.data[2 * index]; + const auto v = concatenated.data[2 * index + 1]; + if (u == v) { + continue; + } + merged.insert(bioimage_cpp::detail::edge_key(u, v)); + } + + std::vector sorted_edges(merged.begin(), merged.end()); + std::sort(sorted_edges.begin(), sorted_edges.end()); + return sorted_edges; +} + +// Fold one block's partial edge statistics into a running whole-volume +// accumulator, returning the updated accumulator. `current_stats` and the return +// value are `(number_of_edges, 5)` rows aligned to `global_graph`'s edge ids; +// `block_edges`/`block_stats` are the `(n, 2)` / `(n, 5)` outputs of a block +// extraction. Each block edge is mapped to its global edge id via +// `global_graph.find_edge`; edges absent from the global graph (id `-1`) are +// skipped. `count/sum/sum_of_squares` add; `min/max` reduce, seeded from the +// first contribution to each edge (so `current_stats` may be zero-initialized). +// +// Block edge endpoints must be valid node ids of `global_graph` +// (`< number_of_nodes`); `find_edge` throws `std::out_of_range` otherwise. +inline std::vector merge_block_edge_stats( + const UndirectedGraph &global_graph, + const ConstArrayView ¤t_stats, + const ConstArrayView &block_edges, + const ConstArrayView &block_stats +) { + if (current_stats.ndim() != 2 || current_stats.shape[1] != 5) { + throw std::invalid_argument("current_stats must have shape (number_of_edges, 5)"); + } + if (block_edges.ndim() != 2 || block_edges.shape[1] != 2) { + throw std::invalid_argument("block_edges must have shape (n_edges, 2)"); + } + if (block_stats.ndim() != 2 || block_stats.shape[1] != 5) { + throw std::invalid_argument("block_stats must have shape (n_edges, 5)"); + } + if (block_stats.shape[0] != block_edges.shape[0]) { + throw std::invalid_argument( + "block_edges and block_stats must have the same number of rows" + ); + } + if (static_cast(current_stats.shape[0]) != global_graph.number_of_edges()) { + throw std::invalid_argument( + "current_stats rows must match global_graph number_of_edges" + ); + } + + const auto number_of_edges = static_cast(current_stats.shape[0]); + std::vector out( + current_stats.data, current_stats.data + number_of_edges * 5 + ); + + const auto number_of_block_edges = static_cast(block_edges.shape[0]); + for (std::size_t index = 0; index < number_of_block_edges; ++index) { + const auto u = block_edges.data[2 * index]; + const auto v = block_edges.data[2 * index + 1]; + const auto edge = global_graph.find_edge(u, v); + if (edge < 0) { + continue; + } + + double *const o = out.data() + static_cast(edge) * 5; + const double *const b = block_stats.data + index * 5; + if (o[0] == 0.0) { + o[3] = b[3]; + o[4] = b[4]; + } else { + o[3] = std::min(o[3], b[3]); + o[4] = std::max(o[4], b[4]); + } + o[0] += b[0]; + o[1] += b[1]; + o[2] += b[2]; + } + return out; +} + +// Turn accumulated partial statistics `(number_of_edges, 5)` into edge features. +// `compute_complex_features` selects the output width: +// false -> `(number_of_edges, 2)`: [mean, size] +// true -> `(number_of_edges, 5)`: [mean, std, min, max, size] +// Edges with zero count produce all-zero rows. These columns equal the in-core +// `accumulate_*_features` results; the complex output is the moment subset of +// the 12-column in-core complex features (median and percentiles are not +// recoverable from block partials). +inline void finalize_edge_features( + const ConstArrayView &stats, + const bool compute_complex_features, + const ArrayView &out +) { + if (stats.ndim() != 2 || stats.shape[1] != 5) { + throw std::invalid_argument("stats must have shape (number_of_edges, 5)"); + } + const auto number_of_edges = static_cast(stats.shape[0]); + const auto number_of_features = compute_complex_features ? 5 : 2; + if (out.ndim() != 2 || + out.shape[0] != stats.shape[0] || + out.shape[1] != number_of_features) { + throw std::invalid_argument("out shape must be (number_of_edges, number_of_features)"); + } + + for (std::size_t edge = 0; edge < number_of_edges; ++edge) { + const double *const s = stats.data + edge * 5; + double *const o = out.data + edge * static_cast(number_of_features); + const auto count = s[0]; + if (count == 0.0) { + for (int feature = 0; feature < number_of_features; ++feature) { + o[feature] = 0.0; + } + continue; + } + + const auto mean = s[1] / count; + if (compute_complex_features) { + const auto variance = std::max(0.0, s[2] / count - mean * mean); + o[0] = mean; + o[1] = std::sqrt(variance); + o[2] = s[3]; + o[3] = s[4]; + o[4] = count; + } else { + o[0] = mean; + o[1] = count; + } + } +} + +} // namespace bioimage_cpp::graph::distributed diff --git a/include/bioimage_cpp/graph/feature_accumulation.hxx b/include/bioimage_cpp/graph/feature_accumulation.hxx index fbdead4..1e13ef2 100644 --- a/include/bioimage_cpp/graph/feature_accumulation.hxx +++ b/include/bioimage_cpp/graph/feature_accumulation.hxx @@ -196,25 +196,10 @@ 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; - } -} +// `valid_axis_range` lives in `detail/grid.hxx` (it is pure grid geometry, +// shared with the distributed block-extraction sweeps). Pull it into this +// namespace so the offset-box sweeps below can call it unqualified. +using bioimage_cpp::detail::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 diff --git a/include/bioimage_cpp/graph/region_adjacency_graph.hxx b/include/bioimage_cpp/graph/region_adjacency_graph.hxx index 58e9f1c..f048208 100644 --- a/include/bioimage_cpp/graph/region_adjacency_graph.hxx +++ b/include/bioimage_cpp/graph/region_adjacency_graph.hxx @@ -2,6 +2,7 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/detail/label_cast.hxx" #include "bioimage_cpp/detail/threading.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" @@ -41,21 +42,12 @@ private: namespace detail { +using bioimage_cpp::detail::checked_label_to_node; using bioimage_cpp::detail::Edge; using bioimage_cpp::detail::EdgeHash; using bioimage_cpp::detail::edge_key; using bioimage_cpp::detail::normalize_thread_count; -template -std::uint64_t checked_label_to_node(const T value) { - if constexpr (std::is_signed_v) { - if (value < 0) { - throw std::invalid_argument("labels must not contain negative values"); - } - } - return static_cast(value); -} - template std::uint64_t max_label(const ConstArrayView &labels) { const auto number_of_pixels = static_cast(std::accumulate( diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 0500b5b..b47df57 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -23,6 +23,8 @@ #include "bioimage_cpp/graph/node_label_projection.hxx" #include "bioimage_cpp/graph/proposal_generator.hxx" #include "bioimage_cpp/graph/proposal_generators/greedy_additive_multicut.hxx" +#include "bioimage_cpp/graph/distributed/block_extraction.hxx" +#include "bioimage_cpp/graph/distributed/merge.hxx" #include "bioimage_cpp/graph/proposal_generators/watershed.hxx" #include "bioimage_cpp/graph/rag_coordinates.hxx" #include "bioimage_cpp/graph/region_adjacency_graph.hxx" @@ -198,6 +200,15 @@ UInt64Array edges_to_uv_array(const std::vector &edg return result; } +// Copy a flat per-edge partial-statistics buffer (row-major, 5 columns) into a +// NumPy `(n_edges, 5)` array. Used by the distributed block-extraction bindings. +DoubleArray block_stats_to_array(const std::vector &stats) { + const auto rows = stats.size() / 5; + auto result = make_double_array({rows, std::size_t{5}}); + std::copy(stats.begin(), stats.end(), result.data()); + return result; +} + template std::vector coordinate_to_vector( const typename graph::GridGraph::Coordinate &coordinate @@ -1423,6 +1434,134 @@ std::vector ndarray_shape(ConstDoubleArray array) { return shape; } +// ---- Distributed region-adjacency-graph + edge-feature primitives ---- + +template +UInt64Array block_region_adjacency_edges_t( + LabelArray labels, + std::vector own_begin, + std::vector own_shape, + const std::size_t number_of_threads +) { + ConstArrayView labels_view{labels.data(), ndarray_shape(labels), {}}; + + std::vector edges; + { + nb::gil_scoped_release release; + edges = graph::distributed::block_region_adjacency_edges( + labels_view, own_begin, own_shape, number_of_threads + ); + } + return edges_to_uv_array(edges); +} + +template +nb::tuple block_edge_map_stats_t( + LabelArray labels, + ConstDoubleArray edge_map, + std::vector own_begin, + std::vector own_shape, + const std::size_t number_of_threads +) { + ConstArrayView labels_view{labels.data(), ndarray_shape(labels), {}}; + ConstArrayView edge_map_view{edge_map.data(), ndarray_shape(edge_map), {}}; + + graph::distributed::BlockEdgeStats result; + { + nb::gil_scoped_release release; + result = graph::distributed::block_edge_map_stats( + labels_view, edge_map_view, own_begin, own_shape, number_of_threads + ); + } + return nb::make_tuple( + edges_to_uv_array(result.edges), block_stats_to_array(result.stats) + ); +} + +template +nb::tuple block_affinity_stats_t( + LabelArray labels, + ConstDoubleArray affinities, + std::vector> offsets, + std::vector own_begin, + std::vector own_shape, + const std::size_t number_of_threads +) { + ConstArrayView labels_view{labels.data(), ndarray_shape(labels), {}}; + ConstArrayView affinities_view{affinities.data(), ndarray_shape(affinities), {}}; + + graph::distributed::BlockEdgeStats result; + { + nb::gil_scoped_release release; + result = graph::distributed::block_affinity_stats( + labels_view, affinities_view, offsets, own_begin, own_shape, number_of_threads + ); + } + return nb::make_tuple( + edges_to_uv_array(result.edges), block_stats_to_array(result.stats) + ); +} + +UInt64Array distributed_merge_edges(ConstUInt64Array edges) { + ConstArrayView edges_view{edges.data(), ndarray_shape(edges), {}}; + + std::vector merged; + { + nb::gil_scoped_release release; + merged = graph::distributed::merge_edges(edges_view); + } + return edges_to_uv_array(merged); +} + +DoubleArray distributed_merge_block_edge_stats( + const Graph &global_graph, + ConstDoubleArray current_stats, + ConstUInt64Array block_edges, + ConstDoubleArray block_stats +) { + ConstArrayView current_view{current_stats.data(), ndarray_shape(current_stats), {}}; + ConstArrayView block_edges_view{block_edges.data(), ndarray_shape(block_edges), {}}; + ConstArrayView block_stats_view{block_stats.data(), ndarray_shape(block_stats), {}}; + + std::vector out; + { + nb::gil_scoped_release release; + out = graph::distributed::merge_block_edge_stats( + global_graph, current_view, block_edges_view, block_stats_view + ); + } + return block_stats_to_array(out); +} + +DoubleArray distributed_finalize_edge_features( + ConstDoubleArray stats, + const bool compute_complex_features +) { + if (stats.ndim() != 2 || stats.shape(1) != 5) { + throw std::invalid_argument("stats must have shape (number_of_edges, 5)"); + } + const auto rows = static_cast(stats.shape(0)); + const std::size_t number_of_features = compute_complex_features ? 5 : 2; + auto result = make_double_array({rows, number_of_features}); + + ConstArrayView stats_view{stats.data(), ndarray_shape(stats), {}}; + ArrayView out_view{ + result.data(), + { + static_cast(rows), + static_cast(number_of_features), + }, + {}, + }; + { + nb::gil_scoped_release release; + graph::distributed::finalize_edge_features( + stats_view, compute_complex_features, out_view + ); + } + return result; +} + template DoubleArray accumulate_edge_map_features_t( const Rag &rag, @@ -2448,6 +2587,99 @@ void bind_graph(nb::module_ &m) { nb::arg("compute_complex_features"), nb::arg("number_of_threads") ); + + // Distributed region-adjacency-graph + edge-feature primitives. + m.def( + "_block_region_adjacency_edges_uint32", + &block_region_adjacency_edges_t, + nb::arg("labels"), nb::arg("own_begin"), nb::arg("own_shape"), + nb::arg("number_of_threads"), + "Extract owned region-adjacency edges from a uint32 label block." + ); + m.def( + "_block_region_adjacency_edges_uint64", + &block_region_adjacency_edges_t, + nb::arg("labels"), nb::arg("own_begin"), nb::arg("own_shape"), + nb::arg("number_of_threads"), + "Extract owned region-adjacency edges from a uint64 label block." + ); + m.def( + "_block_region_adjacency_edges_int32", + &block_region_adjacency_edges_t, + nb::arg("labels"), nb::arg("own_begin"), nb::arg("own_shape"), + nb::arg("number_of_threads"), + "Extract owned region-adjacency edges from an int32 label block." + ); + m.def( + "_block_region_adjacency_edges_int64", + &block_region_adjacency_edges_t, + nb::arg("labels"), nb::arg("own_begin"), nb::arg("own_shape"), + nb::arg("number_of_threads"), + "Extract owned region-adjacency edges from an int64 label block." + ); + + m.def( + "_block_edge_map_stats_uint32", + &block_edge_map_stats_t, + nb::arg("labels"), nb::arg("edge_map"), nb::arg("own_begin"), + nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_edge_map_stats_uint64", + &block_edge_map_stats_t, + nb::arg("labels"), nb::arg("edge_map"), nb::arg("own_begin"), + nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_edge_map_stats_int32", + &block_edge_map_stats_t, + nb::arg("labels"), nb::arg("edge_map"), nb::arg("own_begin"), + nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_edge_map_stats_int64", + &block_edge_map_stats_t, + nb::arg("labels"), nb::arg("edge_map"), nb::arg("own_begin"), + nb::arg("own_shape"), nb::arg("number_of_threads") + ); + + m.def( + "_block_affinity_stats_uint32", + &block_affinity_stats_t, + nb::arg("labels"), nb::arg("affinities"), nb::arg("offsets"), + nb::arg("own_begin"), nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_affinity_stats_uint64", + &block_affinity_stats_t, + nb::arg("labels"), nb::arg("affinities"), nb::arg("offsets"), + nb::arg("own_begin"), nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_affinity_stats_int32", + &block_affinity_stats_t, + nb::arg("labels"), nb::arg("affinities"), nb::arg("offsets"), + nb::arg("own_begin"), nb::arg("own_shape"), nb::arg("number_of_threads") + ); + m.def( + "_block_affinity_stats_int64", + &block_affinity_stats_t, + nb::arg("labels"), nb::arg("affinities"), nb::arg("offsets"), + nb::arg("own_begin"), nb::arg("own_shape"), nb::arg("number_of_threads") + ); + + m.def("_merge_edges", &distributed_merge_edges, nb::arg("edges")); + m.def( + "_merge_block_edge_stats", + &distributed_merge_block_edge_stats, + nb::arg("global_graph"), nb::arg("current_stats"), + nb::arg("block_edges"), nb::arg("block_stats") + ); + m.def( + "_finalize_edge_features", + &distributed_finalize_edge_features, + nb::arg("stats"), nb::arg("compute_complex_features") + ); m.def( "_lifted_edges_from_affinities_uint32", &lifted_edges_from_affinities_t, diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 7a3ffad..3d01e55 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -19,6 +19,9 @@ (with and without semantic constraints). - :mod:`bioimage_cpp.graph.features` — edge-feature accumulation on RAGs and grid graphs. +- :mod:`bioimage_cpp.graph.distributed` — low-level per-block extraction and + whole-volume merge primitives for building RAGs and edge features on volumes + too large to hold in memory (block iteration and I/O are left to the caller). Thread safety ------------- @@ -663,6 +666,7 @@ def project_node_labels_to_pixels( from . import agglomeration # noqa: E402 (must follow class/function definitions) +from . import distributed # noqa: E402 from . import features # noqa: E402 from . import lifted_multicut # noqa: E402 from . import multicut # noqa: E402 @@ -678,6 +682,7 @@ def project_node_labels_to_pixels( "agglomeration", "breadth_first_search", "connected_components", + "distributed", "edge_weighted_watershed", "features", "grid_graph", diff --git a/src/bioimage_cpp/graph/distributed.py b/src/bioimage_cpp/graph/distributed.py new file mode 100644 index 0000000..583f1dd --- /dev/null +++ b/src/bioimage_cpp/graph/distributed.py @@ -0,0 +1,300 @@ +"""Low-level primitives for distributed region-adjacency graphs and features. + +These functions compute a region adjacency graph (RAG) and its edge features +one block at a time and merge the per-block results into a whole-volume result. +They are the building blocks for processing volumes too large to hold in memory; +the surrounding orchestration — iterating blocks, sizing halos, serializing and +reading back the per-block artifacts (e.g. to zarr/N5/HDF5), and scheduling the +merges — lives outside this module and is **not** provided here. + +The pipeline has three stages: + +1. **Per block** (with a halo): :func:`block_region_adjacency_edges` extracts the + edges the block owns; :func:`block_edge_map_stats` / :func:`block_affinity_stats` + extract the partial edge statistics it owns. A block owns the pixel-pairs + whose reference pixel lies in its inner (non-halo) box ``[own_begin, + own_begin + own_shape)``; the neighbor pixel is read from the passed (outer, + haloed) block. Because inner boxes tile the volume, every contribution is + counted exactly once. The caller must supply a halo large enough to reach the + neighbors (≥1 on the forward faces for nearest neighbors; ≥ ``max |offset|`` + per side for affinities) — too small a halo silently drops owned pairs. + +2. **Merge the graph**: :func:`merge_edges` unions the per-block edges into the + whole-volume edge set; build the global graph with + :meth:`bioimage_cpp.graph.UndirectedGraph.from_unique_edges`. + +3. **Merge the features**: fold each block's partial statistics onto the global + edges with :func:`merge_block_edge_stats` (starting from + :func:`empty_edge_stats`), then convert to features with + :func:`finalize_edge_features`. + +Exactly-recoverable features are ``size``, ``mean``, ``std``, ``min`` and +``max``. Median and percentiles cannot be reconstructed from block partials, so +the distributed complex output is the moment subset ``[mean, std, min, max, +size]`` — it equals the corresponding columns of the in-core complex features. +""" + +from __future__ import annotations + +import numpy as np + +from .. import _core +from ._shared import _normalize_labels, _normalize_number_of_threads + + +_BLOCK_REGION_ADJACENCY_EDGES_BY_DTYPE = { + np.dtype("uint32"): _core._block_region_adjacency_edges_uint32, + np.dtype("uint64"): _core._block_region_adjacency_edges_uint64, + np.dtype("int32"): _core._block_region_adjacency_edges_int32, + np.dtype("int64"): _core._block_region_adjacency_edges_int64, +} + +_BLOCK_EDGE_MAP_STATS_BY_DTYPE = { + np.dtype("uint32"): _core._block_edge_map_stats_uint32, + np.dtype("uint64"): _core._block_edge_map_stats_uint64, + np.dtype("int32"): _core._block_edge_map_stats_int32, + np.dtype("int64"): _core._block_edge_map_stats_int64, +} + +_BLOCK_AFFINITY_STATS_BY_DTYPE = { + np.dtype("uint32"): _core._block_affinity_stats_uint32, + np.dtype("uint64"): _core._block_affinity_stats_uint64, + np.dtype("int32"): _core._block_affinity_stats_int32, + np.dtype("int64"): _core._block_affinity_stats_int64, +} + + +def _as_index_vector(values, ndim: int, name: str) -> list[int]: + array = np.asarray(values) + if array.ndim != 1 or array.shape[0] != ndim: + raise ValueError(f"{name} must be a 1D sequence of length {ndim}") + return [int(value) for value in array] + + +def _as_stats_array(values, name: str) -> np.ndarray: + array = np.asarray(values, dtype=np.float64) + if array.ndim != 2 or array.shape[1] != 5: + raise ValueError(f"{name} must have shape (number_of_edges, 5)") + return np.ascontiguousarray(array) + + +def _dispatch_labels(labels, table, name: str): + label_array = _normalize_labels(labels) + try: + run = table[label_array.dtype] + except KeyError as error: + supported = ", ".join(str(dtype) for dtype in table) + raise TypeError( + f"{name} labels must have one of dtypes ({supported}), " + f"got dtype={label_array.dtype}" + ) from error + return label_array, run + + +def block_region_adjacency_edges( + labels: np.ndarray, + own_begin, + own_shape, + *, + number_of_threads: int = 0, +) -> np.ndarray: + """Extract the region-adjacency edges a block owns. + + ``labels`` is a (haloed) 2D or 3D label block with global label ids; + ``own_begin`` / ``own_shape`` delimit the owned inner box in block-local + coordinates. Returns an ``(n, 2)`` ``uint64`` array of ``(u, v)`` edges with + ``u < v``, sorted lexicographically. Concatenate these across blocks and + pass them to :func:`merge_edges` to obtain the whole-volume edge set. + """ + label_array, run = _dispatch_labels( + labels, _BLOCK_REGION_ADJACENCY_EDGES_BY_DTYPE, "region adjacency" + ) + own_begin = _as_index_vector(own_begin, label_array.ndim, "own_begin") + own_shape = _as_index_vector(own_shape, label_array.ndim, "own_shape") + return run( + label_array, own_begin, own_shape, _normalize_number_of_threads(number_of_threads) + ) + + +def block_edge_map_stats( + labels: np.ndarray, + edge_map: np.ndarray, + own_begin, + own_shape, + *, + number_of_threads: int = 0, +) -> tuple[np.ndarray, np.ndarray]: + """Accumulate the partial edge-map statistics a block owns. + + The value on an edge is the average of the two endpoint pixel values, + matching :func:`bioimage_cpp.graph.features.edge_map_features`. ``edge_map`` + must have the same shape as ``labels``. Returns ``(edges, stats)`` where + ``edges`` is ``(n, 2)`` ``uint64`` (matching + :func:`block_region_adjacency_edges` for the same block) and ``stats`` is + ``(n, 5)`` ``float64`` with columns ``[count, sum, sum_of_squares, min, + max]`` aligned row-by-row to ``edges``. + """ + label_array, run = _dispatch_labels( + labels, _BLOCK_EDGE_MAP_STATS_BY_DTYPE, "edge-map" + ) + edge_map_array = np.asarray(edge_map, dtype=np.float64) + if edge_map_array.shape != label_array.shape: + raise ValueError( + "edge_map shape must match labels shape, got " + f"edge_map shape={edge_map_array.shape}, labels shape={label_array.shape}" + ) + own_begin = _as_index_vector(own_begin, label_array.ndim, "own_begin") + own_shape = _as_index_vector(own_shape, label_array.ndim, "own_shape") + return run( + label_array, + np.ascontiguousarray(edge_map_array), + own_begin, + own_shape, + _normalize_number_of_threads(number_of_threads), + ) + + +def block_affinity_stats( + labels: np.ndarray, + affinities: np.ndarray, + offsets, + own_begin, + own_shape, + *, + number_of_threads: int = 0, +) -> tuple[np.ndarray, np.ndarray]: + """Accumulate the partial affinity statistics a block owns. + + ``affinities`` has shape ``(len(offsets), *labels.shape)``; the value on an + edge is the affinity stored at the reference node, matching + :func:`bioimage_cpp.graph.features.affinity_features`. Values from all + offsets are aggregated per ``(u, v)``. Returns ``(edges, stats)`` as in + :func:`block_edge_map_stats`; ``edges`` may include long-range-only pairs, + which are dropped at merge time if absent from the global graph. + """ + label_array, run = _dispatch_labels( + labels, _BLOCK_AFFINITY_STATS_BY_DTYPE, "affinity" + ) + affinity_array = np.asarray(affinities, dtype=np.float64) + if affinity_array.ndim != label_array.ndim + 1: + raise ValueError("affinities must have shape (channels, *labels.shape)") + if affinity_array.shape[1:] != label_array.shape: + raise ValueError( + "affinities spatial shape must match labels shape, got " + f"affinities shape={affinity_array.shape}, labels shape={label_array.shape}" + ) + + normalized_offsets = [tuple(int(value) for value in offset) for offset in offsets] + if len(normalized_offsets) != affinity_array.shape[0]: + raise ValueError( + "offsets length must match affinities channel count, got " + f"offsets length={len(normalized_offsets)}, channels={affinity_array.shape[0]}" + ) + if any(len(offset) != label_array.ndim for offset in normalized_offsets): + raise ValueError("each offset must have length matching labels ndim") + + own_begin = _as_index_vector(own_begin, label_array.ndim, "own_begin") + own_shape = _as_index_vector(own_shape, label_array.ndim, "own_shape") + return run( + label_array, + np.ascontiguousarray(affinity_array), + normalized_offsets, + own_begin, + own_shape, + _normalize_number_of_threads(number_of_threads), + ) + + +def merge_edges(edges) -> np.ndarray: + """Merge per-block edge arrays into the whole-volume edge set. + + ``edges`` is either a single ``(N, 2)`` array or a sequence of ``(n_i, 2)`` + arrays (which are concatenated). Edges are canonicalized to ``u < v``, + self-edges are dropped, and the result is deduplicated and sorted + lexicographically. The returned ``(E, 2)`` ``uint64`` array feeds + :meth:`bioimage_cpp.graph.UndirectedGraph.from_unique_edges` directly. + """ + if isinstance(edges, np.ndarray): + stacked = edges + else: + parts = [np.asarray(part, dtype=np.uint64) for part in edges] + if not parts: + return np.empty((0, 2), dtype=np.uint64) + stacked = np.concatenate(parts, axis=0) + + stacked = np.asarray(stacked, dtype=np.uint64) + if stacked.ndim != 2 or stacked.shape[1] != 2: + raise ValueError("edges must have shape (n_edges, 2)") + return _core._merge_edges(np.ascontiguousarray(stacked)) + + +def empty_edge_stats(number_of_edges: int) -> np.ndarray: + """Return a zero-initialized ``(number_of_edges, 5)`` ``float64`` accumulator. + + Use this as the starting ``current_stats`` for :func:`merge_block_edge_stats`. + """ + return np.zeros((int(number_of_edges), 5), dtype=np.float64) + + +def merge_block_edge_stats( + global_graph, + current_stats: np.ndarray, + block_edges: np.ndarray, + block_stats: np.ndarray, +) -> np.ndarray: + """Fold one block's partial statistics onto the global edges. + + ``current_stats`` is the running ``(E, 5)`` accumulator (rows aligned to + ``global_graph`` edge ids; start from :func:`empty_edge_stats`). + ``block_edges`` / ``block_stats`` are a block extraction's ``(n, 2)`` / + ``(n, 5)`` outputs. Each block edge is mapped to its global edge id via + ``global_graph.find_edge``; edges absent from the graph are skipped. + ``count/sum/sum_of_squares`` add and ``min/max`` reduce. Returns the updated + ``(E, 5)`` accumulator. + + Block edge endpoints must be valid node ids of ``global_graph``. + """ + current = _as_stats_array(current_stats, "current_stats") + if int(current.shape[0]) != int(global_graph.number_of_edges): + raise ValueError("current_stats rows must match global_graph number_of_edges") + + block_edge_array = np.asarray(block_edges, dtype=np.uint64) + if block_edge_array.ndim != 2 or block_edge_array.shape[1] != 2: + raise ValueError("block_edges must have shape (n_edges, 2)") + block_stat_array = _as_stats_array(block_stats, "block_stats") + if block_stat_array.shape[0] != block_edge_array.shape[0]: + raise ValueError("block_edges and block_stats must have the same number of rows") + + return _core._merge_block_edge_stats( + global_graph, + current, + np.ascontiguousarray(block_edge_array), + block_stat_array, + ) + + +def finalize_edge_features( + stats: np.ndarray, + *, + compute_complex_features: bool = False, +) -> np.ndarray: + """Convert accumulated partial statistics into edge features. + + ``stats`` is an ``(E, 5)`` accumulator (from :func:`merge_block_edge_stats`). + With ``compute_complex_features=False`` returns ``(E, 2)`` columns + ``[mean, size]``; with ``True`` returns ``(E, 5)`` columns + ``[mean, std, min, max, size]``. Edges with zero count give all-zero rows. + """ + stats_array = _as_stats_array(stats, "stats") + return _core._finalize_edge_features(stats_array, bool(compute_complex_features)) + + +__all__ = [ + "block_affinity_stats", + "block_edge_map_stats", + "block_region_adjacency_edges", + "empty_edge_stats", + "finalize_edge_features", + "merge_block_edge_stats", + "merge_edges", +] diff --git a/tests/graph/test_distributed_rag.py b/tests/graph/test_distributed_rag.py new file mode 100644 index 0000000..9ce2cb4 --- /dev/null +++ b/tests/graph/test_distributed_rag.py @@ -0,0 +1,379 @@ +"""Tests for the distributed region-adjacency-graph and edge-feature primitives. + +The core property is that computing per block (with a halo) and merging +reproduces the whole-volume result: exactly for the region graph and for the +``size/min/max`` feature columns, and to floating-point tolerance for +``mean/std``. Blocking uses :class:`bioimage_cpp.utils.Blocking` to tile the +array; a block owns the pixels of its inner (non-halo) box. +""" + +import numpy as np +import pytest + +import bioimage_cpp as bic + +dist = bic.graph.distributed + +LABEL_DTYPES = [np.uint32, np.uint64, np.int32, np.int64] + +# In-core complex feature columns are +# (mean, median, std, min, max, p5, p10, p25, p75, p90, p95, size); the +# distributed complex output is the moment subset (mean, std, min, max, size). +_COMPLEX_MOMENT_COLUMNS = [0, 2, 3, 4, 11] + + +def _tile_owned(labels, block_shape, halo): + """Yield ``(sub, slices, own_begin, own_shape)`` for every haloed block.""" + ndim = labels.ndim + blocking = bic.utils.Blocking([0] * ndim, list(labels.shape), list(block_shape)) + for block_id in range(blocking.number_of_blocks): + bwh = blocking.get_block_with_halo(block_id, list(halo)) + outer_begin, outer_end = bwh.outer_block.begin, bwh.outer_block.end + slices = tuple(slice(int(outer_begin[a]), int(outer_end[a])) for a in range(ndim)) + inner_begin, inner_end = bwh.inner_block_local.begin, bwh.inner_block_local.end + own_begin = [int(inner_begin[a]) for a in range(ndim)] + own_shape = [int(inner_end[a]) - int(inner_begin[a]) for a in range(ndim)] + yield np.ascontiguousarray(labels[slices]), slices, own_begin, own_shape + + +def _blocked_edges(labels, block_shape, halo, *, number_of_threads=0): + edges = [ + dist.block_region_adjacency_edges( + sub, own_begin, own_shape, number_of_threads=number_of_threads + ) + for (sub, _, own_begin, own_shape) in _tile_owned(labels, block_shape, halo) + ] + return dist.merge_edges(edges) + + +def _global_graph(rag, merged_edges): + return bic.graph.UndirectedGraph.from_unique_edges( + int(rag.number_of_nodes), merged_edges + ) + + +def _reduce_stats(global_graph, per_block): + acc = dist.empty_edge_stats(global_graph.number_of_edges) + for block_edges, block_stats in per_block: + acc = dist.merge_block_edge_stats(global_graph, acc, block_edges, block_stats) + return acc + + +# -------------------------------------------------------------------------- +# Region graph: block extraction + merge +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", LABEL_DTYPES) +def test_single_block_equals_whole_2d(dtype): + rng = np.random.default_rng(0) + labels = rng.integers(0, 6, size=(9, 11)).astype(dtype) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + single = dist.block_region_adjacency_edges(labels, [0, 0], list(labels.shape)) + np.testing.assert_array_equal(dist.merge_edges(single), whole) + + +def test_single_block_equals_whole_3d(): + rng = np.random.default_rng(1) + labels = rng.integers(0, 5, size=(5, 6, 7)).astype(np.uint64) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + single = dist.block_region_adjacency_edges(labels, [0, 0, 0], list(labels.shape)) + np.testing.assert_array_equal(dist.merge_edges(single), whole) + + +@pytest.mark.parametrize("dtype", LABEL_DTYPES) +def test_blocked_vs_whole_graph_2d(dtype): + rng = np.random.default_rng(2) + labels = rng.integers(0, 7, size=(10, 13)).astype(dtype) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + merged = _blocked_edges(labels, [4, 5], [1, 1]) + np.testing.assert_array_equal(merged, whole) + + +def test_blocked_vs_whole_graph_3d(): + rng = np.random.default_rng(3) + labels = rng.integers(0, 5, size=(7, 8, 9)).astype(np.uint32) + whole = bic.graph.region_adjacency_graph(labels).uv_ids() + merged = _blocked_edges(labels, [3, 4, 5], [1, 1, 1]) + np.testing.assert_array_equal(merged, whole) + + +def test_blocked_graph_thread_determinism(): + rng = np.random.default_rng(4) + labels = rng.integers(0, 8, size=(12, 15)).astype(np.uint32) + one = _blocked_edges(labels, [5, 6], [1, 1], number_of_threads=1) + many = _blocked_edges(labels, [5, 6], [1, 1], number_of_threads=4) + np.testing.assert_array_equal(one, many) + + +def test_block_edges_are_sorted_unique_and_ordered(): + # A single block's edges match the in-core RAG ordering (sorted, u < v). + rng = np.random.default_rng(5) + labels = rng.integers(0, 6, size=(8, 8)).astype(np.uint32) + edges = dist.block_region_adjacency_edges(labels, [0, 0], list(labels.shape)) + assert np.all(edges[:, 0] < edges[:, 1]) + # sorted lexicographically + order = np.lexsort((edges[:, 1], edges[:, 0])) + np.testing.assert_array_equal(edges, edges[order]) + + +# -------------------------------------------------------------------------- +# Edge-map features +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize("dtype", [np.uint32, np.int64]) +def test_edge_map_features_blocked_vs_whole_2d(dtype): + rng = np.random.default_rng(10) + labels = rng.integers(0, 6, size=(10, 13)).astype(dtype) + edge_map = rng.standard_normal((10, 13)).astype(np.float64) + + rag = bic.graph.region_adjacency_graph(labels) + merged = _blocked_edges(labels, [4, 5], [1, 1]) + np.testing.assert_array_equal(merged, rag.uv_ids()) + graph = _global_graph(rag, merged) + + per_block = [ + dist.block_edge_map_stats( + sub, np.ascontiguousarray(edge_map[slices]), own_begin, own_shape + ) + for (sub, slices, own_begin, own_shape) in _tile_owned(labels, [4, 5], [1, 1]) + ] + acc = _reduce_stats(graph, per_block) + + simple = dist.finalize_edge_features(acc, compute_complex_features=False) + complex_ = dist.finalize_edge_features(acc, compute_complex_features=True) + + whole_simple = bic.graph.features.edge_map_features(rag, labels, edge_map) + whole_complex = bic.graph.features.edge_map_features_complex(rag, labels, edge_map) + + np.testing.assert_array_equal(simple[:, 1], whole_simple[:, 1]) + np.testing.assert_allclose(simple[:, 0], whole_simple[:, 0], rtol=1e-6, atol=1e-9) + # complex moment subset + np.testing.assert_array_equal(complex_[:, 2], whole_complex[:, 3]) # min + np.testing.assert_array_equal(complex_[:, 3], whole_complex[:, 4]) # max + np.testing.assert_array_equal(complex_[:, 4], whole_complex[:, 11]) # size + np.testing.assert_allclose(complex_[:, 0], whole_complex[:, 0], rtol=1e-6, atol=1e-9) # mean + np.testing.assert_allclose(complex_[:, 1], whole_complex[:, 2], rtol=1e-6, atol=1e-9) # std + + +def test_edge_map_features_blocked_vs_whole_3d(): + rng = np.random.default_rng(11) + labels = rng.integers(0, 5, size=(6, 7, 8)).astype(np.uint64) + edge_map = rng.standard_normal((6, 7, 8)).astype(np.float64) + + rag = bic.graph.region_adjacency_graph(labels) + merged = _blocked_edges(labels, [3, 4, 4], [1, 1, 1]) + graph = _global_graph(rag, merged) + + per_block = [ + dist.block_edge_map_stats( + sub, np.ascontiguousarray(edge_map[slices]), own_begin, own_shape + ) + for (sub, slices, own_begin, own_shape) in _tile_owned(labels, [3, 4, 4], [1, 1, 1]) + ] + acc = _reduce_stats(graph, per_block) + complex_ = dist.finalize_edge_features(acc, compute_complex_features=True) + whole_complex = bic.graph.features.edge_map_features_complex(rag, labels, edge_map) + for dist_col, whole_col in zip(range(5), _COMPLEX_MOMENT_COLUMNS): + np.testing.assert_allclose( + complex_[:, dist_col], whole_complex[:, whole_col], rtol=1e-6, atol=1e-9 + ) + + +# -------------------------------------------------------------------------- +# Affinity features +# -------------------------------------------------------------------------- + + +@pytest.mark.parametrize( + "offsets, halo", + [ + ([[1, 0], [0, 1]], [1, 1]), + ([[1, 0], [0, 1], [3, 0], [0, 3]], [3, 3]), + ], +) +def test_affinity_features_blocked_vs_whole_2d(offsets, halo): + rng = np.random.default_rng(20) + labels = rng.integers(0, 6, size=(11, 12)).astype(np.uint32) + affinities = rng.standard_normal((len(offsets),) + labels.shape).astype(np.float64) + + rag = bic.graph.region_adjacency_graph(labels) + merged = _blocked_edges(labels, [5, 5], halo) + graph = _global_graph(rag, merged) + + per_block = [] + for sub, slices, own_begin, own_shape in _tile_owned(labels, [5, 5], halo): + sub_aff = np.ascontiguousarray(affinities[(slice(None),) + slices]) + per_block.append(dist.block_affinity_stats(sub, sub_aff, offsets, own_begin, own_shape)) + acc = _reduce_stats(graph, per_block) + + complex_ = dist.finalize_edge_features(acc, compute_complex_features=True) + whole_complex = bic.graph.features.affinity_features_complex(rag, labels, affinities, offsets) + for dist_col, whole_col in zip(range(5), _COMPLEX_MOMENT_COLUMNS): + np.testing.assert_allclose( + complex_[:, dist_col], whole_complex[:, whole_col], rtol=1e-6, atol=1e-9 + ) + + +def test_affinity_features_single_block_equals_whole_3d(): + rng = np.random.default_rng(21) + labels = rng.integers(0, 5, size=(5, 6, 7)).astype(np.uint64) + offsets = [[1, 0, 0], [0, 1, 0], [0, 0, 1], [0, 2, 0]] + affinities = rng.standard_normal((len(offsets),) + labels.shape).astype(np.float64) + + rag = bic.graph.region_adjacency_graph(labels) + graph = _global_graph(rag, rag.uv_ids()) + be, bs = dist.block_affinity_stats( + labels, affinities, offsets, [0, 0, 0], list(labels.shape) + ) + acc = dist.merge_block_edge_stats(graph, dist.empty_edge_stats(graph.number_of_edges), be, bs) + complex_ = dist.finalize_edge_features(acc, compute_complex_features=True) + whole_complex = bic.graph.features.affinity_features_complex(rag, labels, affinities, offsets) + for dist_col, whole_col in zip(range(5), _COMPLEX_MOMENT_COLUMNS): + np.testing.assert_allclose( + complex_[:, dist_col], whole_complex[:, whole_col], rtol=1e-6, atol=1e-9 + ) + + +# -------------------------------------------------------------------------- +# merge_edges +# -------------------------------------------------------------------------- + + +def test_merge_edges_dedup_canonicalize_and_self_loops(): + edges = np.array( + [[3, 1], [1, 3], [0, 2], [2, 0], [4, 4], [5, 6]], dtype=np.uint64 + ) + merged = dist.merge_edges(edges) + expected = np.array([[0, 2], [1, 3], [5, 6]], dtype=np.uint64) + np.testing.assert_array_equal(merged, expected) + + +def test_merge_edges_accepts_list_and_builds_graph(): + a = np.array([[0, 1], [1, 2]], dtype=np.uint64) + b = np.array([[1, 2], [2, 3]], dtype=np.uint64) + merged = dist.merge_edges([a, b]) + np.testing.assert_array_equal(merged, np.array([[0, 1], [1, 2], [2, 3]], dtype=np.uint64)) + graph = bic.graph.UndirectedGraph.from_unique_edges(4, merged) + assert graph.number_of_edges == 3 + assert graph.find_edge(2, 1) == 1 # find_edge canonicalizes + + +def test_merge_edges_empty(): + assert dist.merge_edges([]).shape == (0, 2) + empty = np.empty((0, 2), dtype=np.uint64) + assert dist.merge_edges(empty).shape == (0, 2) + + +# -------------------------------------------------------------------------- +# merge_block_edge_stats / finalize / empty_edge_stats +# -------------------------------------------------------------------------- + + +def test_empty_edge_stats(): + stats = dist.empty_edge_stats(4) + assert stats.shape == (4, 5) + assert stats.dtype == np.float64 + assert np.all(stats == 0.0) + + +def test_merge_block_edge_stats_reduction_and_skipping(): + graph = bic.graph.UndirectedGraph.from_unique_edges( + 3, np.array([[0, 1], [1, 2]], dtype=np.uint64) + ) + acc = dist.empty_edge_stats(graph.number_of_edges) + + # Block A touches edge (0,1) with values summarized as count=2,sum=3,sq=5,min=1,max=2. + edges_a = np.array([[0, 1]], dtype=np.uint64) + stats_a = np.array([[2.0, 3.0, 5.0, 1.0, 2.0]], dtype=np.float64) + acc = dist.merge_block_edge_stats(graph, acc, edges_a, stats_a) + + # Block B touches edge (0,1) again plus (0,2) which is absent -> skipped. + edges_b = np.array([[0, 1], [0, 2]], dtype=np.uint64) + stats_b = np.array( + [[1.0, 10.0, 100.0, 10.0, 10.0], [5.0, 5.0, 5.0, -1.0, 9.0]], dtype=np.float64 + ) + acc = dist.merge_block_edge_stats(graph, acc, edges_b, stats_b) + + # edge (0,1): count 2+1, sum 3+10, sq 5+100, min(1,10)=1, max(2,10)=10 + np.testing.assert_array_equal(acc[0], [3.0, 13.0, 105.0, 1.0, 10.0]) + # edge (1,2): never touched + np.testing.assert_array_equal(acc[1], [0.0, 0.0, 0.0, 0.0, 0.0]) + + +def test_finalize_edge_features_formulas_and_zero_count(): + # edge 0: count=4, sum=8 -> mean 2; sum_sq=24 -> var = 24/4 - 4 = 2 -> std sqrt(2) + # edge 1: empty + stats = np.array( + [[4.0, 8.0, 24.0, -1.0, 5.0], [0.0, 0.0, 0.0, 0.0, 0.0]], dtype=np.float64 + ) + simple = dist.finalize_edge_features(stats, compute_complex_features=False) + np.testing.assert_allclose(simple[0], [2.0, 4.0]) + np.testing.assert_array_equal(simple[1], [0.0, 0.0]) + + complex_ = dist.finalize_edge_features(stats, compute_complex_features=True) + np.testing.assert_allclose(complex_[0], [2.0, np.sqrt(2.0), -1.0, 5.0, 4.0]) + np.testing.assert_array_equal(complex_[1], [0.0, 0.0, 0.0, 0.0, 0.0]) + + +# -------------------------------------------------------------------------- +# Validation +# -------------------------------------------------------------------------- + + +def test_owned_box_out_of_bounds_raises(): + labels = np.zeros((6, 6), dtype=np.uint32) + with pytest.raises(ValueError, match="owned box must lie within the block"): + dist.block_region_adjacency_edges(labels, [0, 0], [7, 6]) + + +def test_owned_box_nonpositive_shape_raises(): + labels = np.zeros((6, 6), dtype=np.uint32) + with pytest.raises(ValueError, match="own_shape values must be positive"): + dist.block_region_adjacency_edges(labels, [0, 0], [0, 6]) + + +def test_owned_box_wrong_length_raises(): + labels = np.zeros((6, 6), dtype=np.uint32) + with pytest.raises(ValueError, match="own_begin must be a 1D sequence of length 2"): + dist.block_region_adjacency_edges(labels, [0, 0, 0], [3, 3]) + + +def test_edge_map_shape_mismatch_raises(): + labels = np.zeros((6, 6), dtype=np.uint32) + edge_map = np.zeros((6, 5), dtype=np.float64) + with pytest.raises(ValueError, match="edge_map shape must match labels shape"): + dist.block_edge_map_stats(labels, edge_map, [0, 0], [6, 6]) + + +def test_affinity_offsets_length_mismatch_raises(): + labels = np.zeros((6, 6), dtype=np.uint32) + affinities = np.zeros((2, 6, 6), dtype=np.float64) + with pytest.raises(ValueError, match="offsets length must match affinities channel count"): + dist.block_affinity_stats(labels, affinities, [[1, 0]], [0, 0], [6, 6]) + + +def test_merge_block_edge_stats_row_mismatch_raises(): + graph = bic.graph.UndirectedGraph.from_unique_edges( + 2, np.array([[0, 1]], dtype=np.uint64) + ) + acc = dist.empty_edge_stats(1) + with pytest.raises(ValueError, match="same number of rows"): + dist.merge_block_edge_stats( + graph, + acc, + np.array([[0, 1]], dtype=np.uint64), + np.zeros((2, 5), dtype=np.float64), + ) + + +def test_finalize_wrong_stats_shape_raises(): + with pytest.raises(ValueError, match=r"stats must have shape \(number_of_edges, 5\)"): + dist.finalize_edge_features(np.zeros((3, 4), dtype=np.float64)) + + +def test_unsupported_label_dtype_raises(): + labels = np.zeros((5, 5), dtype=np.uint8) + with pytest.raises(TypeError): + dist.block_region_adjacency_edges(labels, [0, 0], [5, 5])