33class Skeleton:
 34    """Class representing a skeleton.
 35
 36    Typically returned as results from a skeletonization.
 37
 38    Attributes
 39    ----------
 40    swc :       pd.DataFrame, optional
 41                SWC table.
 42    vertices :  (N, 3) array
 43                Vertex (node) positions.
 44    edges :     (M, 2) array
 45                Indices of connected vertex pairs.
 46    radii :     (N, ) array, optional
 47                Radii for each vertex (node) in the skeleton.
 48    mesh :      trimesh, optional
 49                The original mesh.
 50    mesh_map :  array, optional
 51                Same length as ``mesh``. Maps mesh vertices to vertices (nodes)
 52                in the skeleton.
 53    skel_map :  array of arrays, optional
 54                Inverse of `mesh_map`: maps skeleton vertices (nodes) to mesh
 55                vertices.
 56    method :    str, optional
 57                Which method was used to generate the skeleton.
 58
 59    """
 60
 61    def __init__(self, swc, mesh=None, mesh_map=None, method=None):
 62        self.swc = swc
 63        self.mesh = mesh
 64        self.mesh_map = mesh_map
 65        self.method = method
 66
 67    def __str__(self):
 68        """Summary."""
 69        return self.__repr__()
 70
 71    def __repr__(self):
 72        """Return quick summary of the skeleton's geometry."""
 73        elements = []
 74        if hasattr(self, "vertices"):
 75            elements.append(f"vertices={self.vertices.shape}")
 76        if hasattr(self, "edges"):
 77            elements.append(f"edges={self.edges.shape}")
 78        if hasattr(self, "method"):
 79            elements.append(f"method={self.method}")
 80        return f'<Skeleton({", ".join(elements)})>'
 81
 82    @property
 83    def edges(self):
 84        """Return skeleton edges."""
 85        return self.swc.loc[self.swc.parent_id >= 0, ["node_id", "parent_id"]].values
 86
 87    @property
 88    def vertices(self):
 89        """Return skeleton vertices (nodes)."""
 90        return self.swc[["x", "y", "z"]].values
 91
 92    @property
 93    def radius(self):
 94        """Return radii."""
 95        if "radius" not in self.swc.columns:
 96            raise ValueError(
 97                "No radius info found. Run `skeletor.post.radii()`" " to get them."
 98            )
 99        return self.swc["radius"].values
100
101    @property
102    def skeleton(self):
103        """Skeleton as trimesh Path3D."""
104        if not hasattr(self, "_skeleton"):
105            lines = [tm.path.entities.Line(e) for e in self.edges]
106
107            self._skeleton = tm.path.Path3D(
108                entities=lines, vertices=self.vertices, process=False
109            )
110        return self._skeleton
111
112    @property
113    def skel_map(self):
114        """Skeleton vertex (nodes) to mesh vertices. Based on `mesh_map`."""
115        if isinstance(self.mesh_map, type(None)):
116            return None
117        return (
118            pd.DataFrame(self.mesh_map)
119            .reset_index(drop=False)
120            .groupby(0)["index"]
121            .apply(np.array)
122            .values
123        )
124
125    @property
126    def roots(self):
127        """Root node(s)."""
128        return self.swc.loc[self.swc.parent_id < 0].index.values
129
130    @property
131    def leafs(self):
132        """Leaf nodes (includes root)."""
133        swc = self.swc
134        leafs = swc[~swc.node_id.isin(swc.parent_id.values) | (swc.parent_id < 0)]
135        return leafs.copy()
136
137    def reindex(self, inplace=False):
138        """Clean up skeleton."""
139        x = self
140        if not inplace:
141            x = x.copy()
142
143        # Re-index to make node IDs continous again
144        x.swc, new_ids = reindex_swc(x.swc)
145
146        # Update mesh map
147        if not isinstance(x.mesh_map, type(None)):
148            x.mesh_map = np.array([new_ids.get(i, i) for i in x.mesh_map])
149
150        if not inplace:
151            return x
152
153    def reroot(self, new_root):
154        """Reroot the skeleton.
155
156        Parameters
157        ----------
158        new_root :  int
159                    Index of node to use as new root. If the skeleton
160                    consists of multiple trees, only the tree containing the
161                    new root will be updated.
162
163        Returns
164        -------
165        Skeleton
166                    Skeleton with new root.
167
168        """
169        assert new_root in self.swc.index.values, f"Node index {new_root} not in skeleton."
170
171        # Make copy of self
172        x = self.copy()
173
174        # Check if the new root is already a root
175        if new_root in x.roots:
176            return x
177
178        # Get graph representation
179        G = x.get_graph()
180
181        # Get the path from the new root to the current root (of the same tree)
182        for r in x.roots:
183            try:
184                path = nx.shortest_path(G, source=new_root, target=r)
185                break
186            except nx.NetworkXNoPath:
187                continue
188
189        # Now we need to invert the path from the old root to the new root
190        new_parents = x.swc.set_index('node_id').parent_id.to_dict()
191        new_parents.update({c: p for p, c in zip(path[:-1], path[1:])})
192        new_parents[new_root] = -1
193
194        # Update the SWC table
195        x.swc["parent_id"] = x.swc.node_id.map(new_parents)
196
197        return x
198
199    def copy(self):
200        """Return copy of the skeleton."""
201        return Skeleton(
202            swc=self.swc.copy() if not isinstance(self.swc, type(None)) else None,
203            mesh=self.mesh.copy() if not isinstance(self.mesh, type(None)) else None,
204            mesh_map=self.mesh_map.copy()
205            if not isinstance(self.mesh_map, type(None))
206            else None,
207            method=self.method,
208        )
209
210    def get_graph(self):
211        """Generate networkX representation of the skeletons.
212
213        Distance between nodes will be used as edge weights.
214
215        Returns
216        -------
217        networkx.DiGraph
218
219        """
220        not_root = self.swc.parent_id >= 0
221        nodes = self.swc.loc[not_root]
222        parents = self.swc.set_index("node_id").loc[
223            self.swc.loc[not_root, "parent_id"].values
224        ]
225
226        dists = nodes[["x", "y", "z"]].values - parents[["x", "y", "z"]].values
227        dists = np.sqrt((dists**2).sum(axis=1))
228
229        G = nx.DiGraph()
230        G.add_nodes_from(self.swc.node_id.values)
231        G.add_weighted_edges_from(
232            zip(nodes.node_id.values, nodes.parent_id.values, dists)
233        )
234
235        return G
236
237    def get_segments(self, weight="weight", return_lengths=False):
238        """Generate a list of linear segments while maximizing segment lengths.
239
240        Parameters
241        ----------
242        weight :    'weight' | None, optional
243                    If ``"weight"`` use physical, geodesic length to determine
244                    segment length. If ``None`` use number of nodes (faster).
245        return_lengths : bool
246                    If True, also return lengths of segments according to ``weight``.
247
248        Returns
249        -------
250        segments :  list
251                    Segments as list of lists containing node IDs. List is
252                    sorted by segment lengths.
253        lengths :   list
254                    Length for each segment according to ``weight``. Only provided
255                    if `return_lengths` is True.
256
257        """
258        assert weight in ("weight", None), f'Unable to use weight "{weight}"'
259
260        # Get graph representation
261        G = self.get_graph()
262
263        # Get distances to root
264        dists = {}
265        for root in self.swc[self.swc.parent_id < 0].node_id.values:
266            dists.update(nx.shortest_path_length(G, target=root, weight=weight))
267
268        # Sort leaf nodes
269        endNodeIDs = self.leafs[self.leafs.parent_id >= 0].node_id.values
270        endNodeIDs = sorted(endNodeIDs, key=lambda x: dists.get(x, 0), reverse=True)
271
272        seen: set = set()
273        sequences = []
274        for nodeID in endNodeIDs:
275            sequence = [nodeID]
276            parents = list(G.successors(nodeID))
277            while True:
278                if not parents:
279                    break
280                parentID = parents[0]
281                sequence.append(parentID)
282                if parentID in seen:
283                    break
284                seen.add(parentID)
285                parents = list(G.successors(parentID))
286
287            if len(sequence) > 1:
288                sequences.append(sequence)
289
290        # Sort sequences by length
291        lengths = [dists[s[0]] - dists[s[-1]] for s in sequences]
292        sequences = [x for _, x in sorted(zip(lengths, sequences), reverse=True)]
293
294        if return_lengths:
295            return sequences, sorted(lengths, reverse=True)
296        else:
297            return sequences
298
299    def save_swc(self, filepath):
300        """Save skeleton in SWC format.
301
302        Parameters
303        ----------
304        filepath :      path-like
305                        Filepath to save SWC to.
306
307        """
308        header = dedent(f"""\
309        # SWC format file
310        # based on specifications at http://www.neuronland.org/NLMorphologyConverter/MorphologyFormats/SWC/Spec.html
311        # Created on {datetime.date.today()} using skeletor (https://github.com/navis-org/skeletor)
312        # PointNo Label X Y Z Radius Parent
313        # Labels:
314        # 0 = undefined, 1 = soma, 5 = fork point, 6 = end point
315        """)
316
317        # Make copy of SWC table
318        swc = self.swc.copy()
319
320        # Set all labels to undefined
321        swc["label"] = 0
322        swc.loc[~swc.node_id.isin(swc.parent_id.values), "label"] = 6
323        n_childs = swc.groupby("parent_id").size()
324        bp = n_childs[n_childs > 1].index.values
325        swc.loc[swc.node_id.isin(bp), "label"] = 5
326
327        # Add radius if missing
328        if "radius" not in swc.columns:
329            swc["radius"] = 0
330
331        # Get things in order
332        swc = swc[["node_id", "label", "x", "y", "z", "radius", "parent_id"]]
333
334        # Adjust column titles
335        swc.columns = ["PointNo", "Label", "X", "Y", "Z", "Radius", "Parent"]
336
337        with open(filepath, "w") as file:
338            # Write header
339            file.write(header)
340
341            # Write data
342            writer = csv.writer(file, delimiter=" ")
343            writer.writerows(swc.astype(str).values)
344
345    def scene(self, mesh=False, **kwargs):
346        """Return a Scene object containing the skeleton.
347
348        Returns
349        -------
350        scene :     trimesh.scene.scene.Scene
351                    Contains the skeleton and optionally the mesh.
352
353        """
354        if mesh:
355            if isinstance(self.mesh, type(None)):
356                raise ValueError("Skeleton has no mesh.")
357
358            self.mesh.visual.face_colors = [100, 100, 100, 100]
359
360            # Note the copy(): without it the transform in show() changes
361            # the original meshes
362            sc = tm.Scene([self.mesh.copy(), self.skeleton.copy()], **kwargs)
363        else:
364            sc = tm.Scene(self.skeleton.copy(), **kwargs)
365
366        return sc
367
368    def show(self, mesh=False, **kwargs):
369        """Render the skeleton in an opengl window. Requires pyglet.
370
371        Parameters
372        ----------
373        mesh :      bool
374                    If True, will render transparent mesh on top of the
375                    skeleton.
376
377        Returns
378        --------
379        scene :     trimesh.scene.Scene
380                    Scene with skeleton in it.
381
382        """
383        scene = self.scene(mesh=mesh)
384
385        # I encountered some issues if object space is big and the easiest
386        # way to work around this is to apply a transform such that the
387        # coordinates have -5 to +5 bounds
388        fac = 5 / np.fabs(self.skeleton.bounds).max()
389        scene.apply_transform(np.diag([fac, fac, fac, 1]))
390
391        return scene.show(**kwargs)
392
393    def mend_breaks(self, dist_mult=5, dist_min=0, dist_max=np.inf):
394        """Mend breaks in the skeleton using the original mesh.
395
396        This works by comparing the connectivity of the original mesh with that
397        of the skeleton. If the shortest path between two adjacent vertices on the mesh
398        is shorter than the distance between the nodes in the skeleton, a new edge
399        is added to the skeleton.
400
401        Parameters
402        ----------
403        dist_mult : float, optional
404                    Factor by which the new edge should be shorter than the
405                    current shortest path between two nodes to be added.
406                    Lower values = fewer false negatives; higher values = fewer
407                    false positive edges.
408        dist_min :  float, optional
409                    Minimum distance between nodes to consider adding an edge.
410                    Use this to avoid adding very short edges.
411        dist_max :  float, optional
412                    Maximum distance between nodes to consider adding an edge.
413                    Use this to avoid adding very long edges.
414
415        Returns
416        -------
417        edges :     (N, 2) array
418                    Edges connecting the skeleton nodes.
419        vertices :  (N, 3) array
420                    Positions of the skeleton nodes.
421
422        """
423        # We need `.mesh_map` and `.mesh` to exist
424        if self.mesh_map is None:
425            raise ValueError("Skeleton must have a `mesh_map` to mend breaks.")
426        if self.mesh is None:
427            raise ValueError("Skeleton must have a `mesh` to mend breaks.")
428
429        # Make a copy of the mesh edges
430        edges = self.mesh.edges.copy()
431        # Map mesh vertices to skeleton vertices
432        edges[:, 0] = self.mesh_map[self.mesh.edges[:, 0]]
433        edges[:, 1] = self.mesh_map[self.mesh.edges[:, 1]]
434        # Deduplicate
435        edges = np.unique(edges, axis=0)
436        # Remove self edges
437        edges = edges[edges[:, 0] != edges[:, 1]]
438
439        G = self.get_graph().to_undirected()
440
441        # Remove edges that are already in the skeleton
442        edges = np.array([e for e in edges if not G.has_edge(*e)])
443
444        # Calculate distance between these new edge candidates
445        dists = np.sqrt(
446            ((self.vertices[edges[:, 0]] - self.vertices[edges[:, 1]]) ** 2).sum(axis=1)
447        )
448
449        # Sort by distance (lowest first)
450        edges = edges[np.argsort(dists)]
451        dists = dists[np.argsort(dists)]
452
453        for e, d in zip(edges, dists):
454            # Check if the new path would be shorter than the current shortest path
455            if (d * dist_mult) < nx.shortest_path_length(G, e[0], e[1]):
456                continue
457            # Check if the distance is within bounds
458            elif d < dist_min:
459                continue
460            elif d > dist_max:
461                continue
462            # Add edge
463            G.add_edge(*e, weight=d)
464
465        # The above may have introduced small triangles which we should try to remove
466        # by removing the longest edge in a triangle. I have also spotted more
467        # complex cases of four or more nodes forming false-positive loops but
468        # these will be harder to detect and remove.
469
470        # First collect neighbors for each node
471        later_nbrs = {}
472        for node, neighbors in G.adjacency():
473            later_nbrs[node] = {
474                n for n in neighbors if n not in later_nbrs and n != node
475            }
476
477        # Go over each node
478        triangles = set()
479        for node1, neighbors in later_nbrs.items():
480            # Go over each neighbor
481            for node2 in neighbors:
482                # Check if there is one or more third nodes that are connected to both
483                third_nodes = neighbors & later_nbrs[node2]
484                for node3 in third_nodes:
485                    # Add triangle (sort to deduplicate)
486                    triangles.add(tuple(sorted([node1, node2, node3])))
487
488        # Remove longest edge in each triangle
489        for t in triangles:
490            e1, e2, e3 = t[:2], t[1:], t[::2]
491            # Make sure all edges still exist (we may have already removed edges
492            # that were part of a previous triangle)
493            if any(not G.has_edge(*e) for e in (e1, e2, e3)):
494                continue
495            # Remove the longest edge
496            G.remove_edge(*max((e1, e2, e3), key=lambda e: G.edges[e]["weight"]))
497
498        return np.array(G.edges), self.vertices.copy()