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/chunkedgraph.py

@@ -216,6 +216,7 @@ def get_parents(
                         if fail_to_zero:
                             parents.append(0)
                         else:
+                            exc.add_note(f"timestamp: {time_stamp}")
                             raise KeyError from exc
                 parents = np.array(parents, dtype=basetypes.NODE_ID)
             else:
@@ -1016,7 +1017,7 @@ def get_chunk_layers_and_coordinates(self, node_or_chunk_ids: typing.Sequence):
         """
         node_or_chunk_ids = np.array(node_or_chunk_ids, dtype=basetypes.NODE_ID)
         layers = self.get_chunk_layers(node_or_chunk_ids)
-        chunk_coords = np.zeros(shape=(len(node_or_chunk_ids), 3))
+        chunk_coords = np.zeros(shape=(len(node_or_chunk_ids), 3), dtype=int)
         for _layer in np.unique(layers):
             mask = layers == _layer
             _nodes = node_or_chunk_ids[mask]

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(
@@ -279,7 +279,7 @@ def _get_l2chunkids_along_boundary(max_layer, coord_a, coord_b):
     def _get_filtered_l2ids(node_a, node_b, chunks_map):
         def _filter(node):
             result = []
-            children = cg.get_children(node)
+            children = np.array([node], dtype=basetypes.NODE_ID)
             while True:
                 chunk_ids = cg.get_chunk_ids_from_node_ids(children)
                 mask = np.isin(chunk_ids, chunks_map[node])
@@ -296,24 +296,24 @@ def _filter(node):
 
         return _filter(node_a), _filter(node_b)
 
-    result = []
+    result = [types.empty_2d]
     chunks_map = {}
     for edge_layer, _edge in zip(edge_layers, stale_edges):
         node_a, node_b = _edge
         mlayer, coord_a, coord_b = _get_normalized_coords(node_a, node_b)
         chunks_a, chunks_b = _get_l2chunkids_along_boundary(mlayer, coord_a, coord_b)
 
-        chunks_map[node_a] = []
-        chunks_map[node_b] = []
+        chunks_map[node_a] = [np.array([cg.get_chunk_id(node_a)])]
+        chunks_map[node_b] = [np.array([cg.get_chunk_id(node_b)])]
         _layer = 2
         while _layer < mlayer:
             chunks_map[node_a].append(chunks_a)
             chunks_map[node_b].append(chunks_b)
             chunks_a = np.unique(cg.get_parent_chunk_id_multiple(chunks_a))
             chunks_b = np.unique(cg.get_parent_chunk_id_multiple(chunks_b))
             _layer += 1
-        chunks_map[node_a] = np.concatenate(chunks_map[node_a])
-        chunks_map[node_b] = np.concatenate(chunks_map[node_b])
+        chunks_map[node_a] = np.concatenate(chunks_map[node_a]).astype(basetypes.NODE_ID)
+        chunks_map[node_b] = np.concatenate(chunks_map[node_b]).astype(basetypes.NODE_ID)
 
         l2ids_a, l2ids_b = _get_filtered_l2ids(node_a, node_b, chunks_map)
         edges_d = cg.get_cross_chunk_edges(
@@ -326,32 +326,57 @@ def _filter(node):
         _edges = np.concatenate(_edges)
         mask = np.isin(_edges[:, 1], l2ids_b)
 
-        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 = np.unique(
-            cg.get_roots(
-                children_a, stop_layer=mlayer, ceil=False, time_stamp=parent_ts
-            )
-        )
-        assert parents_a.size == 1 and parents_a[0] == node_a, (
-            node_a,
-            parents_a,
-            children_a,
-        )
 
+        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_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 = np.isin(_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,
+        )
+    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
@@ -104,14 +105,17 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
                 assert not exists_as_parent(cg, node, edges[:, 0]), f"{node}, {node_ts}"
                 return rows
 
+    row_id = serializers.serialize_uint64(node)
     for ts, cx_edges_d in CX_EDGES[node].items():
-        edges = np.concatenate([empty_2d] + list(cx_edges_d.values()))
+        if node_ts > ts:
+            continue
+        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(
@@ -121,7 +125,6 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
             layer_edges = np.unique(layer_edges, axis=0)
             col = Connectivity.CrossChunkEdge[_layer]
             val_dict[col] = layer_edges
-        row_id = serializers.serialize_uint64(node)
         rows.append(cg.client.mutate_row(row_id, val_dict, time_stamp=ts))
     return rows
 

pychunkedgraph/ingest/utils.py

@@ -202,6 +202,7 @@ def queue_layer_helper(parent_layer: int, imanager: IngestionManager, fn):
                 )
             )
         q.enqueue_many(job_datas)
+        logging.info(f"Queued {len(job_datas)} chunks.")
 
 
 def job_type_guard(job_type: str):