wip: find stale edges and their latest nodes

Author: akhileshh

pychunkedgraph/graph/chunkedgraph.py

@@ -940,6 +940,11 @@ def get_parent_chunk_id(
             self.meta, node_or_chunk_id, parent_layer
         )
 
+    def get_parent_chunk_id_multiple(self, node_or_chunk_ids: typing.Sequence):
+        return chunk_hierarchy.get_parent_chunk_id_multiple(
+            self.meta, node_or_chunk_ids
+        )
+
     def get_parent_chunk_ids(self, node_or_chunk_id: basetypes.NODE_ID):
         return chunk_hierarchy.get_parent_chunk_ids(self.meta, node_or_chunk_id)
 
@@ -984,3 +989,38 @@ def get_operation_ids(self, node_ids: typing.Sequence):
             except KeyError:
                 ...
         return result
+
+    def get_supervoxels(self, node_ids):
+        """Returns the first supervoxel found for each node_id."""
+        result = {}
+        node_ids_copy = np.copy(node_ids)
+        children = np.copy(node_ids)
+        children_d = self.get_children(node_ids)
+        while True:
+            children = [children_d[k][0] for k in children]
+            children = np.array(children, dtype=basetypes.NODE_ID)
+            mask = self.get_chunk_layers(children) == 1
+            result.update(
+                [(node, sv) for node, sv in zip(node_ids[mask], children[mask])]
+            )
+            node_ids = node_ids[~mask]
+            children = children[~mask]
+            if children.size == 0:
+                break
+            children_d = self.get_children(children)
+        return np.array([result[k] for k in node_ids_copy], dtype=basetypes.NODE_ID)
+
+    def get_chunk_layers_and_coordinates(self, node_or_chunk_ids: typing.Sequence):
+        """
+        Helper function that wraps get chunk layer and coordinates for nodes at any layer.
+        """
+        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))
+        for _layer in np.unique(layers):
+            mask = layers == _layer
+            _nodes = node_or_chunk_ids[mask]
+            chunk_coords[mask] = chunk_utils.get_chunk_coordinates_multiple(
+                self.meta, _nodes
+            )
+        return layers, chunk_coords

pychunkedgraph/graph/chunks/hierarchy.py

@@ -43,7 +43,7 @@ def get_children_chunk_ids(
     else:
         children_coords = get_children_chunk_coords(meta, layer, (x, y, z))
         children_chunk_ids = []
-        for (x, y, z) in children_coords:
+        for x, y, z in children_coords:
             children_chunk_ids.append(
                 utils.get_chunk_id(meta, layer=layer - 1, x=x, y=y, z=z)
             )
@@ -62,6 +62,19 @@ def get_parent_chunk_id(
     return utils.get_chunk_id(meta, layer=parent_layer, x=x, y=y, z=z)
 
 
+def get_parent_chunk_id_multiple(
+    meta: ChunkedGraphMeta, node_or_chunk_ids: np.ndarray
+) -> np.ndarray:
+    """Parent chunk IDs for multiple nodes. Assumes nodes at same layer."""
+
+    node_layers = utils.get_chunk_layers(meta, node_or_chunk_ids)
+    assert np.unique(node_layers).size == 1, np.unique(node_layers)
+    parent_layer = node_layers[0] + 1
+    coords = utils.get_chunk_coordinates_multiple(meta, node_or_chunk_ids)
+    coords = coords // meta.graph_config.FANOUT
+    return utils.get_chunk_ids_from_coords(meta, layer=parent_layer, coords=coords)
+
+
 def get_parent_chunk_ids(
     meta: ChunkedGraphMeta, node_or_chunk_id: np.uint64
 ) -> np.ndarray:

pychunkedgraph/graph/edges/__init__.py

@@ -3,14 +3,20 @@
 """
 
 from collections import namedtuple
+import datetime
 from os import environ
-from typing import Optional
+from copy import copy
+from typing import Iterable, Optional
 
 import numpy as np
 import tensorstore as ts
 import zstandard as zstd
 from graph_tool import Graph
 
+from pychunkedgraph.graph import types
+from pychunkedgraph.graph.chunks import utils as chunk_utils
+from pychunkedgraph.graph.utils import basetypes
+
 from ..utils import basetypes
 
 
@@ -189,3 +195,125 @@ def get_edges(source: str, nodes: np.ndarray) -> Edges:
         affinities=np.concatenate(affinities),
         areas=np.concatenate(areas),
     )
+
+
+def get_stale_nodes(
+    cg, edge_nodes: Iterable[basetypes.NODE_ID], parent_ts: datetime.datetime = None
+):
+    """
+    Checks to see if partner nodes in edges (edges[:,1]) are stale.
+    This is done by getting a supervoxel of the node and check
+    if it has a new parent at the same layer as the node.
+    """
+    edge_supervoxels = cg.get_supervoxels(edge_nodes)
+    # nodes can be at different layers due to skip connections
+    edge_nodes_layers = cg.get_chunk_layers(edge_nodes)
+    stale_nodes = [types.empty_1d]
+    for layer in np.unique(edge_nodes_layers):
+        layer_nodes = edge_nodes[edge_nodes_layers == layer]
+        _nodes = cg.get_roots(
+            edge_supervoxels,
+            stop_layer=layer,
+            ceil=False,
+            time_stamp=parent_ts,
+        )
+        stale_mask = layer_nodes != _nodes
+        stale_nodes.append(layer_nodes[stale_mask])
+    return np.concatenate(stale_nodes), edge_supervoxels
+
+
+def get_latest_edges(cg, stale_edges: Iterable, edge_layers: Iterable) -> dict:
+    """
+    For each of stale_edges [[`node`, `partner`]], get their L2 edge equivalent.
+    Then get supervoxels of those L2 IDs and get parent(s) at `node` level.
+    These parents would be the new identities for the stale `partner`.
+    """
+
+    _nodes = np.unique(stale_edges)
+    nodes_ts_map = dict(zip(_nodes, cg.get_node_timestamps(_nodes, return_numpy=False)))
+    layers, coords = cg.get_chunk_layers_and_coordinates(_nodes)
+    layers_d = dict(zip(_nodes, layers))
+    coords_d = dict(zip(_nodes, coords))
+
+    def _get_normalized_coords(node_a, node_b) -> tuple:
+        max_layer = layers_d[node_a]
+        coord_a, coord_b = coords_d[node_a], coords_d[node_b]
+        if layers_d[node_a] != layers_d[node_b]:
+            # normalize if nodes are not from the same layer
+            max_layer = max(layers_d[node_a], layers_d[node_b])
+            chunk_a = cg.get_parent_chunk_id(node_a, parent_layer=max_layer)
+            chunk_b = cg.get_parent_chunk_id(node_b, parent_layer=max_layer)
+            coord_a, coord_b = cg.get_chunk_coordinates_multiple([chunk_a, chunk_b])
+        return max_layer, coord_a, coord_b
+
+    def _get_l2chunkids_along_boundary(max_layer, coord_a, coord_b):
+        direction = coord_a - coord_b
+        axis = np.flatnonzero(direction)
+        assert len(axis) == 1, f"{direction}, {coord_a}, {coord_b}"
+        axis = axis[0]
+        children_a = chunk_utils.get_bounding_children_chunks(
+            cg.meta, max_layer, coord_a, children_layer=2
+        )
+        children_b = chunk_utils.get_bounding_children_chunks(
+            cg.meta, max_layer, coord_b, children_layer=2
+        )
+        if direction[axis] > 0:
+            mid = coord_a[axis] * 2 ** (max_layer - 2)
+            l2chunks_a = children_a[children_a[:, axis] == mid]
+            l2chunks_b = children_b[children_b[:, axis] == mid - 1]
+        else:
+            mid = coord_b[axis] * 2 ** (max_layer - 2)
+            l2chunks_a = children_a[children_a[:, axis] == mid - 1]
+            l2chunks_b = children_b[children_b[:, axis] == mid]
+
+        l2chunk_ids_a = chunk_utils.get_chunk_ids_from_coords(cg.meta, 2, l2chunks_a)
+        l2chunk_ids_b = chunk_utils.get_chunk_ids_from_coords(cg.meta, 2, l2chunks_b)
+        return l2chunk_ids_a, l2chunk_ids_b
+
+    def _get_filtered_l2ids(node_a, node_b, chunks_map):
+        def _filter(node):
+            result = []
+            children = cg.get_children(node)
+            while True:
+                chunk_ids = cg.get_chunk_ids_from_node_ids(children)
+                mask = np.isin(chunk_ids, chunks_map[node])
+                children = children[mask]
+
+                mask = cg.get_chunk_layers(children) == 2
+                result.append(children[mask])
+
+                mask = cg.get_chunk_layers(children) > 2
+                if children[mask].size == 0:
+                    break
+                children = cg.get_children(children[mask], flatten=True)
+            return np.concatenate(result)
+
+        return _filter(node_a), _filter(node_b)
+
+    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] = []
+        _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])
+
+        l2ids_a, l2ids_b = _get_filtered_l2ids(node_a, node_b, chunks_map)
+        edges_d = cg.get_cross_chunk_edges(node_ids=l2ids_a, time_stamp=nodes_ts_map[node_b])
+
+        _edges = []
+        for v in edges_d.values():
+            _edges.append(v.get(edge_layer, types.empty_2d))
+        _edges = np.concatenate(_edges)
+        mask = np.isin(_edges[:,1], l2ids_b)
+        print(_edges[mask])

pychunkedgraph/graph/edges/utils.py

@@ -135,7 +135,7 @@ def categorize_edges_v2(
 
 
 def get_cross_chunk_edges_layer(meta: ChunkedGraphMeta, cross_edges: Iterable):
-    """Computes the layer in which a cross chunk edge becomes relevant.
+    """Computes the layer in which an atomic cross chunk edge becomes relevant.
     I.e. if a cross chunk edge links two nodes in layer 4 this function
     returns 3.
     :param cross_edges: n x 2 array

pychunkedgraph/graph/edits.py

@@ -15,6 +15,7 @@
 from . import types
 from . import attributes
 from . import cache as cache_utils
+from .edges import get_stale_nodes
 from .edges.utils import concatenate_cross_edge_dicts
 from .edges.utils import merge_cross_edge_dicts
 from .utils import basetypes
@@ -248,6 +249,7 @@ def add_edges(
     new_roots = create_parents.run()
     sanity_check(cg, new_roots, operation_id)
     create_parents.create_new_entries()
+    raise ValueError("success merge")
     return new_roots, new_l2_ids, create_parents.new_entries
 
 
@@ -378,6 +380,7 @@ def remove_edges(
     new_roots = create_parents.run()
     sanity_check(cg, new_roots, operation_id)
     create_parents.create_new_entries()
+    raise ValueError("success split")
     return new_roots, new_l2_ids, create_parents.new_entries
 
 
@@ -497,25 +500,6 @@ def _update_neighbor_cross_edges(
     return updated_entries
 
 
-def get_supervoxels(cg, node_ids):
-    """Returns the first supervoxel found for each node_id."""
-    result = {}
-    node_ids_copy = np.copy(node_ids)
-    children = np.copy(node_ids)
-    children_d = cg.get_children(node_ids)
-    while True:
-        children = [children_d[k][0] for k in children]
-        children = np.array(children, dtype=basetypes.NODE_ID)
-        mask = cg.get_chunk_layers(children) == 1
-        result.update([(node, sv) for node, sv in zip(node_ids[mask], children[mask])])
-        node_ids = node_ids[~mask]
-        children = children[~mask]
-        if children.size == 0:
-            break
-        children_d = cg.get_children(children)
-    return np.array([result[k] for k in node_ids_copy], dtype=basetypes.NODE_ID)
-
-
 class CreateParentNodes:
     def __init__(
         self,
@@ -605,8 +589,28 @@ 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())
-        edge_nodes = np.unique(np.concatenate([*cx_edges_d.values(), types.empty_2d]))
-        edge_supervoxels = get_supervoxels(self.cg, edge_nodes)
+
+        _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
+        )
+        stale_mask = np.isin(edge_nodes, stale_nodes)
+
+        if np.any(stale_mask):
+            mask = _cx_edges[:, 1] == stale_nodes
+            stale_edges = _cx_edges[mask]
+            stalte_edge_layers = _edge_layers[mask]
+            # latest_parents = get_latest_nodes(self.cg, stale_edges)
+            raise ValueError()
+
         edge_parents = self.cg.get_roots(
             edge_supervoxels,
             stop_layer=parent_layer,

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -110,7 +110,7 @@ def update_cross_edges(cg: ChunkedGraph, layer, node, node_ts, earliest_ts) -> l
         if edges.size == 0:
             continue
         nodes = np.unique(edges[:, 1])
-        svs = get_supervoxels(cg, nodes)
+        svs = cg.get_supervoxels(nodes)
         parents = cg.get_roots(svs, time_stamp=ts, stop_layer=layer, ceil=False)
         edge_parents_d = dict(zip(nodes, parents))
         val_dict = {}