allow Morphology objects to be either mut or immut (#1049)

* Maintain mut/immut after transformation
* Add tests

Co-authored-by: Eleftherios Zisis <eleftherios.zisis@epfl.ch>

neurom/core/morphology.py

@@ -30,6 +30,7 @@
 
 import warnings
 from collections import deque
+from pathlib import Path
 
 import morphio
 import numpy as np
@@ -530,7 +531,11 @@ def __init__(self, filename, name=None):
             filename (str|Path): a filename or morphio.{mut}.Morphology object
             name (str): an optional morphology name
         """
-        self._morphio_morph = morphio.mut.Morphology(filename).as_immutable()
+        self._morphio_morph = morphio.mut.Morphology(filename)
+
+        if isinstance(filename, (str, Path, morphio.Morphology)):
+            self._morphio_morph = self._morphio_morph.as_immutable()
+
         self.name = name if name else 'Morphology'
         self.soma = make_soma(self._morphio_morph.soma)
 
@@ -559,13 +564,24 @@ def points(self):
 
     def transform(self, trans):
         """Return a copy of this morphology with a 3D transformation applied."""
-        mut = self._morphio_morph.as_mutable()
-        mut.soma.points = trans(mut.soma.points)
+        morph = self._morphio_morph
+
+        is_immutable = hasattr(morph, 'as_mutable')
+
+        # make copy or convert to mutable if immutable
+        if is_immutable:
+            morph = morph.as_mutable()
+        else:
+            morph = morphio.mut.Morphology(morph)
+
+        morph.soma.points = trans(morph.soma.points)
 
-        for section in mut.iter():
+        for section in morph.iter():
             section.points = trans(section.points)
 
-        return Morphology(mut)
+        if is_immutable:
+            return Morphology(morph.as_immutable())
+        return Morphology(morph)
 
     def __copy__(self):
         """Creates a deep copy of Morphology instance."""

tests/core/test_neuron.py

@@ -131,3 +131,20 @@ def test_str():
     n = nm.load_morphology(SWC_PATH / 'simple.swc')
     assert 'Morphology' in str(n)
     assert 'Section' in str(n.neurites[0].root_node)
+
+
+def test_mut_nonmut_constructor():
+
+    path = SWC_PATH / 'simple.swc'
+
+    m = Morphology(path)
+    assert isinstance(m.to_morphio(), morphio.Morphology)
+
+    m = Morphology(str(path))
+    assert isinstance(m.to_morphio(), morphio.Morphology)
+
+    m = Morphology(morphio.Morphology(path))
+    assert isinstance(m.to_morphio(), morphio.Morphology)
+
+    m = Morphology(morphio.mut.Morphology(path))
+    assert isinstance(m.to_morphio(), morphio.mut.Morphology)

tests/features/test_morphology.py

@@ -62,7 +62,10 @@
 
 def _add_neurite_trunk(morph, elevation, azimuth, neurite_type=SectionType.basal_dendrite):
     """Add a neurite from the elevation and azimuth to a given morphology."""
-    mut = morph.to_morphio().as_mutable()
+    mut = morph.to_morphio()
+    if hasattr(mut, 'as_mutable'):
+        mut = mut.as_mutable()
+
     new_pts = np.array(morphmath.vector_from_spherical(elevation, azimuth), ndmin=2)
 
     point_lvl = PointLevel(new_pts, [1])

tests/geom/test_transform.py

@@ -27,6 +27,7 @@
 # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 import math
+import morphio
 from pathlib import Path
 
 import neurom.geom.transform as gtr
@@ -222,6 +223,29 @@ def test_transform_translate_morphology_h5():
     _check_morphology_translate(m, tm, t)
 
 
+def test_transform__mut_immut():
+
+    t = np.array([100.0, 100.0, 100.0])
+
+    morph = morphio.Morphology(H5_NRN_PATH)
+
+    m1 = load_morphology(morph)
+    m2 = m1.transform(gtr.Translation(t))
+
+    assert isinstance(m2.to_morphio(), morphio.Morphology), type(m2.to_morphio())
+
+    _check_morphology_translate(m1, m2, t)
+
+    morph = morphio.mut.Morphology(H5_NRN_PATH)
+
+    m3 = load_morphology(morph)
+    m4 = m3.transform(gtr.Translation(t))
+
+    assert isinstance(m4.to_morphio(), morphio.mut.Morphology), type(m4.to_morphio())
+
+    _check_morphology_translate(m3, m4, t)
+
+
 def _apply_rot(points, rot_mat):
     return np.dot(rot_mat, np.array(points).T).T