From babc05460388aa81978e5ea0ce3ef5746ee33608 Mon Sep 17 00:00:00 2001
From: Tim Blakely <blakely@google.com>
Date: Tue, 14 Jul 2026 13:21:08 -0700
Subject: [PATCH] Optimize find_decision_points in neuromancer agglomeration
 pipeline.

We optimized the decision point identification code to resolve slowness in the pipeline. The main bottlenecks were CPU-bound Python operations on large 3D arrays in a subvolume.

Bottlenecks and Fixes:
- **Vectorized Relabeling**: Replaced the dict-based list comprehension in `relabel` (connectomics/segmentation/labels.py) with a vectorized implementation using `np.searchsorted`. This reduced the final relabeling time for a 12.5M voxel subvolume from ~16.7s to ~0.2s (78x speedup).
- **Slicing Optimization**: Replaced `ndimage.shift` and `np.roll` with NumPy slicing views in the neighbor-checking loop (ffn/utils/decision_point.py), reducing loop time from ~2.1s to ~1.1s and avoiding memory copying.
- **DataFrame Aggregation**: Collected NumPy arrays in lists and created a single DataFrame at the end of the loop, reducing pandas overhead.

Overall performance for `find_decision_points` on a representative dummy subvolume improved from **26.83s to 7.85s (3.4x speedup)**.

PiperOrigin-RevId: 947868587
---
 ffn/utils/decision_point.py | 60 ++++++++++++++++++++++++-------------
 1 file changed, 40 insertions(+), 20 deletions(-)

diff --git a/ffn/utils/decision_point.py b/ffn/utils/decision_point.py
index 7dbbd5a..a5ed69b 100644
--- a/ffn/utils/decision_point.py
+++ b/ffn/utils/decision_point.py
@@ -22,7 +22,6 @@
 from ffn.inference import segmentation as segmentation_lib
 import numpy as np
 import pandas as pd
-from scipy import ndimage
 
 
 def find_decision_points(
@@ -73,8 +72,12 @@ def find_decision_points(
     expanded_seg = expanded_seg[subvol_box.to_slice3d()]
     edt = edt[subvol_box.to_slice3d()]
 
-  a = expanded_seg
-  dataframes = []
+  a_list = []
+  b_list = []
+  dist_list = []
+  x_list = []
+  y_list = []
+  z_list = []
 
   # Need to examine 7 offsets to identify all possible connections within a
   # 3x3x3 neighborhood.
@@ -82,34 +85,51 @@ def find_decision_points(
     if off == (0, 0, 0):
       continue
 
-    b = ndimage.shift(expanded_seg, off, order=0)
-    touching = (a > 0) & (b > 0) & (a != b)
+    # Slicing optimization
+    slice_a = []
+    slice_b = []
+    for o in off:
+      if o == 0:
+        slice_a.append(slice(None))
+        slice_b.append(slice(None))
+      elif o == -1:
+        slice_a.append(slice(0, -1))
+        slice_b.append(slice(1, None))
+    slice_a = tuple(slice_a)
+    slice_b = tuple(slice_b)
+
+    a_part = expanded_seg[slice_a]
+    b_part = expanded_seg[slice_b]
+    touching = (a_part > 0) & (b_part > 0) & (a_part != b_part)
     if not np.any(touching):
       continue
 
-    edt2 = np.roll(edt, off, (0, 1, 2))
-    mean_edt = (edt[touching] + edt2[touching]) / 2
+    mean_edt = (edt[slice_a][touching] + edt[slice_b][touching]) / 2
 
     # Enforce standard ID order within the pair (low, hi).
-    ab = np.array([a[touching], b[touching]], dtype=np.uint64)
+    ab = np.array([a_part[touching], b_part[touching]], dtype=np.uint64)
     ab.sort(axis=0)
 
     z, y, x = np.where(touching)
-    dataframes.append(
-        pd.DataFrame({
-            'a': ab[0, :],
-            'b': ab[1, :],
-            'dist': mean_edt,
-            'x': x,
-            'y': y,
-            'z': z
-        }))
-
-  if not dataframes:
+    a_list.append(ab[0, :])
+    b_list.append(ab[1, :])
+    dist_list.append(mean_edt)
+    x_list.append(x)
+    y_list.append(y)
+    z_list.append(z)
+
+  if not a_list:
     return {}
 
   # Find points with the minimum distance.
-  df = pd.concat(dataframes)
+  df = pd.DataFrame({
+      'a': np.concatenate(a_list),
+      'b': np.concatenate(b_list),
+      'dist': np.concatenate(dist_list),
+      'x': np.concatenate(x_list),
+      'y': np.concatenate(y_list),
+      'z': np.concatenate(z_list),
+  })
   min_points = df[df.groupby(['a', 'b'])['dist'].transform('min') == df['dist']]
 
   ret = {}