Merge pull request #524 from juanchopanza/issue-523

Move iter_neurites to neurom.core and neurom modules.

Author: mgeplf

neurom/__init__.py

@@ -59,7 +59,7 @@
 '''
 
 from .version import VERSION as __version__
-from .core import iter_neurites
+from .core import iter_neurites, iter_sections
 from .fst import load_neuron, load_neurons, NeuriteType, get, NEURITE_TYPES
 
 

neurom/core/__init__.py

@@ -28,7 +28,7 @@
 
 ''' Core functionality and data types of NeuroM '''
 
-from itertools import ifilter, imap
+from itertools import ifilter, imap, chain
 from .tree import Tree
 from .types import NeuriteType
 from ._soma import Soma, make_soma, SomaError
@@ -67,3 +67,28 @@ def iter_neurites(obj, mapfun=None, filt=None):
 
     neurite_iter = iter(neurites) if filt is None else ifilter(filt, neurites)
     return neurite_iter if mapfun is None else imap(mapfun, neurite_iter)
+
+
+def iter_sections(neurites, iterator_type=Tree.ipreorder, neurite_filter=None):
+    '''Iterator to the sections in a neurite, neuron or neuron population.
+
+    Parameters:
+        neurites: neuron, population, neurite, or iterable containing neurite objects
+        iterator_type: type of the iteration (ipreorder, iupstream, ibifurcation_point)
+        neurite_filter: optional top level filter on properties of neurite neurite objects.
+
+    Examples:
+
+        Get the number of points in each section of all the axons in a neuron population
+
+        >>> import neurom as nm
+        >>> from neurom.core import ites_sections
+        >>> filter = lambda n : n.type == nm.AXON
+        >>> n_points = [len(s.points) for s in iter_sections(pop,  neurite_filter=filter)]
+
+    '''
+    def _mapfun(neurite):
+        '''Map an iterator type to the root node of a neurite'''
+        return iterator_type(neurite.root_node)
+
+    return chain.from_iterable(imap(_mapfun, iter_neurites(neurites, filt=neurite_filter)))

neurom/core/tests/test_core.py

@@ -28,12 +28,14 @@
 
 import os
 from os.path import join as joinp
+from itertools import ifilter
 
 from nose import tools as nt
 import neurom as nm
 from neurom.core.population import Population
 from neurom import load_neuron
 from neurom import core
+from neurom.core import Tree
 
 _path = os.path.dirname(os.path.abspath(__file__))
 DATA_PATH = joinp(_path, '../../../test_data')
@@ -74,3 +76,47 @@ def test_iter_neurites_filter_mapping():
 
     ref = [500, 500, 500]
     nt.assert_sequence_equal(n, ref)
+
+
+def test_iter_sections_default():
+
+    ref = [s for n in POP.neurites for s in n.iter_sections()]
+    nt.assert_sequence_equal(ref,
+                             [n for n in core.iter_sections(POP)])
+
+
+def test_iter_sections_filter():
+
+    for ntyp in nm.NEURITE_TYPES:
+        a = [s for n in ifilter(lambda nn: nn.type == ntyp, POP.neurites)
+             for s in n.iter_sections()]
+        b = [n for n in core.iter_sections(POP, neurite_filter=lambda n : n.type == ntyp)]
+        nt.assert_sequence_equal(a, b)
+
+
+def test_iter_sections_ipostorder():
+
+    ref = [s for n in POP.neurites for s in n.iter_sections(Tree.ipostorder)]
+    nt.assert_sequence_equal(ref,
+                             [n for n in core.iter_sections(POP, iterator_type=Tree.ipostorder)])
+
+
+def test_iter_sections_ibifurcation():
+
+    ref = [s for n in POP.neurites for s in n.iter_sections(Tree.ibifurcation_point)]
+    nt.assert_sequence_equal(ref,
+                             [n for n in core.iter_sections(POP, iterator_type=Tree.ibifurcation_point)])
+
+
+def test_iter_sections_iforking():
+
+    ref = [s for n in POP.neurites for s in n.iter_sections(Tree.iforking_point)]
+    nt.assert_sequence_equal(ref,
+                             [n for n in core.iter_sections(POP, iterator_type=Tree.iforking_point)])
+
+
+def test_iter_sections_ileaf():
+
+    ref = [s for n in POP.neurites for s in n.iter_sections(Tree.ileaf)]
+    nt.assert_sequence_equal(ref,
+                             [n for n in core.iter_sections(POP, iterator_type=Tree.ileaf)])

neurom/fst/_neuritefunc.py

@@ -29,9 +29,9 @@
 '''Neurite functions'''
 
 from functools import partial
-from itertools import imap, chain, izip
+from itertools import chain, izip
 import numpy as np
-from neurom.core import Tree, iter_neurites
+from neurom.core import Tree, iter_neurites, iter_sections
 from neurom.core.types import tree_type_checker as is_type
 from neurom.core.types import NeuriteType
 from neurom.geom import convex_hull
@@ -42,21 +42,6 @@
                                bifurcation_partition)
 
 
-def iter_sections(neurites, iterator_type=Tree.ipreorder, neurite_filter=None):
-    '''Returns an iterator to the nodes in a iterable of neurite objects
-
-    Parameters:
-        neurites: neuron, population, neurite, or iterable containing neurite objects
-        iterator_type: type of the iteration (ipreorder, iupstream, ibifurcation_point)
-        neurite_filter: optional top level filter on properties of neurite neurite objects.
-    '''
-    def _mapfun(neurite):
-        '''Map an iterator type to the root node of a neurite'''
-        return iterator_type(neurite.root_node)
-
-    return chain.from_iterable(imap(_mapfun, iter_neurites(neurites, filt=neurite_filter)))
-
-
 def iter_segments(neurites, neurite_filter=None):
     '''Return an iterator to the segments in a collection of neurites