feat: long range proposals, profiles (#170)

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

src/deep_neurographs/generate_proposals.py

@@ -15,7 +15,9 @@
 ENDPOINT_DIST = 10
 
 
-def run(neurograph, search_radius, complex_bool=True):
+def run(
+    neurograph, search_radius, complex_bool=True, long_range_proposals=False
+):
     """
     Generates proposals emanating from "leaf".
 
@@ -28,6 +30,9 @@ def run(neurograph, search_radius, complex_bool=True):
     complex_bool : bool, optional
         Indication of whether to generate complex proposals. The default is
         True.
+    long_range_proposals : bool
+        Indication of whether to generate simple proposals within distance of
+        2 * search_radius of leaf.
 
     Returns
     -------
@@ -38,13 +43,24 @@ def run(neurograph, search_radius, complex_bool=True):
     connections = dict()
     for leaf in neurograph.leafs:
         neurograph, connections = run_on_leaf(
-            neurograph, connections, leaf, search_radius, complex_bool
+            neurograph,
+            connections,
+            leaf,
+            search_radius,
+            complex_bool,
+            long_range_proposals,
         )
-    # neurograph.filter_nodes()
     return neurograph
 
 
-def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool):
+def run_on_leaf(
+    neurograph,
+    connections,
+    leaf,
+    search_radius,
+    complex_bool,
+    long_range_proposals,
+):
     """
     Generates proposals emanating from "leaf".
 
@@ -62,6 +78,9 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool):
         Maximum Euclidean distance between endpoints of proposal.
     complex_bool : bool
         Indication of whether to generate complex proposals.
+    long_range_proposals : bool
+        Indication of whether to generate simple proposals within distance of
+        2 * search_radius of leaf.
 
     Returns
     -------
@@ -72,10 +91,17 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool):
         were added to "neurograph".
 
     """
+    # Get candidates
     leaf_swc_id = neurograph.nodes[leaf]["swc_id"]
-    for xyz in get_candidates(neurograph, leaf, search_radius):
+    candidates = get_candidates(neurograph, leaf, search_radius)
+    if len(candidates) == 0 and long_range_proposals:
+        candidates = get_candidates(neurograph, leaf, 2 * search_radius)
+        candidates = parse_long_range(neurograph, candidates, leaf)
+
+    # Parse candidates
+    for xyz in candidates:
         # Get connection
-        neurograph, node = get_conection(neurograph, leaf, xyz, search_radius)
+        neurograph, node = get_conection(neurograph, leaf, xyz)
         if not complex_bool and neurograph.degree[node] > 1:
             continue
 
@@ -100,6 +126,19 @@ def run_on_leaf(neurograph, connections, leaf, search_radius, complex_bool):
     return neurograph, connections
 
 
+def parse_long_range(neurograph, candidates, leaf):
+    hit_swc_ids = set()
+    filtered_candidates = []
+    for xyz in candidates:
+        neurograph, i = get_conection(neurograph, leaf, xyz)
+        if neurograph.degree[i] > 1:
+            continue
+        else:
+            filtered_candidates.append(xyz)
+            hit_swc_ids.add(neurograph.nodes[i]["swc_id"])
+    return filtered_candidates if len(hit_swc_ids) == 1 else []
+
+
 def get_candidates(neurograph, leaf, search_radius):
     """
     Generates proposals for node "leaf" in "neurograph" by finding candidate
@@ -150,11 +189,10 @@ def get_best_candidates(neurograph, candidates, dists):
         return list(candidates.values())
 
 
-def get_conection(neurograph, leaf, xyz, search_radius):
+def get_conection(neurograph, leaf, xyz):
     edge = neurograph.xyz_to_edge[xyz]
     node, d = get_closer_endpoint(neurograph, edge, xyz)
     if d > ENDPOINT_DIST:
-        # or neurograph.dist(leaf, node) > search_radius:
         attrs = neurograph.get_edge_data(*edge)
         idx = np.where(np.all(attrs["xyz"] == xyz, axis=1))[0][0]
         node = neurograph.split_edge(edge, attrs, idx)

src/deep_neurographs/geometry.py

@@ -16,7 +16,7 @@
 
 
 # Directional Vectors
-def get_directional(neurograph, i, origin, window_size):
+def get_directional(neurograph, i, origin, depth):
     """
     Computes the directional vector of a branch or bifurcation in a neurograph
     relative to a specified origin.
@@ -30,9 +30,9 @@ def get_directional(neurograph, i, origin, window_size):
     origin : numpy.ndarray
         The origin point xyz relative to which the directional vector is
         computed.
-    window_size : numpy.ndarry
-        The size of the window around the branch or bifurcation to consider
-        for computing the directional vector.
+    depth : numpy.ndarry
+        The size of the window in microns around the branch or bifurcation to
+        consider for computing the directional vector.
 
     Returns
     -------
@@ -44,26 +44,26 @@ def get_directional(neurograph, i, origin, window_size):
     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))
+        return compute_tangent(get_subarray(branches[0], depth))
     elif len(branches) == 2:
-        branch_1 = get_subarray(branches[0], window_size)
-        branch_2 = get_subarray(branches[1], window_size)
+        branch_1 = get_subarray(branches[0], depth)
+        branch_2 = get_subarray(branches[1], depth)
         branch = np.concatenate((branch_1, branch_2))
         return compute_tangent(branch)
     else:
         return np.array([0, 0, 0])
 
 
-def get_subarray(arr, window_size):
+def get_subarray(arr, depth):
     """
     Extracts a sub-array of a specified window size from a given input array.
 
     Parameters
     ----------
     branch : numpy.ndarray
         Array from which the sub-branch will be extracted.
-    window_size : int
-        Size of the window to extract from "arr".
+    depth : int
+        Size of the window in microns to extract from "arr".
 
     Returns
     -------
@@ -72,10 +72,12 @@ def get_subarray(arr, window_size):
         smaller than the window size, the entire branch array is returned.
 
     """
-    if arr.shape[0] < window_size:
-        return arr
-    else:
-        return arr[0:window_size, :]
+    length = 0
+    for i in range(1, arr.shape[0]):
+        length += dist(arr[i - 1], arr[i])
+        if length > depth:
+            return arr[0:i, :]
+    return arr
 
 
 def compute_svd(xyz):

src/deep_neurographs/machine_learning/feature_generation.py

@@ -273,10 +273,6 @@ def proposal_profiles(neurograph, proposals, img):
     return profiles
 
 
-def generate_edge_profiles(neurograph, img):
-    pass
-
-
 def get_profile(img, coords, thread_id):
     """
     Gets the image intensity profile for a given proposal.
@@ -301,12 +297,24 @@ def get_profile(img, coords, thread_id):
     coords["bbox"]["max"] = [coords["bbox"]["max"][i] + 1 for i in range(3)]
     chunk = img_utils.read_tensorstore_with_bbox(img, coords["bbox"])
     chunk = img_utils.normalize(chunk)
-    profile = [chunk[tuple(xyz)] for xyz in coords["profile_path"]]
+    profile = read_intensities(chunk, coords)
     avg, std = utils.get_avg_std(profile)
     profile.extend([avg, std])
     return thread_id, profile
 
 
+def read_intensities(img, coords):
+    profile = []
+    for xyz in coords["profile_path"]:
+        start = xyz - 1
+        end = xyz + 2
+        val = np.max(
+            img[start[0]: end[0], start[1]: end[1], start[2]: end[2]]
+        )
+        profile.append(val)
+    return profile
+
+
 def get_proposal_profile_coords(neurograph, proposal):
     """
     Gets coordinates needed to compute an image intensity profile.
@@ -330,10 +338,14 @@ def get_proposal_profile_coords(neurograph, proposal):
     coord_0 = utils.to_voxels(xyz_0)
     coord_1 = utils.to_voxels(xyz_1)
 
-    # Store coordinates
+    # Store local coordinates
     bbox = utils.get_minimal_bbox(coord_0, coord_1)
-    start = [coord_0[i] - bbox["min"][i] for i in range(3)]
-    end = [coord_1[i] - bbox["min"][i] for i in range(3)]
+    start = [coord_0[i] - bbox["min"][i] + 1 for i in range(3)]
+    end = [coord_1[i] - bbox["min"][i] + 1 for i in range(3)]
+
+    # Shift bbox
+    bbox["min"] = [bbox["min"][i] - 1 for i in range(3)]
+    bbox["max"] = [bbox["max"][i] + 2 for i in range(3)]
     coords = {
         "bbox": bbox,
         "profile_path": geometry.make_line(start, end, N_PROFILE_PTS),

src/deep_neurographs/machine_learning/graph_trainer.py

@@ -27,6 +27,7 @@
 
 # Training
 LR = 1e-3
+MODEL_TYPE = "GraphNeuralNet"
 N_EPOCHS = 200
 SCHEDULER_GAMMA = 0.5
 SCHEDULER_STEP_SIZE = 1000
@@ -224,7 +225,7 @@ def forward(self, data):
 
         """
         self.optimizer.zero_grad()
-        x, edge_index = toGPU(data)
+        x, edge_index = gnn_utils.get_inputs(data, MODEL_TYPE)
         hat_y = self.model(x, edge_index)
         y = data.y  # .to("cuda:0", dtype=torch.float32)
         return y, truncate(hat_y, y)

src/deep_neurographs/neurograph.py

@@ -275,6 +275,7 @@ def generate_proposals(
         self,
         search_radius,
         complex_bool=True,
+        long_range_proposals=False,
         proposals_per_leaf=3,
         optimize=False,
         optimization_depth=10,
@@ -294,7 +295,10 @@ def generate_proposals(
 
         # Main
         self = generate_proposals.run(
-            self, search_radius, complex_bool=complex_bool
+            self,
+            search_radius,
+            complex_bool=complex_bool,
+            long_range_proposals=long_range_proposals,
         )
 
         # Finish