Some optimization of grid graph creation and features

Author: constantinpape

development/graph/_grid_affinity_compatibility.py

@@ -104,6 +104,12 @@ def bioimage_cpp_local(affinities: np.ndarray, offsets: list[tuple[int, ...]]):
     import bioimage_cpp as bic
 
     graph = bic.graph.grid_graph(affinities.shape[1:])
+    return bioimage_cpp_local_on_graph(graph, affinities, offsets)
+
+
+def bioimage_cpp_local_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]):
+    import bioimage_cpp as bic
+
     weights, 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")
@@ -114,6 +120,12 @@ def bioimage_cpp_lifted(affinities: np.ndarray, offsets: list[tuple[int, ...]]):
     import bioimage_cpp as bic
 
     graph = bic.graph.grid_graph(affinities.shape[1:])
+    return bioimage_cpp_lifted_on_graph(graph, affinities, offsets)
+
+
+def bioimage_cpp_lifted_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]):
+    import bioimage_cpp as bic
+
     local_weights, valid_edges, lifted_uvs, lifted_weights, _ = (
         bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets)
     )
@@ -126,6 +138,10 @@ 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],
@@ -137,6 +153,10 @@ 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],

development/graph/check_grid_affinity_edges.py

@@ -5,17 +5,22 @@
 from _grid_affinity_compatibility import (
     add_common_arguments,
     affogato_edges,
+    affogato_edges_on_graph,
     bioimage_cpp_local,
+    bioimage_cpp_local_on_graph,
     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
+
 
 def run_check(args: argparse.Namespace) -> None:
     affinities, offsets = load_problem(args.data_prefix)
@@ -39,6 +44,27 @@ def run_check(args: argparse.Namespace) -> None:
         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:])
+    bic_affinities = np.ascontiguousarray(affinities, dtype=np.float64)
+    nifty_graph = ng.undirectedGridGraph(list(affinities.shape[1:]))
+    affogato_graph = MWSGridGraph(list(affinities.shape[1:]))
+    bic_feature_timings, _ = time_call(
+        lambda: bioimage_cpp_local_on_graph(bic_graph, bic_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,
+    )
+
     nifty_summary = compare_edge_sets(
         "nifty local", bic_uvs, bic_weights, nifty_uvs, nifty_weights
     )
@@ -56,8 +82,23 @@ def run_check(args: argparse.Namespace) -> None:
         f"affogato edges: {affogato_summary['number_of_edges']}, "
         f"max abs weight diff: {affogato_summary['max_abs_weight_diff']:.6g}"
     )
-    print_timing("local edges", "bioimage-cpp", bic_timings, "nifty", nifty_timings)
-    print_timing("local edges", "bioimage-cpp", bic_timings, "affogato", affogato_timings)
+    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",
+        "bioimage-cpp",
+        bic_feature_timings,
+        "nifty",
+        nifty_feature_timings,
+    )
+    print_timing(
+        "local edges prebuilt",
+        "bioimage-cpp",
+        bic_feature_timings,
+        "affogato",
+        affogato_feature_timings,
+    )
+    print("prebuilt bioimage-cpp timing excludes float32 -> float64 conversion")
 
 
 def main() -> None:

development/graph/check_grid_affinity_lifted_edges.py

@@ -5,7 +5,9 @@
 from _grid_affinity_compatibility import (
     add_common_arguments,
     affogato_edges,
+    affogato_edges_on_graph,
     bioimage_cpp_lifted,
+    bioimage_cpp_lifted_on_graph,
     compare_edge_sets,
     load_problem,
     prepare_2d_problem,
@@ -16,6 +18,8 @@
     time_call,
 )
 
+import numpy as np
+
 
 def run_check(args: argparse.Namespace) -> None:
     affinities, offsets = load_problem(args.data_prefix)
@@ -35,6 +39,21 @@ def run_check(args: argparse.Namespace) -> None:
     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:])
+    bic_affinities = np.ascontiguousarray(affinities, dtype=np.float64)
+    affogato_graph = MWSGridGraph(list(affinities.shape[1:]))
+    bic_feature_timings, _ = time_call(
+        lambda: bioimage_cpp_lifted_on_graph(bic_graph, bic_affinities, offsets),
+        args.repeats,
+    )
+    affogato_feature_timings, _ = time_call(
+        lambda: affogato_edges_on_graph(affogato_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)
     )
@@ -64,7 +83,21 @@ def run_check(args: argparse.Namespace) -> None:
         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", "bioimage-cpp", bic_timings, "affogato", affogato_timings)
+    print_timing(
+        "local+lifted edges total",
+        "bioimage-cpp",
+        bic_timings,
+        "affogato",
+        affogato_timings,
+    )
+    print_timing(
+        "local+lifted edges prebuilt",
+        "bioimage-cpp",
+        bic_feature_timings,
+        "affogato",
+        affogato_feature_timings,
+    )
+    print("prebuilt bioimage-cpp timing excludes float32 -> float64 conversion")
 
 
 def main() -> None: