Test and fix SWC DFS sorting

navis/io/swc_io.py

@@ -450,7 +450,7 @@ def write_swc(x: 'core.NeuronObject',
 
     export_connectors : bool, optional
                         If True, will label nodes with pre- ("7") and
-                        postsynapse ("8"). Because only one label can be given
+                        postsynapse ("8"), or both ("9"). Because only one label can be given
                         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.
@@ -613,7 +613,7 @@ def _sort_swc_dfs(df: pd.DataFrame, roots, sort_children=True, inplace=False):
     while to_visit:
         node_id = to_visit.pop()
         order[node_id_to_orig_idx[node_id]] = count
-        cs = children.pop(order[-1], [])
+        cs = children.pop(node_id, [])
         if sort_children:
             to_visit.extend(sorted(cs, reverse=True))
         else:
@@ -627,7 +627,7 @@ def _sort_swc_dfs(df: pd.DataFrame, roots, sort_children=True, inplace=False):
 
     df["_order"] = order
     df.sort_values("_order", inplace=True)
-    df.drop(columns=["_order"])
+    df.drop(columns=["_order"], inplace=True)
     return df
 
 
@@ -680,6 +680,7 @@ def make_swc_table(x: 'core.TreeNeuron',
         x = x.to_skeleton()
 
     # Work on a copy sorted in depth-first order
+    # swc = _sort_swc_parent(x.nodes, inplace=False)
     swc = _sort_swc_dfs(x.nodes, x.root, inplace=False)
 
     # Add labels

tests/test_io.py

@@ -2,6 +2,7 @@
 import pytest
 import tempfile
 import numpy as np
+import pandas as pd
 
 from pathlib import Path
 
@@ -29,6 +30,48 @@ def test_swc_io(filename):
         assert len(n) == len(n2)
 
 
+@pytest.fixture
+def simple_neuron():
+    """Neuron with 1 branch and no connectors.
+
+    [3]
+     |
+     2 -- 4
+     |
+     1
+    """
+    nrn = navis.TreeNeuron(None)
+    dtypes = {
+        "node_id": np.uint64,
+        "parent_id": np.int64,
+        "x": float,
+        "y": float,
+        "z": float,
+    }
+    df = pd.DataFrame([
+        [1, 2, 0, 2, 0],
+        [2, 3, 0, 1, 0],
+        [3, -1, 0, 0, 0],  # root
+        [4, 2, 1, 1, 0]
+    ], columns=list(dtypes)).astype(dtypes)
+    nrn.nodes = df
+    return nrn
+
+
+def assert_parent_defined(df: pd.DataFrame):
+    defined = set()
+    for node, _structure, _x, _y, _z, _r, parent in df.itertuples(index=False):
+        defined.add(node)
+        if parent == -1:
+            continue
+        assert parent in defined, f"Child {node} has undefined parent {parent}"
+
+
+def test_swc_io_order(simple_neuron):
+    df = navis.io.swc_io.make_swc_table(simple_neuron, True)
+    assert_parent_defined(df)
+
+
 @pytest.mark.parametrize("filename", ['',
                                       'neurons.zip',
                                       '{neuron.id}@neurons.zip'])