feat(edits): batch get descendents, make sanity checks optional

Author: akhileshh

pychunkedgraph/graph/cache.py

@@ -3,7 +3,7 @@
 Cache nodes, parents, children and cross edges.
 """
 import traceback
-from collections import defaultdict
+from collections import defaultdict as defaultd
 from sys import maxsize
 from datetime import datetime
 
@@ -51,7 +51,7 @@ def __init__(self, cg):
             "cross_chunk_edges": {"hits": 0, "misses": 0, "calls": 0},
         }
         # Track where calls/misses come from
-        self.call_sources = defaultdict(lambda: defaultdict(lambda: {"calls": 0, "misses": 0}))
+        self.sources = defaultd(lambda: defaultd(lambda: {"calls": 0, "misses": 0}))
 
     def _get_caller(self, skip_frames=2):
         """Get caller info (filename:line:function)."""
@@ -65,8 +65,8 @@ def _get_caller(self, skip_frames=2):
     def _record_call(self, cache_type, misses=0):
         """Record a call and its source."""
         caller = self._get_caller(skip_frames=3)
-        self.call_sources[cache_type][caller]["calls"] += 1
-        self.call_sources[cache_type][caller]["misses"] += misses
+        self.sources[cache_type][caller]["calls"] += 1
+        self.sources[cache_type][caller]["misses"] += misses
 
     def __len__(self):
         return (
@@ -90,7 +90,7 @@ def get_stats(self):
                 **s,
                 "total": total,
                 "hit_rate": f"{hit_rate:.1%}",
-                "sources": dict(self.call_sources[name]),
+                "sources": dict(self.sources[name]),
             }
         return result
 
@@ -99,7 +99,7 @@ def reset_stats(self):
             s["hits"] = 0
             s["misses"] = 0
             s["calls"] = 0
-        self.call_sources.clear()
+        self.sources.clear()
 
     def parent(self, node_id: np.uint64, *, time_stamp: datetime = None):
         self.stats["parents"]["calls"] += 1
@@ -154,7 +154,13 @@ def cross_edges_decorated(node_id):
 
         return cross_edges_decorated(node_id)
 
-    def parents_multiple(self, node_ids: np.ndarray, *, time_stamp: datetime = None):
+    def parents_multiple(
+        self,
+        node_ids: np.ndarray,
+        *,
+        time_stamp: datetime = None,
+        fail_to_zero: bool = False,
+    ):
         node_ids = np.asarray(node_ids, dtype=NODE_ID)
         if not node_ids.size:
             return node_ids
@@ -168,7 +174,10 @@ def parents_multiple(self, node_ids: np.ndarray, *, time_stamp: datetime = None)
         parents = node_ids.copy()
         parents[mask] = self._parent_vec(node_ids[mask])
         parents[~mask] = self._cg.get_parents(
-            node_ids[~mask], raw_only=True, time_stamp=time_stamp
+            node_ids[~mask],
+            raw_only=True,
+            time_stamp=time_stamp,
+            fail_to_zero=fail_to_zero,
         )
         update(self.parents_cache, node_ids[~mask], parents[~mask])
         return parents

pychunkedgraph/graph/chunkedgraph.py

@@ -241,7 +241,9 @@ def get_parents(
                         else:
                             raise KeyError from exc
             return parents
-        return self.cache.parents_multiple(node_ids, time_stamp=time_stamp)
+        return self.cache.parents_multiple(
+            node_ids, time_stamp=time_stamp, fail_to_zero=fail_to_zero
+        )
 
     def get_parent(
         self,
@@ -807,6 +809,7 @@ def add_edges(
         source_coords: typing.Sequence[int] = None,
         sink_coords: typing.Sequence[int] = None,
         allow_same_segment_merge: typing.Optional[bool] = False,
+        do_sanity_check: typing.Optional[bool] = False,
     ) -> operation.GraphEditOperation.Result:
         """
         Adds an edge to the chunkedgraph
@@ -823,6 +826,7 @@ def add_edges(
             source_coords=source_coords,
             sink_coords=sink_coords,
             allow_same_segment_merge=allow_same_segment_merge,
+            do_sanity_check=do_sanity_check,
         ).execute()
 
     def remove_edges(
@@ -838,6 +842,7 @@ def remove_edges(
         path_augment: bool = True,
         disallow_isolating_cut: bool = True,
         bb_offset: typing.Tuple[int, int, int] = (240, 240, 24),
+        do_sanity_check: typing.Optional[bool] = False,
     ) -> operation.GraphEditOperation.Result:
         """
         Removes edges - either directly or after applying a mincut
@@ -862,6 +867,7 @@ def remove_edges(
                 bbox_offset=bb_offset,
                 path_augment=path_augment,
                 disallow_isolating_cut=disallow_isolating_cut,
+                do_sanity_check=do_sanity_check,
             ).execute()
 
         if not atomic_edges:

pychunkedgraph/graph/edges/__init__.py

@@ -207,6 +207,29 @@ def flip_ids(id_map, node_ids):
     return np.concatenate(ids).astype(basetypes.NODE_ID)
 
 
+def _get_new_nodes(
+    cg, nodes: np.ndarray, layer: int, parent_ts: datetime.datetime = None
+):
+    unique_nodes, inverse = np.unique(nodes, return_inverse=True)
+    node_root_map = {n: n for n in unique_nodes}
+    lookup = np.ones(len(unique_nodes), dtype=unique_nodes.dtype)
+    while np.any(lookup):
+        roots = np.fromiter(node_root_map.values(), dtype=basetypes.NODE_ID)
+        roots = cg.get_parents(roots, time_stamp=parent_ts, fail_to_zero=True)
+        layers = cg.get_chunk_layers(roots)
+        lookup[layers > layer] = 0
+        lookup[roots == 0] = 0
+
+        layer_mask = layers <= layer
+        non_zero_mask = roots != 0
+        mask = layer_mask & non_zero_mask
+        for node, root in zip(unique_nodes[mask], roots[mask]):
+            node_root_map[node] = root
+
+    unique_results = np.fromiter(node_root_map.values(), dtype=basetypes.NODE_ID)
+    return unique_results[inverse]
+
+
 def get_stale_nodes(
     cg, nodes: Iterable[basetypes.NODE_ID], parent_ts: datetime.datetime = None
 ):
@@ -215,20 +238,17 @@ def get_stale_nodes(
     This is done by getting a supervoxel of a node and checking
     if it has a new parent at the same layer as the node.
     """
-    nodes = np.array(nodes, dtype=basetypes.NODE_ID)
+    nodes = np.unique(np.array(nodes, dtype=basetypes.NODE_ID))
+    new_ids = set() if cg.cache is None else cg.cache.new_ids
+    nodes = nodes[~np.isin(nodes, new_ids)]
     supervoxels = cg.get_single_leaf_multiple(nodes)
     # nodes can be at different layers due to skip connections
     node_layers = cg.get_chunk_layers(nodes)
     stale_nodes = [types.empty_1d]
     for layer in np.unique(node_layers):
         _mask = node_layers == layer
         layer_nodes = nodes[_mask]
-        _nodes = cg.get_roots(
-            supervoxels[_mask],
-            stop_layer=layer,
-            ceil=False,
-            time_stamp=parent_ts,
-        )
+        _nodes = _get_new_nodes(cg, supervoxels[_mask], layer, parent_ts)
         stale_mask = layer_nodes != _nodes
         stale_nodes.append(layer_nodes[stale_mask])
     return np.concatenate(stale_nodes)
@@ -544,10 +564,7 @@ def _get_new_edge(edge, edge_layer, parent_ts, padding, fallback: bool = False):
                     if fallback:
                         parents_b = _get_parents_b(_edges, parent_ts, edge_layer, True)
 
-        parents_b = np.unique(
-            cg.get_roots(parents_b, stop_layer=mlayer, ceil=False, time_stamp=parent_ts)
-        )
-
+        parents_b = np.unique(_get_new_nodes(cg, parents_b, mlayer, parent_ts))
         parents_a = np.array([node_a] * parents_b.size, dtype=basetypes.NODE_ID)
         return np.column_stack((parents_a, parents_b))
 
@@ -607,8 +624,6 @@ def get_latest_edges_wrapper(
                 stale_edge_layers,
                 parent_ts=parent_ts,
             )
-            stale_nodes = get_stale_nodes(cg, latest_edges.ravel(), parent_ts=parent_ts)
-            assert stale_nodes.size == 0, f"latest_edges failed, stale: {stale_nodes}"
             logging.debug(f"{stale_edges} -> {latest_edges}; {parent_ts}")
             _new_cx_edges.append(latest_edges)
         new_cx_edges_d[layer] = np.concatenate(_new_cx_edges)