diff --git a/MIGRATION_GUIDE.md b/MIGRATION_GUIDE.md index 4b09083..c8d8729 100644 --- a/MIGRATION_GUIDE.md +++ b/MIGRATION_GUIDE.md @@ -107,6 +107,13 @@ Important differences: - Bulk methods accept array-like inputs and return NumPy arrays. - Python-style names are preferred. A few nifty-style aliases are still present on `UndirectedGraph` for convenience, but new code should use snake_case. +- `graph.clone()` returns an independent deep copy. The C++ class is + move-only (it owns a CSR adjacency buffer), so prefer this over + reassignment-by-value. +- `graph.freeze()` eagerly builds the internal adjacency. Call it after a + batch of `insert_edge` calls if you intend to hand the graph to multiple + reader threads, or if you want to ensure subsequent `node_adjacency` + reads carry no first-call rebuild cost. Common method/property mapping: @@ -126,6 +133,49 @@ Common method/property mapping: | `extractSubgraphFromNodes` | `extract_subgraph_from_nodes` | | `edgesFromNodeList` | `edges_from_node_list` | +## Grid Graphs + +Nifty-style regular grid graphs map to explicit 2D or 3D grid graph classes: + +```python +graph = bic.graph.GridGraph2D((height, width)) +graph = bic.graph.GridGraph3D((depth, height, width)) +graph = bic.graph.grid_graph((height, width)) +``` + +Grid graph nodes use NumPy C-order ids. For a 2D shape `(y, x)`, node +`(row, col)` has id `row * x + col`; for 3D `(z, y, x)`, ids follow the same +row-major convention. `GridGraph2D` and `GridGraph3D` inherit the regular +`UndirectedGraph` API, so solvers, connected components, breadth-first search, +and `uv_ids()` work unchanged. + +Important differences: + +- Only nearest-neighbor 2D and 3D grids are exposed for now. +- Edge ids are deterministic axis blocks: axis 0 edges first, then axis 1, and + axis 2 for 3D. +- Scalar boundary maps can be converted to edge weights with + `grid_boundary_features(graph, boundary_map)`. +- Local affinity channels can be converted to edge-aligned weights with + `grid_affinity_features(graph, affinities, offsets)`. +- Mixed local and long-range affinity offsets can be converted with + `grid_affinity_features_with_lifted(...)`, which returns local graph weights + plus explicit long-range `uv_ids` and weights for lifted multicut or mutex + watershed style workflows. +- The three grid feature functions preserve `float32` and `float64` input + dtype end-to-end (no internal copy to `float64`); other dtypes are + promoted to `float64`. Output weight arrays match the input dtype. +- Grid graph construction does not materialize the per-node adjacency + list. If you only need `uv_ids()` and edge features (the common case) + you pay nothing for adjacency. The first call to `node_adjacency`, + `connected_components`, `breadth_first_search`, or + `extract_subgraph_from_nodes` on a grid graph triggers a one-shot + rebuild; call `graph.freeze()` on the construction thread before + fan-out if you intend to use those from multiple threads. +- Affogato-style masks and seed edges are not part of the public grid feature + API yet; the implementation is structured so these filters/extra edges can + be added later. + ## Region Adjacency Graphs Nifty: diff --git a/README.md b/README.md index a466c73..5f72694 100644 --- a/README.md +++ b/README.md @@ -45,6 +45,15 @@ segmentation = bic.segmentation.mutex_watershed( graph = bic.graph.UndirectedGraph.from_edges(4, [[0, 1], [1, 2], [2, 3]]) graph.find_edge(2, 1) # 1 + +grid = bic.graph.grid_graph((64, 64)) +grid.number_of_edges +# 8064 + +edge_weights = bic.graph.grid_boundary_features(grid, boundary_map) +local_weights, valid_edges, lifted_uvs, lifted_weights, offset_ids = ( + bic.graph.grid_affinity_features_with_lifted(grid, affinities, offsets) +) ``` ## Scope diff --git a/development/graph/_grid_affinity_compatibility.py b/development/graph/_grid_affinity_compatibility.py new file mode 100644 index 0000000..5dfcad7 --- /dev/null +++ b/development/graph/_grid_affinity_compatibility.py @@ -0,0 +1,274 @@ +from __future__ import annotations + +import argparse +from pathlib import Path +from statistics import median +from time import perf_counter +from typing import Callable + +import numpy as np + + +PROJECT_ROOT = Path(__file__).resolve().parents[2] +DEFAULT_DATA_PREFIX = PROJECT_ROOT / "examples" / "segmentation" / "isbi-data-" + + +def load_problem(data_prefix: Path | str = DEFAULT_DATA_PREFIX): + from elf.segmentation.utils import load_mutex_watershed_problem + + affinities, offsets = load_mutex_watershed_problem(prefix=str(data_prefix)) + return np.ascontiguousarray(affinities), [tuple(offset) for offset in offsets] + + +def prepare_2d_problem( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + z: int, + yx_shape: tuple[int, int], +): + channels_2d = [index for index, offset in enumerate(offsets) if offset[0] == 0] + y, x = yx_shape + affinities_2d = affinities[channels_2d, z, :y, :x] + offsets_2d = [offsets[index][1:] for index in channels_2d] + return np.ascontiguousarray(affinities_2d, dtype=np.float32), offsets_2d + + +def prepare_3d_problem( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], + zyx_shape: tuple[int, int, int], +): + z, y, x = zyx_shape + cropped = affinities[:, :z, :y, :x] + return np.ascontiguousarray(cropped, dtype=np.float32), offsets + + +def select_local_offsets( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], +): + local_channels = [ + index for index, offset in enumerate(offsets) + if sum(abs(value) for value in offset) == 1 + ] + local_affinities = np.ascontiguousarray(affinities[local_channels], dtype=np.float32) + local_offsets = [tuple(offsets[index]) for index in local_channels] + return local_affinities, local_offsets + + +def select_mixed_offsets( + affinities: np.ndarray, + offsets: list[tuple[int, ...]], +): + selected = [ + index for index, offset in enumerate(offsets) + if sum(abs(value) for value in offset) >= 1 + ] + mixed_affinities = np.ascontiguousarray(affinities[selected], dtype=np.float32) + mixed_offsets = [tuple(offsets[index]) for index in selected] + return mixed_affinities, mixed_offsets + + +def time_call(function: Callable[[], tuple[np.ndarray, np.ndarray]], repeats: int): + # One untimed warm-up before the measured loop. The first call typically + # pays for code-page faults and one-time library initialization + # (nanobind tuple shapes, numpy ufunc caches, ...) that aren't part of + # the steady-state cost we care about. Without warm-up these costs leak + # into the first sample and skew the median for low `repeats` values. + function() + timings = [] + result = None + for _ in range(repeats): + start = perf_counter() + result = function() + timings.append(perf_counter() - start) + assert result is not None + return timings, result + + +def sorted_edges_and_weights(uvs: np.ndarray, weights: np.ndarray): + uvs = np.asarray(uvs, dtype=np.uint64) + weights = np.asarray(weights, dtype=np.float64) + if uvs.shape[0] == 0: + return uvs.reshape(0, 2), weights.reshape(0) + normalized = np.sort(uvs, axis=1) + order = np.lexsort((normalized[:, 1], normalized[:, 0])) + return normalized[order], weights[order] + + +def split_affogato_edges(uvs: np.ndarray, weights: np.ndarray, graph): + local_mask = np.asarray(graph.find_edges(uvs), dtype=np.int64) >= 0 + return ( + uvs[local_mask], + weights[local_mask], + uvs[~local_mask], + weights[~local_mask], + ) + + +def bioimage_cpp_local(affinities: np.ndarray, offsets: list[tuple[int, ...]]): + import bioimage_cpp as bic + + graph = bic.graph.grid_graph(affinities.shape[1:]) + weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets) + return graph.uv_ids(), weights + + +def bioimage_cpp_local_weights_only( + graph, affinities: np.ndarray, offsets: list[tuple[int, ...]] +): + """Compute edge weights only — no (uvs, weights) materialization. + + This isolates the cost of the feature kernel from the cost of returning + the canonical uv_ids array. Use this when comparing against libraries + that already cache uvs in the graph object. + """ + import bioimage_cpp as bic + + weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets) + return weights + + +def bioimage_cpp_local_with_uvs( + graph, affinities: np.ndarray, offsets: list[tuple[int, ...]] +): + """Compute weights AND materialize uvs — apples-to-apples with nifty's + ``affinitiesToEdgeMapWithOffsets`` and affogato's + ``compute_nh_and_weights``, both of which return uvs in their output.""" + import bioimage_cpp as bic + + weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets) + return graph.uv_ids(), weights + + +def bioimage_cpp_lifted(affinities: np.ndarray, offsets: list[tuple[int, ...]]): + import bioimage_cpp as bic + + graph = bic.graph.grid_graph(affinities.shape[1:]) + local_weights, _, lifted_uvs, lifted_weights, _ = ( + bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets) + ) + return graph, graph.uv_ids(), local_weights, lifted_uvs, lifted_weights + + +def bioimage_cpp_lifted_features_only( + graph, affinities: np.ndarray, offsets: list[tuple[int, ...]] +): + """Lifted features without graph.uv_ids() — see the local variant.""" + import bioimage_cpp as bic + + local_weights, _, lifted_uvs, lifted_weights, _ = ( + bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets) + ) + return local_weights, lifted_uvs, lifted_weights + + +def bioimage_cpp_lifted_with_uvs( + graph, affinities: np.ndarray, offsets: list[tuple[int, ...]] +): + """Lifted features WITH local uvs (apples-to-apples with affogato).""" + import bioimage_cpp as bic + + local_weights, _, lifted_uvs, lifted_weights, _ = ( + bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets) + ) + return graph.uv_ids(), local_weights, lifted_uvs, lifted_weights + + +def assert_local_offsets_cover_all_edges(graph, affinities, offsets) -> None: + """One-shot correctness check called outside of the timing loop.""" + import bioimage_cpp as bic + + _, valid_edges = bic.graph.grid_affinity_features(graph, affinities, offsets) + if not np.all(valid_edges): + raise AssertionError("local offsets did not cover all grid edges") + + +def nifty_local(affinities: np.ndarray, offsets: list[tuple[int, ...]]): + import nifty.graph as ng + + graph = ng.undirectedGridGraph(list(affinities.shape[1:])) + return nifty_local_on_graph(graph, affinities, offsets) + + +def nifty_local_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]): + n_edges, uvs, weights = graph.affinitiesToEdgeMapWithOffsets( + affinities, + [list(offset) for offset in offsets], + ) + return np.asarray(uvs[:n_edges], dtype=np.uint64), np.asarray(weights[:n_edges]) + + +def affogato_edges(affinities: np.ndarray, offsets: list[tuple[int, ...]]): + from affogato.segmentation import MWSGridGraph + + graph = MWSGridGraph(list(affinities.shape[1:])) + return affogato_edges_on_graph(graph, affinities, offsets) + + +def affogato_edges_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]): + uvs, weights = graph.compute_nh_and_weights( + affinities, + [list(offset) for offset in offsets], + strides=[1] * (affinities.ndim - 1), + randomize_strides=False, + ) + return np.asarray(uvs, dtype=np.uint64), np.asarray(weights) + + +def compare_edge_sets( + name: str, + candidate_uvs: np.ndarray, + candidate_weights: np.ndarray, + reference_uvs: np.ndarray, + reference_weights: np.ndarray, +): + candidate_uvs, candidate_weights = sorted_edges_and_weights( + candidate_uvs, candidate_weights + ) + reference_uvs, reference_weights = sorted_edges_and_weights( + reference_uvs, reference_weights + ) + np.testing.assert_array_equal(candidate_uvs, reference_uvs) + np.testing.assert_allclose(candidate_weights, reference_weights, rtol=1.0e-6, atol=1.0e-6) + max_abs_diff = ( + float(np.max(np.abs(candidate_weights - reference_weights))) + if candidate_weights.size + else 0.0 + ) + return { + "name": name, + "number_of_edges": int(candidate_uvs.shape[0]), + "max_abs_weight_diff": max_abs_diff, + } + + +def print_timing(name: str, first_name: str, first_timings: list[float], + second_name: str, second_timings: list[float]): + first_median = median(first_timings) + second_median = median(second_timings) + ratio = second_median / first_median if first_median > 0 else float("inf") + print(f"{name} {first_name} median runtime: {first_median:.6f} s") + print(f"{name} {second_name} median runtime: {second_median:.6f} s") + print(f"{name} {second_name} / {first_name} runtime ratio: {ratio:.3f}x") + + +def add_common_arguments(parser: argparse.ArgumentParser) -> None: + parser.add_argument("--ndim", type=int, choices=(2, 3), default=2) + parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX) + # Default bumped from 3 to 5 — median of 3 is the middle sample and is + # noisy if anything (GC, cache eviction) lands inside one of the three + # runs. With `time_call` doing one warm-up before this, 5 samples gives + # a usable median without much added cost. + parser.add_argument("--repeats", type=int, default=5) + parser.add_argument("--z", type=int, default=20) + parser.add_argument("--yx-shape", type=int, nargs=2, default=(512, 512)) + parser.add_argument("--zyx-shape", type=int, nargs=3, default=(16, 512, 512)) + # Affinity dtype that every library receives. nifty and affogato accept + # both float32 and float64 at near-identical speed (verified separately), + # and bioimage-cpp now templates on the value type, so feeding all three + # the same dtype removes the previous implicit float32 -> float64 copy + # that was charged only to bioimage-cpp. + parser.add_argument( + "--dtype", choices=("float32", "float64"), default="float32" + ) diff --git a/development/graph/check_grid_affinity_edges.py b/development/graph/check_grid_affinity_edges.py new file mode 100644 index 0000000..d382c3c --- /dev/null +++ b/development/graph/check_grid_affinity_edges.py @@ -0,0 +1,154 @@ +from __future__ import annotations + +import argparse + +from _grid_affinity_compatibility import ( + add_common_arguments, + affogato_edges, + affogato_edges_on_graph, + assert_local_offsets_cover_all_edges, + bioimage_cpp_local, + bioimage_cpp_local_weights_only, + bioimage_cpp_local_with_uvs, + compare_edge_sets, + load_problem, + nifty_local, + nifty_local_on_graph, + prepare_2d_problem, + prepare_3d_problem, + print_timing, + select_local_offsets, + time_call, +) + +import numpy as np + + +# Notes on remaining apples-to-apples caveats (left in-code rather than fixed): +# * affogato's `MWSGridGraph` carries MWS-specific state and is heavier to +# construct than a pure undirected grid graph. The "total" timings include +# that cost on the affogato side. This isn't a bug in the comparison, it +# reflects affogato's intended workload. +# * nifty / affogato accept the chosen dtype directly (verified separately); +# feeding all three libraries the same dtype removes the previous implicit +# float32 -> float64 copy that was charged only to bioimage-cpp. +# * `time_call` does an untimed warm-up call before the measured loop so the +# first sample doesn't carry one-shot init costs. + + +def run_check(args: argparse.Namespace) -> None: + affinities, offsets = load_problem(args.data_prefix) + if args.ndim == 2: + affinities, offsets = prepare_2d_problem( + affinities, offsets, z=args.z, yx_shape=tuple(args.yx_shape) + ) + else: + affinities, offsets = prepare_3d_problem( + affinities, offsets, zyx_shape=tuple(args.zyx_shape) + ) + affinities, offsets = select_local_offsets(affinities, offsets) + # One conversion to the chosen common dtype, BEFORE any timed work. + # bioimage-cpp, nifty, and affogato all consume this same array. + affinities = np.ascontiguousarray(affinities, dtype=np.dtype(args.dtype)) + + bic_timings, (bic_uvs, bic_weights) = time_call( + lambda: bioimage_cpp_local(affinities, offsets), args.repeats + ) + nifty_timings, (nifty_uvs, nifty_weights) = time_call( + lambda: nifty_local(affinities, offsets), args.repeats + ) + affogato_timings, (affogato_uvs, affogato_weights) = time_call( + lambda: affogato_edges(affinities, offsets), args.repeats + ) + + import bioimage_cpp as bic + import nifty.graph as ng + from affogato.segmentation import MWSGridGraph + + bic_graph = bic.graph.grid_graph(affinities.shape[1:]) + nifty_graph = ng.undirectedGridGraph(list(affinities.shape[1:])) + affogato_graph = MWSGridGraph(list(affinities.shape[1:])) + # One-shot correctness check, OUTSIDE the timing loop. + assert_local_offsets_cover_all_edges(bic_graph, affinities, offsets) + + # Apples-to-apples timing #1: each library returns (uvs, weights) for a + # pre-built graph. nifty and affogato bundle uvs in their return value, + # bioimage-cpp materializes them via graph.uv_ids(). + bic_with_uvs_timings, _ = time_call( + lambda: bioimage_cpp_local_with_uvs(bic_graph, affinities, offsets), + args.repeats, + ) + nifty_feature_timings, _ = time_call( + lambda: nifty_local_on_graph(nifty_graph, affinities, offsets), + args.repeats, + ) + affogato_feature_timings, _ = time_call( + lambda: affogato_edges_on_graph(affogato_graph, affinities, offsets), + args.repeats, + ) + # Apples-to-apples timing #2: bioimage-cpp ONLY computes weights (no + # uvs materialization). This isolates the cost of the feature kernel + # itself. There is no nifty/affogato equivalent that returns just + # weights, so this is reported on its own. + bic_weights_only_timings, _ = time_call( + lambda: bioimage_cpp_local_weights_only(bic_graph, affinities, offsets), + args.repeats, + ) + + nifty_summary = compare_edge_sets( + "nifty local", bic_uvs, bic_weights, nifty_uvs, nifty_weights + ) + affogato_summary = compare_edge_sets( + "affogato local", bic_uvs, bic_weights, affogato_uvs, affogato_weights + ) + + print(f"Grid local affinity edge check ({args.ndim}D)") + print( + f"affinities shape: {affinities.shape}, dtype: {affinities.dtype}, " + f"size: {affinities.nbytes / 1e6:.2f} MB, offsets: {len(offsets)}, " + f"repeats: {args.repeats}" + ) + print(f"offsets: {offsets}") + print( + f"nifty edges: {nifty_summary['number_of_edges']}, " + f"max abs weight diff: {nifty_summary['max_abs_weight_diff']:.6g}" + ) + print( + f"affogato edges: {affogato_summary['number_of_edges']}, " + f"max abs weight diff: {affogato_summary['max_abs_weight_diff']:.6g}" + ) + print_timing("local edges total", "bioimage-cpp", bic_timings, "nifty", nifty_timings) + print_timing("local edges total", "bioimage-cpp", bic_timings, "affogato", affogato_timings) + print_timing( + "local edges prebuilt (with uvs)", + "bioimage-cpp", + bic_with_uvs_timings, + "nifty", + nifty_feature_timings, + ) + print_timing( + "local edges prebuilt (with uvs)", + "bioimage-cpp", + bic_with_uvs_timings, + "affogato", + affogato_feature_timings, + ) + # Weights-only kernel timing (no comparator — informational). + from statistics import median + + print( + "local edges prebuilt (weights only, no uvs materialization) " + f"bioimage-cpp median runtime: {median(bic_weights_only_timings):.6f} s" + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare local grid affinity topology and weights." + ) + add_common_arguments(parser) + run_check(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/development/graph/check_grid_affinity_lifted_edges.py b/development/graph/check_grid_affinity_lifted_edges.py new file mode 100644 index 0000000..e32c424 --- /dev/null +++ b/development/graph/check_grid_affinity_lifted_edges.py @@ -0,0 +1,144 @@ +from __future__ import annotations + +import argparse + +from _grid_affinity_compatibility import ( + add_common_arguments, + affogato_edges, + affogato_edges_on_graph, + bioimage_cpp_lifted, + bioimage_cpp_lifted_features_only, + bioimage_cpp_lifted_with_uvs, + compare_edge_sets, + load_problem, + prepare_2d_problem, + prepare_3d_problem, + print_timing, + select_mixed_offsets, + split_affogato_edges, + time_call, +) + +import numpy as np + + +# Notes on remaining apples-to-apples caveats (left in-code rather than fixed): +# * affogato bundles local + lifted edges into a single (uvs, weights) array, +# bioimage-cpp returns local and lifted in separate arrays. The two output +# formats are intrinsic to each library's design; the bookkeeping cost of +# producing them is included in the respective timings. +# * affogato's `MWSGridGraph` is heavier to construct than a plain grid graph +# (MWS-specific state). The "total" timings include that cost. +# * `time_call` does an untimed warm-up before the measured loop. + + +def run_check(args: argparse.Namespace) -> None: + affinities, offsets = load_problem(args.data_prefix) + if args.ndim == 2: + affinities, offsets = prepare_2d_problem( + affinities, offsets, z=args.z, yx_shape=tuple(args.yx_shape) + ) + else: + affinities, offsets = prepare_3d_problem( + affinities, offsets, zyx_shape=tuple(args.zyx_shape) + ) + affinities, offsets = select_mixed_offsets(affinities, offsets) + # One conversion to the chosen common dtype, BEFORE any timed work. + affinities = np.ascontiguousarray(affinities, dtype=np.dtype(args.dtype)) + + bic_timings, (graph, bic_local_uvs, bic_local_weights, bic_lifted_uvs, bic_lifted_weights) = ( + time_call(lambda: bioimage_cpp_lifted(affinities, offsets), args.repeats) + ) + affogato_timings, (affogato_uvs, affogato_weights) = time_call( + lambda: affogato_edges(affinities, offsets), args.repeats + ) + + import bioimage_cpp as bic + from affogato.segmentation import MWSGridGraph + + bic_graph = bic.graph.grid_graph(affinities.shape[1:]) + affogato_graph = MWSGridGraph(list(affinities.shape[1:])) + # Apples-to-apples #1: both return (uvs, local_weights, lifted_uvs, + # lifted_weights) — bundle of arrays sized for downstream multicut. + bic_with_uvs_timings, _ = time_call( + lambda: bioimage_cpp_lifted_with_uvs(bic_graph, affinities, offsets), + args.repeats, + ) + affogato_feature_timings, _ = time_call( + lambda: affogato_edges_on_graph(affogato_graph, affinities, offsets), + args.repeats, + ) + # Apples-to-apples #2: bioimage-cpp ONLY returns weight arrays and the + # lifted uvs (which must be returned because they aren't grid-indexed). + # Isolates the cost of the feature kernel from local uv_ids() materialization. + bic_features_only_timings, _ = time_call( + lambda: bioimage_cpp_lifted_features_only(bic_graph, affinities, offsets), + args.repeats, + ) + + affogato_local_uvs, affogato_local_weights, affogato_lifted_uvs, affogato_lifted_weights = ( + split_affogato_edges(affogato_uvs, affogato_weights, graph) + ) + + local_summary = compare_edge_sets( + "affogato local", + bic_local_uvs, + bic_local_weights, + affogato_local_uvs, + affogato_local_weights, + ) + lifted_summary = compare_edge_sets( + "affogato lifted", + bic_lifted_uvs, + bic_lifted_weights, + affogato_lifted_uvs, + affogato_lifted_weights, + ) + + print(f"Grid lifted affinity edge check ({args.ndim}D)") + print( + f"affinities shape: {affinities.shape}, dtype: {affinities.dtype}, " + f"size: {affinities.nbytes / 1e6:.2f} MB, offsets: {len(offsets)}, " + f"repeats: {args.repeats}" + ) + print(f"offsets: {offsets}") + print( + f"local edges: {local_summary['number_of_edges']}, " + f"max abs weight diff: {local_summary['max_abs_weight_diff']:.6g}" + ) + print( + f"lifted edges: {lifted_summary['number_of_edges']}, " + f"max abs weight diff: {lifted_summary['max_abs_weight_diff']:.6g}" + ) + print_timing( + "local+lifted edges total", + "bioimage-cpp", + bic_timings, + "affogato", + affogato_timings, + ) + print_timing( + "local+lifted edges prebuilt (with uvs)", + "bioimage-cpp", + bic_with_uvs_timings, + "affogato", + affogato_feature_timings, + ) + from statistics import median + + print( + "local+lifted edges prebuilt (features only, no local uvs materialization) " + f"bioimage-cpp median runtime: {median(bic_features_only_timings):.6f} s" + ) + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Compare local and lifted grid affinity topology and weights." + ) + add_common_arguments(parser) + run_check(parser.parse_args()) + + +if __name__ == "__main__": + main() diff --git a/include/bioimage_cpp/graph/grid_features.hxx b/include/bioimage_cpp/graph/grid_features.hxx new file mode 100644 index 0000000..554a385 --- /dev/null +++ b/include/bioimage_cpp/graph/grid_features.hxx @@ -0,0 +1,396 @@ +#pragma once + +#include "bioimage_cpp/array_view.hxx" +#include "bioimage_cpp/detail/edge_hash.hxx" +#include "bioimage_cpp/graph/grid_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +namespace detail_grid_features { + +template +std::array offset_to_array( + const std::vector &offset +) { + if (offset.size() != D) { + throw std::invalid_argument("each offset must have length matching graph ndim"); + } + std::array result{}; + for (std::size_t axis = 0; axis < D; ++axis) { + result[axis] = offset[axis]; + } + return result; +} + +inline std::ptrdiff_t l1_norm(const std::vector &offset) { + std::ptrdiff_t result = 0; + for (const auto value : offset) { + result += value < 0 ? -value : value; + } + return result; +} + +template +std::ptrdiff_t l1_norm(const std::array &offset) { + std::ptrdiff_t result = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + const auto value = offset[axis]; + result += value < 0 ? -value : value; + } + return result; +} + +template +void require_graph_shape( + const GridGraph &graph, + const std::vector &shape, + const char *argument_name +) { + if (shape.size() != D) { + throw std::invalid_argument( + std::string(argument_name) + " ndim must match graph ndim" + ); + } + const auto &graph_shape = graph.shape(); + for (std::size_t axis = 0; axis < D; ++axis) { + if (shape[axis] != static_cast(graph_shape[axis])) { + throw std::invalid_argument( + std::string(argument_name) + " shape must match graph shape" + ); + } + } +} + +template +void require_affinity_shape( + const GridGraph &graph, + const ConstArrayView &affinities, + const std::vector> &offsets +) { + if (affinities.ndim() != static_cast(D + 1)) { + throw std::invalid_argument("affinities must have shape (channels, *graph.shape)"); + } + if (static_cast(affinities.shape[0]) != offsets.size()) { + throw std::invalid_argument("offsets length must match affinities channel count"); + } + const auto &graph_shape = graph.shape(); + for (std::size_t axis = 0; axis < D; ++axis) { + if (affinities.shape[axis + 1] != static_cast(graph_shape[axis])) { + throw std::invalid_argument("affinities spatial shape must match graph shape"); + } + } + for (const auto &offset : offsets) { + if (offset.size() != D) { + throw std::invalid_argument("each offset must have length matching graph ndim"); + } + if (l1_norm(offset) == 0) { + throw std::invalid_argument("offsets must not contain the zero offset"); + } + } +} + +// Two long-range offsets `a` and `b` produce the same set of edges iff +// `b = -a` — each offset, restricted to its valid source region, is +// internally injective, and `edge_key` canonicalizes endpoint order, so +// only sign-flipped offset pairs collide. Checking pairs is O(n_offsets^2) +// — a small constant in practice — and avoids sorting the full lifted +// edge list at the end of `grid_affinity_features_with_lifted`. +template +void check_no_negated_long_range_offsets( + const std::vector> &offsets +) { + for (std::size_t i = 0; i < offsets.size(); ++i) { + if (l1_norm(offsets[i]) <= 1) continue; + for (std::size_t j = i + 1; j < offsets.size(); ++j) { + if (l1_norm(offsets[j]) <= 1) continue; + bool negated = true; + for (std::size_t axis = 0; axis < D; ++axis) { + if (offsets[i][axis] != -offsets[j][axis]) { + negated = false; + break; + } + } + if (negated) { + throw std::invalid_argument( + "offsets produce duplicate long-range grid edges" + ); + } + } + } +} + +template +std::array uint64_strides_from_shape( + const std::array &shape +) { + std::array strides{}; + strides[D - 1] = 1; + for (std::size_t axis = D - 1; axis > 0; --axis) { + strides[axis - 1] = strides[axis] * shape[axis]; + } + return strides; +} + +template +std::size_t local_offset_axis(const std::array &offset) { + for (std::size_t axis = 0; axis < D; ++axis) { + if (offset[axis] != 0) { + return axis; + } + } + return D; +} + +// Template-recursive enumeration of (source_node, target_node, coordinate) +// triples for a fixed grid offset. Flat node ids are computed incrementally +// by addition only. With `Axis` known at compile time, the compiler unrolls +// the entire nest for small D and inlines `callback`. +template +void for_each_valid_offset_link_impl( + const std::array &begin, + const std::array &end, + const std::array &strides, + const std::int64_t delta, + std::array &coordinate, + const std::uint64_t node, + Callback &&callback +) { + if constexpr (Axis == D) { + const auto target = + static_cast(static_cast(node) + delta); + callback(node, target, coordinate); + } else { + std::uint64_t axis_node = node + begin[Axis] * strides[Axis]; + for (std::uint64_t coord = begin[Axis]; coord < end[Axis]; ++coord) { + coordinate[Axis] = coord; + for_each_valid_offset_link_impl( + begin, end, strides, delta, coordinate, axis_node, callback + ); + axis_node += strides[Axis]; + } + } +} + +template +void for_each_valid_offset_link( + const GridGraph &graph, + const std::array &offset, + Callback &&callback +) { + const auto &shape = graph.shape(); + const auto &strides = graph.strides(); + std::array begin{}; + std::array end{}; + std::int64_t delta = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + const auto axis_offset = offset[axis]; + if (axis_offset < 0) { + begin[axis] = static_cast(-axis_offset); + end[axis] = shape[axis]; + } else { + begin[axis] = 0; + end[axis] = shape[axis] - static_cast(axis_offset); + } + delta += static_cast(axis_offset) * + static_cast(strides[axis]); + } + std::array coordinate{}; + for_each_valid_offset_link_impl( + begin, end, strides, delta, coordinate, 0, callback + ); +} + +// Compute the local edge id from a source coordinate. The per-axis pivot +// adjustment for negative offsets is hoisted into `pivot_adjustment` — a +// constant per (offset, axis) the caller pre-computes once, so the inner +// loop is a branch-free dot product. +template +std::uint64_t local_edge_id_from_source_coordinate( + const std::uint64_t edge_offset_base, + const std::array &source_coordinate, + const std::array &edge_strides, + const std::uint64_t pivot_adjustment +) { + std::uint64_t local_edge = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + local_edge += source_coordinate[axis] * edge_strides[axis]; + } + return edge_offset_base + local_edge - pivot_adjustment; +} + +} // namespace detail_grid_features + +template +void grid_boundary_features( + const GridGraph &graph, + const ConstArrayView &boundary_map, + const ArrayView &out +) { + detail_grid_features::require_graph_shape(graph, boundary_map.shape, "boundary_map"); + if (out.shape != std::vector{ + static_cast(graph.number_of_edges())}) { + throw std::invalid_argument("out shape must be (number_of_edges,)"); + } + + const auto &uv_ids = graph.uv_ids(); + for (std::size_t edge = 0; edge < uv_ids.size(); ++edge) { + const auto &uv = uv_ids[edge]; + out.data[edge] = static_cast(0.5) * + (boundary_map.data[uv.first] + boundary_map.data[uv.second]); + } +} + +template +void grid_local_affinity_features( + const GridGraph &graph, + const ConstArrayView &affinities, + const std::vector> &offsets, + const ArrayView &weights, + const ArrayView &valid_edges +) { + detail_grid_features::require_affinity_shape(graph, affinities, offsets); + if (weights.shape != std::vector{ + static_cast(graph.number_of_edges())}) { + throw std::invalid_argument("weights shape must be (number_of_edges,)"); + } + if (valid_edges.shape != weights.shape) { + throw std::invalid_argument("valid_edges shape must match weights shape"); + } + for (const auto &offset : offsets) { + if (detail_grid_features::l1_norm(offset) != 1) { + throw std::invalid_argument("grid_affinity_features accepts only local offsets"); + } + } + + const auto number_of_nodes = graph.number_of_nodes(); + for (std::uint64_t edge = 0; edge < graph.number_of_edges(); ++edge) { + weights.data[edge] = static_cast(0); + valid_edges.data[edge] = 0; + } + + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + const auto offset_array = + detail_grid_features::offset_to_array(offsets[channel]); + const auto axis = detail_grid_features::local_offset_axis(offset_array); + auto edge_shape = graph.shape(); + --edge_shape[axis]; + const auto edge_strides = + detail_grid_features::uint64_strides_from_shape(edge_shape); + const auto edge_offset_base = graph.edge_offset(axis); + const std::uint64_t pivot_adjustment = + offset_array[axis] < 0 ? edge_strides[axis] : 0; + const auto channel_offset = + static_cast(channel) * number_of_nodes; + detail_grid_features::for_each_valid_offset_link( + graph, + offset_array, + [&](const std::uint64_t node, const std::uint64_t, const auto &coordinate) { + const auto edge = + detail_grid_features::local_edge_id_from_source_coordinate( + edge_offset_base, coordinate, edge_strides, pivot_adjustment + ); + if (valid_edges.data[edge] != 0) { + throw std::invalid_argument( + "offsets produce duplicate local grid edges" + ); + } + weights.data[edge] = affinities.data[channel_offset + node]; + valid_edges.data[edge] = 1; + } + ); + } +} + +template +struct GridLiftedAffinityFeatures { + std::vector local_weights; + std::vector valid_local_edges; + std::vector lifted_uvs; + std::vector lifted_weights; + std::vector lifted_offset_ids; +}; + +template +GridLiftedAffinityFeatures grid_affinity_features_with_lifted( + const GridGraph &graph, + const ConstArrayView &affinities, + const std::vector> &offsets +) { + detail_grid_features::require_affinity_shape(graph, affinities, offsets); + detail_grid_features::check_no_negated_long_range_offsets(offsets); + + const auto number_of_nodes = graph.number_of_nodes(); + GridLiftedAffinityFeatures result; + result.local_weights.assign(static_cast(graph.number_of_edges()), static_cast(0)); + result.valid_local_edges.assign(static_cast(graph.number_of_edges()), 0); + + std::uint64_t lifted_capacity = 0; + for (const auto &offset : offsets) { + if (detail_grid_features::l1_norm(offset) > 1) { + lifted_capacity += number_of_nodes; + } + } + result.lifted_uvs.reserve(static_cast(lifted_capacity)); + result.lifted_weights.reserve(static_cast(lifted_capacity)); + result.lifted_offset_ids.reserve(static_cast(lifted_capacity)); + + for (std::size_t channel = 0; channel < offsets.size(); ++channel) { + const auto offset_array = + detail_grid_features::offset_to_array(offsets[channel]); + const auto l1 = detail_grid_features::l1_norm(offset_array); + const auto channel_offset = + static_cast(channel) * number_of_nodes; + if (l1 == 1) { + const auto axis = detail_grid_features::local_offset_axis(offset_array); + auto edge_shape = graph.shape(); + --edge_shape[axis]; + const auto edge_strides = + detail_grid_features::uint64_strides_from_shape(edge_shape); + const auto edge_offset_base = graph.edge_offset(axis); + const std::uint64_t pivot_adjustment = + offset_array[axis] < 0 ? edge_strides[axis] : 0; + detail_grid_features::for_each_valid_offset_link( + graph, + offset_array, + [&](const std::uint64_t node, const std::uint64_t, const auto &coordinate) { + const auto edge = + detail_grid_features::local_edge_id_from_source_coordinate( + edge_offset_base, coordinate, edge_strides, pivot_adjustment + ); + if (result.valid_local_edges[static_cast(edge)] != 0) { + throw std::invalid_argument( + "offsets produce duplicate local grid edges" + ); + } + result.local_weights[static_cast(edge)] = + affinities.data[channel_offset + node]; + result.valid_local_edges[static_cast(edge)] = 1; + } + ); + } else { + detail_grid_features::for_each_valid_offset_link( + graph, + offset_array, + [&](const std::uint64_t node, const std::uint64_t target, const auto &) { + const auto uv = bioimage_cpp::detail::edge_key(node, target); + result.lifted_uvs.push_back(uv); + result.lifted_weights.push_back(affinities.data[channel_offset + node]); + result.lifted_offset_ids.push_back(static_cast(channel)); + } + ); + } + } + + return result; +} + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/grid_graph.hxx b/include/bioimage_cpp/graph/grid_graph.hxx new file mode 100644 index 0000000..2d838b8 --- /dev/null +++ b/include/bioimage_cpp/graph/grid_graph.hxx @@ -0,0 +1,352 @@ +#pragma once + +#include "bioimage_cpp/graph/undirected_graph.hxx" + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +namespace bioimage_cpp::graph { + +template +class GridGraph : public UndirectedGraph { +public: + using UndirectedGraph::number_of_edges; + + using Coordinate = std::array; + + explicit GridGraph(const Coordinate &shape) + : UndirectedGraph(checked_node_count(shape), checked_edge_count(shape), 0), + shape_(shape), + strides_(compute_strides(shape)), + edge_offsets_(compute_edge_offsets(shape)), + edge_strides_(compute_edge_strides(shape)) { + build_edges(); + } + + explicit GridGraph(const std::vector &shape) + : GridGraph(vector_to_coordinate(shape)) { + } + + [[nodiscard]] const Coordinate &shape() const { + return shape_; + } + + [[nodiscard]] const Coordinate &strides() const { + return strides_; + } + + [[nodiscard]] std::size_t ndim() const { + return D; + } + + [[nodiscard]] std::uint64_t node_id(const Coordinate &coordinate) const { + std::uint64_t node = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + if (coordinate[axis] >= shape_[axis]) { + throw std::out_of_range( + "coordinate[" + std::to_string(axis) + "] must be < shape[" + + std::to_string(axis) + "]" + ); + } + node += coordinate[axis] * strides_[axis]; + } + return node; + } + + [[nodiscard]] Coordinate coordinates(std::uint64_t node) const { + validate_node(node); + Coordinate coordinate{}; + for (std::size_t axis = 0; axis < D; ++axis) { + coordinate[axis] = node / strides_[axis]; + node -= coordinate[axis] * strides_[axis]; + } + return coordinate; + } + + [[nodiscard]] std::size_t edge_axis(const std::uint64_t edge) const { + validate_edge(edge); + const auto found = std::upper_bound( + edge_offsets_.begin(), + edge_offsets_.end(), + edge + ); + return static_cast(found - edge_offsets_.begin() - 1); + } + + [[nodiscard]] std::uint64_t edge_offset(const std::size_t axis) const { + if (axis >= D) { + throw std::out_of_range("axis must be < ndim"); + } + return edge_offsets_[axis]; + } + + [[nodiscard]] std::uint64_t number_of_edges(const std::size_t axis) const { + if (axis >= D) { + throw std::out_of_range("axis must be < ndim"); + } + return edge_offsets_[axis + 1] - edge_offsets_[axis]; + } + + [[nodiscard]] std::uint64_t edge_id( + const std::size_t axis, + const Coordinate &pivot_coordinate + ) const { + if (axis >= D) { + throw std::out_of_range("axis must be < ndim"); + } + const auto shape = edge_shape(axis); + std::uint64_t local_edge = 0; + for (std::size_t coord_axis = 0; coord_axis < D; ++coord_axis) { + if (pivot_coordinate[coord_axis] >= shape[coord_axis]) { + throw std::out_of_range( + "pivot_coordinate[" + std::to_string(coord_axis) + + "] is outside the edge grid" + ); + } + local_edge += pivot_coordinate[coord_axis] * edge_strides_[axis][coord_axis]; + } + return edge_offsets_[axis] + local_edge; + } + + EdgeId insert_edge(const NodeId u, const NodeId v) override { + if (u == v) { + throw std::invalid_argument("self edges are not supported"); + } + const auto existing = find_edge(u, v); + if (existing >= 0) { + return static_cast(existing); + } + // `find_edge` already validated the node ids and missed both the grid + // analytical check and the parent hash map, so we can insert directly + // without paying for another hash lookup inside `UndirectedGraph::insert_edge`. + const auto key = detail::edge_key(u, v); + return UndirectedGraph::insert_new_edge(key.first, key.second); + } + + [[nodiscard]] std::int64_t find_edge(const NodeId u, const NodeId v) const override { + validate_node(u); + validate_node(v); + if (u == v) { + return -1; + } + const auto lower = std::min(u, v); + const auto upper = std::max(u, v); + const auto diff = upper - lower; + for (std::size_t axis = 0; axis < D; ++axis) { + if (diff != strides_[axis]) { + continue; + } + const auto axis_coordinate = (lower / strides_[axis]) % shape_[axis]; + if (axis_coordinate + 1 >= shape_[axis]) { + return -1; + } + std::uint64_t local_edge = 0; + for (std::size_t coord_axis = 0; coord_axis < D; ++coord_axis) { + const auto coordinate = (lower / strides_[coord_axis]) % shape_[coord_axis]; + local_edge += coordinate * edge_strides_[axis][coord_axis]; + } + return static_cast(edge_offsets_[axis] + local_edge); + } + return UndirectedGraph::find_edge(u, v); + } + + [[nodiscard]] std::pair + edge_coordinates(const std::uint64_t edge) const { + const auto axis = edge_axis(edge); + const auto local_edge = edge - edge_offsets_[axis]; + return {edge_pivot_coordinates(axis, local_edge), axis}; + } + + [[nodiscard]] bool valid_offset_target( + const std::uint64_t node, + const std::array &offset, + std::uint64_t &target_out + ) const { + const auto coordinate = coordinates(node); + std::int64_t signed_delta = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + const auto coord = static_cast(coordinate[axis]); + const auto neighbor = coord + offset[axis]; + if (neighbor < 0 || neighbor >= static_cast(shape_[axis])) { + return false; + } + signed_delta += offset[axis] * static_cast(strides_[axis]); + } + target_out = static_cast(static_cast(node) + signed_delta); + return true; + } + +private: + static Coordinate vector_to_coordinate(const std::vector &shape) { + if (shape.size() != D) { + throw std::invalid_argument( + "shape must have length " + std::to_string(D) + + ", got length=" + std::to_string(shape.size()) + ); + } + Coordinate coordinate{}; + std::copy(shape.begin(), shape.end(), coordinate.begin()); + return coordinate; + } + + static std::uint64_t checked_multiply( + const std::uint64_t a, + const std::uint64_t b, + const char *name + ) { + if (a != 0 && b > std::numeric_limits::max() / a) { + throw std::overflow_error(std::string(name) + " exceeds uint64 range"); + } + return a * b; + } + + static std::uint64_t checked_node_count(const Coordinate &shape) { + std::uint64_t count = 1; + for (std::size_t axis = 0; axis < D; ++axis) { + if (shape[axis] == 0) { + throw std::invalid_argument("shape dimensions must be greater than zero"); + } + count = checked_multiply(count, shape[axis], "number_of_nodes"); + } + return count; + } + + static Coordinate compute_strides(const Coordinate &shape) { + Coordinate strides{}; + strides[D - 1] = 1; + for (std::size_t axis = D - 1; axis > 0; --axis) { + strides[axis - 1] = checked_multiply(strides[axis], shape[axis], "stride"); + } + return strides; + } + + static std::uint64_t axis_edge_count(const Coordinate &shape, const std::size_t axis) { + if (shape[axis] <= 1) { + return 0; + } + std::uint64_t count = shape[axis] - 1; + for (std::size_t other_axis = 0; other_axis < D; ++other_axis) { + if (other_axis != axis) { + count = checked_multiply(count, shape[other_axis], "number_of_edges"); + } + } + return count; + } + + static std::uint64_t checked_edge_count(const Coordinate &shape) { + std::uint64_t count = 0; + for (std::size_t axis = 0; axis < D; ++axis) { + const auto axis_count = axis_edge_count(shape, axis); + if (axis_count > std::numeric_limits::max() - count) { + throw std::overflow_error("number_of_edges exceeds uint64 range"); + } + count += axis_count; + } + return count; + } + + static std::array compute_edge_offsets(const Coordinate &shape) { + std::array offsets{}; + for (std::size_t axis = 0; axis < D; ++axis) { + offsets[axis + 1] = offsets[axis] + axis_edge_count(shape, axis); + } + return offsets; + } + + static std::array compute_edge_strides(const Coordinate &shape) { + std::array strides{}; + for (std::size_t axis = 0; axis < D; ++axis) { + auto edge_shape = shape; + --edge_shape[axis]; + strides[axis] = compute_strides(edge_shape); + } + return strides; + } + + [[nodiscard]] Coordinate edge_shape(const std::size_t axis) const { + auto result = shape_; + --result[axis]; + return result; + } + + [[nodiscard]] Coordinate edge_pivot_coordinates( + const std::size_t axis, + std::uint64_t local_edge + ) const { + const auto shape = edge_shape(axis); + Coordinate coordinate{}; + for (std::size_t coord_axis = 0; coord_axis < D; ++coord_axis) { + coordinate[coord_axis] = local_edge / edge_strides_[axis][coord_axis]; + local_edge -= coordinate[coord_axis] * edge_strides_[axis][coord_axis]; + } + return coordinate; + } + + // Walk a sub-shape of the node grid in C-order and invoke `callback(flat)` + // for every node id in it. Flat node ids are computed incrementally by + // addition only — no per-step divisions, no per-step bounds checks. The + // template recursion lets the compiler unroll the entire nest for D = 2/3. + template + static void enumerate_subshape_in_c_order( + const Coordinate &subshape, + const Coordinate &node_strides, + std::uint64_t base_flat, + Callback &&callback + ) { + if constexpr (Axis == D) { + callback(base_flat); + } else { + std::uint64_t flat = base_flat; + for (std::uint64_t i = 0; i < subshape[Axis]; ++i) { + enumerate_subshape_in_c_order( + subshape, node_strides, flat, callback + ); + flat += node_strides[Axis]; + } + } + } + + void build_edges() { + // Emit every edge into `edges_` in canonical order (axis-major, + // C-order within each axis). Sequential `emplace_back` into a + // single pre-reserved vector is cache-friendly. `edges_` was + // already reserved by the `UndirectedGraph` ctor. + // + // We deliberately do NOT call `rebuild_adjacency_from_edges()` + // here. The dominant grid-graph workflow (build graph → compute + // edge features via `uv_ids()`) never touches adjacency, so the + // ~400 ms rebuild on a 12 M-edge 3D grid is pure waste in that + // case. Callers that need adjacency (BFS, connected components, + // `extract_subgraph_from_nodes`) trigger a lazy rebuild via the + // base-class `node_adjacency` path on first use, paying the cost + // exactly once. Multi-threaded readers should call `freeze()` on + // the construction thread before fan-out — the same convention + // any insert-built `UndirectedGraph` already follows. + auto &edges = access_edges(); + for (std::size_t axis = 0; axis < D; ++axis) { + auto pivot_shape = shape_; + --pivot_shape[axis]; + const auto axis_step = strides_[axis]; + enumerate_subshape_in_c_order<0>( + pivot_shape, strides_, 0, + [&edges, axis_step](const std::uint64_t u) { + edges.emplace_back(u, u + axis_step); + } + ); + } + } + + Coordinate shape_{}; + Coordinate strides_{}; + std::array edge_offsets_{}; + std::array edge_strides_{}; +}; + +} // namespace bioimage_cpp::graph diff --git a/include/bioimage_cpp/graph/undirected_graph.hxx b/include/bioimage_cpp/graph/undirected_graph.hxx index 7aa5b20..8eb8f25 100644 --- a/include/bioimage_cpp/graph/undirected_graph.hxx +++ b/include/bioimage_cpp/graph/undirected_graph.hxx @@ -5,8 +5,11 @@ #include #include #include +#include +#include #include #include +#include #include #include #include @@ -19,34 +22,72 @@ struct Adjacency { std::uint64_t edge; }; +// Undirected graph storing adjacency in CSR-style (compressed sparse row) +// layout: a single contiguous `adjacency_data_` array of `2 * E` `Adjacency` +// entries, plus an `adjacency_offsets_` array of length `N + 1` giving the +// start of each node's adjacency slice. This replaces the previous +// `std::vector>` representation, which paid for one +// allocator call per node and scattered the per-node buffers throughout the +// heap — the dominant cost in `GridGraph` construction on large 3D problems. +// +// CSR is rebuilt lazily: incremental `insert_edge*` only appends to `edges_` +// (and `edge_lookup_`) and marks the adjacency dirty; the first subsequent +// `node_adjacency` read triggers a single bulk rebuild. Bulk construction +// paths (`from_sorted_unique_edges`, subclass `build_edges` overrides) call +// `rebuild_adjacency_from_edges()` explicitly, which keeps reads cheap and +// thread-safe. +// +// Thread safety: as long as a graph is "frozen" before being shared with +// reader threads (no concurrent inserts, no first-read-from-dirty-state +// race), reads of `node_adjacency` are safe to share across threads. The +// lazy rebuild is not internally synchronized — call +// `rebuild_adjacency_from_edges()` once on the construction thread before +// fan-out if you built the graph via `insert_edge*`. class UndirectedGraph { public: using NodeId = std::uint64_t; using EdgeId = std::uint64_t; using Edge = detail::Edge; - using AdjacencyList = std::vector; + // Non-owning view over a node's adjacency slice in the CSR buffer. + // Backward-compatible with the previous `std::vector` type + // for the uses we have today (range-for and `.size()`); not assignable + // and not extendable. + using AdjacencyList = std::span; explicit UndirectedGraph( const NodeId number_of_nodes = 0, const EdgeId reserve_number_of_edges = 0 ) - : number_of_nodes_(number_of_nodes), - adjacency_(static_cast(number_of_nodes)) { - edges_.reserve(static_cast(reserve_number_of_edges)); + : UndirectedGraph(number_of_nodes, reserve_number_of_edges, reserve_number_of_edges) { } virtual ~UndirectedGraph() = default; + // CSR data lives in a `unique_ptr`, so the type is move-only. + // The user-declared destructor above suppresses implicit move generation, + // so we re-default the moves explicitly. Copies are deleted on purpose: + // the previous vector-of-vectors layout was implicitly copyable but every + // such copy paid for a deep clone of millions of small adjacency vectors, + // and no caller in this codebase actually needs that. Add an explicit + // deep-copy method if one becomes necessary. + UndirectedGraph(const UndirectedGraph &) = delete; + UndirectedGraph &operator=(const UndirectedGraph &) = delete; + UndirectedGraph(UndirectedGraph &&) noexcept = default; + UndirectedGraph &operator=(UndirectedGraph &&) noexcept = default; + void assign( const NodeId number_of_nodes = 0, const EdgeId reserve_number_of_edges = 0 ) { number_of_nodes_ = number_of_nodes; edges_.clear(); - edge_lookup_.clear(); - adjacency_.clear(); - adjacency_.resize(static_cast(number_of_nodes)); edges_.reserve(static_cast(reserve_number_of_edges)); + edge_lookup_.clear(); + edge_lookup_.reserve(static_cast(reserve_number_of_edges)); + adjacency_offsets_.clear(); + adjacency_data_.reset(); + adjacency_data_size_ = 0; + adjacency_dirty_ = true; } [[nodiscard]] NodeId number_of_nodes() const { @@ -98,12 +139,18 @@ public: return edges_; } - [[nodiscard]] const AdjacencyList &node_adjacency(const NodeId node) const { + [[nodiscard]] AdjacencyList node_adjacency(const NodeId node) const { validate_node(node); - return adjacency_[static_cast(node)]; + ensure_adjacency_built(); + const auto begin = adjacency_offsets_[static_cast(node)]; + const auto end = adjacency_offsets_[static_cast(node) + 1]; + return AdjacencyList( + adjacency_data_.get() + begin, + static_cast(end - begin) + ); } - EdgeId insert_edge(const NodeId u, const NodeId v) { + virtual EdgeId insert_edge(const NodeId u, const NodeId v) { validate_node(u); validate_node(v); if (u == v) { @@ -118,7 +165,7 @@ public: return insert_new_edge(key.first, key.second); } - [[nodiscard]] std::int64_t find_edge(const NodeId u, const NodeId v) const { + [[nodiscard]] virtual std::int64_t find_edge(const NodeId u, const NodeId v) const { validate_node(u); validate_node(v); if (u == v) { @@ -137,6 +184,47 @@ public: return 2 + 2 * number_of_edges(); } + // Force the CSR adjacency to be rebuilt if it's currently stale. Use this + // after a batch of `insert_edge*` calls, before handing the graph to + // multiple reader threads or before holding a span returned from + // `node_adjacency` across other graph operations. Lazy rebuild from a + // dirty state happens via `mutable` writes inside a `const` method and is + // not internally synchronized; once a graph is "frozen" via `freeze` + // (or via an explicit `rebuild_adjacency_from_edges` from a subclass + // constructor) it is safe to share by `const&` across threads. + void freeze() const { + ensure_adjacency_built(); + } + + // Explicit deep copy. The previous vector-of-vectors layout made + // `UndirectedGraph` implicitly copyable; switching to CSR with a + // `unique_ptr` buffer made the class move-only, so callers that need + // a clone use this method. If the source graph's CSR is already built + // we copy the buffer so the clone doesn't pay another rebuild on + // first read. + [[nodiscard]] UndirectedGraph clone() const { + UndirectedGraph copy( + number_of_nodes_, + static_cast(edges_.size()), + edge_lookup_.empty() ? 0 : static_cast(edges_.size()) + ); + copy.edges_ = edges_; + copy.edge_lookup_ = edge_lookup_; + if (adjacency_dirty_) { + copy.adjacency_dirty_ = true; + } else { + copy.adjacency_offsets_ = adjacency_offsets_; + copy.adjacency_data_size_ = adjacency_data_size_; + copy.adjacency_data_ = + std::make_unique_for_overwrite(adjacency_data_size_); + std::copy_n( + adjacency_data_.get(), adjacency_data_size_, copy.adjacency_data_.get() + ); + copy.adjacency_dirty_ = false; + } + return copy; + } + // Fast construction from a pre-sorted, deduplicated edge list. // // Precondition: `edges` is sorted ascending by `(u, v)` with `u < v` in @@ -155,36 +243,21 @@ public: std::vector edges, const bool populate_lookup = true ) { - UndirectedGraph graph(number_of_nodes, static_cast(edges.size())); - - std::vector degree(static_cast(number_of_nodes), 0); - for (const auto &edge : edges) { - ++degree[static_cast(edge.first)]; - ++degree[static_cast(edge.second)]; - } - for (NodeId node = 0; node < number_of_nodes; ++node) { - graph.adjacency_[static_cast(node)].reserve( - degree[static_cast(node)] - ); - } - + const auto n_edges = static_cast(edges.size()); + UndirectedGraph graph( + number_of_nodes, + n_edges, + populate_lookup ? n_edges : 0 + ); + graph.edges_ = std::move(edges); + graph.rebuild_adjacency_from_edges(); if (populate_lookup) { - graph.edge_lookup_.reserve(edges.size()); - } - for (std::size_t index = 0; index < edges.size(); ++index) { - const auto &edge = edges[index]; - const auto edge_id = static_cast(index); - graph.adjacency_[static_cast(edge.first)].push_back( - Adjacency{edge.second, edge_id} - ); - graph.adjacency_[static_cast(edge.second)].push_back( - Adjacency{edge.first, edge_id} - ); - if (populate_lookup) { - graph.edge_lookup_.emplace(edge, edge_id); + for (std::size_t index = 0; index < graph.edges_.size(); ++index) { + graph.edge_lookup_.emplace( + graph.edges_[index], static_cast(index) + ); } } - graph.edges_ = std::move(edges); return graph; } @@ -216,15 +289,98 @@ public: } protected: + // Internal constructor that lets subclasses opt out of pre-reserving the + // `(u, v) -> edge_id` hash map. Subclasses (e.g. `GridGraph`) that build + // their edges analytically and never populate `edge_lookup_` for their + // intrinsic edges should pass `reserve_lookup = 0` so we don't allocate + // millions of empty hash buckets that are never written to. + UndirectedGraph( + const NodeId number_of_nodes, + const EdgeId reserve_edges, + const EdgeId reserve_lookup + ) + : number_of_nodes_(number_of_nodes) { + edges_.reserve(static_cast(reserve_edges)); + edge_lookup_.reserve(static_cast(reserve_lookup)); + } + EdgeId insert_new_edge(const NodeId u, const NodeId v) { const auto edge = static_cast(edges_.size()); edges_.emplace_back(u, v); edge_lookup_.emplace(Edge{u, v}, edge); - adjacency_[static_cast(u)].push_back(Adjacency{v, edge}); - adjacency_[static_cast(v)].push_back(Adjacency{u, edge}); + adjacency_dirty_ = true; return edge; } + EdgeId insert_new_edge_without_lookup(const NodeId u, const NodeId v) { + const auto edge = static_cast(edges_.size()); + edges_.emplace_back(u, v); + adjacency_dirty_ = true; + return edge; + } + + // Mutable access to the underlying edge list for subclasses that emit + // edges in bulk. Marks adjacency dirty; the next `node_adjacency` read + // (or an explicit `rebuild_adjacency_from_edges` call) will rebuild + // the CSR buffers. + std::vector &access_edges() { + adjacency_dirty_ = true; + return edges_; + } + + // Build the CSR adjacency from `edges_`. After this call, + // `adjacency_offsets_` has length `N + 1`, `adjacency_data_` has length + // `2 * E`, and the slice `adjacency_data_[offsets[u]:offsets[u+1]]` is + // exactly the adjacency of node `u`. + // + // Builds in three sequential passes: + // 1. Count degree per node (one pass over `edges_`). + // 2. Prefix-sum into `adjacency_offsets_` (one pass over nodes). + // 3. Place each edge's two adjacency entries via a fill cursor copied + // from the offsets (one pass over `edges_`). + // + // All writes in pass (3) land in a single contiguous buffer, so cache + // behavior is good; there are no per-node allocator calls. This is + // typically 4-5x faster than the previous vector-of-vectors fill on + // large grids. + void rebuild_adjacency_from_edges() { + const auto n_nodes = static_cast(number_of_nodes_); + const auto data_size = 2 * edges_.size(); + // Pass 1: count per-node degree into the offsets buffer (shifted by 1). + adjacency_offsets_.assign(n_nodes + 1, 0); + for (const auto &edge : edges_) { + ++adjacency_offsets_[static_cast(edge.first) + 1]; + ++adjacency_offsets_[static_cast(edge.second) + 1]; + } + // Pass 2: inclusive prefix sum turns degree-counts into slice starts. + for (std::size_t node = 1; node <= n_nodes; ++node) { + adjacency_offsets_[node] += adjacency_offsets_[node - 1]; + } + // Pass 3: allocate the contiguous adjacency buffer. + // `make_unique_for_overwrite` leaves trivially-default-constructible + // `Adjacency` elements UNINITIALIZED — every slot is overwritten in + // the fill pass below, so the zero-init that `std::vector::resize` + // would do is pure overhead. On a 12M-edge grid this skips ~125 ms + // of memset, though the underlying page-fault cost shifts to the + // fill pass; net win is modest but real. + static_assert(std::is_trivially_default_constructible_v); + adjacency_data_ = std::make_unique_for_overwrite(data_size); + adjacency_data_size_ = data_size; + // Pass 4: `cursor[u]` is the next free slot in node `u`'s adjacency + // range. Initialized from the slice starts and incremented per write. + std::vector cursor(adjacency_offsets_); + Adjacency *const data = adjacency_data_.get(); + for (std::size_t index = 0; index < edges_.size(); ++index) { + const auto &edge = edges_[index]; + const auto edge_id = static_cast(index); + data[cursor[static_cast(edge.first)]++] = + Adjacency{edge.second, edge_id}; + data[cursor[static_cast(edge.second)]++] = + Adjacency{edge.first, edge_id}; + } + adjacency_dirty_ = false; + } + void validate_node(const NodeId node) const { if (node >= number_of_nodes_) { throw std::out_of_range( @@ -246,9 +402,29 @@ protected: } private: + void ensure_adjacency_built() const { + if (adjacency_dirty_) { + const_cast(this)->rebuild_adjacency_from_edges(); + } + } + NodeId number_of_nodes_; std::vector edges_; - std::vector adjacency_; + // CSR adjacency. `adjacency_offsets_` has length `N + 1`; + // `adjacency_data_` has length `2 * number_of_edges()` after a rebuild. + // The `mutable` qualifiers allow the lazy rebuild from a const reader + // path — see `ensure_adjacency_built`. Writers (insert / bulk-edit + // helpers) set `adjacency_dirty_ = true` and the rebuild happens on the + // next read. + mutable std::vector adjacency_offsets_; + // Heap-allocated CSR data buffer. Using `unique_ptr` instead of + // `std::vector` lets `make_unique_for_overwrite` skip the + // zero-initialization that `vector::resize` would force — a substantial + // saving for 12 M-edge graphs where the buffer is ~384 MB and every + // slot is overwritten in the fill pass anyway. + mutable std::unique_ptr adjacency_data_; + mutable std::size_t adjacency_data_size_ = 0; + mutable bool adjacency_dirty_ = true; std::unordered_map edge_lookup_; }; diff --git a/src/bindings/graph.cxx b/src/bindings/graph.cxx index 62653e5..82d516b 100644 --- a/src/bindings/graph.cxx +++ b/src/bindings/graph.cxx @@ -5,6 +5,8 @@ #include "bioimage_cpp/graph/connected_components.hxx" #include "bioimage_cpp/graph/edge_weighted_watershed.hxx" #include "bioimage_cpp/graph/feature_accumulation.hxx" +#include "bioimage_cpp/graph/grid_features.hxx" +#include "bioimage_cpp/graph/grid_graph.hxx" #include "bioimage_cpp/graph/lifted_from_affinities.hxx" #include "bioimage_cpp/graph/lifted_multicut.hxx" #include "bioimage_cpp/graph/multicut.hxx" @@ -39,17 +41,28 @@ namespace bioimage_cpp::bindings { namespace { using Graph = graph::UndirectedGraph; +using GridGraph2D = graph::GridGraph<2>; +using GridGraph3D = graph::GridGraph<3>; using Rag = graph::RegionAdjacencyGraph; +using UInt8Array = nb::ndarray; using UInt64Array = nb::ndarray; using ConstUInt8Array = nb::ndarray; using ConstUInt64Array = nb::ndarray; using Int64Array = nb::ndarray; +using ConstInt64Array = nb::ndarray; using DoubleArray = nb::ndarray; using ConstDoubleArray = nb::ndarray; +using FloatArray = nb::ndarray; +using ConstFloatArray = nb::ndarray; template using LabelArray = nb::ndarray; +template +using FloatingArray = nb::ndarray; +template +using ConstFloatingArray = nb::ndarray; + void require_uv_array(const ConstUInt64Array &uvs, const char *argument_name) { if (uvs.ndim() != 2 || uvs.shape(1) != 2) { throw std::invalid_argument( @@ -68,6 +81,16 @@ UInt64Array make_uint64_array(const std::vector &shape) { return UInt64Array(data, shape.size(), shape.data(), owner); } +UInt8Array make_uint8_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new std::uint8_t[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return UInt8Array(data, shape.size(), shape.data(), owner); +} + Int64Array make_int64_array(const std::vector &shape) { std::size_t size = 1; for (const auto axis_size : shape) { @@ -88,12 +111,236 @@ DoubleArray make_double_array(const std::vector &shape) { return DoubleArray(data, shape.size(), shape.data(), owner); } +template +FloatingArray make_floating_array(const std::vector &shape) { + std::size_t size = 1; + for (const auto axis_size : shape) { + size *= axis_size; + } + auto *data = new T[size](); + nb::capsule owner(data, [](void *p) noexcept { delete[] static_cast(p); }); + return FloatingArray(data, shape.size(), shape.data(), owner); +} + +template +FloatingArray vector_to_floating_array(const std::vector &values) { + auto result = make_floating_array({values.size()}); + std::copy(values.begin(), values.end(), result.data()); + return result; +} + +template +std::vector const_floating_shape(ConstFloatingArray array) { + std::vector shape(array.ndim()); + for (std::size_t axis = 0; axis < array.ndim(); ++axis) { + shape[axis] = static_cast(array.shape(axis)); + } + return shape; +} + +std::vector const_double_shape(ConstDoubleArray array) { + std::vector shape(array.ndim()); + for (std::size_t axis = 0; axis < array.ndim(); ++axis) { + shape[axis] = static_cast(array.shape(axis)); + } + return shape; +} + UInt64Array vector_to_uint64_array(const std::vector &values) { auto result = make_uint64_array({values.size()}); std::copy(values.begin(), values.end(), result.data()); return result; } +UInt8Array vector_to_uint8_array(const std::vector &values) { + auto result = make_uint8_array({values.size()}); + std::copy(values.begin(), values.end(), result.data()); + return result; +} + +DoubleArray vector_to_double_array(const std::vector &values) { + auto result = make_double_array({values.size()}); + std::copy(values.begin(), values.end(), result.data()); + return result; +} + +UInt64Array edges_to_uv_array(const std::vector &edges) { + auto result = make_uint64_array({edges.size(), 2}); + auto *data = result.data(); + for (std::size_t index = 0; index < edges.size(); ++index) { + data[2 * index] = edges[index].first; + data[2 * index + 1] = edges[index].second; + } + return result; +} + +template +std::vector coordinate_to_vector( + const typename graph::GridGraph::Coordinate &coordinate +) { + return std::vector(coordinate.begin(), coordinate.end()); +} + +template +typename graph::GridGraph::Coordinate coordinate_from_array( + ConstUInt64Array coordinate, + const char *argument_name +) { + if (coordinate.ndim() != 1 || coordinate.shape(0) != D) { + throw std::invalid_argument( + std::string(argument_name) + " must be a 1D uint64 array of length " + + std::to_string(D) + ); + } + typename graph::GridGraph::Coordinate result{}; + std::copy(coordinate.data(), coordinate.data() + D, result.begin()); + return result; +} + +template +std::array offset_from_array( + ConstInt64Array offset, + const char *argument_name +) { + if (offset.ndim() != 1 || offset.shape(0) != D) { + throw std::invalid_argument( + std::string(argument_name) + " must be a 1D int64 array of length " + + std::to_string(D) + ); + } + std::array result{}; + std::copy(offset.data(), offset.data() + D, result.begin()); + return result; +} + +template +std::uint64_t grid_node_id( + const graph::GridGraph &graph, + ConstUInt64Array coordinate +) { + return graph.node_id(coordinate_from_array(coordinate, "coordinate")); +} + +template +UInt64Array grid_coordinates(const graph::GridGraph &graph, const std::uint64_t node) { + return vector_to_uint64_array(coordinate_to_vector(graph.coordinates(node))); +} + +template +std::vector grid_shape(const graph::GridGraph &graph) { + return coordinate_to_vector(graph.shape()); +} + +template +std::vector grid_strides(const graph::GridGraph &graph) { + return coordinate_to_vector(graph.strides()); +} + +template +UInt64Array grid_edge_coordinates( + const graph::GridGraph &graph, + const std::uint64_t edge +) { + return vector_to_uint64_array(coordinate_to_vector(graph.edge_coordinates(edge).first)); +} + +template +std::int64_t grid_offset_target( + const graph::GridGraph &graph, + const std::uint64_t node, + ConstInt64Array offset +) { + std::uint64_t target = 0; + if (!graph.valid_offset_target(node, offset_from_array(offset, "offset"), target)) { + return -1; + } + return static_cast(target); +} + +template +FloatingArray grid_boundary_features_t( + const graph::GridGraph &graph, + ConstFloatingArray boundary_map +) { + auto result = make_floating_array({static_cast(graph.number_of_edges())}); + ConstArrayView boundary_view{ + boundary_map.data(), + const_floating_shape(boundary_map), + {}, + }; + ArrayView out_view{ + result.data(), + {static_cast(graph.number_of_edges())}, + {}, + }; + + nb::gil_scoped_release release; + graph::grid_boundary_features(graph, boundary_view, out_view); + return result; +} + +template +nb::tuple grid_affinity_features_t( + const graph::GridGraph &graph, + ConstFloatingArray affinities, + const std::vector> &offsets +) { + auto weights = make_floating_array({static_cast(graph.number_of_edges())}); + auto valid_edges = make_uint8_array({static_cast(graph.number_of_edges())}); + ConstArrayView affinities_view{ + affinities.data(), + const_floating_shape(affinities), + {}, + }; + ArrayView weights_view{ + weights.data(), + {static_cast(graph.number_of_edges())}, + {}, + }; + ArrayView valid_view{ + valid_edges.data(), + {static_cast(graph.number_of_edges())}, + {}, + }; + + { + nb::gil_scoped_release release; + graph::grid_local_affinity_features( + graph, affinities_view, offsets, weights_view, valid_view + ); + } + return nb::make_tuple(weights, valid_edges); +} + +template +nb::tuple grid_affinity_features_with_lifted_t( + const graph::GridGraph &graph, + ConstFloatingArray affinities, + const std::vector> &offsets +) { + ConstArrayView affinities_view{ + affinities.data(), + const_floating_shape(affinities), + {}, + }; + + graph::GridLiftedAffinityFeatures features; + { + nb::gil_scoped_release release; + features = graph::grid_affinity_features_with_lifted( + graph, affinities_view, offsets + ); + } + + return nb::make_tuple( + vector_to_floating_array(features.local_weights), + vector_to_uint8_array(features.valid_local_edges), + edges_to_uv_array(features.lifted_uvs), + vector_to_floating_array(features.lifted_weights), + vector_to_uint64_array(features.lifted_offset_ids) + ); +} + std::vector double_array_to_vector( ConstDoubleArray array, const char *argument_name, @@ -873,6 +1120,8 @@ void bind_graph(nb::module_ &m) { nb::arg("nodes") ) .def("edges_from_node_list", &graph_edges_from_node_list, nb::arg("nodes")) + .def("freeze", &Graph::freeze) + .def("clone", &Graph::clone) .def_static( "from_edges", &graph_from_edges, @@ -908,9 +1157,92 @@ void bind_graph(nb::module_ &m) { ) .def("edgesFromNodeList", &graph_edges_from_node_list, nb::arg("nodes")); + nb::class_(m, "GridGraph2D") + .def(nb::init &>(), nb::arg("shape")) + .def_prop_ro("shape", &grid_shape<2>) + .def_prop_ro("strides", &grid_strides<2>) + .def_prop_ro("ndim", &GridGraph2D::ndim) + .def("node_id", &grid_node_id<2>, nb::arg("coordinate")) + .def("coordinates", &grid_coordinates<2>, nb::arg("node")) + .def("edge_axis", &GridGraph2D::edge_axis, nb::arg("edge")) + .def("edge_coordinates", &grid_edge_coordinates<2>, nb::arg("edge")) + .def("offset_target", &grid_offset_target<2>, nb::arg("node"), nb::arg("offset")) + .def_prop_ro("numberOfDimensions", &GridGraph2D::ndim) + .def("nodeId", &grid_node_id<2>, nb::arg("coordinate")) + .def("edgeAxis", &GridGraph2D::edge_axis, nb::arg("edge")) + .def("edgeCoordinates", &grid_edge_coordinates<2>, nb::arg("edge")) + .def("offsetTarget", &grid_offset_target<2>, nb::arg("node"), nb::arg("offset")); + + nb::class_(m, "GridGraph3D") + .def(nb::init &>(), nb::arg("shape")) + .def_prop_ro("shape", &grid_shape<3>) + .def_prop_ro("strides", &grid_strides<3>) + .def_prop_ro("ndim", &GridGraph3D::ndim) + .def("node_id", &grid_node_id<3>, nb::arg("coordinate")) + .def("coordinates", &grid_coordinates<3>, nb::arg("node")) + .def("edge_axis", &GridGraph3D::edge_axis, nb::arg("edge")) + .def("edge_coordinates", &grid_edge_coordinates<3>, nb::arg("edge")) + .def("offset_target", &grid_offset_target<3>, nb::arg("node"), nb::arg("offset")) + .def_prop_ro("numberOfDimensions", &GridGraph3D::ndim) + .def("nodeId", &grid_node_id<3>, nb::arg("coordinate")) + .def("edgeAxis", &GridGraph3D::edge_axis, nb::arg("edge")) + .def("edgeCoordinates", &grid_edge_coordinates<3>, nb::arg("edge")) + .def("offsetTarget", &grid_offset_target<3>, nb::arg("node"), nb::arg("offset")); + nb::class_(m, "RegionAdjacencyGraph") .def_prop_ro("shape", &Rag::shape); + const auto register_grid_boundary = [&m](const char *name) { + m.def( + name, + &grid_boundary_features_t, + nb::arg("graph"), + nb::arg("boundary_map") + ); + }; + register_grid_boundary.operator()("_grid_boundary_features_2d_float32"); + register_grid_boundary.operator()("_grid_boundary_features_2d_float64"); + register_grid_boundary.operator()("_grid_boundary_features_3d_float32"); + register_grid_boundary.operator()("_grid_boundary_features_3d_float64"); + + const auto register_grid_affinity = [&m](const char *name) { + m.def( + name, + &grid_affinity_features_t, + nb::arg("graph"), + nb::arg("affinities"), + nb::arg("offsets") + ); + }; + register_grid_affinity.operator()("_grid_affinity_features_2d_float32"); + register_grid_affinity.operator()("_grid_affinity_features_2d_float64"); + register_grid_affinity.operator()("_grid_affinity_features_3d_float32"); + register_grid_affinity.operator()("_grid_affinity_features_3d_float64"); + + const auto register_grid_affinity_lifted = [&m]( + const char *name + ) { + m.def( + name, + &grid_affinity_features_with_lifted_t, + nb::arg("graph"), + nb::arg("affinities"), + nb::arg("offsets") + ); + }; + register_grid_affinity_lifted.operator()( + "_grid_affinity_features_with_lifted_2d_float32" + ); + register_grid_affinity_lifted.operator()( + "_grid_affinity_features_with_lifted_2d_float64" + ); + register_grid_affinity_lifted.operator()( + "_grid_affinity_features_with_lifted_3d_float32" + ); + register_grid_affinity_lifted.operator()( + "_grid_affinity_features_with_lifted_3d_float64" + ); + m.def( "_breadth_first_search", &graph_breadth_first_search, diff --git a/src/bioimage_cpp/graph/__init__.py b/src/bioimage_cpp/graph/__init__.py index 3ebd6b1..bdc310a 100644 --- a/src/bioimage_cpp/graph/__init__.py +++ b/src/bioimage_cpp/graph/__init__.py @@ -139,6 +139,77 @@ def deserialize(cls, serialization): return cls.from_edges(number_of_nodes, uvs) +def _as_shape(shape, ndim: int, name: str = "shape") -> list[int]: + array = np.asarray(shape) + if array.ndim != 1 or array.shape[0] != ndim: + raise ValueError(f"{name} must be a 1D sequence of length {ndim}") + if not np.issubdtype(array.dtype, np.integer): + raise TypeError(f"{name} must contain integers") + if np.any(array <= 0): + raise ValueError(f"{name} dimensions must be greater than zero") + return [int(axis_size) for axis_size in array] + + +class GridGraph2D(_core.GridGraph2D): + """Regular 2D nearest-neighbor grid graph. + + Nodes use C-order ids for ``shape=(y, x)``. Edge ids are deterministic: + all y-axis edges first, then all x-axis edges. + """ + + def __init__(self, shape): + super().__init__(_as_shape(shape, 2)) + + def node_id(self, coordinate): + return super().node_id(_as_coordinate_array(coordinate, 2, "coordinate")) + + def nodeId(self, coordinate): + return self.node_id(coordinate) + + def offset_target(self, node: int, offset): + return super().offset_target(int(node), _as_offset_array(offset, 2, "offset")) + + def offsetTarget(self, node: int, offset): + return self.offset_target(node, offset) + + +class GridGraph3D(_core.GridGraph3D): + """Regular 3D nearest-neighbor grid graph. + + Nodes use C-order ids for ``shape=(z, y, x)``. Edge ids are deterministic: + all z-axis edges first, then y-axis edges, then x-axis edges. + """ + + def __init__(self, shape): + super().__init__(_as_shape(shape, 3)) + + def node_id(self, coordinate): + return super().node_id(_as_coordinate_array(coordinate, 3, "coordinate")) + + def nodeId(self, coordinate): + return self.node_id(coordinate) + + def offset_target(self, node: int, offset): + return super().offset_target(int(node), _as_offset_array(offset, 3, "offset")) + + def offsetTarget(self, node: int, offset): + return self.offset_target(node, offset) + + +def _as_coordinate_array(coordinate, ndim: int, name: str) -> np.ndarray: + array = np.asarray(coordinate, dtype=np.uint64) + if array.ndim != 1 or array.shape[0] != ndim: + raise ValueError(f"{name} must be a 1D sequence of length {ndim}") + return np.ascontiguousarray(array) + + +def _as_offset_array(offset, ndim: int, name: str) -> np.ndarray: + array = np.asarray(offset, dtype=np.int64) + if array.ndim != 1 or array.shape[0] != ndim: + raise ValueError(f"{name} must be a 1D sequence of length {ndim}") + return np.ascontiguousarray(array) + + def _as_uv_array(uvs, name: str) -> np.ndarray: array = np.asarray(uvs, dtype=np.uint64) if array.ndim != 2 or array.shape[1] != 2: @@ -165,7 +236,7 @@ def _as_serialization_array(serialization) -> np.ndarray: return np.ascontiguousarray(array) -def _copy_graph(graph: UndirectedGraph | RegionAdjacencyGraph) -> _core.UndirectedGraph: +def _copy_graph(graph) -> _core.UndirectedGraph: # `uv_ids()` always returns a unique list (graphs deduplicate on insert), # so we can use the bulk constructor that skips per-edge hash dedup — # significantly faster than `insert_edges` for large graphs. The result @@ -228,6 +299,169 @@ def undirected_graph(number_of_nodes: int) -> UndirectedGraph: return UndirectedGraph(number_of_nodes) +def grid_graph(shape): + """Create a regular 2D or 3D nearest-neighbor grid graph.""" + ndim = np.asarray(shape).ndim + if ndim != 1: + raise ValueError("shape must be a 1D sequence") + n_axes = len(shape) + if n_axes == 2: + return GridGraph2D(shape) + if n_axes == 3: + return GridGraph3D(shape) + raise ValueError(f"shape must have length 2 or 3, got length={n_axes}") + + +def _grid_ndim(graph) -> int: + if isinstance(graph, GridGraph2D): + return 2 + if isinstance(graph, GridGraph3D): + return 3 + raise TypeError("graph must be a GridGraph2D or GridGraph3D") + + +def _grid_shape(graph) -> tuple[int, ...]: + return tuple(int(size) for size in graph.shape) + + +_GRID_FLOAT_DTYPES = (np.dtype(np.float32), np.dtype(np.float64)) + + +def _as_grid_data(values, graph, name: str, *, with_channels: bool) -> np.ndarray: + array = np.asarray(values) + if array.dtype not in _GRID_FLOAT_DTYPES: + # Integer / non-float input falls back to float64 — the previous default. + # float32 and float64 inputs are passed through end-to-end, no copy. + array = array.astype(np.float64) + shape = _grid_shape(graph) + if with_channels: + if array.ndim != len(shape) + 1 or array.shape[1:] != shape: + raise ValueError( + f"{name} must have shape (channels, *graph.shape), got " + f"shape={array.shape}, graph shape={shape}" + ) + elif array.shape != shape: + raise ValueError( + f"{name} shape must match graph shape, got " + f"{name} shape={array.shape}, graph shape={shape}" + ) + return np.ascontiguousarray(array) + + +def _normalize_grid_offsets(offsets, ndim: int, n_channels: int) -> list[tuple[int, ...]]: + normalized = [tuple(int(value) for value in offset) for offset in offsets] + if len(normalized) != n_channels: + raise ValueError( + "offsets length must match affinities channel count, got " + f"offsets length={len(normalized)}, channels={n_channels}" + ) + if any(len(offset) != ndim for offset in normalized): + raise ValueError("each offset must have length matching graph ndim") + if any(all(value == 0 for value in offset) for offset in normalized): + raise ValueError("offsets must not contain the zero offset") + return normalized + + +def _grid_dtype_suffix(array: np.ndarray) -> str: + if array.dtype == np.float32: + return "float32" + return "float64" + + +_GRID_BOUNDARY_DISPATCH = { + (2, "float32"): _core._grid_boundary_features_2d_float32, + (2, "float64"): _core._grid_boundary_features_2d_float64, + (3, "float32"): _core._grid_boundary_features_3d_float32, + (3, "float64"): _core._grid_boundary_features_3d_float64, +} +_GRID_AFFINITY_DISPATCH = { + (2, "float32"): _core._grid_affinity_features_2d_float32, + (2, "float64"): _core._grid_affinity_features_2d_float64, + (3, "float32"): _core._grid_affinity_features_3d_float32, + (3, "float64"): _core._grid_affinity_features_3d_float64, +} +_GRID_AFFINITY_LIFTED_DISPATCH = { + (2, "float32"): _core._grid_affinity_features_with_lifted_2d_float32, + (2, "float64"): _core._grid_affinity_features_with_lifted_2d_float64, + (3, "float32"): _core._grid_affinity_features_with_lifted_3d_float32, + (3, "float64"): _core._grid_affinity_features_with_lifted_3d_float64, +} + + +def grid_boundary_features(graph, boundary_map) -> np.ndarray: + """Compute scalar nearest-neighbor grid edge weights from a boundary map. + + The output is a 1D array aligned to ``graph.edges()``. Output dtype matches + the input: ``float32`` and ``float64`` inputs are processed without copying, + other dtypes are promoted to ``float64``. Each edge receives the average of + the two endpoint boundary-map values. + """ + ndim = _grid_ndim(graph) + boundary_array = _as_grid_data( + boundary_map, graph, "boundary_map", with_channels=False + ) + return _GRID_BOUNDARY_DISPATCH[(ndim, _grid_dtype_suffix(boundary_array))]( + graph, boundary_array + ) + + +def grid_affinity_features(graph, affinities, offsets) -> tuple[np.ndarray, np.ndarray]: + """Map local affinity channels to grid graph edge weights. + + Only nearest-neighbor offsets with L1 norm 1 are accepted. The returned + tuple is ``(edge_weights, valid_edges)``, both aligned to ``graph.edges()``. + ``edge_weights`` has the same dtype as ``affinities`` (``float32`` or + ``float64``; other dtypes are promoted to ``float64``). ``valid_edges`` is + boolean and marks local graph edges covered by offsets. + """ + ndim = _grid_ndim(graph) + affinity_array = _as_grid_data( + affinities, graph, "affinities", with_channels=True + ) + normalized_offsets = _normalize_grid_offsets( + offsets, ndim, int(affinity_array.shape[0]) + ) + weights, valid = _GRID_AFFINITY_DISPATCH[ + (ndim, _grid_dtype_suffix(affinity_array)) + ](graph, affinity_array, normalized_offsets) + return weights, valid.astype(bool, copy=False) + + +def grid_affinity_features_with_lifted( + graph, + affinities, + offsets, +) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray]: + """Map local affinities and emit explicit long-range grid edges. + + Returns ``(local_weights, valid_local_edges, lifted_uvs, lifted_weights, + lifted_offset_ids)``. Local arrays are aligned to ``graph.edges()``. + Long-range arrays have one row/value per valid offset hit and are suitable + for lifted multicut or mutex-watershed style callers. Weight arrays share + the dtype of ``affinities`` (``float32`` or ``float64``; other dtypes are + promoted to ``float64``). + """ + ndim = _grid_ndim(graph) + affinity_array = _as_grid_data( + affinities, graph, "affinities", with_channels=True + ) + normalized_offsets = _normalize_grid_offsets( + offsets, ndim, int(affinity_array.shape[0]) + ) + local_weights, valid, lifted_uvs, lifted_weights, lifted_offset_ids = ( + _GRID_AFFINITY_LIFTED_DISPATCH[ + (ndim, _grid_dtype_suffix(affinity_array)) + ](graph, affinity_array, normalized_offsets) + ) + return ( + local_weights, + valid.astype(bool, copy=False), + lifted_uvs, + lifted_weights, + lifted_offset_ids, + ) + + RegionAdjacencyGraph = _core.RegionAdjacencyGraph @@ -1545,6 +1779,8 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "GreedyAdditiveMulticut", "GreedyAdditiveProposalGenerator", "GreedyFixationMulticut", + "GridGraph2D", + "GridGraph3D", "KernighanLinMulticut", "LiftedChainedSolvers", "LiftedGreedyAdditiveMulticut", @@ -1567,6 +1803,10 @@ def _normalize_number_of_threads(number_of_threads: int) -> int: "edge_map_features", "edge_map_features_complex", "edge_weighted_watershed", + "grid_graph", + "grid_affinity_features", + "grid_affinity_features_with_lifted", + "grid_boundary_features", "lifted_affinity_features", "lifted_affinity_features_complex", "lifted_edges_from_affinities", diff --git a/tests/graph/test_grid_affinity_multicut_integration.py b/tests/graph/test_grid_affinity_multicut_integration.py new file mode 100644 index 0000000..59962df --- /dev/null +++ b/tests/graph/test_grid_affinity_multicut_integration.py @@ -0,0 +1,89 @@ +import numpy as np + +import bioimage_cpp as bic + + +def _same_partition(labels, expected): + labels = np.asarray(labels) + expected = np.asarray(expected) + assert labels.shape == expected.shape + np.testing.assert_array_equal( + labels[:, None] == labels[None, :], + expected[:, None] == expected[None, :], + ) + + +def _node_labels_from_grid_partition(partition): + return np.asarray(partition, dtype=np.uint64).reshape(-1) + + +def _affinities_from_partition(partition, offsets, *, high=0.95, low=0.05): + partition = np.asarray(partition) + affinities = np.zeros((len(offsets),) + partition.shape, dtype=np.float64) + for channel, offset in enumerate(offsets): + offset = tuple(int(v) for v in offset) + for source in np.ndindex(partition.shape): + target = tuple(source[axis] + offset[axis] for axis in range(partition.ndim)) + if any( + target[axis] < 0 or target[axis] >= partition.shape[axis] + for axis in range(partition.ndim) + ): + continue + affinities[(channel,) + source] = ( + high if partition[source] == partition[target] else low + ) + return affinities + + +def test_grid_local_affinities_drive_multicut_partition(): + partition = np.zeros((4, 6), dtype=np.uint64) + partition[:, 3:] = 1 + graph = bic.graph.GridGraph2D(partition.shape) + offsets = [(1, 0), (0, 1)] + affinities = _affinities_from_partition(partition, offsets) + + weights, valid_edges = bic.graph.grid_affinity_features(graph, affinities, offsets) + np.testing.assert_array_equal(valid_edges, np.ones(graph.number_of_edges, dtype=bool)) + edge_costs = weights - 0.5 + + objective = bic.graph.MulticutObjective(graph, edge_costs) + labels = bic.graph.ChainedMulticutSolvers( + [ + bic.graph.GreedyAdditiveMulticut(), + bic.graph.KernighanLinMulticut(number_of_outer_iterations=5), + ] + ).optimize(objective) + + _same_partition(labels, _node_labels_from_grid_partition(partition)) + assert np.unique(labels).size == 2 + + +def test_grid_long_range_affinities_drive_lifted_multicut_partition(): + partition = np.zeros((4, 8), dtype=np.uint64) + partition[:, 4:] = 1 + graph = bic.graph.GridGraph2D(partition.shape) + offsets = [(1, 0), (0, 1), (0, 2), (0, 3)] + affinities = _affinities_from_partition(partition, offsets) + + local_weights, valid_edges, lifted_uvs, lifted_weights, lifted_offset_ids = ( + bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets) + ) + np.testing.assert_array_equal(valid_edges, np.ones(graph.number_of_edges, dtype=bool)) + assert lifted_uvs.shape[0] > 0 + assert set(np.unique(lifted_offset_ids).tolist()) == {2, 3} + + objective = bic.graph.LiftedMulticutObjective( + graph, + local_weights - 0.5, + lifted_uvs=lifted_uvs, + lifted_costs=lifted_weights - 0.5, + ) + labels = bic.graph.LiftedChainedSolvers( + [ + bic.graph.LiftedGreedyAdditiveMulticut(), + bic.graph.LiftedKernighanLinMulticut(number_of_outer_iterations=5), + ] + ).optimize(objective) + + _same_partition(labels, _node_labels_from_grid_partition(partition)) + assert np.unique(labels).size == 2 diff --git a/tests/graph/test_grid_features.py b/tests/graph/test_grid_features.py new file mode 100644 index 0000000..285db49 --- /dev/null +++ b/tests/graph/test_grid_features.py @@ -0,0 +1,203 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_grid_boundary_features_2d_align_to_grid_edges(): + graph = bic.graph.GridGraph2D((2, 3)) + boundary_map = np.array([[0.0, 2.0, 4.0], [10.0, 12.0, 14.0]]) + + weights = bic.graph.grid_boundary_features(graph, boundary_map) + + assert weights.dtype == np.float64 + np.testing.assert_allclose( + weights, + np.array([5.0, 7.0, 9.0, 1.0, 3.0, 11.0, 13.0]), + ) + + +def test_grid_affinity_features_2d_local_offsets(): + graph = bic.graph.GridGraph2D((2, 3)) + affinities = np.zeros((2, 2, 3), dtype=np.float64) + affinities[0] = np.array([[10.0, 11.0, 12.0], [0.0, 0.0, 0.0]]) + affinities[1] = np.array([[20.0, 21.0, 0.0], [23.0, 24.0, 0.0]]) + + weights, valid = bic.graph.grid_affinity_features( + graph, affinities, offsets=[(1, 0), (0, 1)] + ) + + np.testing.assert_allclose(weights, np.array([10.0, 11.0, 12.0, 20.0, 21.0, 23.0, 24.0])) + np.testing.assert_array_equal(valid, np.ones(graph.number_of_edges, dtype=bool)) + + +def test_grid_affinity_features_partial_local_offsets(): + graph = bic.graph.GridGraph2D((2, 3)) + affinities = np.arange(6, dtype=np.float64).reshape(1, 2, 3) + + weights, valid = bic.graph.grid_affinity_features( + graph, affinities, offsets=[(1, 0)] + ) + + np.testing.assert_allclose(weights, np.array([0.0, 1.0, 2.0, 0.0, 0.0, 0.0, 0.0])) + np.testing.assert_array_equal(valid, np.array([True, True, True, False, False, False, False])) + + +def test_grid_affinity_features_with_lifted_2d(): + graph = bic.graph.GridGraph2D((3, 3)) + affinities = np.zeros((2, 3, 3), dtype=np.float64) + affinities[0] = np.arange(9, dtype=np.float64).reshape(3, 3) + affinities[1] = 100.0 + np.arange(9, dtype=np.float64).reshape(3, 3) + + local_weights, valid, lifted_uvs, lifted_weights, lifted_offset_ids = ( + bic.graph.grid_affinity_features_with_lifted( + graph, affinities, offsets=[(1, 0), (0, 2)] + ) + ) + + np.testing.assert_allclose( + local_weights, + np.array([0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]), + ) + np.testing.assert_array_equal( + valid, + np.array([True, True, True, True, True, True, False, False, False, False, False, False]), + ) + np.testing.assert_array_equal( + lifted_uvs, + np.array([[0, 2], [3, 5], [6, 8]], dtype=np.uint64), + ) + np.testing.assert_allclose(lifted_weights, np.array([100.0, 103.0, 106.0])) + np.testing.assert_array_equal(lifted_offset_ids, np.array([1, 1, 1], dtype=np.uint64)) + + +def test_grid_affinity_features_with_lifted_3d(): + graph = bic.graph.GridGraph3D((2, 1, 3)) + affinities = np.zeros((2, 2, 1, 3), dtype=np.float64) + affinities[0, 0, 0, :] = [1.0, 2.0, 3.0] + affinities[1, 0, 0, 0] = 9.0 + affinities[1, 1, 0, 0] = 10.0 + + local_weights, valid, lifted_uvs, lifted_weights, lifted_offset_ids = ( + bic.graph.grid_affinity_features_with_lifted( + graph, affinities, offsets=[(1, 0, 0), (0, 0, 2)] + ) + ) + + np.testing.assert_allclose(local_weights[:3], np.array([1.0, 2.0, 3.0])) + np.testing.assert_array_equal(valid[:3], np.array([True, True, True])) + np.testing.assert_array_equal( + lifted_uvs, + np.array([[0, 2], [3, 5]], dtype=np.uint64), + ) + np.testing.assert_allclose(lifted_weights, np.array([9.0, 10.0])) + np.testing.assert_array_equal(lifted_offset_ids, np.array([1, 1], dtype=np.uint64)) + + +def test_grid_affinity_features_rejects_long_range_offsets(): + graph = bic.graph.GridGraph2D((3, 3)) + affinities = np.zeros((1, 3, 3), dtype=np.float64) + + with pytest.raises(ValueError, match="only local offsets"): + bic.graph.grid_affinity_features(graph, affinities, offsets=[(0, 2)]) + + +def test_grid_affinity_features_rejects_duplicate_local_edges(): + graph = bic.graph.GridGraph2D((2, 3)) + affinities = np.zeros((2, 2, 3), dtype=np.float64) + + with pytest.raises(ValueError, match="duplicate local"): + bic.graph.grid_affinity_features(graph, affinities, offsets=[(0, 1), (0, -1)]) + + +def test_grid_affinity_features_rejects_duplicate_lifted_edges(): + graph = bic.graph.GridGraph2D((1, 3)) + affinities = np.zeros((2, 1, 3), dtype=np.float64) + + with pytest.raises(ValueError, match="duplicate long-range"): + bic.graph.grid_affinity_features_with_lifted( + graph, affinities, offsets=[(0, 2), (0, -2)] + ) + + +def test_grid_boundary_features_preserves_float32_dtype(): + graph = bic.graph.GridGraph2D((2, 3)) + boundary_map = np.array( + [[0.0, 2.0, 4.0], [10.0, 12.0, 14.0]], dtype=np.float32 + ) + + weights = bic.graph.grid_boundary_features(graph, boundary_map) + + assert weights.dtype == np.float32 + np.testing.assert_allclose( + weights, np.array([5.0, 7.0, 9.0, 1.0, 3.0, 11.0, 13.0], dtype=np.float32) + ) + + +def test_grid_affinity_features_preserves_float32_dtype(): + graph = bic.graph.GridGraph2D((2, 3)) + affinities = np.zeros((2, 2, 3), dtype=np.float32) + affinities[0] = np.array([[10.0, 11.0, 12.0], [0.0, 0.0, 0.0]], dtype=np.float32) + affinities[1] = np.array([[20.0, 21.0, 0.0], [23.0, 24.0, 0.0]], dtype=np.float32) + + weights, valid = bic.graph.grid_affinity_features( + graph, affinities, offsets=[(1, 0), (0, 1)] + ) + + assert weights.dtype == np.float32 + np.testing.assert_allclose( + weights, + np.array([10.0, 11.0, 12.0, 20.0, 21.0, 23.0, 24.0], dtype=np.float32), + ) + np.testing.assert_array_equal(valid, np.ones(graph.number_of_edges, dtype=bool)) + + +def test_grid_affinity_features_with_lifted_preserves_float32_dtype(): + graph = bic.graph.GridGraph2D((3, 3)) + affinities = np.zeros((2, 3, 3), dtype=np.float32) + affinities[0] = np.arange(9, dtype=np.float32).reshape(3, 3) + affinities[1] = 100.0 + np.arange(9, dtype=np.float32).reshape(3, 3) + + local_weights, valid, lifted_uvs, lifted_weights, _ = ( + bic.graph.grid_affinity_features_with_lifted( + graph, affinities, offsets=[(1, 0), (0, 2)] + ) + ) + + assert local_weights.dtype == np.float32 + assert lifted_weights.dtype == np.float32 + np.testing.assert_allclose( + local_weights, + np.array( + [0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + dtype=np.float32, + ), + ) + np.testing.assert_allclose( + lifted_weights, np.array([100.0, 103.0, 106.0], dtype=np.float32) + ) + np.testing.assert_array_equal( + lifted_uvs, np.array([[0, 2], [3, 5], [6, 8]], dtype=np.uint64) + ) + + +def test_grid_features_integer_input_promotes_to_float64(): + graph = bic.graph.GridGraph2D((2, 3)) + boundary_map = np.array([[0, 2, 4], [10, 12, 14]], dtype=np.int32) + + weights = bic.graph.grid_boundary_features(graph, boundary_map) + + assert weights.dtype == np.float64 + + +def test_grid_feature_validation(): + graph = bic.graph.GridGraph2D((2, 3)) + + with pytest.raises(ValueError, match="boundary_map shape"): + bic.graph.grid_boundary_features(graph, np.zeros((3, 2))) + with pytest.raises(ValueError, match="affinities must have shape"): + bic.graph.grid_affinity_features(graph, np.zeros((2, 3)), offsets=[(0, 1)]) + with pytest.raises(ValueError, match="offsets length"): + bic.graph.grid_affinity_features(graph, np.zeros((2, 2, 3)), offsets=[(0, 1)]) + with pytest.raises(ValueError, match="zero offset"): + bic.graph.grid_affinity_features(graph, np.zeros((1, 2, 3)), offsets=[(0, 0)]) diff --git a/tests/graph/test_grid_graph.py b/tests/graph/test_grid_graph.py new file mode 100644 index 0000000..8f25994 --- /dev/null +++ b/tests/graph/test_grid_graph.py @@ -0,0 +1,189 @@ +import numpy as np +import pytest + +import bioimage_cpp as bic + + +def test_grid_graph_2d_topology_and_coordinates(): + graph = bic.graph.GridGraph2D((2, 3)) + + assert graph.number_of_nodes == 6 + assert graph.number_of_edges == 7 + assert graph.ndim == 2 + np.testing.assert_array_equal(graph.shape, np.array([2, 3], dtype=np.uint64)) + np.testing.assert_array_equal(graph.strides, np.array([3, 1], dtype=np.uint64)) + np.testing.assert_array_equal( + graph.uv_ids(), + np.array( + [[0, 3], [1, 4], [2, 5], [0, 1], [1, 2], [3, 4], [4, 5]], + dtype=np.uint64, + ), + ) + + assert graph.node_id((1, 2)) == 5 + np.testing.assert_array_equal(graph.coordinates(4), np.array([1, 1], dtype=np.uint64)) + assert graph.find_edge(0, 3) == 0 + assert graph.find_edge(2, 5) == 2 + assert graph.find_edge(0, 1) == 3 + assert graph.find_edge(0, 2) == -1 + assert graph.edge_axis(0) == 0 + assert graph.edge_axis(3) == 1 + np.testing.assert_array_equal( + graph.edge_coordinates(6), np.array([1, 1], dtype=np.uint64) + ) + + +def test_grid_graph_3d_topology_and_offset_targets(): + graph = bic.graph.GridGraph3D((2, 2, 2)) + + assert graph.number_of_nodes == 8 + assert graph.number_of_edges == 12 + np.testing.assert_array_equal(graph.shape, np.array([2, 2, 2], dtype=np.uint64)) + np.testing.assert_array_equal(graph.strides, np.array([4, 2, 1], dtype=np.uint64)) + np.testing.assert_array_equal( + graph.uv_ids(), + np.array( + [ + [0, 4], + [1, 5], + [2, 6], + [3, 7], + [0, 2], + [1, 3], + [4, 6], + [5, 7], + [0, 1], + [2, 3], + [4, 5], + [6, 7], + ], + dtype=np.uint64, + ), + ) + + assert graph.node_id((1, 1, 1)) == 7 + np.testing.assert_array_equal( + graph.coordinates(6), np.array([1, 1, 0], dtype=np.uint64) + ) + assert graph.offset_target(0, (1, 0, 0)) == 4 + assert graph.offset_target(3, (0, 0, -1)) == 2 + assert graph.offset_target(3, (0, 0, 1)) == -1 + + +def test_grid_graph_factory_and_algorithm_interop(): + graph = bic.graph.grid_graph((3, 3)) + + assert isinstance(graph, bic.graph.GridGraph2D) + components = bic.graph.connected_components(graph) + np.testing.assert_array_equal(components, np.zeros(9, dtype=np.uint64)) + + nodes, distances = bic.graph.breadth_first_search( + graph, 0, max_distance=1, include_source=True + ) + np.testing.assert_array_equal(nodes, np.array([0, 3, 1], dtype=np.uint64)) + np.testing.assert_array_equal(distances, np.array([0, 1, 1], dtype=np.uint64)) + + +def _adjacency_pairs(adjacency: np.ndarray) -> set[tuple[int, int]]: + return {(int(adjacency[i, 0]), int(adjacency[i, 1])) for i in range(adjacency.shape[0])} + + +def test_grid_graph_node_adjacency_returns_grid_neighbors(): + # Triggers the lazy CSR rebuild path: GridGraph::build_edges only emits + # edges, the first node_adjacency call populates the CSR buffer. + graph = bic.graph.GridGraph2D((3, 3)) + + # Corner (0, 0): two neighbors — (1, 0)=node 3 via edge 0, (0, 1)=node 1 via edge 6. + assert _adjacency_pairs(graph.node_adjacency(0)) == {(3, 0), (1, 6)} + # Center (1, 1)=node 4: degree 4. + assert _adjacency_pairs(graph.node_adjacency(4)) == {(1, 1), (7, 4), (3, 8), (5, 9)} + # Opposite corner (2, 2)=node 8. + assert _adjacency_pairs(graph.node_adjacency(8)) == {(5, 5), (7, 11)} + + +def test_grid_graph_3d_node_adjacency_returns_six_neighbors_for_interior_node(): + graph = bic.graph.GridGraph3D((3, 3, 3)) + # Center node (1, 1, 1) — flat id = 1*9 + 1*3 + 1 = 13. Should have 6 neighbors. + adj = graph.node_adjacency(13) + assert adj.shape == (6, 2) + neighbors = {int(adj[i, 0]) for i in range(adj.shape[0])} + # Expected neighbors: 4 (z-), 22 (z+), 10 (y-), 16 (y+), 12 (x-), 14 (x+). + assert neighbors == {4, 22, 10, 16, 12, 14} + + +def test_grid_graph_adjacency_calls_are_idempotent(): + # The lazy rebuild must run exactly once and survive subsequent reads. + graph = bic.graph.GridGraph2D((3, 3)) + first = _adjacency_pairs(graph.node_adjacency(4)) + second = _adjacency_pairs(graph.node_adjacency(4)) + assert first == second + # Reading a different node after the rebuild must still work. + assert _adjacency_pairs(graph.node_adjacency(0)) == {(3, 0), (1, 6)} + + +def test_grid_graph_breadth_first_search_with_distance_2(): + graph = bic.graph.GridGraph2D((3, 3)) + nodes, distances = bic.graph.breadth_first_search( + graph, 4, max_distance=2, include_source=True + ) + # BFS from center reaches every node within distance 2. + assert set(int(n) for n in nodes) == set(range(9)) + # Center has distance 0; the four cardinal neighbors have distance 1; + # the four diagonals have distance 2 (no diagonal grid edges). + distance_of_node = {int(n): int(d) for n, d in zip(nodes, distances)} + assert distance_of_node[4] == 0 + for neighbor in (1, 3, 5, 7): + assert distance_of_node[neighbor] == 1 + for corner in (0, 2, 6, 8): + assert distance_of_node[corner] == 2 + + +def test_grid_graph_connected_components_with_edge_mask(): + # Mask out the axis-0 edges that bridge row 1 ↔ row 2, splitting the + # grid into two components. + graph = bic.graph.GridGraph2D((3, 3)) + edge_mask = np.ones(graph.number_of_edges, dtype=bool) + edge_mask[3] = False # (3, 6) + edge_mask[4] = False # (4, 7) + edge_mask[5] = False # (5, 8) + labels = bic.graph.connected_components(graph, edge_mask=edge_mask) + assert labels[0] == labels[5] # rows 0..1 form one component + assert labels[6] == labels[8] # row 2 forms another + assert labels[0] != labels[6] + + +def test_grid_graph_extract_subgraph_from_nodes(): + # Pick the top-left 2x2 block of a 3x3 grid. + graph = bic.graph.GridGraph2D((3, 3)) + nodes = np.array([0, 1, 3, 4], dtype=np.uint64) + inner, outer = graph.extract_subgraph_from_nodes(nodes) + # Inner edges: (0,1), (0,3), (1,4), (3,4) — ids 6, 0, 1, 8. + assert set(int(e) for e in inner) == {0, 1, 6, 8} + # Outer edges: edges crossing the boundary of the 2x2 block. + # (1,2)=7, (3,6)=3, (4,5)=9, (4,7)=4. + assert set(int(e) for e in outer) == {3, 4, 7, 9} + + +def test_grid_graph_freeze_is_idempotent_and_safe_to_call_repeatedly(): + graph = bic.graph.GridGraph2D((3, 3)) + graph.freeze() # eager rebuild + graph.freeze() # already frozen; no-op + # Adjacency must still be correct after explicit freeze. + assert _adjacency_pairs(graph.node_adjacency(4)) == {(1, 1), (7, 4), (3, 8), (5, 9)} + + +def test_grid_graph_rejects_invalid_shapes_and_coordinates(): + with pytest.raises(ValueError, match="length 2"): + bic.graph.GridGraph2D((2, 3, 4)) + with pytest.raises(ValueError, match="greater than zero"): + bic.graph.GridGraph3D((2, 0, 4)) + with pytest.raises(ValueError, match="length 2 or 3"): + bic.graph.grid_graph((5,)) + + graph = bic.graph.GridGraph2D((2, 3)) + with pytest.raises(ValueError, match="coordinate must be"): + graph.node_id((1, 2, 3)) + with pytest.raises(IndexError, match="coordinate\\[0\\]"): + graph.node_id((2, 0)) + with pytest.raises(ValueError, match="offset must be"): + graph.offset_target(0, (1, 0, 0)) diff --git a/tests/graph/test_undirected_graph.py b/tests/graph/test_undirected_graph.py index b85912f..f9198ac 100644 --- a/tests/graph/test_undirected_graph.py +++ b/tests/graph/test_undirected_graph.py @@ -117,3 +117,41 @@ def test_undirected_graph_rejects_invalid_edges(): graph.insert_edge(0, 2) with pytest.raises(ValueError, match="uvs must have shape"): graph.insert_edges([0, 1]) + + +def test_undirected_graph_clone_is_independent_deep_copy(): + original = bic.graph.UndirectedGraph(4) + original.insert_edge(0, 1) + original.insert_edge(1, 2) + original.insert_edge(2, 3) + + copy = original.clone() + assert copy.number_of_nodes == 4 + assert copy.number_of_edges == 3 + np.testing.assert_array_equal(copy.uv_ids(), original.uv_ids()) + np.testing.assert_array_equal(copy.node_adjacency(1), original.node_adjacency(1)) + + # Mutating the copy must not affect the original. + copy.insert_edge(0, 3) + assert copy.number_of_edges == 4 + assert original.number_of_edges == 3 + assert original.find_edge(0, 3) == -1 + assert copy.find_edge(0, 3) == 3 + + +def test_undirected_graph_freeze_is_callable_after_inserts_and_after_reads(): + graph = bic.graph.UndirectedGraph(3) + graph.insert_edge(0, 1) + graph.insert_edge(1, 2) + # Eagerly rebuild adjacency before any read. + graph.freeze() + np.testing.assert_array_equal( + graph.node_adjacency(1), np.array([[0, 0], [2, 1]], dtype=np.uint64) + ) + # Frozen graph still accepts new edges; lazy rebuild fires again on next read. + graph.insert_edge(0, 2) + graph.freeze() + assert graph.number_of_edges == 3 + np.testing.assert_array_equal( + np.sort(graph.node_adjacency(0)[:, 0]), np.array([1, 2], dtype=np.uint64) + )