fix(edits): more precise filter for latest edges; error on chunk_id mismatch (ingest); bump version

Author: akhileshh

.bumpversion.cfg

@@ -1,5 +1,5 @@
 [bumpversion]
-current_version = 3.0.13
+current_version = 3.1.0
 commit = True
 tag = True
 

pychunkedgraph/__init__.py

@@ -1 +1 @@
-__version__ = "3.0.13"
+__version__ = "3.1.0"

pychunkedgraph/graph/edges/__init__.py

@@ -220,7 +220,7 @@ def get_stale_nodes(
         )
         stale_mask = layer_nodes != _nodes
         stale_nodes.append(layer_nodes[stale_mask])
-    return np.concatenate(stale_nodes), edge_supervoxels
+    return np.concatenate(stale_nodes)
 
 
 def get_latest_edges(
@@ -328,13 +328,18 @@ def _filter(node):
 
         children_a = cg.get_children(_edges[mask][:, 0], flatten=True)
         children_b = cg.get_children(_edges[mask][:, 1], flatten=True)
-        if 85431849467249595 in children_a and 85502218144317440 in children_b:
-            print("woohoo0")
-            continue
 
-        if 85502218144317440 in children_a and 85431849467249595 in children_b:
-            print("woohoo1")
-            continue
+        parents_a = _edges[mask][:, 0]
+        parents_b = np.unique(cg.get_parents(children_b, time_stamp=parent_ts))
+        _cx_edges_d = cg.get_cross_chunk_edges(parents_b)
+        parents_b = []
+        for _node, _edges_d in _cx_edges_d.items():
+            for _edges in _edges_d.values():
+                _mask = np.isin(_edges[:,1], parents_a)
+                if np.any(_mask):
+                    parents_b.append(_node)
+
+        parents_b = np.array(parents_b, dtype=basetypes.NODE_ID)
         parents_a = np.unique(
             cg.get_roots(
                 children_a, stop_layer=mlayer, ceil=False, time_stamp=parent_ts
@@ -348,10 +353,43 @@ def _filter(node):
 
         parents_b = np.unique(
             cg.get_roots(
-                children_b, stop_layer=mlayer, ceil=False, time_stamp=parent_ts
+                parents_b, stop_layer=mlayer, ceil=False, time_stamp=parent_ts
             )
         )
 
         parents_a = np.array([node_a] * parents_b.size, dtype=basetypes.NODE_ID)
         result.append(np.column_stack((parents_a, parents_b)))
     return np.concatenate(result)
+
+
+def get_latest_edges_wrapper(
+    cg,
+    cx_edges_d: dict,
+    parent_ts: datetime.datetime = None,
+) -> np.ndarray:
+    """Helper function to filter stale edges and replace with latest edges."""
+    _cx_edges = [types.empty_2d]
+    _edge_layers = [types.empty_1d]
+    for k, v in cx_edges_d.items():
+        _cx_edges.append(v)
+        _edge_layers.append([k] * len(v))
+    _cx_edges = np.concatenate(_cx_edges)
+    _edge_layers = np.concatenate(_edge_layers, dtype=int)
+
+    edge_nodes = np.unique(_cx_edges)
+    stale_nodes = get_stale_nodes(cg, edge_nodes, parent_ts=parent_ts)
+    stale_nodes_mask = np.isin(edge_nodes, stale_nodes)
+
+    latest_edges = types.empty_2d.copy()
+    if np.any(stale_nodes_mask):
+        stalte_edges_mask = _cx_edges[:, 1] == stale_nodes
+        stale_edges = _cx_edges[stalte_edges_mask]
+        stale_edge_layers = _edge_layers[stalte_edges_mask]
+        latest_edges = get_latest_edges(
+            cg,
+            stale_edges,
+            stale_edge_layers,
+            parent_ts=parent_ts,
+        )
+        print(list(stale_edges[:, 1]), list(latest_edges[:, 1]))
+    return np.concatenate([_cx_edges, latest_edges])

pychunkedgraph/graph/edits.py

@@ -15,7 +15,7 @@
 from . import types
 from . import attributes
 from . import cache as cache_utils
-from .edges import get_latest_edges, get_stale_nodes
+from .edges import get_latest_edges, get_latest_edges_wrapper, get_stale_nodes
 from .edges.utils import concatenate_cross_edge_dicts
 from .edges.utils import merge_cross_edge_dicts
 from .utils import basetypes
@@ -587,36 +587,10 @@ def _update_cross_edge_cache(self, parent, children):
             children, time_stamp=self._last_successful_ts
         )
         cx_edges_d = concatenate_cross_edge_dicts(cx_edges_d.values())
-
-        _cx_edges = [types.empty_2d]
-        _edge_layers = [types.empty_1d]
-        for k, v in cx_edges_d.items():
-            _cx_edges.append(v)
-            _edge_layers.append([k] * len(v))
-        _cx_edges = np.concatenate(_cx_edges)
-        _edge_layers = np.concatenate(_edge_layers, dtype=int)
-
-        edge_nodes = np.unique(_cx_edges)
-        stale_nodes, edge_supervoxels = get_stale_nodes(
-            self.cg, edge_nodes, parent_ts=self._last_successful_ts
+        _cx_edges = get_latest_edges_wrapper(
+            self.cg, cx_edges_d, parent_ts=self._last_successful_ts
         )
-        stale_nodes_mask = np.isin(edge_nodes, stale_nodes)
-
-        latest_edges = types.empty_2d.copy()
-        if np.any(stale_nodes_mask):
-            stalte_edges_mask = _cx_edges[:, 1] == stale_nodes
-            stale_edges = _cx_edges[stalte_edges_mask]
-            stale_edge_layers = _edge_layers[stalte_edges_mask]
-            latest_edges = get_latest_edges(
-                self.cg,
-                stale_edges,
-                stale_edge_layers,
-                parent_ts=self._last_successful_ts,
-            )
-
-        _cx_edges = np.concatenate([_cx_edges, latest_edges])
-        edge_nodes  = np.unique(_cx_edges)
-
+        edge_nodes = np.unique(_cx_edges)
         edge_parents = self.cg.get_roots(
             edge_nodes,
             stop_layer=parent_layer,

pychunkedgraph/ingest/cluster.py

@@ -108,15 +108,8 @@ def _check_edges_direction(
     chunk_id = cg.get_chunk_id(layer=1, x=x, y=y, z=z)
     for edge_type in [EDGE_TYPES.between_chunk, EDGE_TYPES.cross_chunk]:
         edges = chunk_edges[edge_type]
-        e1 = edges.node_ids1
-        e2 = edges.node_ids2
-
-        e2_chunk_ids = cg.get_chunk_ids_from_node_ids(e2)
-        mask = e2_chunk_ids == chunk_id
-        e1[mask], e2[mask] = e2[mask], e1[mask]
-
-        e1_chunk_ids = cg.get_chunk_ids_from_node_ids(e1)
-        mask = e1_chunk_ids == chunk_id
+        chunk_ids = cg.get_chunk_ids_from_node_ids(edges.node_ids1)
+        mask = chunk_ids == chunk_id
         assert np.all(mask), "all IDs must belong to same chunk"
 
 

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -10,6 +10,7 @@
 
 from pychunkedgraph.graph import ChunkedGraph
 from pychunkedgraph.graph.attributes import Connectivity, Hierarchy
+from pychunkedgraph.graph.edges import get_latest_edges_wrapper
 from pychunkedgraph.graph.utils import serializers
 from pychunkedgraph.graph.types import empty_2d
 from pychunkedgraph.utils.general import chunked
@@ -105,19 +106,21 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
                 return rows
 
     for ts, cx_edges_d in CX_EDGES[node].items():
-        edges = np.concatenate([empty_2d] + list(cx_edges_d.values()))
+        edges = get_latest_edges_wrapper(cg, cx_edges_d, parent_ts=ts)
         if edges.size == 0:
             continue
-        nodes = np.unique(edges[:, 1])
-        svs = cg.get_single_leaf_multiple(nodes)
-        parents = cg.get_roots(svs, time_stamp=ts, stop_layer=layer, ceil=False)
-        edge_parents_d = dict(zip(nodes, parents))
+
+        edge_nodes = np.unique(edges)
+        parents = cg.get_roots(edge_nodes, time_stamp=ts, stop_layer=layer, ceil=False)
+        edge_parents_d = dict(zip(edge_nodes, parents))
         val_dict = {}
         for _layer, layer_edges in cx_edges_d.items():
             layer_edges = fastremap.remap(
                 layer_edges, edge_parents_d, preserve_missing_labels=True
             )
-            layer_edges[:, 0] = node
+            mask = layer_edges[:, 0] == node
+            assert np.all(mask), (node, layer_edges[~mask])
+            # layer_edges[:, 0] = node
             layer_edges = np.unique(layer_edges, axis=0)
             col = Connectivity.CrossChunkEdge[_layer]
             val_dict[col] = layer_edges