Mc stress test (#15)

* Add script to build and serialize grid LMC problem, add more sample data

* Update sample data loading

* Add comprehensive solver checks

Author: constantinpape

MIGRATION_GUIDE.md

@@ -696,8 +696,8 @@ shims default to sample A, size small and continue to honor the
 `BIOIMAGE_CPP_EXTERNAL_MULTICUT_PATH` and
 `BIOIMAGE_CPP_EXTERNAL_MULTICUT_CACHE` environment variables.
 
-Lifted multicut problems (2D ISBI slice and full 3D volume, built by
-`examples/segmentation/serialize_lifted_problem.py`):
+Lifted multicut problems (2D ISBI slice, RAG-based 3D volume, and grid-graph
+volume):
 
 ```python
 problem = bic.graph.load_lifted_multicut_problem(size="2d")
@@ -711,6 +711,8 @@ objective = bic.graph.LiftedMulticutObjective(
 )
 ```
 
+Valid sizes are `"2d"`, `"3d"`, and `"grid"`.
+
 Notes:
 
 - Every download is integrity-checked against a SHA256 in the registry; a

development/graph/_grid_affinity_compatibility.py

@@ -1,22 +1,16 @@
 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
 
+def load_problem():
+    from bioimage_cpp._data import load_isbi_affinities
 
-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))
+    affinities, offsets = load_isbi_affinities()
     return np.ascontiguousarray(affinities), [tuple(offset) for offset in offsets]
 
 
@@ -255,7 +249,6 @@ def print_timing(name: str, first_name: str, first_timings: list[float],
 
 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

development/graph/_rag_compatibility.py

@@ -1,7 +1,6 @@
 from __future__ import annotations
 
 import argparse
-from pathlib import Path
 from statistics import median
 from time import perf_counter
 
@@ -10,15 +9,10 @@
 from skimage.measure import label as label_components
 from skimage.segmentation import watershed
 
+def load_problem():
+    from bioimage_cpp._data import load_isbi_affinities
 
-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))
+    affinities, offsets = load_isbi_affinities()
     return np.ascontiguousarray(affinities), [tuple(offset) for offset in offsets]
 
 
@@ -227,7 +221,6 @@ def run_compatibility_check(
     ndim: int,
     repeats: int,
     threads: int,
-    data_prefix: Path,
     z: int,
     yx_shape: tuple[int, int],
     zyx_shape: tuple[int, int, int],
@@ -238,7 +231,7 @@ def run_compatibility_check(
     import bioimage_cpp as bic
     import nifty.graph.rag as nrag
 
-    affinities, offsets = load_problem(data_prefix)
+    affinities, offsets = load_problem()
     if ndim == 2:
         direct_affinities, direct_offsets = prepare_2d_problem(affinities, offsets, z, yx_shape)
     elif ndim == 3:
@@ -310,7 +303,6 @@ def _print_timing(name: str, bic_timings: list[float], nifty_timings: list[float
 
 
 def add_common_arguments(parser: argparse.ArgumentParser) -> None:
-    parser.add_argument("--data-prefix", type=Path, default=DEFAULT_DATA_PREFIX)
     parser.add_argument("--repeats", type=int, default=3)
     parser.add_argument("--threads", type=int, default=1)
     parser.add_argument("--watershed-min-distance", type=int, default=5)

development/graph/check_grid_affinity_edges.py

@@ -37,7 +37,7 @@
 
 
 def run_check(args: argparse.Namespace) -> None:
-    affinities, offsets = load_problem(args.data_prefix)
+    affinities, offsets = load_problem()
     if args.ndim == 2:
         affinities, offsets = prepare_2d_problem(
             affinities, offsets, z=args.z, yx_shape=tuple(args.yx_shape)

development/graph/check_grid_affinity_lifted_edges.py

@@ -33,7 +33,7 @@
 
 
 def run_check(args: argparse.Namespace) -> None:
-    affinities, offsets = load_problem(args.data_prefix)
+    affinities, offsets = load_problem()
     if args.ndim == 2:
         affinities, offsets = prepare_2d_problem(
             affinities, offsets, z=args.z, yx_shape=tuple(args.yx_shape)

development/graph/check_rag_2d.py

@@ -18,7 +18,6 @@ def main() -> None:
         ndim=2,
         repeats=args.repeats,
         threads=args.threads,
-        data_prefix=args.data_prefix,
         z=args.z,
         yx_shape=tuple(args.shape),
         zyx_shape=(0, 0, 0),

development/graph/check_rag_3d.py

@@ -17,7 +17,6 @@ def main() -> None:
         ndim=3,
         repeats=args.repeats,
         threads=args.threads,
-        data_prefix=args.data_prefix,
         z=0,
         yx_shape=(0, 0),
         zyx_shape=tuple(args.shape),

development/graph/lifted_multicut/_compatibility.py

@@ -12,7 +12,7 @@ def parser(description: str) -> argparse.ArgumentParser:
     arg_parser = argparse.ArgumentParser(description=description)
     arg_parser.add_argument(
         "--size",
-        choices=("2d", "3d"),
+        choices=("2d", "3d", "grid"),
         default="3d",
         help="Lifted multicut problem instance to load (default: 3d).",
     )

development/graph/lifted_multicut/evaluate_solvers.py

@@ -0,0 +1,204 @@
+from __future__ import annotations
+
+import argparse
+from dataclasses import dataclass
+from statistics import mean
+from time import perf_counter
+from typing import Callable
+
+import numpy as np
+
+
+PROBLEMS = ("2d", "3d", "grid")
+
+
+@dataclass(frozen=True)
+class SolverConfig:
+    make_bic_solver: Callable[[], object]
+    make_nifty_factory: Callable[[object], object]
+
+
+def solver_configs():
+    import bioimage_cpp as bic
+
+    return {
+        "lifted_greedy_additive": SolverConfig(
+            make_bic_solver=lambda: bic.graph.LiftedGreedyAdditiveMulticut(),
+            make_nifty_factory=lambda objective: objective.liftedMulticutGreedyAdditiveFactory(),
+        ),
+        "lifted_kernighan_lin": SolverConfig(
+            make_bic_solver=lambda: bic.graph.LiftedKernighanLinMulticut(
+                number_of_outer_iterations=10
+            ),
+            make_nifty_factory=lambda objective: objective.chainedSolversFactory(
+                [
+                    objective.liftedMulticutGreedyAdditiveFactory(),
+                    objective.liftedMulticutKernighanLinFactory(
+                        numberOfOuterIterations=10
+                    ),
+                ]
+            ),
+        ),
+    }
+
+
+def load_problem(size: str, *, timeout: float):
+    import bioimage_cpp as bic
+    import nifty
+    import nifty.graph.opt.lifted_multicut as nlmc
+
+    problem = bic.graph.load_lifted_multicut_problem(size, timeout=timeout)
+    bic_graph = bic.graph.UndirectedGraph.from_edges(problem.n_nodes, problem.local_uvs)
+    nifty_graph = nifty.graph.undirectedGraph(int(problem.n_nodes))
+    nifty_graph.insertEdges(problem.local_uvs.astype(np.uint64, copy=False))
+
+    def make_nifty_objective():
+        objective = nlmc.liftedMulticutObjective(nifty_graph)
+        objective.setGraphEdgesCosts(problem.local_costs)
+        if problem.lifted_uvs.shape[0] > 0:
+            objective.setCosts(
+                problem.lifted_uvs.astype(np.uint64, copy=False),
+                problem.lifted_costs.astype(np.float64, copy=False),
+            )
+        return objective
+
+    return bic_graph, nifty_graph, make_nifty_objective, problem
+
+
+def make_bic_objective(bic_graph, problem):
+    import bioimage_cpp as bic
+
+    return bic.graph.LiftedMulticutObjective(
+        bic_graph,
+        problem.local_costs,
+        lifted_uvs=problem.lifted_uvs,
+        lifted_costs=problem.lifted_costs,
+    )
+
+
+def bic_energy(bic_graph, problem, labels: np.ndarray) -> float:
+    return float(make_bic_objective(bic_graph, problem).energy(labels))
+
+
+def nifty_energy(nifty_objective, labels: np.ndarray) -> float:
+    return float(nifty_objective.evalNodeLabels(labels.astype(np.uint64, copy=False)))
+
+
+def evaluate(problem_name: str, solver_name: str, config: SolverConfig, args):
+    bic_graph, _, make_nifty_objective, problem = load_problem(
+        problem_name, timeout=args.timeout
+    )
+    bic_energies = []
+    nifty_energies = []
+    bic_runtimes = []
+    nifty_runtimes = []
+
+    for _ in range(args.n_repeats):
+        bic_objective = make_bic_objective(bic_graph, problem)
+        start = perf_counter()
+        bic_labels = config.make_bic_solver().optimize(bic_objective)
+        bic_runtimes.append(perf_counter() - start)
+        bic_energies.append(bic_energy(bic_graph, problem, bic_labels))
+
+        nifty_objective = make_nifty_objective()
+        start = perf_counter()
+        nifty_labels = (
+            config.make_nifty_factory(nifty_objective)
+            .create(nifty_objective)
+            .optimize()
+        )
+        nifty_runtimes.append(perf_counter() - start)
+        nifty_energies.append(nifty_energy(nifty_objective, np.asarray(nifty_labels)))
+
+    bic_runtime = mean(bic_runtimes)
+    nifty_runtime = mean(nifty_runtimes)
+    return {
+        "problem": problem_name,
+        "solver": solver_name,
+        "nodes": int(problem.n_nodes),
+        "local_edges": int(problem.local_uvs.shape[0]),
+        "lifted_edges": int(problem.lifted_uvs.shape[0]),
+        "bic_energy": mean(bic_energies),
+        "nifty_energy": mean(nifty_energies),
+        "energy_diff": mean(bic_energies) - mean(nifty_energies),
+        "bic_runtime_s": bic_runtime,
+        "nifty_runtime_s": nifty_runtime,
+        "runtime_ratio": nifty_runtime / bic_runtime if bic_runtime > 0 else float("inf"),
+    }
+
+
+def format_float(value: float) -> str:
+    return f"{value:.6g}"
+
+
+def print_markdown_table(rows: list[dict]) -> None:
+    headers = [
+        "problem",
+        "solver",
+        "nodes",
+        "local_edges",
+        "lifted_edges",
+        "bic_energy",
+        "nifty_energy",
+        "energy_diff",
+        "bic_runtime_s",
+        "nifty_runtime_s",
+        "runtime_ratio_nifty_over_bic",
+    ]
+    print("| " + " | ".join(headers) + " |")
+    print("| " + " | ".join(["---"] * len(headers)) + " |")
+    for row in rows:
+        values = [
+            row["problem"],
+            row["solver"],
+            str(row["nodes"]),
+            str(row["local_edges"]),
+            str(row["lifted_edges"]),
+            format_float(row["bic_energy"]),
+            format_float(row["nifty_energy"]),
+            format_float(row["energy_diff"]),
+            format_float(row["bic_runtime_s"]),
+            format_float(row["nifty_runtime_s"]),
+            format_float(row["runtime_ratio"]),
+        ]
+        print("| " + " | ".join(values) + " |")
+
+
+def main() -> None:
+    configs = solver_configs()
+    parser = argparse.ArgumentParser(
+        description=(
+            "Evaluate matched bioimage-cpp and nifty lifted multicut solvers "
+            "on registered lifted multicut problems."
+        )
+    )
+    parser.add_argument(
+        "--solvers",
+        nargs="+",
+        choices=tuple(configs.keys()),
+        default=tuple(configs.keys()),
+        help="Solvers to evaluate. Defaults to all.",
+    )
+    parser.add_argument(
+        "--problems",
+        nargs="+",
+        choices=PROBLEMS,
+        default=PROBLEMS,
+        help="Problems to evaluate. Defaults to all.",
+    )
+    parser.add_argument("--n-repeats", type=int, default=1)
+    parser.add_argument("--timeout", type=float, default=60.0)
+    args = parser.parse_args()
+
+    if args.n_repeats < 1:
+        raise ValueError("--n-repeats must be at least 1")
+
+    rows = []
+    for problem_name in args.problems:
+        for solver_name in args.solvers:
+            rows.append(evaluate(problem_name, solver_name, configs[solver_name], args))
+    print_markdown_table(rows)
+
+
+if __name__ == "__main__":
+    main()