From ad1e1434082f92faf2050341c32b7a1a801aa8a9 Mon Sep 17 00:00:00 2001
From: anna-grim <anna.grim@alleninstitute.org>
Date: Tue, 16 Jan 2024 17:54:48 +0000
Subject: [PATCH] bug: add irreducibles to graph

---
 src/deep_neurographs/intake.py     | 12 ++++----
 src/deep_neurographs/neurograph.py | 48 ++++++++++++++++--------------
 2 files changed, 32 insertions(+), 28 deletions(-)

diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py
index 65ed526..5db4933 100644
--- a/src/deep_neurographs/intake.py
+++ b/src/deep_neurographs/intake.py
@@ -51,11 +51,9 @@ def build_neurograph_from_local(
     smooth=SMOOTH,
 ):
     # Process swc files
-    t0 = time()
     assert swc_dir or swc_paths, "Provide swc_dir or swc_paths!"
     bbox = utils.get_bbox(img_patch_origin, img_patch_shape)
     paths = get_paths(swc_dir) if swc_dir else swc_paths
-    t0 = time()
     swc_dicts = process_local_paths(paths, min_size, bbox=bbox)
 
     # Build neurograph
@@ -243,19 +241,21 @@ def build_neurograph(
     print("# nodes:", utils.reformat_number(n_nodes))
     print("# edges:", utils.reformat_number(n_edges))
     neurograph = NeuroGraph(bbox=bbox, img_path=img_path, swc_paths=swc_paths)
-    start_ids = get_start_ids(swc_dicts)
     t0, t1 = utils.init_timers()
-    for key in swc_dicts.keys():
+    n_components = len(irreducibles.keys())
+    chunk_size = int(n_components) * 0.05
+    cnt = 1
+    for i, key in enumerate(irreducibles.keys()):
         neurograph.add_immutables(
             irreducibles[key], key
         )
-        del swc_dicts[key]
         if i > cnt * chunk_size:
-            cnt, t1 = report_progress(i, len(swc_dicts), chunk_size, cnt, t0, t1)
+            cnt, t1 = report_progress(i, n_components, chunk_size, cnt, t0, t1)
     print(f"add_irreducibles(): {time() - t0} seconds")
 
     """
     t0 = time()
+    start_ids = get_start_ids(swc_dicts)
     with ThreadPoolExecutor() as executor:
         futures = {
             executor.submit(
diff --git a/src/deep_neurographs/neurograph.py b/src/deep_neurographs/neurograph.py
index 4a9f628..ed5d536 100644
--- a/src/deep_neurographs/neurograph.py
+++ b/src/deep_neurographs/neurograph.py
@@ -83,40 +83,28 @@ def init_densegraph(self):
 
     # --- Add nodes or edges ---
     def add_immutables(self, irreducibles, swc_id, start_id=None):
-        # Initializations
-        node_id = dict()
+        # Nodes
+        node_ids = dict()
         cur_id = start_id if start_id else len(self.nodes)
-        leaf_ids = irreducibles["leafs"].keys()
-        junction_ids = irreducibles["junctions"].keys()
-        for i in leaf_ids + junction_ids:
-            node_id[i] = cur_id
-            self.add_node(
-                node_id[i],
-                xyz=irreducibles[i]["xyz"],
-                radius=irreducibles[i]["radius"],
-                swc_id=swc_id,
-            )
-            cur_id += 1
-
-        # Update leafs and junctions
-        for leaf in irreducibles["leafs"].keys():
-            self.leafs.add(node_id[leaf])
-
-        for junction in irreducibles["junctions"].keys():
-            self.junctions.add(node_id[junction])
+        node_ids, cur_id = self.__add_nodes(
+            irreducibles, "leafs", node_ids, cur_id, swc_id
+        )
+        node_ids, cur_id = self.__add_nodes(
+            irreducibles, "junctions", node_ids, cur_id, swc_id
+        )
 
         # Add edges
         edges = irreducibles["edges"]
         for i, j in edges.keys():
             # Get edge
-            edge = (node_id[i], node_id[j])
+            edge = (node_ids[i], node_ids[j])
             xyz = np.array(edges[(i, j)]["xyz"])
             radii = np.array(edges[(i, j)]["radius"])
 
             # Add edge
             self.immutable_edges.add(frozenset(edge))
             self.add_edge(
-                node_id[i], node_id[j], xyz=xyz, radius=radii, swc_id=swc_id
+                node_ids[i], node_ids[j], xyz=xyz, radius=radii, swc_id=swc_id
             )
             xyz_to_edge = dict((tuple(xyz), edge) for xyz in xyz)
             check_xyz = set(xyz_to_edge.keys())
@@ -125,6 +113,22 @@ def add_immutables(self, irreducibles, swc_id, start_id=None):
                 for xyz in collisions:
                     del xyz_to_edge[xyz]
             self.xyz_to_edge.update(xyz_to_edge)
+    
+    def __add_nodes(self, nodes, key, node_ids, cur_id, swc_id):
+        for i in nodes[key].keys():
+            node_ids[i] = cur_id
+            self.add_node(
+                node_ids[i],
+                xyz=nodes[key][i]["xyz"],
+                radius=nodes[key][i]["radius"],
+                swc_id=swc_id,
+            )
+            if key == "leafs":
+                self.leafs.add(cur_id)
+            else:
+                self.junctions.add(cur_id)
+            cur_id += 1
+        return node_ids, cur_id
 
     # --- Proposal Generation ---
     def generate_proposals(