Skip to content

Commit 750272d

Browse files
constantinpapeconstantinpape
committed
More optimization
1 parent 9e6f47c commit 750272d

9 files changed

Lines changed: 480 additions & 164 deletions

File tree

‎development/graph/_grid_affinity_compatibility.py‎

Lines changed: 72 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -70,6 +70,12 @@ def select_mixed_offsets(
7070

7171

7272
def time_call(function: Callable[[], tuple[np.ndarray, np.ndarray]], repeats: int):
73+
# One untimed warm-up before the measured loop. The first call typically
74+
# pays for code-page faults and one-time library initialization
75+
# (nanobind tuple shapes, numpy ufunc caches, ...) that aren't part of
76+
# the steady-state cost we care about. Without warm-up these costs leak
77+
# into the first sample and skew the median for low `repeats` values.
78+
function()
7379
timings = []
7480
result = None
7581
for _ in range(repeats):
@@ -104,34 +110,78 @@ def bioimage_cpp_local(affinities: np.ndarray, offsets: list[tuple[int, ...]]):
104110
import bioimage_cpp as bic
105111

106112
graph = bic.graph.grid_graph(affinities.shape[1:])
107-
return bioimage_cpp_local_on_graph(graph, affinities, offsets)
113+
weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets)
114+
return graph.uv_ids(), weights
115+
108116

117+
def bioimage_cpp_local_weights_only(
118+
graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]
119+
):
120+
"""Compute edge weights only — no (uvs, weights) materialization.
109121
110-
def bioimage_cpp_local_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]):
122+
This isolates the cost of the feature kernel from the cost of returning
123+
the canonical uv_ids array. Use this when comparing against libraries
124+
that already cache uvs in the graph object.
125+
"""
111126
import bioimage_cpp as bic
112127

113-
weights, valid_edges = bic.graph.grid_affinity_features(graph, affinities, offsets)
114-
if not np.all(valid_edges):
115-
raise AssertionError("local offsets did not cover all grid edges")
128+
weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets)
129+
return weights
130+
131+
132+
def bioimage_cpp_local_with_uvs(
133+
graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]
134+
):
135+
"""Compute weights AND materialize uvs — apples-to-apples with nifty's
136+
``affinitiesToEdgeMapWithOffsets`` and affogato's
137+
``compute_nh_and_weights``, both of which return uvs in their output."""
138+
import bioimage_cpp as bic
139+
140+
weights, _ = bic.graph.grid_affinity_features(graph, affinities, offsets)
116141
return graph.uv_ids(), weights
117142

118143

119144
def bioimage_cpp_lifted(affinities: np.ndarray, offsets: list[tuple[int, ...]]):
120145
import bioimage_cpp as bic
121146

122147
graph = bic.graph.grid_graph(affinities.shape[1:])
123-
return bioimage_cpp_lifted_on_graph(graph, affinities, offsets)
148+
local_weights, _, lifted_uvs, lifted_weights, _ = (
149+
bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets)
150+
)
151+
return graph, graph.uv_ids(), local_weights, lifted_uvs, lifted_weights
124152

125153

126-
def bioimage_cpp_lifted_on_graph(graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]):
154+
def bioimage_cpp_lifted_features_only(
155+
graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]
156+
):
157+
"""Lifted features without graph.uv_ids() — see the local variant."""
127158
import bioimage_cpp as bic
128159

129-
local_weights, valid_edges, lifted_uvs, lifted_weights, _ = (
160+
local_weights, _, lifted_uvs, lifted_weights, _ = (
130161
bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets)
131162
)
163+
return local_weights, lifted_uvs, lifted_weights
164+
165+
166+
def bioimage_cpp_lifted_with_uvs(
167+
graph, affinities: np.ndarray, offsets: list[tuple[int, ...]]
168+
):
169+
"""Lifted features WITH local uvs (apples-to-apples with affogato)."""
170+
import bioimage_cpp as bic
171+
172+
local_weights, _, lifted_uvs, lifted_weights, _ = (
173+
bic.graph.grid_affinity_features_with_lifted(graph, affinities, offsets)
174+
)
175+
return graph.uv_ids(), local_weights, lifted_uvs, lifted_weights
176+
177+
178+
def assert_local_offsets_cover_all_edges(graph, affinities, offsets) -> None:
179+
"""One-shot correctness check called outside of the timing loop."""
180+
import bioimage_cpp as bic
181+
182+
_, valid_edges = bic.graph.grid_affinity_features(graph, affinities, offsets)
132183
if not np.all(valid_edges):
133184
raise AssertionError("local offsets did not cover all grid edges")
134-
return graph, graph.uv_ids(), local_weights, lifted_uvs, lifted_weights
135185

136186

137187
def nifty_local(affinities: np.ndarray, offsets: list[tuple[int, ...]]):
@@ -206,7 +256,19 @@ def print_timing(name: str, first_name: str, first_timings: list[float],
206256
def add_common_arguments(parser: argparse.ArgumentParser) -> None:
207257
parser.add_argument("--ndim", type=int, choices=(2, 3), default=2)
208258
parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX)
209-
parser.add_argument("--repeats", type=int, default=3)
259+
# Default bumped from 3 to 5 — median of 3 is the middle sample and is
260+
# noisy if anything (GC, cache eviction) lands inside one of the three
261+
# runs. With `time_call` doing one warm-up before this, 5 samples gives
262+
# a usable median without much added cost.
263+
parser.add_argument("--repeats", type=int, default=5)
210264
parser.add_argument("--z", type=int, default=20)
211265
parser.add_argument("--yx-shape", type=int, nargs=2, default=(512, 512))
212266
parser.add_argument("--zyx-shape", type=int, nargs=3, default=(16, 512, 512))
267+
# Affinity dtype that every library receives. nifty and affogato accept
268+
# both float32 and float64 at near-identical speed (verified separately),
269+
# and bioimage-cpp now templates on the value type, so feeding all three
270+
# the same dtype removes the previous implicit float32 -> float64 copy
271+
# that was charged only to bioimage-cpp.
272+
parser.add_argument(
273+
"--dtype", choices=("float32", "float64"), default="float32"
274+
)

‎development/graph/check_grid_affinity_edges.py‎

Lines changed: 51 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -6,8 +6,10 @@
66
add_common_arguments,
77
affogato_edges,
88
affogato_edges_on_graph,
9+
assert_local_offsets_cover_all_edges,
910
bioimage_cpp_local,
10-
bioimage_cpp_local_on_graph,
11+
bioimage_cpp_local_weights_only,
12+
bioimage_cpp_local_with_uvs,
1113
compare_edge_sets,
1214
load_problem,
1315
nifty_local,
@@ -22,6 +24,18 @@
2224
import numpy as np
2325

2426

27+
# Notes on remaining apples-to-apples caveats (left in-code rather than fixed):
28+
# * affogato's `MWSGridGraph` carries MWS-specific state and is heavier to
29+
# construct than a pure undirected grid graph. The "total" timings include
30+
# that cost on the affogato side. This isn't a bug in the comparison, it
31+
# reflects affogato's intended workload.
32+
# * nifty / affogato accept the chosen dtype directly (verified separately);
33+
# feeding all three libraries the same dtype removes the previous implicit
34+
# float32 -> float64 copy that was charged only to bioimage-cpp.
35+
# * `time_call` does an untimed warm-up call before the measured loop so the
36+
# first sample doesn't carry one-shot init costs.
37+
38+
2539
def run_check(args: argparse.Namespace) -> None:
2640
affinities, offsets = load_problem(args.data_prefix)
2741
if args.ndim == 2:
@@ -33,6 +47,9 @@ def run_check(args: argparse.Namespace) -> None:
3347
affinities, offsets, zyx_shape=tuple(args.zyx_shape)
3448
)
3549
affinities, offsets = select_local_offsets(affinities, offsets)
50+
# One conversion to the chosen common dtype, BEFORE any timed work.
51+
# bioimage-cpp, nifty, and affogato all consume this same array.
52+
affinities = np.ascontiguousarray(affinities, dtype=np.dtype(args.dtype))
3653

3754
bic_timings, (bic_uvs, bic_weights) = time_call(
3855
lambda: bioimage_cpp_local(affinities, offsets), args.repeats
@@ -49,11 +66,16 @@ def run_check(args: argparse.Namespace) -> None:
4966
from affogato.segmentation import MWSGridGraph
5067

5168
bic_graph = bic.graph.grid_graph(affinities.shape[1:])
52-
bic_affinities = np.ascontiguousarray(affinities, dtype=np.float64)
5369
nifty_graph = ng.undirectedGridGraph(list(affinities.shape[1:]))
5470
affogato_graph = MWSGridGraph(list(affinities.shape[1:]))
55-
bic_feature_timings, _ = time_call(
56-
lambda: bioimage_cpp_local_on_graph(bic_graph, bic_affinities, offsets),
71+
# One-shot correctness check, OUTSIDE the timing loop.
72+
assert_local_offsets_cover_all_edges(bic_graph, affinities, offsets)
73+
74+
# Apples-to-apples timing #1: each library returns (uvs, weights) for a
75+
# pre-built graph. nifty and affogato bundle uvs in their return value,
76+
# bioimage-cpp materializes them via graph.uv_ids().
77+
bic_with_uvs_timings, _ = time_call(
78+
lambda: bioimage_cpp_local_with_uvs(bic_graph, affinities, offsets),
5779
args.repeats,
5880
)
5981
nifty_feature_timings, _ = time_call(
@@ -64,6 +86,14 @@ def run_check(args: argparse.Namespace) -> None:
6486
lambda: affogato_edges_on_graph(affogato_graph, affinities, offsets),
6587
args.repeats,
6688
)
89+
# Apples-to-apples timing #2: bioimage-cpp ONLY computes weights (no
90+
# uvs materialization). This isolates the cost of the feature kernel
91+
# itself. There is no nifty/affogato equivalent that returns just
92+
# weights, so this is reported on its own.
93+
bic_weights_only_timings, _ = time_call(
94+
lambda: bioimage_cpp_local_weights_only(bic_graph, affinities, offsets),
95+
args.repeats,
96+
)
6797

6898
nifty_summary = compare_edge_sets(
6999
"nifty local", bic_uvs, bic_weights, nifty_uvs, nifty_weights
@@ -73,7 +103,12 @@ def run_check(args: argparse.Namespace) -> None:
73103
)
74104

75105
print(f"Grid local affinity edge check ({args.ndim}D)")
76-
print(f"affinities shape: {affinities.shape}, offsets: {offsets}")
106+
print(
107+
f"affinities shape: {affinities.shape}, dtype: {affinities.dtype}, "
108+
f"size: {affinities.nbytes / 1e6:.2f} MB, offsets: {len(offsets)}, "
109+
f"repeats: {args.repeats}"
110+
)
111+
print(f"offsets: {offsets}")
77112
print(
78113
f"nifty edges: {nifty_summary['number_of_edges']}, "
79114
f"max abs weight diff: {nifty_summary['max_abs_weight_diff']:.6g}"
@@ -85,20 +120,26 @@ def run_check(args: argparse.Namespace) -> None:
85120
print_timing("local edges total", "bioimage-cpp", bic_timings, "nifty", nifty_timings)
86121
print_timing("local edges total", "bioimage-cpp", bic_timings, "affogato", affogato_timings)
87122
print_timing(
88-
"local edges prebuilt",
123+
"local edges prebuilt (with uvs)",
89124
"bioimage-cpp",
90-
bic_feature_timings,
125+
bic_with_uvs_timings,
91126
"nifty",
92127
nifty_feature_timings,
93128
)
94129
print_timing(
95-
"local edges prebuilt",
130+
"local edges prebuilt (with uvs)",
96131
"bioimage-cpp",
97-
bic_feature_timings,
132+
bic_with_uvs_timings,
98133
"affogato",
99134
affogato_feature_timings,
100135
)
101-
print("prebuilt bioimage-cpp timing excludes float32 -> float64 conversion")
136+
# Weights-only kernel timing (no comparator — informational).
137+
from statistics import median
138+
139+
print(
140+
"local edges prebuilt (weights only, no uvs materialization) "
141+
f"bioimage-cpp median runtime: {median(bic_weights_only_timings):.6f} s"
142+
)
102143

103144

104145
def main() -> None:

‎development/graph/check_grid_affinity_lifted_edges.py‎

Lines changed: 40 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,8 @@
77
affogato_edges,
88
affogato_edges_on_graph,
99
bioimage_cpp_lifted,
10-
bioimage_cpp_lifted_on_graph,
10+
bioimage_cpp_lifted_features_only,
11+
bioimage_cpp_lifted_with_uvs,
1112
compare_edge_sets,
1213
load_problem,
1314
prepare_2d_problem,
@@ -21,6 +22,16 @@
2122
import numpy as np
2223

2324

25+
# Notes on remaining apples-to-apples caveats (left in-code rather than fixed):
26+
# * affogato bundles local + lifted edges into a single (uvs, weights) array,
27+
# bioimage-cpp returns local and lifted in separate arrays. The two output
28+
# formats are intrinsic to each library's design; the bookkeeping cost of
29+
# producing them is included in the respective timings.
30+
# * affogato's `MWSGridGraph` is heavier to construct than a plain grid graph
31+
# (MWS-specific state). The "total" timings include that cost.
32+
# * `time_call` does an untimed warm-up before the measured loop.
33+
34+
2435
def run_check(args: argparse.Namespace) -> None:
2536
affinities, offsets = load_problem(args.data_prefix)
2637
if args.ndim == 2:
@@ -32,6 +43,8 @@ def run_check(args: argparse.Namespace) -> None:
3243
affinities, offsets, zyx_shape=tuple(args.zyx_shape)
3344
)
3445
affinities, offsets = select_mixed_offsets(affinities, offsets)
46+
# One conversion to the chosen common dtype, BEFORE any timed work.
47+
affinities = np.ascontiguousarray(affinities, dtype=np.dtype(args.dtype))
3548

3649
bic_timings, (graph, bic_local_uvs, bic_local_weights, bic_lifted_uvs, bic_lifted_weights) = (
3750
time_call(lambda: bioimage_cpp_lifted(affinities, offsets), args.repeats)
@@ -44,16 +57,25 @@ def run_check(args: argparse.Namespace) -> None:
4457
from affogato.segmentation import MWSGridGraph
4558

4659
bic_graph = bic.graph.grid_graph(affinities.shape[1:])
47-
bic_affinities = np.ascontiguousarray(affinities, dtype=np.float64)
4860
affogato_graph = MWSGridGraph(list(affinities.shape[1:]))
49-
bic_feature_timings, _ = time_call(
50-
lambda: bioimage_cpp_lifted_on_graph(bic_graph, bic_affinities, offsets),
61+
# Apples-to-apples #1: both return (uvs, local_weights, lifted_uvs,
62+
# lifted_weights) — bundle of arrays sized for downstream multicut.
63+
bic_with_uvs_timings, _ = time_call(
64+
lambda: bioimage_cpp_lifted_with_uvs(bic_graph, affinities, offsets),
5165
args.repeats,
5266
)
5367
affogato_feature_timings, _ = time_call(
5468
lambda: affogato_edges_on_graph(affogato_graph, affinities, offsets),
5569
args.repeats,
5670
)
71+
# Apples-to-apples #2: bioimage-cpp ONLY returns weight arrays and the
72+
# lifted uvs (which must be returned because they aren't grid-indexed).
73+
# Isolates the cost of the feature kernel from local uv_ids() materialization.
74+
bic_features_only_timings, _ = time_call(
75+
lambda: bioimage_cpp_lifted_features_only(bic_graph, affinities, offsets),
76+
args.repeats,
77+
)
78+
5779
affogato_local_uvs, affogato_local_weights, affogato_lifted_uvs, affogato_lifted_weights = (
5880
split_affogato_edges(affogato_uvs, affogato_weights, graph)
5981
)
@@ -74,7 +96,12 @@ def run_check(args: argparse.Namespace) -> None:
7496
)
7597

7698
print(f"Grid lifted affinity edge check ({args.ndim}D)")
77-
print(f"affinities shape: {affinities.shape}, offsets: {offsets}")
99+
print(
100+
f"affinities shape: {affinities.shape}, dtype: {affinities.dtype}, "
101+
f"size: {affinities.nbytes / 1e6:.2f} MB, offsets: {len(offsets)}, "
102+
f"repeats: {args.repeats}"
103+
)
104+
print(f"offsets: {offsets}")
78105
print(
79106
f"local edges: {local_summary['number_of_edges']}, "
80107
f"max abs weight diff: {local_summary['max_abs_weight_diff']:.6g}"
@@ -91,13 +118,18 @@ def run_check(args: argparse.Namespace) -> None:
91118
affogato_timings,
92119
)
93120
print_timing(
94-
"local+lifted edges prebuilt",
121+
"local+lifted edges prebuilt (with uvs)",
95122
"bioimage-cpp",
96-
bic_feature_timings,
123+
bic_with_uvs_timings,
97124
"affogato",
98125
affogato_feature_timings,
99126
)
100-
print("prebuilt bioimage-cpp timing excludes float32 -> float64 conversion")
127+
from statistics import median
128+
129+
print(
130+
"local+lifted edges prebuilt (features only, no local uvs materialization) "
131+
f"bioimage-cpp median runtime: {median(bic_features_only_timings):.6f} s"
132+
)
101133

102134

103135
def main() -> None:

0 commit comments

Comments
 (0)