SWC io: DFS-sort nodes, optionally do not reindex

Previously, nodes were sorted by parent index, which would ensure that
root nodes appear before non-roots but did not guarantee that all
children are defined after their parent.
Now, nodes are sorted in depth-first order (addressing their children in
order of their index).

Additionally, added the option of return_node_map=None, which does not
re-index nodes before writing.
This is useful when exporting a SWC files of several neurons along with
where node IDs must persist their global uniqueness.

navis/io/swc_io.py

@@ -453,9 +453,12 @@ def write_swc(x: 'core.NeuronObject',
                         this might drop synapses (i.e. in case of multiple
                         pre- and/or postsynapses on a single node)! ``labels``
                         must be ``True`` for this to have any effect.
-    return_node_map :   bool
+    return_node_map :   bool, optional
                         If True, will return a dictionary mapping the old node
                         ID to the new reindexed node IDs in the file.
+                        If False (default), will reindex nodes but not return the mapping.
+                        If None, will not reindex nodes (i.e. their current IDs will be written).
+
 
     Returns
     -------
@@ -535,7 +538,7 @@ def _write_swc(x: 'core.TreeNeuron',
                write_meta: Union[bool, List[str], dict] = True,
                labels: Union[str, dict, bool] = True,
                export_connectors: bool = False,
-               return_node_map: bool = False) -> None:
+               return_node_map: Optional[bool] = False) -> None:
     """Write single TreeNeuron to file."""
     # Generate SWC table
     res = make_swc_table(x,
@@ -589,6 +592,45 @@ def _write_swc(x: 'core.TreeNeuron',
         return node_map
 
 
+def sort_swc(df: pd.DataFrame, roots, sort_children=True, inplace=False):
+    """Depth-first search tree to ensure parents are always defined before children."""
+    children = defaultdict(list)
+    node_id_to_orig_idx = dict()
+    for row in df.itertuples():
+        child = row.node_id
+        parent = row.parent_id
+        children[parent].append(child)
+        node_id_to_orig_idx[child] = row.index
+
+    if sort_children:
+        to_visit = sorted(roots, reverse=True)
+    else:
+        to_visit = list(roots)[::-1]
+
+    idx = 0
+    order = np.full(len(df), np.nan)
+    count = 0
+    while to_visit:
+        node_id = to_visit.pop()
+        order[node_id_to_orig_idx[node_id]] = count
+        cs = children.pop(order[-1], [])
+        if sort_children:
+            to_visit.append(sorted(sort_children, reverse=True))
+        else:
+            to_visit.append(cs[::-1])
+        count += 1
+
+    # undefined behaviour if any nodes are not reachable from the given roots
+
+    if not inplace:
+        df = df.copy()
+
+    df["_order"] = order
+    df.sort_values("_order", inplace=True)
+    df.drop(columns=["_order"])
+    return df
+
+
 def make_swc_table(x: 'core.TreeNeuron',
                    labels: Union[str, dict, bool] = None,
                    export_connectors: bool = False,
@@ -606,17 +648,20 @@ def make_swc_table(x: 'core.TreeNeuron',
 
                         str : column name in node table
                         dict: must be of format {node_id: 'label', ...}.
-                        bool: if True, will generate automatic labels, if False all nodes have label "0".
+                        bool: if True, will generate automatic labels for branches ("5") and ends ("6"),
+                        soma where labelled ("1"), and optionally connectors (see below).
+                        If False (or for all nodes not labelled as above) all nodes have label "0".
 
     export_connectors : bool, optional
-                        If True, will label nodes with pre- ("7") and
-                        postsynapse ("8"). Because only one label can be given
-                        this might drop synapses (i.e. in case of multiple
-                        pre- or postsynapses on a single node)! ``labels``
-                        must be ``True`` for this to have any effect.
-    return_node_map :   bool
+                        If True, will label nodes with only presynapses ("7"),
+                        only postsynapses ("8"), or both ("9").
+                        This overrides branch/end/soma labels.
+                        ``labels`` must be ``True`` for this to have any effect.
+    return_node_map :   bool, optional
                         If True, will return a dictionary mapping the old node
                         ID to the new reindexed node IDs in the file.
+                        If False, will remap IDs but not return the mapping.
+                        If None, will not remap IDs.
 
     Returns
     -------
@@ -625,8 +670,8 @@ def make_swc_table(x: 'core.TreeNeuron',
                         Only if ``return_node_map=True``.
 
     """
-    # Work on a copy
-    swc = x.nodes.copy()
+    # Work on a copy sorted in depth-first order
+    swc = sort_swc(x.nodes, x.root, inplace=False)
 
     # Add labels
     swc['label'] = 0
@@ -642,31 +687,41 @@ def make_swc_table(x: 'core.TreeNeuron',
         if not isinstance(x.soma, type(None)):
             soma = utils.make_iterable(x.soma)
             swc.loc[swc.node_id.isin(soma), 'label'] = 1
+
         if export_connectors:
             # Add synapse label
             pre_ids = x.presynapses.node_id.values
             post_ids = x.postsynapses.node_id.values
-            swc.loc[swc.node_id.isin(pre_ids), 'label'] = 7
-            swc.loc[swc.node_id.isin(post_ids), 'label'] = 8
-
-    # Sort such that the parent is always before the child
-    swc.sort_values('parent_id', ascending=True, inplace=True)
 
-    # Reset index
-    swc.reset_index(drop=True, inplace=True)
+            is_pre = swc["node_id"].isin(pre_ids)
+            swc.loc[is_pre, 'label'] = 7
 
-    # Generate mapping
-    new_ids = dict(zip(swc.node_id.values, swc.index.values + 1))
+            is_post = swc["node"].isin(post_ids)
+            swc.loc[is_post, 'label'] = 8
 
-    swc['node_id'] = swc.node_id.map(new_ids)
-    # Lambda prevents potential issue with missing parents
-    swc['parent_id'] = swc.parent_id.map(lambda x: new_ids.get(x, -1))
+            is_both = np.logical_and(is_pre, is_post)
+            swc.loc[is_both, 'label'] = 9
 
-    # Get things in order
+    # Order columns
     swc = swc[['node_id', 'label', 'x', 'y', 'z', 'radius', 'parent_id']]
 
-    # Make sure radius has no `None`
+    # Make sure radius has no `None` or negative
     swc['radius'] = swc.radius.fillna(0)
+    swc['radius'][swc['radius'] < 0] = 0
+
+    if return_node_map is not None:
+        # remap IDs
+
+        # Reset index
+        swc.reset_index(drop=True, inplace=True)
+
+        # Generate mapping
+        new_ids = dict(zip(swc.node_id.values, swc.index.values + 1))
+
+        swc['node_id'] = swc.node_id.map(new_ids)
+        # Lambda prevents potential issue with missing parents
+        swc['parent_id'] = swc.parent_id.map(lambda x: new_ids.get(x, -1))
+
 
     # Adjust column titles
     swc.columns = ['PointNo', 'Label', 'X', 'Y', 'Z', 'Radius', 'Parent']