Implement Multicut (#4)

* Implement multicut functionality

* Optimization round 1

* Refactor to avoid duplicate functionality

* Kernighan lin updates

* Implement PQ and update KL

* Multicut performance improvements

* Update agent instructions

Author: constantinpape

AGENTS.md

@@ -23,5 +23,8 @@ nanobind_add_module(_core
 
 target_include_directories(_core PRIVATE include)
 target_compile_features(_core PRIVATE cxx_std_20)
+target_compile_options(_core PRIVATE
+    $<$<NOT:$<CONFIG:Debug>>:-O3>
+)
 
 install(TARGETS _core LIBRARY DESTINATION bioimage_cpp)

CMakeLists.txt

@@ -0,0 +1,134 @@
+from __future__ import annotations
+
+import argparse
+from statistics import median
+from time import perf_counter
+from typing import Callable
+
+import numpy as np
+
+
+def parser(description: str) -> argparse.ArgumentParser:
+    arg_parser = argparse.ArgumentParser(description=description)
+    arg_parser.add_argument(
+        "--path",
+        default=None,
+        help="Path to the external multicut problem. Defaults to the package cache.",
+    )
+    arg_parser.add_argument(
+        "--repeats",
+        type=int,
+        default=3,
+        help="Number of timed repeats per implementation.",
+    )
+    arg_parser.add_argument(
+        "--threads",
+        type=int,
+        default=1,
+        help="Number of threads for solvers that support it.",
+    )
+    arg_parser.add_argument(
+        "--timeout",
+        type=float,
+        default=30.0,
+        help="Download timeout in seconds if the external problem is not cached.",
+    )
+    arg_parser.add_argument(
+        "--energy-bound",
+        type=float,
+        default=-76900.0,
+        help="Maximum accepted energy for both implementations.",
+    )
+    return arg_parser
+
+
+def load_problem(path: str | None, *, timeout: float):
+    import bioimage_cpp as bic
+    import nifty
+
+    uv_ids, costs = bic.graph.load_external_multicut_problem_data(
+        path,
+        timeout=timeout,
+    )
+    bic_graph = bic.graph.UndirectedGraph.from_edges(int(uv_ids.max()) + 1, uv_ids)
+    nifty_graph = nifty.graph.undirectedGraph(int(uv_ids.max()) + 1)
+    nifty_graph.insertEdges(uv_ids)
+    return bic_graph, nifty_graph, costs
+
+
+def time_call(function: Callable[[], np.ndarray], repeats: int):
+    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 optimize_bic_solver(make_bic_solver, objective):
+    objective.reset_labels()
+    return make_bic_solver().optimize(objective)
+
+
+def bic_energy(graph, costs: np.ndarray, labels: np.ndarray) -> float:
+    import bioimage_cpp as bic
+
+    return bic.graph.MulticutObjective(graph, costs).energy(labels)
+
+
+def nifty_energy(graph, costs: np.ndarray, labels: np.ndarray) -> float:
+    import nifty.graph.opt.multicut as nmc
+
+    return float(nmc.multicutObjective(graph, costs).evalNodeLabels(labels))
+
+
+def run_comparison(
+    name: str,
+    make_bic_solver,
+    make_nifty_solver,
+    args: argparse.Namespace,
+) -> dict[str, float]:
+    import bioimage_cpp as bic
+    import nifty.graph.opt.multicut as nmc
+
+    bic_graph, nifty_graph, costs = load_problem(args.path, timeout=args.timeout)
+    bic_objective = bic.graph.MulticutObjective(bic_graph, costs)
+    nifty_objective = nmc.multicutObjective(nifty_graph, costs)
+
+    bic_timings, bic_labels = time_call(
+        lambda: optimize_bic_solver(make_bic_solver, bic_objective),
+        args.repeats,
+    )
+    nifty_timings, nifty_labels = time_call(
+        lambda: make_nifty_solver(nifty_objective).create(nifty_objective).optimize(),
+        args.repeats,
+    )
+
+    bic_score = bic_energy(bic_graph, costs, bic_labels)
+    nifty_score = nifty_energy(nifty_graph, costs, nifty_labels)
+    if bic_score > args.energy_bound:
+        raise AssertionError(
+            f"bioimage-cpp {name} energy {bic_score:.6f} exceeds bound {args.energy_bound:.6f}"
+        )
+    if nifty_score > args.energy_bound:
+        raise AssertionError(
+            f"nifty {name} energy {nifty_score:.6f} exceeds bound {args.energy_bound:.6f}"
+        )
+
+    result = {
+        "bioimage_cpp_energy": bic_score,
+        "nifty_energy": nifty_score,
+        "energy_difference": bic_score - nifty_score,
+        "bioimage_cpp_median_runtime": median(bic_timings),
+        "nifty_median_runtime": median(nifty_timings),
+    }
+    print(f"solver: {name}")
+    print(f"nodes: {bic_graph.number_of_nodes}, edges: {bic_graph.number_of_edges}")
+    print(f"bioimage-cpp energy: {bic_score:.6f}")
+    print(f"nifty energy:        {nifty_score:.6f}")
+    print(f"energy difference:   {bic_score - nifty_score:.6f}")
+    print(f"bioimage-cpp median runtime [s]: {median(bic_timings):.6f}")
+    print(f"nifty median runtime [s]:        {median(nifty_timings):.6f}")
+    return result

development/graph/multicut/_compatibility.py

@@ -0,0 +1,29 @@
+from __future__ import annotations
+
+import bioimage_cpp as bic
+
+from _compatibility import parser, run_comparison
+
+
+def main() -> None:
+    args = parser("Compare bioimage-cpp and nifty chained multicut solvers.").parse_args()
+    run_comparison(
+        "chained",
+        lambda: bic.graph.ChainedMulticutSolvers(
+            [
+                bic.graph.GreedyAdditiveMulticut(),
+                bic.graph.KernighanLinMulticut(number_of_outer_iterations=5),
+            ]
+        ),
+        lambda objective: objective.chainedSolversFactory(
+            [
+                objective.greedyAdditiveFactory(),
+                objective.kernighanLinFactory(numberOfOuterIterations=5),
+            ]
+        ),
+        args,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/graph/multicut/check_chained.py

@@ -0,0 +1,23 @@
+from __future__ import annotations
+
+import bioimage_cpp as bic
+
+from _compatibility import parser, run_comparison
+
+
+def main() -> None:
+    args = parser("Compare bioimage-cpp and nifty decomposer multicut.").parse_args()
+    run_comparison(
+        "decomposer",
+        lambda: bic.graph.MulticutDecomposer(bic.graph.GreedyAdditiveMulticut()),
+        lambda objective: objective.multicutDecomposerFactory(
+            submodelFactory=objective.greedyAdditiveFactory(),
+            fallthroughFactory=objective.greedyAdditiveFactory(),
+            numberOfThreads=args.threads,
+        ),
+        args,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/graph/multicut/check_decomposer.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import bioimage_cpp as bic
+
+from _compatibility import parser, run_comparison
+
+
+def main() -> None:
+    args = parser("Compare bioimage-cpp and nifty greedy-additive multicut.").parse_args()
+    run_comparison(
+        "greedy_additive",
+        lambda: bic.graph.GreedyAdditiveMulticut(),
+        lambda objective: objective.greedyAdditiveFactory(),
+        args,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/graph/multicut/check_greedy_additive.py

@@ -0,0 +1,19 @@
+from __future__ import annotations
+
+import bioimage_cpp as bic
+
+from _compatibility import parser, run_comparison
+
+
+def main() -> None:
+    args = parser("Compare bioimage-cpp and nifty greedy-fixation multicut.").parse_args()
+    run_comparison(
+        "greedy_fixation",
+        lambda: bic.graph.GreedyFixationMulticut(),
+        lambda objective: objective.greedyFixationFactory(),
+        args,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/graph/multicut/check_greedy_fixation.py

@@ -0,0 +1,22 @@
+from __future__ import annotations
+
+import bioimage_cpp as bic
+
+from _compatibility import parser, run_comparison
+
+
+def main() -> None:
+    args = parser("Compare bioimage-cpp and nifty Kernighan-Lin multicut.").parse_args()
+    run_comparison(
+        "kernighan_lin",
+        lambda: bic.graph.KernighanLinMulticut(number_of_outer_iterations=5),
+        lambda objective: objective.kernighanLinFactory(
+            warmStartGreedy=True,
+            numberOfOuterIterations=5,
+        ),
+        args,
+    )
+
+
+if __name__ == "__main__":
+    main()

development/graph/multicut/check_kernighan_lin.py

@@ -0,0 +1,27 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <utility>
+
+namespace bioimage_cpp::detail {
+
+using NodeId = std::uint64_t;
+using Edge = std::pair<NodeId, NodeId>;
+
+inline Edge edge_key(NodeId u, NodeId v) {
+    if (v < u) {
+        std::swap(u, v);
+    }
+    return {u, v};
+}
+
+struct EdgeHash {
+    std::size_t operator()(const Edge &edge) const {
+        const auto first = static_cast<std::size_t>(edge.first);
+        const auto second = static_cast<std::size_t>(edge.second);
+        return first ^ (second + 0x9e3779b97f4a7c15ULL + (first << 6U) + (first >> 2U));
+    }
+};
+
+} // namespace bioimage_cpp::detail

include/bioimage_cpp/detail/edge_hash.hxx

@@ -0,0 +1,59 @@
+#pragma once
+
+#include <cstddef>
+#include <cstdint>
+#include <vector>
+
+namespace bioimage_cpp::detail {
+
+// C-order strides for a row-major array of the given shape, in units of array
+// elements (not bytes). The innermost (last) axis has stride 1.
+inline std::vector<std::ptrdiff_t> c_order_strides(const std::vector<std::ptrdiff_t> &shape) {
+    std::vector<std::ptrdiff_t> strides(shape.size(), 1);
+    for (std::ptrdiff_t axis = static_cast<std::ptrdiff_t>(shape.size()) - 2; axis >= 0; --axis) {
+        strides[static_cast<std::size_t>(axis)] =
+            strides[static_cast<std::size_t>(axis + 1)] *
+            shape[static_cast<std::size_t>(axis + 1)];
+    }
+    return strides;
+}
+
+// Translate a flat node index by a per-axis offset on a row-major grid.
+//
+// Returns true when the offset keeps the result inside the grid, in which case
+// `target_out` is set to the neighbor's flat index. Returns false otherwise and
+// leaves `target_out` unchanged. `strides` must match `shape` and is typically
+// `c_order_strides(shape)`.
+inline bool valid_offset_target(
+    const std::uint64_t node,
+    const std::vector<std::ptrdiff_t> &offset,
+    const std::vector<std::ptrdiff_t> &shape,
+    const std::vector<std::ptrdiff_t> &strides,
+    std::uint64_t &target_out
+) {
+    std::int64_t target_signed = static_cast<std::int64_t>(node);
+    for (std::size_t axis = 0; axis < shape.size(); ++axis) {
+        const auto coord =
+            static_cast<std::ptrdiff_t>(node / static_cast<std::uint64_t>(strides[axis])) %
+            shape[axis];
+        const auto neighbor = coord + offset[axis];
+        if (neighbor < 0 || neighbor >= shape[axis]) {
+            return false;
+        }
+        target_signed += static_cast<std::int64_t>(offset[axis] * strides[axis]);
+    }
+    target_out = static_cast<std::uint64_t>(target_signed);
+    return true;
+}
+
+inline bool is_valid_grid_edge(
+    const std::uint64_t node,
+    const std::vector<std::ptrdiff_t> &offset,
+    const std::vector<std::ptrdiff_t> &shape,
+    const std::vector<std::ptrdiff_t> &strides
+) {
+    std::uint64_t unused = 0;
+    return valid_offset_target(node, offset, shape, strides, unused);
+}
+
+} // namespace bioimage_cpp::detail