Feat: Use MorphIO by default and fallbacks to internal loader in case of failure

Author: adrien-berchet

tests/test_neuron_conversion.py

@@ -8,8 +8,8 @@
 from numpy import testing as npt
 
 from tmd.io import conversion as tested
+from tmd.io.io import _load_neuron_morphio
 from tmd.io.io import load_neuron
-from tmd.io.io import load_neuron_from_morphio
 
 _path = os.path.dirname(os.path.abspath(__file__))
 DATA_PATH = os.path.join(_path, "data")
@@ -166,11 +166,11 @@ def test_neuron_building_consistency__h5():
     path = f"{DATA_PATH}/valid/C010398B-P2.h5"
 
     neuron1 = load_neuron(path)
-    neuron2 = load_neuron_from_morphio(path)
+    neuron2 = _load_neuron_morphio(path)
 
     _assert_neurons_equal(neuron1, neuron2)
 
-    neuron2 = load_neuron_from_morphio(morphio.Morphology(path))
+    neuron2 = _load_neuron_morphio(morphio.Morphology(path))
 
     _assert_neurons_equal(neuron1, neuron2)
 
@@ -180,10 +180,10 @@ def test_neuron_building_consistency__swc():
     path = f"{DATA_PATH}/valid/C010398B-P2.CNG.swc"
 
     neuron1 = load_neuron(path)
-    neuron2 = load_neuron_from_morphio(path)
+    neuron2 = _load_neuron_morphio(path)
 
     _assert_neurons_equal(neuron1, neuron2)
 
-    neuron2 = load_neuron_from_morphio(morphio.Morphology(path))
+    neuron2 = _load_neuron_morphio(morphio.Morphology(path))
 
     _assert_neurons_equal(neuron1, neuron2)

tmd/Neuron/Neuron.py

@@ -91,8 +91,9 @@ def append_tree(self, new_tree, tree_types):
         correct list of trees in neuron.
         """
         if isinstance(new_tree, Tree.Tree):
-            if int(np.median(new_tree.t)) in tree_types.keys():
-                neurite_type = tree_types[int(np.median(new_tree.t))]
+            tree_type_key = int(np.median(new_tree.t))
+            if tree_type_key in tree_types.keys():
+                neurite_type = tree_types[tree_type_key]
             else:
                 neurite_type = "undefined"
             getattr(self, neurite_type).append(new_tree)

tmd/io/conversion.py

@@ -19,6 +19,7 @@
 
 import numpy as np
 
+from tmd.Neuron import Neuron
 from tmd.Soma.Soma import Soma
 from tmd.Tree import Tree
 
@@ -128,3 +129,14 @@ def convert_morphio_trees(morphio_neuron):
             t=t[tree_beg:tree_end],
             p=p[tree_beg:tree_end],
         )
+
+
+def convert_morphio_neuron(morph, tree_types, name=""):
+    """Convert a MorphIO morphology into a Neuron object."""
+    neuron = Neuron.Neuron()
+    neuron.name = name
+    neuron.set_soma(convert_morphio_soma(morph.soma))
+    for tree in convert_morphio_trees(morph):
+        neuron.append_tree(tree, tree_types)
+
+    return neuron

tmd/io/io.py

@@ -20,11 +20,12 @@
 from pathlib import Path
 
 import numpy as _np
+from morphio import Morphology
+from morphio import Option
 from scipy import sparse as sp
 from scipy.sparse import csgraph as cs
 
-from tmd.io.conversion import convert_morphio_soma
-from tmd.io.conversion import convert_morphio_trees
+from tmd.io.conversion import convert_morphio_neuron
 from tmd.io.h5 import read_h5
 from tmd.io.swc import SWC_DCT
 from tmd.io.swc import read_swc
@@ -76,7 +77,7 @@ def redefine_types(user_types=None):
     return final_tree_types
 
 
-def load_neuron(
+def _load_neuron_internal(
     input_file, line_delimiter="\n", soma_type=None, user_tree_types=None, remove_duplicates=True
 ):
     """I/O method to load an swc or h5 file into a Neuron object."""
@@ -146,7 +147,7 @@ def load_neuron(
     return neuron
 
 
-def load_neuron_from_morphio(path_or_obj, user_tree_types=None):
+def _load_neuron_morphio(path_or_obj, user_tree_types=None):
     """Create Neuron object from morphio object or from path loaded via morphio.
 
     Supported file formats: h5, swc, asc.
@@ -158,32 +159,76 @@ def load_neuron_from_morphio(path_or_obj, user_tree_types=None):
     Returns:
         neuron (Neuron): tmd Neuron object
     """
-    from morphio import Morphology  # pylint: disable=import-outside-toplevel
-
     tree_types = redefine_types(user_tree_types)
 
     if isinstance(path_or_obj, (str, Path)):
-        obj = Morphology(path_or_obj)
+        obj = Morphology(
+            path_or_obj,
+            Option.allow_root_bifurcations
+            | Option.allow_soma_bifurcations
+            | Option.allow_custom_root_id
+            | Option.allow_multiple_somata,
+        )
         filename = path_or_obj
     else:
         obj = path_or_obj
         # MorphIO does not support naming of objects yet.
         filename = ""
 
-    neuron = Neuron.Neuron()
-    neuron.name = filename
-    neuron.set_soma(convert_morphio_soma(obj.soma))
-    for tree in convert_morphio_trees(obj):
-        neuron.append_tree(tree, tree_types)
+    return convert_morphio_neuron(obj, tree_types, filename)
 
-    return neuron
 
+def load_neuron(
+    input_file, user_tree_types=None, *, line_delimiter="\n", soma_type=None, remove_duplicates=True
+):
+    """I/O method to load an 'asc', 'h5' or 'swc' file into a Neuron object.
+
+    Args:
+        input_file (Union[str, morphio.Morphology]):
+            Filepath or morphio object
 
-def load_population(neurons, user_tree_types=None, name=None, use_morphio=False):
+    Returns:
+        neuron (Neuron): tmd Neuron object
+
+    """
+    ext = os.path.splitext(input_file)[-1][1:]
+    if ext not in ("h5", "swc", "asc"):
+        raise ValueError("The file extension must be in ['asc', 'h5', 'swc']")
+
+    try:
+        return _load_neuron_morphio(input_file, user_tree_types=user_tree_types)
+    except Exception as morphio_exc:  # pylint: disable=broad-except
+        try:
+            if ext not in ("h5", "swc"):
+                raise ValueError(
+                    "The internal loader can only read '*.h5' and '*.swc' files."
+                ) from morphio_exc
+            neuron = _load_neuron_internal(
+                input_file,
+                line_delimiter=line_delimiter,
+                soma_type=soma_type,
+                user_tree_types=user_tree_types,
+                remove_duplicates=remove_duplicates,
+            )
+            warnings.warn(
+                f"The file {input_file} was loaded using the internal loader because of a MorphIO "
+                "failure."
+            )
+            return neuron
+        except Exception as exc:
+            raise exc from morphio_exc
+
+
+def load_population(neurons, user_tree_types=None, name=None, use_morphio=None):
     """Load all data of recognised format (swc, h5) into a Population object.
 
     Takes as input a directory or a list of files to load.
     """
+    if use_morphio is not None:
+        warnings.warn(
+            "The 'use_morphio' parameter is deprecated as the internal loader is only used "
+            "when MorphIO fails."
+        )
     if isinstance(neurons, (list, tuple)):
         files = neurons
         name = name if name is not None else "Population"
@@ -204,16 +249,7 @@ def load_population(neurons, user_tree_types=None, name=None, use_morphio=False)
 
     for filename in files:
         try:
-            ext = os.path.splitext(filename)[-1][1:].lower()
-            if not use_morphio:
-                assert ext in ("h5", "swc")
-                pop.append_neuron(load_neuron(filename, user_tree_types=user_tree_types))
-            else:
-                assert ext in ("h5", "swc", "asc")
-                pop.append_neuron(
-                    load_neuron_from_morphio(filename, user_tree_types=user_tree_types)
-                )
-
+            pop.append_neuron(load_neuron(filename, user_tree_types=user_tree_types))
         except AssertionError as exc:
             raise Warning(
                 "{filename} is not a valid h5, swc or asc file. If asc set use_morphio to True."