bug: proposal generation (#45)

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

src/deep_neurographs/geometry_utils.py

@@ -79,8 +79,8 @@ def fit_spline(xyz, s=None):
     return spline_x, spline_y, spline_z
 
 
-def sample_path(path, num_points):
-    t = np.linspace(0, 1, num_points)
+def sample_path(path, n_points):
+    t = np.linspace(0, 1, n_points)
     if len(path) > 5:
         spline_x, spline_y, spline_z = fit_spline(path, s=10)
         path = np.column_stack((spline_x(t), spline_y(t), spline_z(t)))
@@ -431,10 +431,40 @@ def dist(v_1, v_2, metric="l2"):
         return distance.euclidean(v_1, v_2)
 
 
-def make_line(xyz_1, xyz_2, num_steps):
+def check_dists(xyz_1, xyz_2, xyz_3, radius):
+    """
+    Checks whether distance between "xyz_1", "xyz_3" and "xyz_2", "xyz_3" is 
+    sufficiently small. Routine is used during edge proposal generation to 
+    determine whether to create new vertex at "xyz_2" or draw proposal between
+    "xyz_1" and existing node at "xyz_3".
+
+    Parameters
+    ----------
+    xyz_1 : np.ndarray
+        xyz coordinate of leaf node (i.e. source of edge proposal).
+    xyz_2 : np.ndarray
+        xyz coordinate queried from kdtree (i.e. dest of edge proposal).
+    xyz_3 : np.ndarray
+        xyz coordinate of existing node in graph that is near "xyz_2".
+    radius : float
+        Maximum Euclidean length of edge proposal.
+
+    Parameters
+    ----------
+    bool
+        Indication of whether to draw edge proposal between "xyz_1" and
+        "xyz_3".
+
+    """
+    d_1 = dist(xyz_1, xyz_3) < radius
+    d_2 = dist(xyz_2, xyz_3) < 5
+    return True if d_1 and d_2 else False
+
+
+def make_line(xyz_1, xyz_2, n_steps):
     xyz_1 = np.array(xyz_1)
     xyz_2 = np.array(xyz_2)
-    t_steps = np.linspace(0, 1, num_steps)
+    t_steps = np.linspace(0, 1, n_steps)
     return np.array([(1 - t) * xyz_1 + t * xyz_2 for t in t_steps], dtype=int)
 
 

src/deep_neurographs/graph_utils.py

@@ -92,11 +92,6 @@ def get_irreducibles(swc_dict, swc_id=None, prune=True, depth=16, smooth=True):
             nbs = utils.append_dict_value(nbs, root, j)
             nbs = utils.append_dict_value(nbs, j, root)
             root = None
-
-            if all(attrs["xyz"][0] == attrs["xyz"][-1]):
-                print(root, j)
-                print(attrs)
-                stop
 
     # Output
     leafs = set_node_attrs(swc_dict, leafs)

src/deep_neurographs/intake.py

@@ -154,7 +154,7 @@ def build_neurograph_from_gcs_zips(
 
     # Generate proposals
     if search_radius > 0:
-        print("Generate edge proposals...")
+        print("Generate Edge Proposals...")
         t0 = time()
         neurograph.generate_proposals(
             search_radius, n_proposals_per_leaf=n_proposals_per_leaf

src/deep_neurographs/neurograph.py

@@ -19,7 +19,7 @@
 from deep_neurographs import graph_utils as gutils
 from deep_neurographs import utils
 from deep_neurographs.densegraph import DenseGraph
-from deep_neurographs.geometry_utils import dist as get_dist
+from deep_neurographs.geometry_utils import dist as get_dist, check_dists
 
 SUPPORTED_LABEL_MASK_TYPES = [dict, np.array, ts.TensorStore]
 
@@ -126,7 +126,7 @@ def __add_nodes(self, nodes, key, node_ids, cur_id, swc_id):
     # --- Proposal Generation ---
     def generate_proposals(
         self,
-        search_radius,
+        radius,
         n_proposals_per_leaf=3,
         optimize=False,
         optimization_depth=10,
@@ -146,23 +146,22 @@ def generate_proposals(
                 continue
             xyz_leaf = self.nodes[leaf]["xyz"]
             proposals = self.__get_proposals(
-                leaf, xyz_leaf, n_proposals_per_leaf, search_radius
+                leaf, xyz_leaf, n_proposals_per_leaf, radius
             )
             for xyz in proposals:
                 # Extract info on mutable connection
                 (i, j) = self.xyz_to_edge[xyz]
                 attrs = self.get_edge_data(i, j)
 
                 # Get connecting node
-                contained_j = self.is_contained(j)
-                if get_dist(xyz, attrs["xyz"][0]) < 10 and self.is_contained(
-                    i
-                ):
+                d1 = check_dists(xyz_leaf, xyz, self.nodes[i]["xyz"], radius)
+                d2 = check_dists(xyz_leaf, xyz, self.nodes[j]["xyz"], radius)
+                if d1 and self.is_contained(i):
+                    xyz = deepcopy(self.nodes[i]["xyz"])
                     node = i
-                    xyz = self.nodes[node]["xyz"]
-                elif get_dist(xyz, attrs["xyz"][-1]) < 10 and contained_j:
+                elif d2 and self.is_contained(j):
+                    xyz = deepcopy(self.nodes[j]["xyz"])
                     node = j
-                    xyz = self.nodes[node]["xyz"]
                 else:
                     idxs = np.where(np.all(attrs["xyz"] == xyz, axis=1))[0]
                     node = self.add_immutable_node((i, j), attrs, idxs[0])
@@ -176,7 +175,7 @@ def generate_proposals(
             self.run_optimization()
 
     def __get_proposals(
-        self, query_id, query_xyz, n_proposals_per_leaf, search_radius
+        self, query_id, query_xyz, n_proposals_per_leaf, radius
     ):
         """
         Generates edge proposals for node "query_id" by finding points on
@@ -190,7 +189,7 @@ def __get_proposals(
             (x,y,z) coordinates of the query node.
         n_proposals_per_leaf : int
             Number of proposals generated from node "query_id".
-        search_radius : float
+        radius : float
             Maximum Euclidean length of edge proposal.
 
         Returns
@@ -203,7 +202,7 @@ def __get_proposals(
         best_xyz = dict()
         best_dist = dict()
         query_swc_id = self.nodes[query_id]["swc_id"]
-        for xyz in self._query_kdtree(query_xyz, search_radius):
+        for xyz in self._query_kdtree(query_xyz, radius):
             # Check whether xyz is contained
             if not self.is_contained(xyz, buffer=36):
                 continue