fix(upgrade): use timestamps of partners at layers > 2

Author: akhileshh

pychunkedgraph/__init__.py

@@ -1 +1 @@
-__version__ = "3.0.7"
+__version__ = "3.0.8"

pychunkedgraph/ingest/upgrade/parent_layer.py

@@ -66,7 +66,14 @@ def _populate_cx_edges_with_timestamps(
     for node, node_ts in zip(nodes, nodes_ts):
         CX_EDGES[node] = {}
         timestamps = timestamps_d[node]
-        timestamps.add(node_ts)
+        cx_edges_d_node_ts = _get_cx_edges_at_timestamp(node, response, node_ts)
+
+        edges = np.concatenate([empty_2d] + list(cx_edges_d_node_ts.values()))
+        partner_parent_ts_d = get_parent_timestamps(cg, edges[:, 1])
+        for v in partner_parent_ts_d.values():
+            timestamps.update(v)
+        CX_EDGES[node][node_ts] = cx_edges_d_node_ts
+
         for ts in sorted(timestamps):
             if ts < earliest_ts:
                 ts = earliest_ts
@@ -152,21 +159,25 @@ def update_chunk(
     nodes_ts = cg.get_node_timestamps(nodes, return_numpy=False, normalize=True)
     _populate_cx_edges_with_timestamps(cg, layer, nodes, nodes_ts, earliest_ts)
 
-    task_size = int(math.ceil(len(nodes) / mp.cpu_count() / 2))
-    chunked_nodes = chunked(nodes, task_size)
-    chunked_nodes_ts = chunked(nodes_ts, task_size)
-    cg_info = cg.get_serialized_info()
-
-    tasks = []
-    for chunk, ts_chunk in zip(chunked_nodes, chunked_nodes_ts):
-        args = (cg_info, layer, chunk, ts_chunk, earliest_ts)
-        tasks.append(args)
-
-    with mp.Pool(min(mp.cpu_count(), len(tasks))) as pool:
-        _ = list(
-            tqdm(
-                pool.imap_unordered(_update_cross_edges_helper, tasks),
-                total=len(tasks),
+    if nodes:
+        for node, node_ts in zip(nodes, nodes_ts):
+            update_cross_edges(cg, layer, node, node_ts, earliest_ts)
+    else:
+        task_size = int(math.ceil(len(nodes) / mp.cpu_count() / 2))
+        chunked_nodes = chunked(nodes, task_size)
+        chunked_nodes_ts = chunked(nodes_ts, task_size)
+        cg_info = cg.get_serialized_info()
+
+        tasks = []
+        for chunk, ts_chunk in zip(chunked_nodes, chunked_nodes_ts):
+            args = (cg_info, layer, chunk, ts_chunk, earliest_ts)
+            tasks.append(args)
+
+        with mp.Pool(min(mp.cpu_count(), len(tasks))) as pool:
+            _ = list(
+                tqdm(
+                    pool.imap_unordered(_update_cross_edges_helper, tasks),
+                    total=len(tasks),
+                )
             )
-        )
     print(f"total elaspsed time: {time.time() - start}")