From 50b51e7ff456c62f6422c866acb986d0110d5ccb Mon Sep 17 00:00:00 2001
From: Forrest Collman <forrest.collman@gmail.com>
Date: Sun, 20 Oct 2024 23:37:06 -0700
Subject: [PATCH] adding first draft edges

---
 navis/core/mesh.py          | 35 ++++++++++++++++++++++++++++++-----
 navis/graph/converters.py   | 16 ++++++++++++----
 navis/interfaces/blender.py |  2 +-
 3 files changed, 43 insertions(+), 10 deletions(-)

diff --git a/navis/core/mesh.py b/navis/core/mesh.py
index 233076e0..0ba86485 100644
--- a/navis/core/mesh.py
+++ b/navis/core/mesh.py
@@ -98,7 +98,7 @@ class MeshNeuron(BaseNeuron):
     TEMP_ATTR = ['_memory_usage', '_trimesh', '_skeleton', '_igraph', '_graph_nx']
 
     #: Core data table(s) used to calculate hash
-    CORE_DATA = ['vertices', 'faces']
+    CORE_DATA = ['vertices', 'faces', 'edges']
 
     def __init__(self,
                  x,
@@ -114,18 +114,29 @@ def __init__(self,
         self._lock = 1
         self._trimesh = None  # this is required to avoid recursion during init
 
+
         if isinstance(x, MeshNeuron):
             self.__dict__.update(x.copy().__dict__)
+            self.vertices, self.faces, self.edges = x.vertices, x.faces, x.edges
+        elif hasattr(x, "faces") and hasattr(x, "vertices"):
             self.vertices, self.faces = x.vertices, x.faces
-        elif hasattr(x, 'faces') and hasattr(x, 'vertices'):
-            self.vertices, self.faces = x.vertices, x.faces
+            if hasattr(x, "edges"):
+                self.edges = x.edges
+            else:
+                self.edges = np.array([[]])
         elif isinstance(x, dict):
-            if 'faces' not in x or 'vertices' not in x:
+            if "faces" not in x or "vertices" not in x:
                 raise ValueError('Dictionary must contain "vertices" and "faces"')
-            self.vertices, self.faces = x['vertices'], x['faces']
+            self.vertices, self.faces = x["vertices"], x["faces"]
+            if "edges" in x:
+                self.edges = x["edges"]
+            else:
+                self.edges = np.array([[]], int)
+
         elif isinstance(x, str) and os.path.isfile(x):
             m = tm.load(x)
             self.vertices, self.faces = m.vertices, m.faces
+            self.edges = np.array([[]], int)
         elif isinstance(x, type(None)):
             # Empty neuron
             self.vertices, self.faces = np.zeros((0, 3)), np.zeros((0, 3))
@@ -282,6 +293,20 @@ def vertices(self, verts):
             raise ValueError('Vertices must be 2-dimensional array')
         self._vertices = verts
         self._clear_temp_attr()
+    
+    @property
+    def edges(self):
+        """Edges making up the neuron."""
+        return self._edges
+
+    @edges.setter
+    def edges(self, edges):
+        if not isinstance(edges, np.ndarray):
+            raise TypeError(f'Edges must be numpy array, got "{type(edges)}"')
+        if edges.ndim != 2:
+            raise ValueError("Edges must be 2-dimensional array")
+        self._edges = edges
+        self._clear_temp_attr()
 
     @property
     def faces(self):
diff --git a/navis/graph/converters.py b/navis/graph/converters.py
index aa0f7edd..7a6daa52 100644
--- a/navis/graph/converters.py
+++ b/navis/graph/converters.py
@@ -300,6 +300,15 @@ def neuron2nx(x: "core.NeuronObject", simplify=False, epsilon=None) -> nx.DiGrap
             for e, l in zip(x.trimesh.edges_unique, x.trimesh.edges_unique_length)
         ]
         G.add_weighted_edges_from(edges)
+        extra_edge_lengths = np.linalg.norm(
+            x.vertices[x.edges[:, 0], :] - x.vertices[x.edges[:, 1], :], axis=1
+        )
+        extra_edges = [
+            (x.edges[i, 0], x.edges[i, 1], extra_edge_lengths[i])
+            for i in range(x.edges.shape[0])
+        ]
+        G.add_weighted_edges_from(extra_edges)
+
     elif isinstance(x, core.Dotprops):
         if epsilon is None:
             epsilon = 5 * x.sampling_resolution
@@ -841,14 +850,13 @@ def edges2neuron(edges, vertices=None, validate=True, **kwargs):
             # (note that we assign -1 as root's parent)
             parents.update({k: v[0] if v else -1 for k, v in this.items()})
 
-    nodes = pd.DataFrame(vertices, columns=['x', 'y', 'z'])
-    nodes.insert(0, 'node_id', nodes.index)
-    nodes.insert(1, 'parent_id', nodes.index.map(parents))
+    nodes = pd.DataFrame(vertices, columns=["x", "y", "z"])
+    nodes.insert(0, "node_id", nodes.index)
+    nodes.insert(1, "parent_id", nodes.index.map(parents))
 
     return core.TreeNeuron(nodes, **kwargs)
 
 
-
 def _find_all_paths(
     g: nx.DiGraph, start, end, mode: str = "OUT", maxlen: Optional[int] = None
 ) -> list:
diff --git a/navis/interfaces/blender.py b/navis/interfaces/blender.py
index b9c154dd..2eeee7d6 100644
--- a/navis/interfaces/blender.py
+++ b/navis/interfaces/blender.py
@@ -399,7 +399,7 @@ def _create_mesh(self, x, mat, collection=None):
         ob['id'] = str(x.id)
 
         blender_verts = verts.tolist()
-        me.from_pydata(list(blender_verts), [], list(x.faces))
+        me.from_pydata(list(blender_verts), list(x.edges), list(x.faces))
         me.update()
 
         me.polygons.foreach_set('use_smooth', [True] * len(me.polygons))