fix(edits): account for fake edges when finding latest edges

Author: akhileshh

pychunkedgraph/graph/chunks/atomic.py

@@ -62,4 +62,6 @@ def get_bounding_atomic_chunks(
     chunkedgraph_meta: ChunkedGraphMeta, layer: int, chunk_coords: Sequence[int]
 ) -> List:
     """Atomic chunk coordinates along the boundary of a chunk"""
-    return get_bounding_children_chunks(chunkedgraph_meta, layer, chunk_coords, 2)
+    return get_bounding_children_chunks(
+        chunkedgraph_meta, layer, tuple(chunk_coords), 2
+    )

pychunkedgraph/graph/chunks/utils.py

@@ -1,11 +1,13 @@
 # pylint: disable=invalid-name, missing-docstring
 
-from typing import List
 from typing import Union
 from typing import Optional
 from typing import Sequence
+from typing import Tuple
 from typing import Iterable
 
+from functools import lru_cache
+
 import numpy as np
 
 
@@ -210,8 +212,9 @@ def _get_chunk_coordinates_from_vol_coordinates(
     return coords.astype(int)
 
 
+@lru_cache()
 def get_bounding_children_chunks(
-    cg_meta, layer: int, chunk_coords: Sequence[int], children_layer, return_unique=True
+    cg_meta, layer: int, chunk_coords: Tuple[int], children_layer, return_unique=True
 ) -> np.ndarray:
     """Children chunk coordinates at given layer, along the boundary of a chunk"""
     chunk_coords = np.array(chunk_coords, dtype=int)

pychunkedgraph/graph/edges/__init__.py

@@ -3,7 +3,7 @@
 """
 
 from collections import namedtuple
-import datetime
+import datetime, logging
 from os import environ
 from copy import copy
 from typing import Iterable, Optional
@@ -14,7 +14,10 @@
 from graph_tool import Graph
 
 from pychunkedgraph.graph import types
-from pychunkedgraph.graph.chunks import utils as chunk_utils
+from pychunkedgraph.graph.chunks.utils import (
+    get_bounding_children_chunks,
+    get_chunk_ids_from_coords,
+)
 from pychunkedgraph.graph.utils import basetypes
 
 from ..utils import basetypes
@@ -235,7 +238,9 @@ def get_latest_edges(
     These parents would be the new identities for the stale `partner`.
     """
     _nodes = np.unique(stale_edges[:, 1])
-    nodes_ts_map = dict(zip(_nodes, cg.get_node_timestamps(_nodes, return_numpy=False)))
+    nodes_ts_map = dict(
+        zip(_nodes, cg.get_node_timestamps(_nodes, return_numpy=False, normalize=True))
+    )
     _nodes = np.unique(stale_edges)
     layers, coords = cg.get_chunk_layers_and_coordinates(_nodes)
     layers_d = dict(zip(_nodes, layers))
@@ -252,31 +257,55 @@ def _get_normalized_coords(node_a, node_b) -> tuple:
             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):
+    def _get_l2chunkids_along_boundary(mlayer: int, coord_a, coord_b, padding: int = 0):
+        """
+        Gets L2 Chunk IDs along opposing faces for larger chunks.
+        If padding is enabled, more faces of L2 chunks are padded on both sides.
+        This is necessary to find fake edges that can span more than 2 L2 chunks.
+        """
         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]
+        major_axis = np.argmax(np.abs(direction))
+        bounds_a = get_bounding_children_chunks(cg.meta, mlayer, tuple(coord_a), 2)
+        bounds_b = get_bounding_children_chunks(cg.meta, mlayer, tuple(coord_b), 2)
+
+        l2chunk_count = 2 ** (mlayer - 2)
+        max_coord = coord_a if direction[major_axis] > 0 else coord_b
+
+        skip = abs(direction[major_axis]) - 1
+        l2_skip = skip * l2chunk_count
 
-        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)
+        mid = max_coord[major_axis] * l2chunk_count
+        face_a = mid if direction[major_axis] > 0 else (mid - l2_skip - 1)
+        face_b = mid if direction[major_axis] < 0 else (mid - l2_skip - 1)
+
+        l2chunks_a = [bounds_a[bounds_a[:, major_axis] == face_a]]
+        l2chunks_b = [bounds_b[bounds_b[:, major_axis] == face_b]]
+
+        step_a, step_b = (1, -1) if direction[major_axis] > 0 else (-1, 1)
+        for _ in range(padding):
+            _l2_chunks_a = copy(l2chunks_a[-1])
+            _l2_chunks_b = copy(l2chunks_b[-1])
+            _l2_chunks_a[:, major_axis] += step_a
+            _l2_chunks_b[:, major_axis] += step_b
+            l2chunks_a.append(_l2_chunks_a)
+            l2chunks_b.append(_l2_chunks_b)
+
+        l2chunks_a = np.concatenate(l2chunks_a)
+        l2chunks_b = np.concatenate(l2chunks_b)
+
+        l2chunk_ids_a = get_chunk_ids_from_coords(cg.meta, 2, l2chunks_a)
+        l2chunk_ids_b = 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 _get_filtered_l2ids(node_a, node_b, padding: int):
+        """
+        Finds L2 IDs along opposing faces for given nodes.
+        Filterting is done by first finding L2 chunks along these faces.
+        Then get their parent chunks iteratively.
+        Then filter children iteratively using these chunks.
+        """
+        chunks_map = {}
+
         def _filter(node):
             result = []
             children = np.array([node], dtype=basetypes.NODE_ID)
@@ -294,59 +323,67 @@ def _filter(node):
                 children = cg.get_children(children[mask], flatten=True)
             return np.concatenate(result)
 
-        return _filter(node_a), _filter(node_b)
-
-    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_a, chunks_b = _get_l2chunkids_along_boundary(
+            mlayer, coord_a, coord_b, padding
+        )
 
-        chunks_map[node_a] = [np.array([cg.get_chunk_id(node_a)])]
-        chunks_map[node_b] = [np.array([cg.get_chunk_id(node_b)])]
+        chunks_map[node_a] = [[cg.get_chunk_id(node_a)]]
+        chunks_map[node_b] = [[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]).astype(basetypes.NODE_ID)
-        chunks_map[node_b] = np.concatenate(chunks_map[node_b]).astype(basetypes.NODE_ID)
+        chunks_map[node_a] = np.concatenate(chunks_map[node_a])
+        chunks_map[node_b] = np.concatenate(chunks_map[node_b])
+        return int(mlayer), _filter(node_a), _filter(node_b)
 
-        l2ids_a, l2ids_b = _get_filtered_l2ids(node_a, node_b, chunks_map)
+    result = [types.empty_2d]
+    for edge_layer, _edge in zip(edge_layers, stale_edges):
+        node_a, node_b = _edge
+        mlayer, l2ids_a, l2ids_b = _get_filtered_l2ids(node_a, node_b, padding=0)
+        if l2ids_a.size == 0 or l2ids_b.size == 0:
+            logging.info(f"{node_a}, {node_b}, expanding search with padding.")
+            mlayer, l2ids_a, l2ids_b = _get_filtered_l2ids(node_a, node_b, padding=2)
+            logging.info(f"Found {l2ids_a} and {l2ids_b}")
+
+        _edges = []
         edges_d = cg.get_cross_chunk_edges(
             node_ids=l2ids_a, time_stamp=nodes_ts_map[node_b], raw_only=True
         )
-
-        _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)
 
-        children_b = cg.get_children(_edges[mask][:, 1], flatten=True)
+        try:
+            _edges = np.concatenate(_edges)
+        except ValueError as exc:
+            logging.warning(f"No edges found for {node_a}, {node_b}")
+            raise ValueError from exc
 
+        mask = np.isin(_edges[:, 1], l2ids_b)
         parents_a = _edges[mask][:, 0]
+        children_b = cg.get_children(_edges[mask][:, 1], flatten=True)
         parents_b = np.unique(cg.get_parents(children_b, time_stamp=parent_ts))
-        _cx_edges_d = cg.get_cross_chunk_edges(parents_b)
+        _cx_edges_d = cg.get_cross_chunk_edges(parents_b, time_stamp=parent_ts)
         parents_b = []
         for _node, _edges_d in _cx_edges_d.items():
             for _edges in _edges_d.values():
-                _mask = np.isin(_edges[:,1], parents_a)
+                _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(
-                parents_b, stop_layer=mlayer, ceil=False, time_stamp=parent_ts
-            )
+            cg.get_roots(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)))
+        _new_edges = np.column_stack((parents_a, parents_b))
+        assert _new_edges.size, f"No edge found for {node_a}, {node_b} at {parent_ts}"
+        result.append(_new_edges)
     return np.concatenate(result)
 
 

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -1,6 +1,6 @@
 # pylint: disable=invalid-name, missing-docstring, c-extension-no-member
 
-import math, random, time
+import logging, math, random, time
 import multiprocessing as mp
 from collections import defaultdict
 
@@ -171,6 +171,7 @@ def update_chunk(
         args = (cg_info, layer, chunk, ts_chunk, earliest_ts)
         tasks.append(args)
 
+    logging.info(f"Processing {len(nodes)} nodes.")
     with mp.Pool(min(mp.cpu_count(), len(tasks))) as pool:
         _ = list(
             tqdm(