Improve agglomeration implementation

Author: constantinpape

MIGRATION_GUIDE.md

@@ -1448,9 +1448,18 @@ Differences from nifty:
 - Both `float32` and `float64` inputs are accepted; computation runs in
   `float64` internally.
 - Tie-breaks follow the deterministic order of edge ids returned by
-  `UndirectedGraph`, which may differ in rare cases from nifty's. Final
-  partitions agree closely; byte-identical label equality is not
-  guaranteed.
+  `UndirectedGraph`, which may differ from nifty's. On inputs where many
+  edges share the same indicator value, this combines with the
+  hierarchical agglomeration's positive feedback loop (each tied merge
+  changes node sizes, which changes the harmonic size factor `sFac`,
+  which changes future priorities) to give cascading divergence. On the
+  external multicut problem sample C/medium, where 86% of indicator
+  values are non-unique, perturbing the indicators of a single bic run by
+  1e-9 random noise can change the final partition's adjusted Rand index
+  vs. its own unperturbed output by ~0.5 (the algorithm is chaotically
+  sensitive to tie-breaking under non-zero `size_regularizer`). Both
+  partitions are valid clusterings; partition agreement (VI, ARI) is the
+  appropriate comparison metric, not label equality.
 
 ### Projecting RAG Node Labels to Pixels
 

development/graph/agglomeration/PERFORMANCE_NOTES.md

@@ -13,10 +13,15 @@
 
 def make_parser() -> argparse.ArgumentParser:
     arg_parser = base_parser(__doc__ or "")
+    # ``abs_max`` is intentionally not compared against nifty: nifty has no
+    # direct sign-aware absolute-maximum linkage. The closest match is
+    # ``MutexWatershedSettings`` but it additionally installs cannot-link
+    # constraints, so the comparison is apples-to-oranges. Run the unit
+    # tests for ``abs_max`` coverage instead.
     arg_parser.add_argument(
         "--linkage",
         default="mean",
-        choices=["sum", "mean", "max", "min", "abs_max", "mutex_watershed"],
+        choices=["sum", "mean", "max", "min", "mutex_watershed"],
         help="GASP linkage rule.",
     )
     return arg_parser
@@ -38,7 +43,6 @@ def main() -> None:
         "sum": nagglo.SumSettings,
         "max": nagglo.MaxSettings,
         "min": nagglo.MinSettings,
-        "abs_max": nagglo.GeneralizedMeanSettings,
         "mutex_watershed": nagglo.MutexWatershedSettings,
     }[args.linkage]
 

development/graph/agglomeration/check_gasp.py

@@ -0,0 +1,173 @@
+"""Diagnose divergences between bioimage-cpp and nifty agglomeration policies.
+
+Targeted at the cases where the benchmark sweep showed ARI < 0.90:
+
+* mala on A/B/C small and C medium,
+* edge_weighted on C medium,
+* gasp_max on A/B small.
+
+For each case, run both implementations, compare partitions, and (for the
+hypothesised root cause) print enough state to confirm or deny it.
+"""
+
+from __future__ import annotations
+
+import argparse
+import numpy as np
+
+import bioimage_cpp as bic
+
+from _compatibility import load_problem
+
+
+def _ari(a: np.ndarray, b: np.ndarray) -> float:
+    try:
+        from sklearn.metrics import adjusted_rand_score
+        return float(adjusted_rand_score(a, b))
+    except ImportError:
+        return float("nan")
+
+
+def _cluster_sizes(labels: np.ndarray, top: int = 8) -> str:
+    _, counts = np.unique(labels, return_counts=True)
+    counts = np.sort(counts)[::-1]
+    head = counts[:top].tolist()
+    return f"n_clusters={len(counts)} top{top}={head} max={int(counts[0])} median={int(np.median(counts))}"
+
+
+def diagnose_mala(sample: str, size: str) -> None:
+    print(f"\n=== MALA  sample={sample} size={size} ===")
+    bic_graph, nifty_graph, indicators, _, _ = load_problem(sample, size, timeout=120.0)
+    n = int(bic_graph.number_of_nodes)
+    e = int(bic_graph.number_of_edges)
+    print(f"  nodes={n} edges={e}")
+
+    bic_labels = bic.graph.agglomeration.MalaClusterPolicy(
+        num_bins=40, bin_min=0.0, bin_max=1.0,
+        num_clusters_stop=1000, threshold=0.5,
+    ).optimize(bic_graph, indicators)
+
+    import nifty.graph.agglo as nagglo
+    policy = nagglo.malaClusterPolicy(
+        graph=nifty_graph,
+        edgeIndicators=indicators.astype(np.float32),
+        nodeSizes=np.ones(n, dtype=np.float32),
+        edgeSizes=np.ones(e, dtype=np.float32),
+        threshold=0.5, numberOfNodesStop=1000,
+    )
+    clustering = nagglo.agglomerativeClustering(policy)
+    clustering.run()
+    nifty_labels = np.asarray(clustering.result(), dtype=np.uint64)
+
+    print(f"  bic   : {_cluster_sizes(bic_labels)}")
+    print(f"  nifty : {_cluster_sizes(nifty_labels)}")
+    print(f"  ARI   : {_ari(bic_labels, nifty_labels):.4f}")
+    # Sample a handful of indicators near 0.5 — the threshold — to highlight
+    # how bin-center vs interpolated median changes the stop decision.
+    mid = indicators[np.abs(indicators - 0.5) < 0.1]
+    print(f"  #indicators within 0.1 of threshold=0.5: {len(mid)} "
+          f"(out of {len(indicators)})")
+
+
+def diagnose_edge_weighted(sample: str, size: str) -> None:
+    print(f"\n=== EDGE_WEIGHTED  sample={sample} size={size} ===")
+    bic_graph, nifty_graph, indicators, _, _ = load_problem(sample, size, timeout=120.0)
+    n = int(bic_graph.number_of_nodes)
+    e = int(bic_graph.number_of_edges)
+    print(f"  nodes={n} edges={e}")
+
+    edge_sizes = np.ones(e, dtype=np.float64)
+    node_sizes = np.ones(n, dtype=np.float64)
+
+    bic_labels = bic.graph.agglomeration.EdgeWeightedClusterPolicy(
+        num_clusters_stop=1000, size_regularizer=0.5,
+    ).optimize(bic_graph, indicators, edge_sizes=edge_sizes, node_sizes=node_sizes)
+
+    import nifty.graph.agglo as nagglo
+    policy = nagglo.edgeWeightedClusterPolicy(
+        graph=nifty_graph,
+        edgeIndicators=indicators.astype(np.float32),
+        edgeSizes=edge_sizes.astype(np.float32),
+        nodeSizes=node_sizes.astype(np.float32),
+        numberOfNodesStop=1000, sizeRegularizer=0.5,
+    )
+    clustering = nagglo.agglomerativeClustering(policy)
+    clustering.run()
+    nifty_labels = np.asarray(clustering.result(), dtype=np.uint64)
+
+    print(f"  bic   : {_cluster_sizes(bic_labels)}")
+    print(f"  nifty : {_cluster_sizes(nifty_labels)}")
+    print(f"  ARI   : {_ari(bic_labels, nifty_labels):.4f}")
+
+    # How many edges share the smallest-bucket priority? (tie-breaking
+    # signal: large equal-priority cohorts let the two impls diverge.)
+    p = np.round(indicators, 6)
+    _, counts = np.unique(p, return_counts=True)
+    top = np.sort(counts)[::-1][:5].tolist()
+    print(f"  unique-priorities up to 6 dp: {len(counts)}  top-5 counts: {top}")
+
+
+def diagnose_gasp_max(sample: str, size: str) -> None:
+    print(f"\n=== GASP max  sample={sample} size={size} ===")
+    bic_graph, nifty_graph, _, signed_weights, _ = load_problem(sample, size, timeout=120.0)
+    n = int(bic_graph.number_of_nodes)
+    e = int(bic_graph.number_of_edges)
+    print(f"  nodes={n} edges={e}")
+    print(f"  signed_weights: min={signed_weights.min():.3f} max={signed_weights.max():.3f} "
+          f"positive_fraction={(signed_weights > 0).mean():.3f}")
+
+    edge_sizes = np.ones(e, dtype=np.float64)
+
+    bic_labels = bic.graph.agglomeration.GaspClusterPolicy(
+        num_clusters_stop=1000, linkage="max",
+    ).optimize(bic_graph, signed_weights, edge_sizes=edge_sizes)
+
+    import nifty.graph.agglo as nagglo
+    policy = nagglo.gaspClusterPolicy(
+        graph=nifty_graph,
+        signedWeights=signed_weights.astype(np.float64),
+        isMergeEdge=np.ones(e, dtype=np.uint8),
+        edgeSizes=edge_sizes.astype(np.float64),
+        nodeSizes=np.ones(n, dtype=np.float64),
+        updateRule0=nagglo.MaxSettings(),
+        numberOfNodesStop=1000,
+    )
+    clustering = nagglo.agglomerativeClustering(policy)
+    clustering.run()
+    nifty_labels = np.asarray(clustering.result(), dtype=np.uint64)
+
+    print(f"  bic   : {_cluster_sizes(bic_labels)}")
+    print(f"  nifty : {_cluster_sizes(nifty_labels)}")
+    print(f"  ARI   : {_ari(bic_labels, nifty_labels):.4f}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--policy", choices=["mala", "edge_weighted", "gasp_max", "all"],
+                        default="all")
+    parser.add_argument("--sample", default=None)
+    parser.add_argument("--size", default=None)
+    args = parser.parse_args()
+
+    cases_mala = [("A", "small"), ("B", "small"), ("C", "small"), ("C", "medium")]
+    cases_ew = [("C", "medium")]
+    cases_gm = [("A", "small"), ("B", "small")]
+
+    if args.sample and args.size:
+        cases_mala = [(args.sample, args.size)]
+        cases_ew = [(args.sample, args.size)]
+        cases_gm = [(args.sample, args.size)]
+
+    if args.policy in ("mala", "all"):
+        for s, z in cases_mala:
+            diagnose_mala(s, z)
+    if args.policy in ("edge_weighted", "all"):
+        for s, z in cases_ew:
+            diagnose_edge_weighted(s, z)
+    if args.policy in ("gasp_max", "all"):
+        for s, z in cases_gm:
+            diagnose_gasp_max(s, z)
+
+
+if __name__ == "__main__":
+    main()

development/graph/agglomeration/diagnose.py

@@ -17,12 +17,20 @@
 namespace bioimage_cpp::graph::agglomeration {
 
 // Linkage criterion for GASP. The criterion determines how parallel edges
-// fold when two clusters merge; the heap priority is always |edge_weight|.
+// fold when two clusters merge and how the agglomeration terminates.
 //
-// `kMutexWatershed` interprets a negative weight on the heap top as a
-// cannot-link constraint (the edge is rejected and a mutex installed
-// between the two clusters), exactly matching the mutex-watershed algorithm
-// applied to a single edge list.
+// For all linkages except ``kMutexWatershed`` the heap priority is
+// ``-edge_weight`` (signed): the most attractive edge pops first, and the
+// agglomeration stops as soon as no positive-weight edges remain. This
+// matches nifty's "stop when no merge candidates are left" behaviour and
+// avoids doing unbounded work after all attractive evidence has been
+// consumed (otherwise critical for ``kSum`` whose combined weight can grow
+// without bound).
+//
+// ``kMutexWatershed`` instead uses priority ``-|edge_weight|`` so the
+// largest-magnitude edge pops first; a negative-weight pop installs a
+// permanent cannot-link constraint between the two clusters and is
+// rejected, exactly matching the mutex-watershed algorithm.
 enum class GaspLinkage {
     kSum = 0,
     kMean = 1,
@@ -96,10 +104,7 @@ public:
             const auto u = static_cast<std::size_t>(uv.first);
             const auto v = static_cast<std::size_t>(uv.second);
             const auto edge_index = static_cast<std::size_t>(edge_id);
-            const double weight = edge_weight_[edge_index];
-            // Min-heap stores ``-|weight|`` so the largest |weight| is the
-            // first to pop (recovers max-heap-on-absolute-value semantics).
-            const double priority = -std::abs(weight);
+            const double priority = priority_of(edge_weight_[edge_index]);
             auto &edge = dynamic_graph.edges[edge_index];
             edge.u = u;
             edge.v = v;
@@ -128,6 +133,16 @@ public:
         if (check_mutex(u, v, cannot_link_)) {
             return Action::kRejectEdge;
         }
+        // Non-mutex-watershed linkages use signed-weight priority and stop
+        // as soon as the top of the heap is non-positive (no attractive
+        // edges remain). Matches `nifty.graph.agglo`'s `isDone` behaviour
+        // and prevents `kSum` / `kAbsMax` from running away when negative
+        // weights flood the queue.
+        if (linkage_ != GaspLinkage::kMutexWatershed) {
+            if (edge_weight_[edge_id] <= 0.0) {
+                return Action::kStop;
+            }
+        }
         if (!is_mergeable_[edge_id]) {
             insert_mutex(u, v, cannot_link_);
             return Action::kRejectEdge;
@@ -195,10 +210,22 @@ public:
         // contribution makes the surviving edge a cannot-link candidate.
         is_mergeable_[existing_id] =
             (is_mergeable_[existing_id] != 0 && is_mergeable_[fold_id] != 0) ? 1 : 0;
-        return -std::abs(combined);
+        return priority_of(combined);
     }
 
 private:
+    // For non-mutex-watershed linkages the priority is the negated signed
+    // weight (min-heap pops the most-positive weight first; the loop
+    // terminates when the top is non-positive). The mutex-watershed
+    // linkage uses ``-|weight|`` so the largest-magnitude edge pops first
+    // regardless of sign.
+    double priority_of(const double weight) const {
+        if (linkage_ == GaspLinkage::kMutexWatershed) {
+            return -std::abs(weight);
+        }
+        return -weight;
+    }
+
     std::vector<double> edge_weight_;
     std::vector<double> edge_size_;
     std::vector<std::uint8_t> is_mergeable_;