Merge branch 'main' into add-features

src/deep_neurographs/graph_utils.py

@@ -7,17 +7,19 @@
 
 Routines that extract the irreducible components of a graph.
 
---define what an irreducible is
-    leafs : set
-        Nodes with degreee 1.
-    junctions : set
-        Nodes with degree > 2.
-    edges : dict
-        Set of edges connecting nodes in leafs and junctions. The keys are
-        pairs of nodes connected by an edge and values are a dictionary of
-        attributes.
 
---define what a branch is
+Terminology
+------------
+
+Leaf: a node with degree 1.
+
+Junction: a node with degree > 2.
+
+Irreducibles: the irreducibles of a graph G=(V,E) consists of (1) leaf nodes
+V_l, (2) junction nodes, and (3) 
+junction nodes along 
+
+Branch: the sequence of nodes between two 
 
 """
 
@@ -26,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):
@@ -58,17 +60,20 @@ def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True):
         all irreducibles from the graph of that type.
 
     """
-    # Initializations
+    # Build dense graph
     dense_graph = swc_utils.to_graph(swc_dict)
     if prune:
         dense_graph = prune_short_branches(dense_graph, depth)
 
-    # Extract irreducibles
+    # Extract nodes
     leafs, junctions = get_irreducible_nodes(dense_graph, swc_dict)
-    assert len(leafs) > 0, "Error: swc with no leaf nodes!"
-    root = None
+    if len(leafs) == 0:
+        return False, None
+
+    # Extract edges
     edges = dict()
     nbs = dict()
+    root = None
     for (i, j) in nx.dfs_edges(dense_graph, source=sample(leafs, 1)[0]):
         # Check if start of path is valid
         if root is None:
@@ -84,8 +89,8 @@ def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True):
                 )
             else:
                 edges[(root, j)] = attrs
-            nbs = append_value(nbs, root, j)
-            nbs = append_value(nbs, j, root)
+            nbs = utils.append_dict_value(nbs, root, j)
+            nbs = utils.append_dict_value(nbs, j, root)
             root = None
 
     # Output
@@ -196,6 +201,7 @@ def get_leafs(graph):
 
 def __smooth_branch(swc_dict, attrs, edges, nbs, root, j):
     attrs["xyz"] = geometry_utils.smooth_branch(np.array(attrs["xyz"]), s=10)
+    attrs["radius"] = np.array(attrs["radius"])
     swc_dict, edges = upd_xyz(swc_dict, attrs, edges, nbs, root, 0)
     swc_dict, edges = upd_xyz(swc_dict, attrs, edges, nbs, j, -1)
     edges[(root, j)] = attrs
@@ -213,14 +219,6 @@ def upd_xyz(swc_dict, attrs, edges, nbs, i, start_or_end):
     return swc_dict, edges
 
 
-def append_value(my_dict, key, value):
-    if key in my_dict.keys():
-        my_dict[key].append(value)
-    else:
-        my_dict[key] = [value]
-    return my_dict
-
-
 def upd_branch_endpoint(edges, key, old_xyz, new_xyz):
     if all(edges[key]["xyz"][0] == old_xyz):
         edges[key]["xyz"][0] = new_xyz
@@ -242,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]
 
 

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)
@@ -243,7 +254,9 @@ def build_neurograph(
         neurograph.add_immutables(irreducible_set, key)
         if i > cnt * chunk_size:
             cnt, t1 = report_progress(i, n_components, chunk_size, cnt, t0, t1)
-    print(f"add_irreducibles(): {time() - t0} seconds")
+        i += 1
+    t, unit = utils.time_writer(time() - t0)
+    print("\n" + f"add_irreducibles(): {round(t, 4)} {unit}")
     return neurograph
 
 
@@ -277,9 +290,10 @@ def get_irreducibles(
         progress_cnt = 1
         for i, process in enumerate(as_completed(processes)):
             process_id, result = process.result()
-            irreducibles[process_id] = result
-            n_nodes += len(result["leafs"]) + len(result["junctions"])
-            n_edges += len(result["edges"])
+            if process_id:
+                irreducibles[process_id] = result
+                n_nodes += len(result["leafs"]) + len(result["junctions"])
+                n_edges += len(result["edges"])
             if i > progress_cnt * chunk_size:
                 progress_cnt, t1 = report_progress(
                     i, n_components, chunk_size, progress_cnt, t0, t1
@@ -289,15 +303,6 @@ def get_irreducibles(
     return irreducibles, n_nodes, n_edges
 
 
-def get_start_ids(swc_dicts):
-    start_id = 0
-    start_ids = dict()
-    for key in swc_dicts.keys():
-        start_ids[key] = start_id
-        start_id += len(swc_dicts[key]["id"])
-    return start_ids
-
-
 # -- Utils --
 def get_paths(swc_dir):
     paths = []
@@ -307,11 +312,8 @@ def get_paths(swc_dir):
 
 
 def report_progress(current, total, chunk_size, cnt, t0, t1):
-    # Compute
     eta = get_eta(current, total, chunk_size, t1)
     runtime = get_runtime(current, total, chunk_size, t0, t1)
-
-    # Write results
     utils.progress_bar(current, total, eta=eta, runtime=runtime)
     return cnt + 1, time()
 

src/deep_neurographs/neurograph.py

@@ -80,10 +80,10 @@ def init_densegraph(self):
         self.densegraph = DenseGraph(self.swc_paths)
 
     # --- Add nodes or edges ---
-    def add_immutables(self, irreducibles, swc_id, start_id=None):
+    def add_immutables(self, irreducibles, swc_id):
         # Nodes
         node_ids = dict()
-        cur_id = start_id if start_id else len(self.nodes)
+        cur_id = len(self.nodes)
         node_ids, cur_id = self.__add_nodes(
             irreducibles, "leafs", node_ids, cur_id, swc_id
         )
@@ -92,25 +92,19 @@ def add_immutables(self, irreducibles, swc_id, start_id=None):
         )
 
         # Add edges
-        edges = irreducibles["edges"]
-        for i, j in edges.keys():
-            # Get edge
-            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))
+        for edge, values in irreducibles["edges"].items():
+            i, j = edge
+            self.immutable_edges.add(frozenset((node_ids[i], node_ids[j])))
             self.add_edge(
-                node_ids[i], node_ids[j], xyz=xyz, radius=radii, swc_id=swc_id
+                node_ids[i],
+                node_ids[j],
+                radius=values["radius"],
+                xyz=values["xyz"],
+                swc_id=swc_id
             )
-            xyz_to_edge = dict((tuple(xyz), edge) for xyz in xyz)
-            check_xyz = set(xyz_to_edge.keys())
-            collisions = check_xyz.intersection(set(self.xyz_to_edge.keys()))
-            if len(collisions) > 0:
-                for xyz in collisions:
-                    del xyz_to_edge[xyz]
-            self.xyz_to_edge.update(xyz_to_edge)
+            for xyz in values["xyz"][::2]:
+                self.xyz_to_edge[tuple(xyz)] = (i, j)
+            self.xyz_to_edge[tuple(values["xyz"][-1])] = (i, j)
 
     def __add_nodes(self, nodes, key, node_ids, cur_id, swc_id):
         for i in nodes[key].keys():

src/deep_neurographs/swc_utils.py

@@ -127,18 +127,17 @@ def fast_parse(contents):
     contents, offset = get_contents(contents)
     min_id = np.inf
     swc_dict = {
-        "id": np.zeros((len(contents)), dtype=int),
-        "radius": np.zeros((len(contents)), dtype=float),
-        "pid": np.zeros((len(contents)), dtype=int),
-        "xyz": [],
+        "id": np.zeros((len(contents)), dtype=np.int32),
+        "radius": np.zeros((len(contents)), dtype=np.float32),
+        "pid": np.zeros((len(contents)), dtype=np.int32),
+        "xyz": np.zeros((len(contents), 3), dtype=np.int32),
     }
     for i, line in enumerate(contents):
         parts = line.split()
-        xyz = read_xyz(parts[2:5], offset=offset)
-        swc_dict["id"][i] = int(parts[0])
-        swc_dict["radius"][i] = float(parts[-2])
-        swc_dict["pid"][i] = int(parts[-1])
-        swc_dict["xyz"].append(xyz)
+        swc_dict["id"][i] = parts[0]
+        swc_dict["radius"][i] = parts[-2]
+        swc_dict["pid"][i] = parts[-1]
+        swc_dict["xyz"][i] = read_xyz(parts[2:5], offset=offset)
 
     # Reindex from zero
     min_id = np.min(swc_dict["id"])

src/deep_neurographs/utils.py

@@ -97,6 +97,32 @@ def remove_key(my_dict, key):
     return my_dict
 
 
+def append_dict_value(my_dict, key, value):
+    """
+    Appends "value" to the list stored at "key".
+
+    Parameters
+    ----------
+    my_dict : dict
+        Dictionary to be queried.
+    key : hashable data type
+        Key to be query.
+    value : list item type
+        Value to append to list stored at "key".
+
+    Returns
+    -------
+    my_dict : dict
+        Updated dictionary.
+
+    """
+    if key in my_dict.keys():
+        my_dict[key].append(value)
+    else:
+        my_dict[key] = [value]
+    return my_dict
+
+
 # --- os utils ---
 def mkdir(path, delete=False):
     """
@@ -500,7 +526,7 @@ def progress_bar(current, total, bar_length=50, eta=None, runtime=None):
     bar = f"[{'=' * progress}{' ' * (bar_length - progress)}]"
     eta = f"Time Remaining: {eta}" if eta else ""
     runtime = f"Estimated Total Runtime: {runtime}" if runtime else ""
-    print(f"\r{bar} {n_completed} | {eta} | {runtime}", end="", flush=True)
+    print(f"\r{bar} {n_completed} | {eta} | {runtime}    ", end="", flush=True)
 
 
 def xor(a, b):