feat: leaf-branch proposals

Author: anna-grim

src/neuron_proofreader/proposal_graph.py

@@ -193,17 +193,23 @@ def add_proposal(self, i, j):
         self.node_proposals[j].add(i)
         self.proposals.add(proposal)
 
-    def generate_proposals(self, search_radius):
+    def generate_proposals(self, search_radius, allow_nonleaf_targets=False):
         """
         Generates proposals from leaf nodes.
 
         Parameters
         ----------
         search_radius : float
             Search radius used to generate proposals.
+        allow_nonleaf_targets : bool, optional
+            Indication of whether to generate proposals between leaf and nodes
+            with degree 2. Default is False.
         """
         # Proposal generation
-        proposals = self.proposal_generator(search_radius)
+        proposals = self.proposal_generator(
+            search_radius, allow_nonleaf_targets=allow_nonleaf_targets
+        )
+
         self.search_radius = search_radius
         self.store_proposals(proposals)
         self.trim_proposals()

src/neuron_proofreader/skeleton_graph.py

@@ -19,7 +19,7 @@
 import numpy as np
 import zipfile
 
-from neuron_proofreader.utils import graph_util as gutil, img_util, util
+from neuron_proofreader.utils import geometry_util, graph_util, img_util, util
 
 
 class SkeletonGraph(nx.Graph):
@@ -82,7 +82,7 @@ def __init__(
 
         # Graph Loader
         anisotropy = anisotropy if use_anisotropy else (1.0, 1.0, 1.0)
-        self.graph_loader = gutil.GraphLoader(
+        self.graph_loader = graph_util.GraphLoader(
             anisotropy=anisotropy,
             min_size=min_size,
             node_spacing=node_spacing,
@@ -686,15 +686,15 @@ def nodes_with_segment_id(self, segment_id):
                 )
         return nodes
 
-    def nodes_within_distance(self, root, max_depth):
+    def nodes_within_distance(self, root, max_dist):
         """
         Gets nodes connected to the given root node up to a certain depth.
 
         Parameters
         ----------
         root : int
             Node ID and root of search.
-        max_depth : float
+        max_dist : float
             Maximum distance (in microns) between returned nodes and root.
 
         Returns
@@ -711,7 +711,7 @@ def nodes_within_distance(self, root, max_depth):
             # Populate queue
             for j in self.neighbors(i):
                 dist_j = dist_i + self.dist(i, j)
-                if dist_j < max_depth and j not in visited:
+                if dist_j < max_dist and j not in visited:
                     queue.append((j, dist_j))
                     visited.add(j)
         return list(visited)
@@ -750,6 +750,25 @@ def path_from_leaf(self, leaf, max_depth=np.inf):
                     path.append(j)
         return path
 
+    def path_length(self, path):
+        """
+        Computes the length of the given path.
+
+        Parameters
+        ----------
+        path : List[int]
+            List of nodes that forms a path.
+
+        Returns
+        -------
+        Length of the given path.
+        """
+        if len(path) > 1:
+            diffs = self.node_xyz[path[1:]] - self.node_xyz[path[:-1]]
+            return np.sqrt(np.sum(diffs ** 2))
+        else:
+            return 0
+
     def rooted_subgraph(self, root, radius):
         """
         Gets a rooted subgraph with the given radius (in microns).
@@ -795,25 +814,6 @@ def rooted_subgraph(self, root, radius):
         subgraph.node_xyz = self.node_xyz[idxs]
         return subgraph
 
-    def path_length(self, path):
-        """
-        Computes the length of the given path.
-
-        Parameters
-        ----------
-        path : List[int]
-            List of nodes that forms a path.
-
-        Returns
-        -------
-        Length of the given path.
-        """
-        if len(path) > 1:
-            diffs = self.node_xyz[path[1:]] - self.node_xyz[path[:-1]]
-            return np.sqrt(np.sum(diffs ** 2))
-        else:
-            return 0
-
     def set_kdtree(self):
         """
         Initializes KD-Tree from node xyz coordinates.
@@ -852,3 +852,20 @@ def summary(self, prefix=""):
         summary.append(f"# Edges: {n_edges}")
         summary.append(f"Memory Consumption: {memory:.2f} GBs")
         return "\n".join(summary)
+
+    def tangent_from_leaf(self, leaf, max_depth=np.inf):
+        """
+        Computes the tangent vector of the path emanating from the given leaf.
+
+        Parameters
+        ----------
+        leaf : int
+            Node ID and starting point of path extracted.
+        max_depth : float
+            Maximum depth (in microns) of path extracted.
+
+        Returns
+        -------
+        """
+        path = self.path_from_leaf(leaf, max_depth=max_depth)
+        return geometry_util.tangent(self.node_xyz[np.array(path)])

src/neuron_proofreader/split_proofreading/proposal_generation.py

@@ -71,8 +71,8 @@ def __call__(self, initial_radius, allow_nonleaf_targets=False):
             with degree 2. Default is False.
         """
         # Initializations
-        self.set_kdtree()
         self.allow_nonleaf_targets = allow_nonleaf_targets
+        self.set_kdtree()
         iterator = self.graph.leaf_nodes()
         if self.graph.verbose:
             iterator = tqdm(iterator, desc="Proposal Generation")
@@ -141,7 +141,7 @@ def find_node_candidates(self, leaf, radius):
         """
         node_candidates = list()
         for node in self.get_nearby_nodes(leaf, radius):
-            # UPDATE - check whether to move to leaf if allowing leaf-branch connections
+            node = self.adjust_position(leaf, node, radius)
             if self.is_valid_proposal(leaf, node):
                 node_candidates.append(node)
         return node_candidates
@@ -178,6 +178,73 @@ def get_nearby_nodes(self, leaf, radius):
         return [val["node"] for val in pts_dict.values()]
 
     # --- Helpers ---
+    def adjust_position(self, leaf, node, radius):
+        """
+        Adjusts a candidate node by searching for a nearby leaf, or falling
+        back to angular alignment if none is found.
+
+        Parameters
+        ----------
+        leaf : int
+            Node ID of the leaf node.
+        node : int
+            Node ID of the candidate node.
+        radius : float
+            Maximum path distance (in microns) for searching neighboring
+            nodes.
+
+        Returns
+        -------
+        node : int
+            Node ID of candidate.
+        """
+        # Check if node is close to another leaf
+        queue = [(node, 0)]
+        visited = {node}
+        while queue:
+            # Visit node
+            i, dist_i = queue.pop()
+            if self.graph.degree[i] == 1:
+                return i
+
+            # Update queue
+            for j in self.graph.neighbors(i):
+                dist_j = dist_i + self.graph.dist(i, j)
+                if dist_j < radius and j not in visited:
+                    queue.append((j, dist_j))
+                    visited.add(j)
+        return self.maximize_branch_angle(leaf, node)
+
+    def maximize_branch_angle(self, leaf, node):
+        """
+        Selects a nearby node that maximizes the angular alignment with a
+        leaf's tangent.
+
+        Parameters
+        ----------
+        leaf : int
+            Node ID of the leaf node.
+        node : int
+            Node ID of the candidate node. The search for better candidates is
+            restricted to nodes within a fixed distance of this node.
+
+        Returns
+        -------
+        node : int
+            Node ID of the node that maximizes angular alignment with the leaf
+            tangent.
+        """
+        max_inner_product = 0
+        leaf_tangent = self.graph.tangent_from_leaf(leaf, max_depth=20)
+        for i in self.graph.nodes_within_distance(node, 30):
+            pts = [self.graph.node_xyz[i], self.graph.node_xyz[leaf]]
+            proposal_tangent = geometry_util.tangent(pts)
+            inner_product = abs(np.sum(leaf_tangent * proposal_tangent))
+            if inner_product > max_inner_product:
+                max_inner_product = inner_product
+                node = i
+        return node
+
     def is_valid_proposal(self, leaf, i):
         """
         Determines whether a pair of nodes satisfies the following:
@@ -226,10 +293,8 @@ def query_nbhd(self, node, radius):
         else:
             nodes = list()
             for idx in self.kdtree.query_ball_point(xyz, radius):
-                xyz = self.kdtree.data[idx]
-                node = self.graph.closest_node(xyz)
+                node = self.graph.closest_node(self.kdtree.data[idx])
                 nodes.append(node)
-                assert self.graph.degree[node] == 1
             return nodes
 
     def select_closest_components(self, pts_dict):

src/neuron_proofreader/utils/ml_util.py

@@ -63,6 +63,14 @@ def build_network(self, input_dim, output_dim, n_layers):
 
     @staticmethod
     def _init_weights(m):
+        """
+        Initializes weights for linear layers using Kaiming initialization.
+
+        Parameters
+        ----------
+        m : torch.nn.Module
+            Module to initialize.
+        """
         if isinstance(m, nn.Linear):
             nn.init.kaiming_normal_(m.weight, nonlinearity="leaky_relu")
             if m.bias is not None:
@@ -116,8 +124,24 @@ def init_mlp(input_dim, hidden_dim, output_dim, dropout=0.1):
 
 # --- Data Structures ---
 class TensorDict(dict):
+    """
+    A class for model inputs in a dictionary.
+    """
 
     def to(self, device):
+        """
+        Moves dictionary values to the specified GPU device.
+
+        Parameters
+        ----------
+        device : str
+            Name of GPU device to move inputs to.
+
+        Returns
+        -------
+        TensorDict
+            Dictionary with values moved to the specified GPU device.
+        """
         return TensorDict({k: self.move(v, device) for k, v in self.items()})
 
     def move(self, v, device):