improved target edge generation (#23)

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

src/deep_neurographs/feature_extraction.py

@@ -8,19 +8,23 @@
 
 """
 
-import numpy as np
 from copy import deepcopy
-from deep_neurographs import utils, geometry_utils
 from random import sample
+
+import numpy as np
 from scipy.linalg import svd
 
+from deep_neurographs import geometry_utils, utils
+
 NUM_IMG_FEATURES = 0
 NUM_SKEL_FEATURES = 9
 NUM_PC_FEATURES = 0
 
 
 # -- Wrappers --
-def generate_mutable_features(neurograph, img=True, pointcloud=True, skel=True):
+def generate_mutable_features(
+    neurograph, img=True, pointcloud=True, skel=True
+):
     features = dict()
     if img:
         features["img"] = generate_img_features(neurograph)
@@ -74,15 +78,15 @@ def generate_mutable_skel_features(neurograph):
         radius_i, radius_j = get_radii(neurograph, edge)
 
         dot1, dot2, dot3 = get_directionals(neurograph, edge, 5)
-        ddot1, ddot2, ddot3 = get_directionals(neurograph, edge, 10)
-        features[edge] = np.concatenate((length, dot1, dot2, dot3), axis=None)
+        ddot1, ddot2, ddot3 = get_directionals(neurograph, edge, 5)
+        features[edge] = np.concatenate((length, radius_i, radius_j, dot1, dot2, dot3), axis=None)
     return features
 
 
 def compute_length(neurograph, edge, metric="l2"):
     i, j = tuple(edge)
     xyz_1, xyz_2 = neurograph.get_edge_attr("xyz", i, j)
-    return utils.dist(xyz_1, xyz_2, metric=metric)
+    return geometry_utils.dist(xyz_1, xyz_2, metric=metric)
 
 
 def get_directionals(neurograph, edge, window_size):
@@ -91,15 +95,19 @@ def get_directionals(neurograph, edge, window_size):
     mutable_xyz_i, mutable_xyz_j = neurograph.get_edge_attr("xyz", i, j)
     mutable_xyz = np.array([mutable_xyz_i, mutable_xyz_j])
     mutable_tangent = geometry_utils.compute_tangent(mutable_xyz)
-    context_tangent_1 = geometry_utils.compute_context_vec(neurograph, i, mutable_tangent, window_size=window_size)
-    context_tangent_2 = geometry_utils.compute_context_vec(neurograph, j, mutable_tangent, window_size=window_size)
-
+    context_tangent_i = geometry_utils.compute_context_vec(
+        neurograph, i, mutable_tangent, window_size=window_size
+    )
+    context_tangent_j = geometry_utils.compute_context_vec(
+        neurograph, j, mutable_tangent, window_size=window_size
+    )
+
     # Compute features
-    inner_product_1 = abs(np.dot(mutable_tangent, context_tangent_1))
-    inner_product_2 = abs(np.dot(mutable_tangent, context_tangent_2))
-    inner_product_3 = np.dot(context_tangent_1, context_tangent_2)
+    inner_product_1 = abs(np.dot(mutable_tangent, context_tangent_i))
+    inner_product_2 = abs(np.dot(mutable_tangent, context_tangent_j))
+    inner_product_3 = np.dot(context_tangent_i, context_tangent_j)
     return inner_product_1, inner_product_2, inner_product_3
-    
+
 
 def get_radii(neurograph, edge):
     i, j = tuple(edge)
@@ -114,15 +122,13 @@ def build_feature_matrix(neurographs, features, blocks):
     X = None
     block_to_idxs = dict()
     idx_to_edge = dict()
-    
+
     # Feature extraction
     for block_id in blocks:
         # Get features
         idx_shift = 0 if X is None else X.shape[0]
         X_i, y_i, idx_to_edge_i = build_feature_submatrix(
-            neurographs[block_id],
-            features[block_id],
-            idx_shift,
+            neurographs[block_id], features[block_id], idx_shift
         )
 
         # Concatenate
@@ -175,7 +181,6 @@ def combine_feature_vecs(features):
     return vec
 
 
-
 """
 
 def generate_node_features(neurograph, img=True, pointcloud=True, skel=True):
@@ -215,4 +220,4 @@ def generate_immutable_skel_features(neurograph):
 def _generate_immutable_skel_features(neurograph, edge):
     mean_radius = np.mean(neurograph.edges[edge]["radius"], axis=0)
     return np.concatenate((mean_radius), axis=None)
-"""
+"""

src/deep_neurographs/geometry_utils.py

@@ -1,17 +1,27 @@
 import numpy as np
-from deep_neurographs import utils
+from scipy.interpolate import CubicSpline, UnivariateSpline
 from scipy.linalg import svd
+from deep_neurographs import utils
 
 
 # Context Tangent Vectors
-def compute_context_vec(neurograph, i, mutable_tangent, window_size=5, return_pts=False, vec_type="tangent"):
+def compute_context_vec(
+    neurograph,
+    i,
+    mutable_tangent,
+    window_size=5,
+    return_pts=False,
+    vec_type="tangent",
+):
     # Compute context vecs
     branches = get_branches(neurograph, i)
     context_vec_list = []
     xyz_list = []
     ref_xyz = neurograph.nodes[i]["xyz"]
     for branch in branches:
-        context_vec, xyz = _compute_context_vec(branch, ref_xyz, window_size, vec_type)
+        context_vec, xyz = _compute_context_vec(
+            branch, ref_xyz, window_size, vec_type
+        )
         context_vec_list.append(context_vec)
         xyz_list.append(xyz)
 
@@ -69,10 +79,63 @@ def compute_svd(xyz):
 
 def compute_tangent(xyz):
     if xyz.shape[0] == 2:
-        tangent = (xyz[1] - xyz[0]) / utils.dist(xyz[1], xyz[0])
+        tangent = (xyz[1] - xyz[0]) / dist(xyz[1], xyz[0])
     else:
+        xyz = smooth_branch(xyz)
         U, S, VT = compute_svd(xyz)
         tangent = VT[0]
     return tangent / np.linalg.norm(tangent)
 
 
+# Smoothing
+def smooth_branch(xyz):
+    t = np.arange(len(xyz[:, 0]) + 12)
+    s = len(t) / 10
+    cs_x = UnivariateSpline(t, extend_boundary(xyz[:, 0]), s=s, k=3)
+    cs_y = UnivariateSpline(t, extend_boundary(xyz[:, 1]), s=s, k=3)
+    cs_z = UnivariateSpline(t, extend_boundary(xyz[:, 2]), s=s, k=3)
+    smoothed_x = trim_boundary(cs_x(t))
+    smoothed_y = trim_boundary(cs_y(t))
+    smoothed_z = trim_boundary(cs_z(t))
+    smoothed = np.column_stack((smoothed_x, smoothed_y, smoothed_z))
+    return smoothed
+
+
+def extend_boundary(x, num_boundary_points=6):
+    extended_x = np.concatenate(
+        (
+            np.linspace(x[0], x[1], num_boundary_points, endpoint=False),
+            x,
+            np.linspace(x[-2], x[-1], num_boundary_points, endpoint=False),
+        )
+    )
+    return extended_x
+
+
+def trim_boundary(x, num_boundary_points=6):
+    return x[num_boundary_points:-num_boundary_points]
+
+
+# Miscellaneous
+def compare_edges(xyx_i, xyz_j, xyz_k):
+    dist_ij = dist(xyx_i, xyz_j)
+    dist_ik = dist(xyx_i, xyz_k)
+    return dist_ij < dist_ik
+
+
+def dist(x, y, metric="l2"):
+    """
+    Computes distance between "x" and "y".
+
+    Parameters
+    ----------
+
+    Returns
+    -------
+    float
+
+    """
+    if metric == "l1":
+        return np.linalg.norm(np.subtract(x, y), ord=1)
+    else:
+        return np.linalg.norm(np.subtract(x, y), ord=2)

src/deep_neurographs/graph_utils.py

@@ -11,12 +11,23 @@
 
 from copy import deepcopy as cp
 
+import os
 import networkx as nx
 import numpy as np
 
 from deep_neurographs import swc_utils, utils
 
 
+def init_dense_graphs(swc_dir):
+    dense_graphs = dict()
+    for f in utils.listdir(swc_dir, ext=".swc"):
+        raw_txt = swc_utils.read_swc(os.path.join(swc_dir, f))
+        swc_dict = swc_utils.parse(raw_txt)
+        graph_id = f.replace(".0.swc", "")
+        graph = swc_utils.file_to_graph(swc_dict, graph_id=graph_id, set_attrs=True)
+        dense_graphs[graph_id] = graph
+    return dense_graphs
+
 def get_irreducibles(graph):
     leafs = []
     junctions = []

src/deep_neurographs/intake.py

@@ -15,7 +15,7 @@
 from torch_geometric.data import Data
 
 from deep_neurographs import neurograph as ng
-from deep_neurographs import s3_utils, swc_utils, utils
+from deep_neurographs import geometry_utils, s3_utils, swc_utils, utils
 
 
 # --- Build graph ---
@@ -25,19 +25,20 @@ def build_neurograph(
     bucket=None,
     access_key_id=None,
     secret_access_key=None,
-    generate_mutables=True,
     max_mutable_degree=5,
     max_mutable_dist=50.0,
     prune=True,
     prune_depth=16,
+    origin=None,
+    shape=None,
 ):
     """
     Builds a neurograph from a directory of swc files, where each swc
     represents a neuron and these neurons are assumed to be near each
     other.
 
     """
-    neurograph = ng.NeuroGraph()
+    neurograph = ng.NeuroGraph(origin=origin, shape=shape)
     if bucket is not None:
         neurograph = init_immutables_from_s3(
             neurograph,
@@ -48,7 +49,6 @@ def build_neurograph(
             secret_access_key=secret_access_key,
             prune=prune,
             prune_depth=prune_depth,
-            smooth=True,
         )
     else:
         neurograph = init_immutables_from_local(
@@ -57,11 +57,10 @@ def build_neurograph(
             anisotropy=anisotropy,
             prune=prune,
             prune_depth=prune_depth,
-            smooth=True,
         )
     neurograph.generate_mutables(
-        max_degree=max_mutable_degree, max_dist=max_mutable_dist
-    )
+            max_degree=max_mutable_degree, max_dist=max_mutable_dist
+        )
     return neurograph
 
 
@@ -74,7 +73,7 @@ def init_immutables_from_s3(
     secret_access_key=None,
     prune=True,
     prune_depth=16,
-    smooth=True,
+    smooth=False,
 ):
     """
     To do...