Merge pull request #1544 from brian-team/load_swc_improvements

Improvements for SWC file loading

brian2/spatialneuron/morphology.py

@@ -1078,38 +1078,33 @@ def _compartments_to_sections(
         if current_compartments is None:
             current_compartments = []
 
-        current_compartments.append(compartment)
-
         # We have to create a new section, if we are either
         # 1. at a leaf of the tree or at a branching point, or
         # 2. if the compartment type changes
-        if (
-            len(compartment.children) != 1
-            or compartment.comp_name != compartment.children[0].comp_name
-        ):
-            parent = current_compartments[0].parent
-            section = Morphology._create_section(
-                current_compartments,
-                compartment.comp_name,
-                parent=parent,
-                sections=sections,
-                spherical_soma=spherical_soma,
-            )
-            sections[current_compartments[-1].index] = section
-            # If we are at a branching point, recurse into all subtrees
-            for child in compartment.children:
-                Morphology._compartments_to_sections(
-                    child,
-                    spherical_soma=spherical_soma,
-                    current_compartments=None,
-                    sections=sections,
-                )
-        else:
-            # A single child of the same type, continue (recursive call)
+        while True:
+            current_compartments.append(compartment)
+            if (
+                len(compartment.children) != 1
+                or compartment.children[0].comp_name != compartment.comp_name
+            ):
+                break
+            compartment = compartment.children[0]
+
+        parent = current_compartments[0].parent
+        section = Morphology._create_section(
+            current_compartments,
+            compartment.comp_name,
+            parent=parent,
+            sections=sections,
+            spherical_soma=spherical_soma,
+        )
+        sections[current_compartments[-1].index] = section
+        # If we are at a branching point, recurse into all subtrees
+        for child in compartment.children:
             Morphology._compartments_to_sections(
-                compartment.children[0],
+                child,
                 spherical_soma=spherical_soma,
-                current_compartments=current_compartments,
+                current_compartments=None,
                 sections=sections,
             )
 
@@ -1124,13 +1119,14 @@ def _replace_three_point_soma(compartment, all_compartments):
         # We are looking for a node with two children of the soma type (and
         # other childen of other types), where the two children don't have any
         # children of their own
+        # To avoid recursion errors we are not searching any further
         soma_children = [c for c in compartment.children if c.comp_name == "soma"]
         if (
             compartment.comp_name == "soma"
             and len(soma_children) == 2
             and all(len(c.children) == 0 for c in soma_children)
         ):
-            # We've found a 3-point soma to replace
+            # This is a 3-point soma to replace
             soma_c = [compartment] + soma_children
             if not all(abs(c.diameter - soma_c[0].diameter) < 1e-15 for c in soma_c):
                 indices = ", ".join(str(c.index) for c in soma_c)
@@ -1146,14 +1142,14 @@ def _replace_three_point_soma(compartment, all_compartments):
             length_1 = np.sqrt(np.sum((point_1 - point_0) ** 2))
             length_2 = np.sqrt(np.sum((point_2 - point_0) ** 2))
             if (
-                np.abs(length_1 - diameter / 2) > 0.01
-                or np.abs(length_2 - diameter / 2) > 0.01
+                np.abs(length_1 - diameter / 2) > 0.1
+                or np.abs(length_2 - diameter / 2) > 0.1
             ):
                 raise ValueError(
                     "Cannot replace '3-point-soma' by a single "
                     "point, the second and third points should "
                     "be positioned one radius away from the "
-                    f"first point. Distances are {length_1:.3d}um and "
+                    f"first point. Distances are {length_1:.3f}um and "
                     f"{length_2:.3f}um, respectively, while the "
                     f"radius is {diameter / 2:.3f}um."
                 )
@@ -1171,11 +1167,15 @@ def _replace_three_point_soma(compartment, all_compartments):
             all_compartments[compartment.index] = compartment
             del all_compartments[soma_children[0].index]
             del all_compartments[soma_children[1].index]
-
-        # Recurse further down the tree
-        all_compartments[compartment.index] = compartment
-        for child in compartment.children:
-            Morphology._replace_three_point_soma(child, all_compartments)
+        else:
+            # Raise a warning if there is more than one soma compartment
+            somata = [c for c in all_compartments.values() if c.comp_name == "soma"]
+            if len(somata) > 1:
+                logger.warning(
+                    f"Found {len(somata)} soma compartments. If you have a 3-point soma, "
+                    "make sure it is at the beginning.",
+                    name_suffix="soma_compartments",
+                )
 
     @staticmethod
     def from_points(points, spherical_soma=True):

brian2/tests/test_morphology.py

@@ -8,6 +8,7 @@
 from brian2.spatialneuron import *
 from brian2.tests.utils import assert_allclose
 from brian2.units import DimensionMismatchError, cm, second, um
+from brian2.utils.logger import catch_logs
 
 
 @pytest.mark.codegen_independent
@@ -718,6 +719,39 @@ def test_tree_soma_from_points_3_point_soma_incorrect():
         Morphology.from_points(points)
 
 
+@pytest.mark.codegen_independent
+def test_tree_soma_from_points_somas_not_at_start():
+    # A single soma in the middle (ununusal but possible)
+    points = [  # dendrite
+        (1, "dend", 100, 0, 0, 10, -1),
+        (2, "dend", 100, 0, 0, 10, 1),
+        (3, "dend", 100, 0, 0, 10, 2),
+        (4, "soma", 100, 0, 0, 30, 3),
+        (5, "dend2", 130, 0, 0, 10, 4),
+        (6, "dend2", 160, 0, 0, 10, 5),
+    ]
+    morpho = Morphology.from_points(points)
+    assert morpho.total_sections == 3
+    assert morpho.total_compartments == 5
+
+    # Several somata (probably an error)
+    points = [  # dendrite
+        (1, "dend", 0, 0, 0, 10, -1),
+        (2, "dend", 30, 0, 0, 10, 1),
+        (3, "dend", 60, 0, 0, 10, 2),
+        (4, "soma", 90, 0, 0, 30, 3),
+        (5, "dend2", 120, 70, 0, 10, 4),
+        (6, "dend2", 150, 70, 0, 10, 5),
+        (7, "soma", 180, 40, 0, 30, 6),
+    ]
+    with catch_logs() as logs:
+        morpho = Morphology.from_points(points)
+    assert len(logs) == 1
+    assert logs[0][1].endswith("soma_compartments")
+    assert morpho.total_sections == 4
+    assert morpho.total_compartments == 6
+
+
 @pytest.mark.codegen_independent
 def test_tree_soma_from_swc():
     swc_content = """
@@ -768,6 +802,58 @@ def test_tree_soma_from_swc_3_point_soma():
     _check_tree_soma(soma, coordinates=True, use_cylinders=False)
 
 
+@pytest.mark.codegen_independent
+def test_tree_soma_from_swc_3_point_soma_incorrect_points():
+    swc_content = """
+# Test file
+1    1  100  0  0  15  -1
+2    1  100  15  0  15  1
+3    1  100  -10  0  15  1
+4   2  114.14213562373095  14.142135623730949  0  4  1
+5   2  128.2842712474619  28.284271247461898  0  3  4
+6   2  142.42640687119285  42.426406871192846  0  2  5
+7   2  156.5685424949238  56.568542494923797  0  1  6
+8   2  170.71067811865476  70.710678118654741  0  0  7
+9   2  107.07106781186548  -7.0710678118654746  0  2.5  1
+10   2  114.14213562373095  -14.142135623730949  0  2.5  9
+11   2  121.21320343559643  -21.213203435596423  0  2.5  10
+12   2  128.2842712474619  -28.284271247461898  0  2.5  11
+13   2  135.35533905932738  -35.35533905932737  0  2.5  12
+"""
+    tmp_filename = tempfile.mktemp("cable_morphology.swc")
+    with open(tmp_filename, "w") as f:
+        f.write(swc_content)
+    with pytest.raises(ValueError, match=r".*radius is 15.000um.*"):
+        Morphology.from_file(tmp_filename)
+    os.remove(tmp_filename)
+
+
+@pytest.mark.codegen_independent
+def test_tree_soma_from_swc_3_point_soma_incorrect_radius():
+    swc_content = """
+# Test file
+1    1  100  0  0  15  -1
+2    1  100  15  0  10  1
+3    1  100  -15  0  15  1
+4   2  114.14213562373095  14.142135623730949  0  4  1
+5   2  128.2842712474619  28.284271247461898  0  3  4
+6   2  142.42640687119285  42.426406871192846  0  2  5
+7   2  156.5685424949238  56.568542494923797  0  1  6
+8   2  170.71067811865476  70.710678118654741  0  0  7
+9   2  107.07106781186548  -7.0710678118654746  0  2.5  1
+10   2  114.14213562373095  -14.142135623730949  0  2.5  9
+11   2  121.21320343559643  -21.213203435596423  0  2.5  10
+12   2  128.2842712474619  -28.284271247461898  0  2.5  11
+13   2  135.35533905932738  -35.35533905932737  0  2.5  12
+"""
+    tmp_filename = tempfile.mktemp("cable_morphology.swc")
+    with open(tmp_filename, "w") as f:
+        f.write(swc_content)
+    with pytest.raises(ValueError, match=r".*not all the diameters are identical.*"):
+        Morphology.from_file(tmp_filename)
+    os.remove(tmp_filename)
+
+
 @pytest.mark.codegen_independent
 def test_construction_incorrect_arguments():
     ### Morphology
@@ -873,6 +959,19 @@ def test_construction_incorrect_arguments():
         )
 
 
+@pytest.mark.codegen_independent
+def test_from_points_long_chain():
+    # in previous versions, this led to a recursion error
+    points = [(1, "soma", 0, 0, 0, 1, -1)]
+    for i in range(2, 10000):
+        points.append((i, "dend", 0 + i * 10, 0, 0, 5, i - 1))
+    morph = Morphology.from_points(points)
+    assert (
+        morph.total_compartments == 10000 - 1
+    )  # compartments are in-between the points
+    assert morph.total_sections == 2
+
+
 @pytest.mark.codegen_independent
 def test_from_points_minimal():
     points = [(1, "soma", 10, 20, 30, 30, -1)]