From 349408e7cb2dd111d3d34b192c422bc251e85c3e Mon Sep 17 00:00:00 2001 From: Constantin Pape Date: Mon, 18 May 2026 07:47:33 +0200 Subject: [PATCH] Expose UFD to python --- CMakeLists.txt | 1 + MIGRATION_GUIDE.md | 62 +++++++- .../graph/connected_components.hxx | 6 +- .../graph/detail/fusion_contract.hxx | 4 +- .../graph/edge_weighted_watershed.hxx | 4 +- .../graph/lifted_multicut/detail.hxx | 6 +- .../graph/lifted_multicut/greedy_additive.hxx | 2 +- .../graph/lifted_multicut/kernighan_lin.hxx | 2 +- .../bioimage_cpp/graph/multicut/detail.hxx | 6 +- .../graph/multicut/greedy_additive.hxx | 2 +- .../graph/multicut/greedy_fixation.hxx | 2 +- .../graph/multicut/kernighan_lin.hxx | 4 +- .../graph/mutex_watershed/clustering.hxx | 4 +- .../graph/mutex_watershed/semantic.hxx | 4 +- .../segmentation/mutex_watershed.hxx | 4 +- .../segmentation/semantic_mutex_watershed.hxx | 4 +- .../{detail => util}/union_find.hxx | 4 +- src/bindings/module.cxx | 2 + src/bindings/util.cxx | 146 ++++++++++++++++++ src/bindings/util.hxx | 9 ++ src/bioimage_cpp/__init__.py | 2 + src/bioimage_cpp/util/__init__.py | 5 + tests/test_util_union_find.py | 138 +++++++++++++++++ 23 files changed, 393 insertions(+), 30 deletions(-) rename include/bioimage_cpp/{detail => util}/union_find.hxx (97%) create mode 100644 src/bindings/util.cxx create mode 100644 src/bindings/util.hxx create mode 100644 src/bioimage_cpp/util/__init__.py create mode 100644 tests/test_util_union_find.py diff --git a/CMakeLists.txt b/CMakeLists.txt index b68202e..9d606b6 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -18,6 +18,7 @@ nanobind_add_module(_core src/bindings/graph.cxx src/bindings/ground_truth.cxx src/bindings/segmentation.cxx + src/bindings/util.cxx src/bindings/utils.cxx src/cpp/segmentation/mutex_watershed.cxx src/cpp/segmentation/semantic_mutex_watershed.cxx diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index e41ddda..fc86b27 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -17,7 +17,8 @@ import bioimage_cpp as bic Graph functionality is under `bic.graph`, segmentation functionality is under `bic.segmentation`, ground-truth comparison functionality is under -`bic.ground_truth`, and small utility functions are under `bic.utils`. +`bic.ground_truth`, small utility functions are under `bic.utils`, and +reusable data structures (e.g. union-find) are under `bic.util`. ## Blocking @@ -1094,6 +1095,65 @@ partition metrics and a semantic-label match fraction so the deviation is measurable. The bioimage-cpp partitions match an independent Python reference implementation of the algorithm. +## Union-Find + +Nifty exposes a disjoint-set / union-find structure as `nifty.ufd.ufd`. +`bioimage-cpp` provides the same primitive under `bic.util`. + +Nifty: + +```python +import nifty.ufd as nufd +import numpy as np + +uf = nufd.ufd(5) +uf.merge(0, 1) +uf.merge(np.array([[2, 3], [3, 4]], dtype="uint64")) +roots = uf.find(np.array([0, 1, 2, 3, 4], dtype="uint64")) +labels = uf.elementLabeling() +``` + +bioimage-cpp: + +```python +import bioimage_cpp as bic +import numpy as np + +uf = bic.util.UnionFind(5) +uf.merge(0, 1) +uf.merge(np.array([[2, 3], [3, 4]], dtype=np.uint64)) +roots = uf.find(np.array([0, 1, 2, 3, 4], dtype=np.uint64)) +labels = uf.element_labeling() +``` + +Method mapping: + +| nifty-style name | bioimage-cpp name | +| --- | --- | +| `find(node)` / `find(array)` | `find(node)` / `find(array)` | +| `merge(u, v)` / `merge(array)` | `merge(u, v)` / `merge(array)` | +| `elementLabeling` | `element_labeling` | +| `numberOfElements` | `size` (property) | + +Notes: + +- The constructor takes a single `size` argument; all elements start as + singletons. +- Scalar `find`/`merge` accept Python integers and return Python integers. +- Bulk `find(nodes)` accepts a 1D `uint64` array and returns a 1D `uint64` + array of roots of the same length. +- Bulk `merge(edges)` accepts an `(N, 2)` `uint64` array of node-pair edges + and applies the merges in row order. +- `element_labeling()` returns a `uint64` array of length `size`, each entry + the (path-compressed) root of that element. Use this when you want the + final labeling as one array rather than via repeated `find` calls. +- `merge_to(stable, removed)` is also available: it forces `stable`'s root + to survive the union regardless of rank. +- `reset(n)` reinitialises the structure to `n` singletons, reusing + capacity where possible. +- The GIL is released around bulk operations, so multiple threads can run + bulk merges on independent `UnionFind` instances in parallel. + ## Dictionary-Based Relabeling If you used a small helper to apply a dictionary to an integer label array, use diff --git a/include/bioimage_cpp/graph/connected_components.hxx b/include/bioimage_cpp/graph/connected_components.hxx index 7f93fc5..7451221 100644 --- a/include/bioimage_cpp/graph/connected_components.hxx +++ b/include/bioimage_cpp/graph/connected_components.hxx @@ -1,6 +1,6 @@ #pragma once -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include @@ -11,7 +11,7 @@ namespace bioimage_cpp::graph { inline std::vector dense_labels_from_union_find( - detail::UnionFind &sets, + bioimage_cpp::util::UnionFind &sets, const std::uint64_t number_of_nodes ) { std::unordered_map relabeling; @@ -31,7 +31,7 @@ inline std::vector connected_components( const UndirectedGraph &graph, const std::uint8_t *edge_mask = nullptr ) { - detail::UnionFind sets(static_cast(graph.number_of_nodes())); + bioimage_cpp::util::UnionFind sets(static_cast(graph.number_of_nodes())); for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { if (edge_mask != nullptr && edge_mask[static_cast(edge)] == 0) { continue; diff --git a/include/bioimage_cpp/graph/detail/fusion_contract.hxx b/include/bioimage_cpp/graph/detail/fusion_contract.hxx index c880aa6..baaee22 100644 --- a/include/bioimage_cpp/graph/detail/fusion_contract.hxx +++ b/include/bioimage_cpp/graph/detail/fusion_contract.hxx @@ -1,6 +1,6 @@ #pragma once -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include @@ -62,7 +62,7 @@ inline AgreementContraction contract_by_agreement( ); } - bioimage_cpp::detail::UnionFind sets(static_cast(number_of_nodes)); + bioimage_cpp::util::UnionFind sets(static_cast(number_of_nodes)); for (std::uint64_t edge = 0; edge < number_of_edges; ++edge) { const auto uv = graph.uv(edge); const auto u = static_cast(uv.first); diff --git a/include/bioimage_cpp/graph/edge_weighted_watershed.hxx b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx index 253a5aa..f21d9da 100644 --- a/include/bioimage_cpp/graph/edge_weighted_watershed.hxx +++ b/include/bioimage_cpp/graph/edge_weighted_watershed.hxx @@ -1,6 +1,6 @@ #pragma once -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include @@ -49,7 +49,7 @@ inline void edge_weighted_watershed_kernel( [](const auto &a, const auto &b) { return a.first < b.first; } ); - bioimage_cpp::detail::UnionFind sets(static_cast(number_of_nodes)); + bioimage_cpp::util::UnionFind sets(static_cast(number_of_nodes)); constexpr SeedT zero{0}; for (const auto &item : scratch.sort_buffer) { diff --git a/include/bioimage_cpp/graph/lifted_multicut/detail.hxx b/include/bioimage_cpp/graph/lifted_multicut/detail.hxx index 051bc41..d269aac 100644 --- a/include/bioimage_cpp/graph/lifted_multicut/detail.hxx +++ b/include/bioimage_cpp/graph/lifted_multicut/detail.hxx @@ -1,7 +1,7 @@ #pragma once #include "bioimage_cpp/detail/indexed_heap.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/connected_components.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" @@ -161,7 +161,7 @@ inline void rename_neighbor( // pops base edges off the heap). inline std::size_t merge_dynamic_nodes( DynamicGraph &dynamic_graph, - bioimage_cpp::detail::UnionFind &sets, + bioimage_cpp::util::UnionFind &sets, EdgeHeap &heap, std::size_t u, std::size_t v @@ -258,7 +258,7 @@ inline std::size_t merge_dynamic_nodes( } inline std::vector labels_from_sets( - bioimage_cpp::detail::UnionFind &sets, + bioimage_cpp::util::UnionFind &sets, const UndirectedGraph &graph ) { return dense_labels_from_union_find(sets, graph.number_of_nodes()); diff --git a/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx index d28e089..f694f1f 100644 --- a/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx +++ b/include/bioimage_cpp/graph/lifted_multicut/greedy_additive.hxx @@ -13,7 +13,7 @@ namespace bioimage_cpp::graph::lifted_multicut { // Reusable scratch state for `lifted_greedy_additive`. struct GreedyAdditiveWorkspace { detail::DynamicGraph dynamic_graph; - bioimage_cpp::detail::UnionFind union_find{0}; + bioimage_cpp::util::UnionFind union_find{0}; detail::EdgeHeap heap; void reset(const UndirectedGraph &lifted_graph) { diff --git a/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx b/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx index 952fed2..0a1b685 100644 --- a/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx +++ b/include/bioimage_cpp/graph/lifted_multicut/kernighan_lin.hxx @@ -3,7 +3,7 @@ #include "bioimage_cpp/detail/edge_hash.hxx" #include "bioimage_cpp/detail/indexed_heap.hxx" #include "bioimage_cpp/detail/profile.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/connected_components.hxx" #include "bioimage_cpp/graph/lifted_multicut/objective.hxx" diff --git a/include/bioimage_cpp/graph/multicut/detail.hxx b/include/bioimage_cpp/graph/multicut/detail.hxx index 7d32d66..5ed3bb9 100644 --- a/include/bioimage_cpp/graph/multicut/detail.hxx +++ b/include/bioimage_cpp/graph/multicut/detail.hxx @@ -1,7 +1,7 @@ #pragma once #include "bioimage_cpp/detail/indexed_heap.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/connected_components.hxx" #include "bioimage_cpp/graph/multicut/objective.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" @@ -150,7 +150,7 @@ inline void initialize_dynamic_graph( } inline std::vector labels_from_sets( - bioimage_cpp::detail::UnionFind &sets, + bioimage_cpp::util::UnionFind &sets, const UndirectedGraph &graph ) { return dense_labels_from_union_find(sets, graph.number_of_nodes()); @@ -199,7 +199,7 @@ inline void rename_neighbor( // edge id appears at most once. inline std::size_t merge_dynamic_nodes( DynamicGraph &dynamic_graph, - bioimage_cpp::detail::UnionFind &sets, + bioimage_cpp::util::UnionFind &sets, EdgeHeap &heap, std::size_t u, std::size_t v, diff --git a/include/bioimage_cpp/graph/multicut/greedy_additive.hxx b/include/bioimage_cpp/graph/multicut/greedy_additive.hxx index ccf50ec..cafc06d 100644 --- a/include/bioimage_cpp/graph/multicut/greedy_additive.hxx +++ b/include/bioimage_cpp/graph/multicut/greedy_additive.hxx @@ -15,7 +15,7 @@ namespace bioimage_cpp::graph::multicut { // internal vectors are reset (not freed) between calls. struct GreedyAdditiveWorkspace { detail::DynamicGraph dynamic_graph; - bioimage_cpp::detail::UnionFind union_find{0}; + bioimage_cpp::util::UnionFind union_find{0}; detail::EdgeHeap heap; void reset(const UndirectedGraph &graph) { diff --git a/include/bioimage_cpp/graph/multicut/greedy_fixation.hxx b/include/bioimage_cpp/graph/multicut/greedy_fixation.hxx index 8795276..948d41a 100644 --- a/include/bioimage_cpp/graph/multicut/greedy_fixation.hxx +++ b/include/bioimage_cpp/graph/multicut/greedy_fixation.hxx @@ -17,7 +17,7 @@ inline std::vector greedy_fixation( ) { validate_costs(graph, costs); detail::DynamicGraph dynamic_graph(graph); - bioimage_cpp::detail::UnionFind sets(static_cast(graph.number_of_nodes())); + bioimage_cpp::util::UnionFind sets(static_cast(graph.number_of_nodes())); detail::EdgeHeap heap; detail::initialize_dynamic_graph(graph, costs, dynamic_graph, heap, true); diff --git a/include/bioimage_cpp/graph/multicut/kernighan_lin.hxx b/include/bioimage_cpp/graph/multicut/kernighan_lin.hxx index d1b62b7..4b8cc9f 100644 --- a/include/bioimage_cpp/graph/multicut/kernighan_lin.hxx +++ b/include/bioimage_cpp/graph/multicut/kernighan_lin.hxx @@ -2,7 +2,7 @@ #include "bioimage_cpp/detail/edge_hash.hxx" #include "bioimage_cpp/detail/indexed_heap.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/multicut/objective.hxx" #include @@ -343,7 +343,7 @@ inline bool apply_joins( if (number_of_clusters == 0) { return false; } - bioimage_cpp::detail::UnionFind sets(static_cast(number_of_clusters)); + bioimage_cpp::util::UnionFind sets(static_cast(number_of_clusters)); bool any_join = false; for (const auto &pair : pairs) { if (pair.cut_weight <= epsilon) { diff --git a/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx b/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx index bf929be..e0fbe0b 100644 --- a/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx +++ b/include/bioimage_cpp/graph/mutex_watershed/clustering.hxx @@ -2,7 +2,7 @@ #include "bioimage_cpp/detail/mutex_storage.hxx" #include "bioimage_cpp/detail/relabel.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include @@ -97,7 +97,7 @@ std::vector mutex_watershed_clustering( return first.index < second.index; }); - bioimage_cpp::detail::UnionFind sets(number_of_nodes); + bioimage_cpp::util::UnionFind sets(number_of_nodes); MutexStorage mutexes(number_of_nodes); for (const auto &edge : edge_order) { diff --git a/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx b/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx index e665fcd..4f2e5c5 100644 --- a/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx +++ b/include/bioimage_cpp/graph/mutex_watershed/semantic.hxx @@ -3,7 +3,7 @@ #include "bioimage_cpp/detail/mutex_storage.hxx" #include "bioimage_cpp/detail/relabel.hxx" #include "bioimage_cpp/detail/semantic_labels.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include "bioimage_cpp/graph/undirected_graph.hxx" #include @@ -144,7 +144,7 @@ SemanticMutexWatershedResult semantic_mutex_watershed_clustering( return first.index < second.index; }); - bioimage_cpp::detail::UnionFind sets(number_of_nodes); + bioimage_cpp::util::UnionFind sets(number_of_nodes); MutexStorage mutexes(number_of_nodes); SemanticLabeling semantic_labels(number_of_nodes, -1); diff --git a/include/bioimage_cpp/segmentation/mutex_watershed.hxx b/include/bioimage_cpp/segmentation/mutex_watershed.hxx index 18f888a..743cf1a 100644 --- a/include/bioimage_cpp/segmentation/mutex_watershed.hxx +++ b/include/bioimage_cpp/segmentation/mutex_watershed.hxx @@ -3,7 +3,7 @@ #include "bioimage_cpp/array_view.hxx" #include "bioimage_cpp/detail/grid.hxx" #include "bioimage_cpp/detail/mutex_storage.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include #include @@ -100,7 +100,7 @@ void mutex_watershed_grid( return first.weight > second.weight; }); - detail::UnionFind sets(static_cast(number_of_nodes)); + bioimage_cpp::util::UnionFind sets(static_cast(number_of_nodes)); MutexStorage mutexes(static_cast(number_of_nodes)); for (const auto &edge : edge_order) { diff --git a/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx b/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx index ff620eb..0db53d2 100644 --- a/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx +++ b/include/bioimage_cpp/segmentation/semantic_mutex_watershed.hxx @@ -4,7 +4,7 @@ #include "bioimage_cpp/detail/grid.hxx" #include "bioimage_cpp/detail/mutex_storage.hxx" #include "bioimage_cpp/detail/semantic_labels.hxx" -#include "bioimage_cpp/detail/union_find.hxx" +#include "bioimage_cpp/util/union_find.hxx" #include #include @@ -132,7 +132,7 @@ void semantic_mutex_watershed_grid( return first.weight > second.weight; }); - detail::UnionFind sets(static_cast(number_of_nodes)); + bioimage_cpp::util::UnionFind sets(static_cast(number_of_nodes)); MutexStorage mutexes(static_cast(number_of_nodes)); SemanticLabeling semantic_labels(static_cast(number_of_nodes), -1); diff --git a/include/bioimage_cpp/detail/union_find.hxx b/include/bioimage_cpp/util/union_find.hxx similarity index 97% rename from include/bioimage_cpp/detail/union_find.hxx rename to include/bioimage_cpp/util/union_find.hxx index 404b9a2..7ea3bb0 100644 --- a/include/bioimage_cpp/detail/union_find.hxx +++ b/include/bioimage_cpp/util/union_find.hxx @@ -6,7 +6,7 @@ #include #include -namespace bioimage_cpp::detail { +namespace bioimage_cpp::util { // Disjoint-set / union-find with path compression and union-by-rank. // @@ -85,4 +85,4 @@ private: std::vector ranks_; }; -} // namespace bioimage_cpp::detail +} // namespace bioimage_cpp::util diff --git a/src/bindings/module.cxx b/src/bindings/module.cxx index 381c15f..7bded51 100644 --- a/src/bindings/module.cxx +++ b/src/bindings/module.cxx @@ -4,6 +4,7 @@ #include "graph.hxx" #include "ground_truth.hxx" #include "segmentation.hxx" +#include "util.hxx" #include "utils.hxx" #include @@ -18,5 +19,6 @@ NB_MODULE(_core, m) { bioimage_cpp::bindings::bind_graph(m); bioimage_cpp::bindings::bind_ground_truth(m); bioimage_cpp::bindings::bind_segmentation(m); + bioimage_cpp::bindings::bind_util(m); bioimage_cpp::bindings::bind_utils(m); } diff --git a/src/bindings/util.cxx b/src/bindings/util.cxx new file mode 100644 index 0000000..90a972e --- /dev/null +++ b/src/bindings/util.cxx @@ -0,0 +1,146 @@ +#include "util.hxx" + +#include "bioimage_cpp/util/union_find.hxx" + +#include + +#include +#include +#include +#include + +namespace nb = nanobind; + +namespace bioimage_cpp::bindings { +namespace { + +using EdgeArray = nb::ndarray; +using NodeArray = nb::ndarray>; +using OutputArray = nb::ndarray; + +void merge_edges( + util::UnionFind &uf, + EdgeArray edges +) { + if (edges.ndim() != 2 || edges.shape(1) != 2) { + std::string shape = "("; + for (std::size_t axis = 0; axis < edges.ndim(); ++axis) { + if (axis > 0) { + shape += ", "; + } + shape += std::to_string(edges.shape(axis)); + } + shape += ")"; + throw std::invalid_argument( + "edges must have shape (N, 2), got shape " + shape + ); + } + + const auto n_edges = edges.shape(0); + const auto *data = edges.data(); + + { + nb::gil_scoped_release release; + for (std::size_t i = 0; i < n_edges; ++i) { + uf.merge(data[2 * i], data[2 * i + 1]); + } + } +} + +OutputArray find_nodes(util::UnionFind &uf, NodeArray nodes) { + const auto n = nodes.shape(0); + auto *out = new std::uint64_t[n](); + nb::capsule owner(out, [](void *p) noexcept { delete[] static_cast(p); }); + + const auto *input = nodes.data(); + { + nb::gil_scoped_release release; + for (std::size_t i = 0; i < n; ++i) { + out[i] = uf.find(input[i]); + } + } + + std::size_t shape[1] = {n}; + return OutputArray(out, 1, shape, owner); +} + +OutputArray element_labeling(util::UnionFind &uf) { + const auto n = uf.size(); + auto *out = new std::uint64_t[n](); + nb::capsule owner(out, [](void *p) noexcept { delete[] static_cast(p); }); + + { + nb::gil_scoped_release release; + for (std::size_t i = 0; i < n; ++i) { + out[i] = uf.find(static_cast(i)); + } + } + + std::size_t shape[1] = {n}; + return OutputArray(out, 1, shape, owner); +} + +} // namespace + +void bind_util(nb::module_ &m) { + nb::module_ util_module = m.def_submodule("util", "Utility data structures."); + + nb::class_(util_module, "UnionFind") + .def( + nb::init(), + nb::arg("size"), + "Create a union-find over `size` singleton elements." + ) + .def_prop_ro( + "size", + &util::UnionFind::size, + "Number of elements in the union-find." + ) + .def( + "find", + &util::UnionFind::find, + nb::arg("node"), + "Return the (path-compressed) root of `node`." + ) + .def( + "find", + &find_nodes, + nb::arg("nodes"), + "Return the roots for a 1-D uint64 array of node indices." + ) + .def( + "merge", + &util::UnionFind::merge, + nb::arg("first"), + nb::arg("second"), + "Union the sets containing `first` and `second`. Returns the new root." + ) + .def( + "merge", + &merge_edges, + nb::arg("edges"), + "Bulk-merge from an (N, 2) uint64 array of node-pair edges." + ) + .def( + "merge_to", + &util::UnionFind::merge_to, + nb::arg("stable"), + nb::arg("removed"), + "Union the sets containing `stable` and `removed`, forcing " + "`stable`'s root to survive. Returns the new root." + ) + .def( + "element_labeling", + &element_labeling, + "Return a uint64 array of length `size` where entry i is the " + "(path-compressed) root of element i." + ) + .def( + "reset", + &util::UnionFind::reset, + nb::arg("size"), + "Reinitialise the union-find to `size` singleton elements." + ); +} + +} // namespace bioimage_cpp::bindings diff --git a/src/bindings/util.hxx b/src/bindings/util.hxx new file mode 100644 index 0000000..5839bcc --- /dev/null +++ b/src/bindings/util.hxx @@ -0,0 +1,9 @@ +#pragma once + +#include + +namespace bioimage_cpp::bindings { + +void bind_util(nanobind::module_ &m); + +} // namespace bioimage_cpp::bindings diff --git a/src/bioimage_cpp/__init__.py b/src/bioimage_cpp/__init__.py index b8b26a1..03cb6c6 100644 --- a/src/bioimage_cpp/__init__.py +++ b/src/bioimage_cpp/__init__.py @@ -7,6 +7,7 @@ from . import graph from . import ground_truth from . import segmentation +from . import util from . import utils __all__ = [ @@ -19,5 +20,6 @@ "graph", "ground_truth", "segmentation", + "util", "utils", ] diff --git a/src/bioimage_cpp/util/__init__.py b/src/bioimage_cpp/util/__init__.py new file mode 100644 index 0000000..85703f4 --- /dev/null +++ b/src/bioimage_cpp/util/__init__.py @@ -0,0 +1,5 @@ +"""Reusable data structures (union-find, ...).""" + +from .._core.util import UnionFind + +__all__ = ["UnionFind"] diff --git a/tests/test_util_union_find.py b/tests/test_util_union_find.py new file mode 100644 index 0000000..2535654 --- /dev/null +++ b/tests/test_util_union_find.py @@ -0,0 +1,138 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_initial_state_is_all_singletons(): + uf = bic.util.UnionFind(5) + assert uf.size == 5 + for i in range(5): + assert uf.find(i) == i + + +def test_size_zero(): + uf = bic.util.UnionFind(0) + assert uf.size == 0 + labels = uf.element_labeling() + assert labels.shape == (0,) + assert labels.dtype == np.uint64 + + +def test_scalar_merge_joins_two_elements(): + uf = bic.util.UnionFind(4) + root = uf.merge(0, 1) + assert uf.find(0) == root + assert uf.find(1) == root + assert uf.find(2) == 2 + assert uf.find(3) == 3 + + +def test_scalar_merge_is_transitive(): + uf = bic.util.UnionFind(5) + uf.merge(0, 1) + uf.merge(2, 3) + uf.merge(1, 2) + assert uf.find(0) == uf.find(3) + assert uf.find(0) != uf.find(4) + + +def test_merge_to_forces_stable_root(): + uf = bic.util.UnionFind(4) + root = uf.merge_to(0, 1) + assert root == 0 + assert uf.find(0) == 0 + assert uf.find(1) == 0 + + +def test_bulk_merge_from_edges(): + uf = bic.util.UnionFind(6) + edges = np.array([[0, 1], [1, 2], [3, 4]], dtype=np.uint64) + uf.merge(edges) + assert uf.find(0) == uf.find(1) == uf.find(2) + assert uf.find(3) == uf.find(4) + assert uf.find(0) != uf.find(3) + assert uf.find(5) == 5 + + +def test_bulk_merge_empty_edges(): + uf = bic.util.UnionFind(3) + edges = np.empty((0, 2), dtype=np.uint64) + uf.merge(edges) + assert uf.find(0) == 0 + assert uf.find(1) == 1 + assert uf.find(2) == 2 + + +def test_bulk_find_returns_array_of_roots(): + uf = bic.util.UnionFind(5) + uf.merge(0, 1) + uf.merge(2, 3) + + nodes = np.array([0, 1, 2, 3, 4], dtype=np.uint64) + roots = uf.find(nodes) + + assert roots.dtype == np.uint64 + assert roots.shape == (5,) + assert roots[0] == roots[1] + assert roots[2] == roots[3] + assert roots[4] == 4 + for i in range(5): + assert roots[i] == uf.find(int(i)) + + +def test_bulk_find_empty_array(): + uf = bic.util.UnionFind(3) + roots = uf.find(np.empty((0,), dtype=np.uint64)) + assert roots.shape == (0,) + assert roots.dtype == np.uint64 + + +def test_element_labeling_matches_scalar_find(): + uf = bic.util.UnionFind(6) + edges = np.array([[0, 1], [2, 3], [3, 4]], dtype=np.uint64) + uf.merge(edges) + + labels = uf.element_labeling() + + assert labels.dtype == np.uint64 + assert labels.shape == (6,) + for i in range(6): + assert labels[i] == uf.find(i) + + +def test_element_labeling_equivalence_classes(): + uf = bic.util.UnionFind(5) + uf.merge(np.array([[0, 2], [1, 3]], dtype=np.uint64)) + + labels = uf.element_labeling() + + assert labels[0] == labels[2] + assert labels[1] == labels[3] + assert labels[0] != labels[1] + assert labels[4] != labels[0] + assert labels[4] != labels[1] + + +def test_reset_reinitialises_to_singletons(): + uf = bic.util.UnionFind(3) + uf.merge(0, 1) + uf.merge(1, 2) + + uf.reset(4) + + assert uf.size == 4 + for i in range(4): + assert uf.find(i) == i + + +def test_bulk_merge_rejects_wrong_shape(): + uf = bic.util.UnionFind(3) + with pytest.raises(Exception, match=r"\(N, 2\)"): + uf.merge(np.zeros((4,), dtype=np.uint64)) + + +def test_bulk_merge_rejects_three_columns(): + uf = bic.util.UnionFind(3) + with pytest.raises(Exception, match=r"\(N, 2\)"): + uf.merge(np.zeros((2, 3), dtype=np.uint64))