bug: fixed branch retrieval (#130)

Co-authored-by: anna-grim <[email protected]>

src/deep_neurographs/geometry.py

@@ -41,7 +41,7 @@ def get_directional(neurograph, i, origin, window_size):
         specified origin.
 
     """
-    branches = neurograph.get_branches(i)
+    branches = neurograph.get_branches(i, ignore_reducibles=True)
     branches = shift_branches(branches, origin)
     if len(branches) == 1:
         return compute_tangent(get_subarray(branches[0], window_size))
@@ -341,8 +341,8 @@ def optimize_simple_alignment(neurograph, img, edge, depth=15):
 
     """
     i, j = tuple(edge)
-    branch_i = neurograph.get_branches(i)[0]
-    branch_j = neurograph.get_branches(j)[0]
+    branch_i = neurograph.get_branches(i, ignore_reducibles=True)[0]
+    branch_j = neurograph.get_branches(j, ignore_reducibles=True,)[0]
     d_i, d_j, _ = align(neurograph, img, branch_i, branch_j, depth)
     return branch_i[d_i], branch_j[d_j]
 

src/deep_neurographs/machine_learning/feature_generation.py

@@ -367,8 +367,8 @@ def get_directionals(neurograph, proposal, window_size):
 
 def get_avg_radii(neurograph, proposal):
     i, j = tuple(proposal)
-    radii_i = neurograph.get_branches(i, key="radius")
-    radii_j = neurograph.get_branches(j, key="radius")
+    radii_i = neurograph.get_branches(i, ignore_reducibles=True, key="radius")
+    radii_j = neurograph.get_branches(j, ignore_reducibles=True, key="radius")
     return np.array([get_avg_radius(radii_i), get_avg_radius(radii_j)])
 
 
@@ -382,8 +382,8 @@ def get_avg_radius(radii_list):
 
 def get_avg_branch_lens(neurograph, proposal):
     i, j = tuple(proposal)
-    branches_i = neurograph.get_branches(i, key="xyz")
-    branches_j = neurograph.get_branches(j, key="xyz")
+    branches_i = neurograph.get_branches(i, ignore_reducibles=True)
+    branches_j = neurograph.get_branches(j, ignore_reducibles=True)
     return np.array([get_branch_len(branches_i), get_branch_len(branches_j)])
 
 

src/deep_neurographs/neurograph.py

@@ -512,12 +512,47 @@ def dist(self, i, j):
         """
         return get_dist(self.nodes[i]["xyz"], self.nodes[j]["xyz"])
 
-    def get_branches(self, i, key="xyz"):
+    def get_branches(self, i, ignore_reducibles=False, key="xyz"):
         branches = []
         for j in self.neighbors(i):
-            branches.append(self.orient_edge((i, j), i, key=key))
+            branch = self.orient_edge((i, j), i, key=key)
+            if ignore_reducibles:
+                root = i
+                while self.degree[j] == 2:
+                    k = self.get_other_nb(j, root)
+                    branch_jk = self.orient_edge((j, k), j, key=key)
+                    if key == "xyz":
+                        branch = np.vstack([branch, branch_jk])
+                    else:
+                        branch = np.concatenate((branch, branch_jk))
+                    root = j
+                    j = k
+            branches.append(branch)
         return branches
 
+    def get_other_nb(self, i, j):
+        """
+        Gets the other neighbor of node "i" which is not "j" such that "j" is
+        a neighbor of node "i".
+
+        Parameters
+        ----------
+        i : int
+            Node with degree 2.
+        j : int
+            Neighbor of node "i"
+
+        Returns
+        -------
+        int
+            Neighbor of node "i" which is not "j".
+
+        """
+        assert self.degree[i] == 2, "Node does not have degree 2."
+        nbs = list(self.neighbors(i))
+        nbs.remove(j)
+        return nbs[0]
+
     def orient_edge(self, edge, i, key="xyz"):
         if (self.edges[edge][key][0] == self.nodes[i][key]).all():
             return self.edges[edge][key]

src/deep_neurographs/visualization.py

@@ -44,13 +44,45 @@ def visualize_connected_components(
         plot(data, title)
 
 
-def visualize_graph(graph, title="Initial Segmentation"):
+def visualize_graph(graph, title=""):
+    """
+    Visualizes the graph with nodes and edges.
+
+    Parameters
+    ----------
+        graph : networkx.Graph
+            Graph to be visualized.
+        title : str, optional
+            Title of the plot. Default is "".
+
+    Returns
+    -------
+        None
+
+    """
     data = plot_edges(graph, graph.edges)
     data.append(plot_nodes(graph))
     plot(data, title)
 
 
-def visualize_proposals(graph, target_graph=None, title="Edge Proposals"):
+def visualize_proposals(graph, target_graph=None, title="Proposals"):
+    """
+    Visualizes a graph with proposals.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        Graph to be visualized.
+    target_graph : networkx.Graph, optional
+        Graph generated from ground truth tracings. The default is None.
+    title : str, optional
+        Title of the plot. Default is "Proposals".
+
+    Returns
+    -------
+    None
+
+    """
     visualize_subset(
         graph,
         graph.proposals,
@@ -60,7 +92,26 @@ def visualize_proposals(graph, target_graph=None, title="Edge Proposals"):
     )
 
 
-def visualize_targets(graph, target_graph=None, title="Target Edges"):
+def visualize_targets(
+    graph, target_graph=None, title="Ground Truth - Accepted Proposals"
+):
+    """
+    Visualizes a graph and its ground truth accept proposals.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        Graph to be visualized.
+    target_graph : networkx.Graph, optional
+        Graph generated from ground truth tracings. The default is None.
+    title : str, optional
+        Title of the plot. Default is "Ground Truth - Accepted Proposals".
+
+    Returns
+    -------
+    None
+
+    """
     visualize_subset(
         graph,
         graph.target_edges,
@@ -72,19 +123,43 @@ def visualize_targets(graph, target_graph=None, title="Target Edges"):
 
 def visualize_subset(
     graph,
-    edges,
+    subset,
     line_width=5,
     proposal_subset=False,
     target_graph=None,
     title="",
 ):
+    """
+    Visualizes a graph and a subset of edges or proposals.
+
+    Parameters
+    ----------
+    graph : networkx.Graph
+        Graph to be visualized.
+    subset : container
+        Subset of edges or proposals to be visualized.
+    line_width : int, optional
+        Line width used to plot "subset". The default is 5.
+    proposals_subset : bool, optional
+        Indication of whether "subset" is a subset of proposals. The default
+        is False.
+    target_graph : networkx.Graph, optional
+        Graph generated from ground truth tracings. The default is None.
+    title : str, optional
+        Title of the plot. Default is "Proposals".
+
+    Returns
+    -------
+    None
+
+    """
     # Plot graph
     data = plot_edges(graph, graph.edges, color="black")
     data.append(plot_nodes(graph))
     if proposal_subset:
-        data.extend(plot_proposals(graph, edges, line_width=line_width))
+        data.extend(plot_proposals(graph, subset, line_width=line_width))
     else:
-        data.extend(plot_edges(graph, edges, line_width=line_width))
+        data.extend(plot_edges(graph, subset, line_width=line_width))
 
     # Add target graph (if applicable)
     if target_graph: