From d4dabd58f95ce691e8eecba484111b7af61bf3d6 Mon Sep 17 00:00:00 2001
From: Anna Grim <108307071+anna-grim@users.noreply.github.com>
Date: Thu, 18 Jan 2024 13:50:39 -0800
Subject: [PATCH] Optimize graph build (#35)

* build: combine whole brain graphs

* bug: edge attrs

---------

Co-authored-by: anna-grim <anna.grim@alleninstitute.org>
---
 src/deep_neurographs/graph_utils.py |  7 +++++-
 src/deep_neurographs/intake.py      | 34 +++++++++++++++++++----------
 2 files changed, 29 insertions(+), 12 deletions(-)

diff --git a/src/deep_neurographs/graph_utils.py b/src/deep_neurographs/graph_utils.py
index f4cfe5a..bc49fd7 100644
--- a/src/deep_neurographs/graph_utils.py
+++ b/src/deep_neurographs/graph_utils.py
@@ -28,7 +28,7 @@
 import networkx as nx
 import numpy as np
 
-from deep_neurographs import geometry_utils, swc_utils
+from deep_neurographs import geometry_utils, swc_utils, utils
 
 
 def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True):
@@ -240,6 +240,11 @@ def upd_edge_attrs(swc_dict, attrs, i):
 
 def get_edge_attr(graph, edge, attr):
     edge_data = graph.get_edge_data(*edge)
+    print("here")
+    print(edge)
+    print(graph.edges[edge])
+    print(edge_data)
+    print(attr)
     return edge_data[attr]
 
 
diff --git a/src/deep_neurographs/intake.py b/src/deep_neurographs/intake.py
index ac3ca16..a6eaa3b 100644
--- a/src/deep_neurographs/intake.py
+++ b/src/deep_neurographs/intake.py
@@ -28,7 +28,7 @@
 OPTIMIZATION_DEPTH = 15
 PRUNE = True
 PRUNE_DEPTH = 16
-SEARCH_RADIUS = 0
+SEARCH_RADIUS = 10
 MIN_SIZE = 35
 SMOOTH = True
 
@@ -135,12 +135,16 @@ def build_neurograph_from_gcs_zips(
 
     """
     # Process swc files
-    t0 = time()
+    print("Process swc files...")
+    total_runtime, t0 = utils.init_timers()
     swc_dicts = download_gcs_zips(bucket_name, cloud_path, min_size)
+    t_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
     t, unit = utils.time_writer(time() - t0)
-    print(f"\ndownload_gcs_zips(): {t} {unit} \n")
+    print(f"\nTime Stamp: {round(t_stamp, 4)} {unit_stamp}")
+    print(f"Module Runtime(): {round(t, 4)} {unit} \n")
 
     # Build neurograph
+    print("Build NeuroGraph...")
     t0 = time()
     neurograph = build_neurograph(
         swc_dicts,
@@ -150,17 +154,25 @@ def build_neurograph_from_gcs_zips(
         smooth=smooth,
     )
     t, unit = utils.time_writer(time() - t0)
-    print(f"build_neurograph(): {t} {unit} \n")
+    t_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
+    print(f"Time Stamp: {round(t_stamp, 4)} {unit_stamp}")
+    print(f"Module Runtime(): {round(t, 4)} {unit} \n")
 
     # Generate proposals
     if search_radius > 0:
-        t0 = time()
         print("Generate edge proposals...")
+        t0 = time()
         neurograph.generate_proposals(
             search_radius, n_proposals_per_leaf=n_proposals_per_leaf
         )
         t, unit = utils.time_writer(time() - t0)
-    print(f"generate_proposals(): {t} {unit} \n")
+        time_stamp, unit_stamp = utils.time_writer(time() - total_runtime)
+        print("\n# proposals:", len(neurograph.mutable_edges))
+        print(f"Time Stamp: {round(t_stamp, 4)} {unit_stamp}")
+        print(f"Module Runtime(): {round(t, 4)} {unit} \n")
+
+    t, unit = utils.time_writer(time() - total_runtime)
+    print(f"Total Runtime: {round(t, 4)} {unit}")
     return neurograph
 
 
@@ -188,13 +200,12 @@ def download_gcs_zips(bucket_name, cloud_path, min_size):
     bucket = storage_client.bucket(bucket_name)
     zip_paths = utils.list_gcs_filenames(bucket, cloud_path, ".zip")
     chunk_size = int(len(zip_paths) * 0.02)
-    print(f"# zip files: {len(zip_paths)} \n")
+    print(f"# zip files: {len(zip_paths)}")
 
     # Parse
     cnt = 1
     t0, t1 = utils.init_timers()
     swc_dicts = dict()
-    print("Process swc files...")
     for i, path in enumerate(zip_paths):
         swc_dicts.update(process_gsc_zip(bucket, path, min_size=min_size))
         if i > cnt * chunk_size:
@@ -224,14 +235,14 @@ def build_neurograph(
 ):
     # Extract irreducibles
     n_components = len(swc_dicts)
-    print("Extract irreducible nodes and edges...")
+    print("(1) Extract irreducible nodes and edges")
     print("# connected components:", utils.reformat_number(n_components))
     irreducibles, n_nodes, n_edges = get_irreducibles(
         swc_dicts, prune=prune, prune_depth=prune_depth, smooth=smooth
     )
 
     # Build neurograph
-    print("Build graph...")
+    print("(2) Combine irreducibles...")
     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)
@@ -244,7 +255,8 @@ def build_neurograph(
         if i > cnt * chunk_size:
             cnt, t1 = report_progress(i, n_components, chunk_size, cnt, t0, t1)
         i += 1
-    print("\n" + f"add_irreducibles(): {time() - t0} seconds")
+    t, unit = utils.time_writer(time() - t0)
+    print("\n" + f"add_irreducibles(): {round(t, 4)} {unit}")
     return neurograph