Improve iter_segments and add segment methods in core objects

neurom/core/morphology.py

@@ -200,6 +200,10 @@ def __repr__(self):
         return (f'Section(id={self.id}, type={self.type}, n_points={len(self.points)})'
                 f'<parent: Section(id={parent_id}), nchildren: {len(self.children)}>')
 
+    def iter_segments(self):
+        """Iterator to the segments of the section."""
+        return iter_segments(self)
+
 
 # NRN simulator iteration order
 # See:
@@ -288,7 +292,10 @@ def iter_sections(neurites,
     )
 
 
-def iter_segments(obj, neurite_filter=None, neurite_order=NeuriteIter.FileOrder):
+def iter_segments(obj,
+                  neurite_filter=None,
+                  neurite_order=NeuriteIter.FileOrder,
+                  section_order=Section.ipreorder):
     """Return an iterator to the segments in a collection of neurites.
 
     Arguments:
@@ -297,6 +304,12 @@ def iter_segments(obj, neurite_filter=None, neurite_order=NeuriteIter.FileOrder)
         neurite_order: order upon which neurite should be iterated. Values:
             - NeuriteIter.FileOrder: order of appearance in the file
             - NeuriteIter.NRN: NRN simulator order: soma -> axon -> basal -> apical
+        section_order: section iteration order within a given neurite. Must be one of:
+            Section.ipreorder: Depth-first pre-order iteration of tree nodes
+            Section.ipostorder: Depth-first post-order iteration of tree nodes
+            Section.iupstream: Iterate from a tree node to the root nodes
+            Section.ibifurcation_point: Iterator to bifurcation points
+            Section.ileaf: Iterator to all leaves of a tree
 
     Note:
         This is a convenience function provided for generic access to
@@ -305,6 +318,7 @@ def iter_segments(obj, neurite_filter=None, neurite_order=NeuriteIter.FileOrder)
     """
     sections = iter((obj,) if isinstance(obj, Section) else
                     iter_sections(obj,
+                                  iterator_type=section_order,
                                   neurite_filter=neurite_filter,
                                   neurite_order=neurite_order))
 
@@ -386,7 +400,7 @@ def volume(self):
         """
         return sum(s.volume for s in self.iter_sections())
 
-    def iter_sections(self, order=Section.ipreorder, neurite_order=NeuriteIter.FileOrder):
+    def iter_sections(self, order=Section.ipreorder):
         """Iteration over section nodes.
 
         Arguments:
@@ -396,12 +410,12 @@ def iter_sections(self, order=Section.ipreorder, neurite_order=NeuriteIter.FileO
                 Section.iupstream: Iterate from a tree node to the root nodes
                 Section.ibifurcation_point: Iterator to bifurcation points
                 Section.ileaf: Iterator to all leaves of a tree
-
-            neurite_order: order upon which neurites should be iterated. Values:
-                - NeuriteIter.FileOrder: order of appearance in the file
-                - NeuriteIter.NRN: NRN simulator order: soma -> axon -> basal -> apical
         """
-        return iter_sections(self, iterator_type=order, neurite_order=neurite_order)
+        return iter_sections(self, iterator_type=order)
+
+    def iter_segments(self, section_order=Section.ipreorder):
+        """Iterator to the segments of the sections."""
+        return flatten(iter_segments(sec) for sec in self.iter_sections(order=section_order))
 
     def __nonzero__(self):
         """If has root node."""
@@ -452,6 +466,11 @@ def sections(self):
         """The array of all sections, excluding the soma."""
         return list(iter_sections(self))
 
+    @property
+    def segments(self):
+        """The array of all segments of the sections."""
+        return list(iter_segments(self))
+
     @property
     def points(self):
         """Returns the list of points."""

tests/core/test_iter.py

@@ -116,8 +116,10 @@ def test_iter_sections_filter():
         assert a == b
 
 def test_iter_sections_inrnorder():
-    assert ([s.id for n in POP.neurites for s in n.iter_sections(neurite_order=NeuriteIter.NRN)] ==
-                          [0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83,  0, 1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20,  0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83,  0,  1,  2])
+    assert (
+        [s.id for s in iter_sections(POP, neurite_order=NeuriteIter.NRN)] ==
+        [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 0, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 1, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 63, 64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 2]
+    )
 
 def test_iter_sections_ipreorder():
     assert ([s.id for n in POP.neurites for s in n.iter_sections(Section.ipreorder)] ==
@@ -159,7 +161,31 @@ def test_iter_section_morph():
 
 def test_iter_segments_morph():
     ref = list(iter_segments(SIMPLE))
-    assert len(ref) == 6
+    assert_array_equal(SIMPLE.segments, ref)
+    assert_array_equal(
+        ref,
+        [
+            [[0.0, 0.0, 0.0, 1.0], [0.0, 5.0, 0.0, 1.0]],
+            [[0.0, 5.0, 0.0, 1.0], [-5.0, 5.0, 0.0, 0.0]],
+            [[0.0, 5.0, 0.0, 1.0], [6.0, 5.0, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0, 1.0], [0.0, -4.0, 0.0, 1.0]],
+            [[0.0, -4.0, 0.0, 1.0], [6.0, -4.0, 0.0, 0.0]],
+            [[0.0, -4.0, 0.0, 1.0], [-5.0, -4.0, 0.0, 0.0]],
+        ]
+    )
+
+    ref = list(iter_segments(SIMPLE, section_order=Section.ipostorder))
+    assert_array_equal(
+        ref,
+        [
+            [[0.0, 5.0, 0.0, 1.0], [-5.0, 5.0, 0.0, 0.0]],
+            [[0.0, 5.0, 0.0, 1.0], [6.0, 5.0, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0, 1.0], [0.0, 5.0, 0.0, 1.0]],
+            [[0.0, -4.0, 0.0, 1.0], [6.0, -4.0, 0.0, 0.0]],
+            [[0.0, -4.0, 0.0, 1.0], [-5.0, -4.0, 0.0, 0.0]],
+            [[0.0, 0.0, 0.0, 1.0], [0.0, -4.0, 0.0, 1.0]],
+        ]
+    )
 
     ref = list(iter_segments(SIMPLE, neurite_filter=lambda n: n.type == nm.AXON))
     assert len(ref) == 3
@@ -214,3 +240,22 @@ def test_iter_segments_section():
     assert_array_equal(ref, [[[1, 2, 3, 4], [5, 6, 7, 8]],
                              [[5, 6, 7, 8], [8, 7, 6, 5]],
                              [[8, 7, 6, 5], [4, 3, 2, 1]]])
+    assert_array_equal(ref, list(sec.iter_segments()))
+
+
+def test_iter_segments_neurite():
+    morph = load_morphology(StringIO(u"""
+                          ((CellBody) (-1 0 0 2) (1 0 0 2))
+
+                          ((Dendrite)
+                          (1 2 3 8)
+                          (5 6 7 16)
+                          (8 7 6 10)
+                          (4 3 2 2))
+                       """), reader='asc')
+    neurite = morph.neurites[0]
+    ref = list(iter_segments(neurite))
+    assert_array_equal(ref, [[[1, 2, 3, 4], [5, 6, 7, 8]],
+                             [[5, 6, 7, 8], [8, 7, 6, 5]],
+                             [[8, 7, 6, 5], [4, 3, 2, 1]]])
+    assert_array_equal(list(neurite.iter_segments()), ref)