Cleanup neurite features

neurom/features/bifurcation.py

@@ -33,7 +33,7 @@
 from neurom.exceptions import NeuroMError
 from neurom.core.dataformat import COLS
 from neurom.core.morphology import Section
-from neurom.features.section import section_mean_radius
+from neurom.features import section as sf
 
 
 def _raise_if_not_bifurcation(section):
@@ -156,8 +156,8 @@ def sibling_ratio(bif_point, method='first'):
         n = bif_point.children[0].points[1, COLS.R]
         m = bif_point.children[1].points[1, COLS.R]
     if method == 'mean':
-        n = section_mean_radius(bif_point.children[0])
-        m = section_mean_radius(bif_point.children[1])
+        n = sf.section_mean_radius(bif_point.children[0])
+        m = sf.section_mean_radius(bif_point.children[1])
     return min(n, m) / max(n, m)
 
 
@@ -182,7 +182,28 @@ def diameter_power_relation(bif_point, method='first'):
         d_child1 = bif_point.children[0].points[1, COLS.R]
         d_child2 = bif_point.children[1].points[1, COLS.R]
     if method == 'mean':
-        d_child = section_mean_radius(bif_point)
-        d_child1 = section_mean_radius(bif_point.children[0])
-        d_child2 = section_mean_radius(bif_point.children[1])
+        d_child = sf.section_mean_radius(bif_point)
+        d_child1 = sf.section_mean_radius(bif_point.children[0])
+        d_child2 = sf.section_mean_radius(bif_point.children[1])
     return (d_child / d_child1)**(1.5) + (d_child / d_child2)**(1.5)
+
+
+def downstream_pathlength_asymmetry(
+    bif_point, normalization_length=1.0, iterator_type=Section.ipreorder
+):
+    """Calculates the downstream pathlength asymmetry at a bifurcation point.
+
+    Args:
+        bif_point: Bifurcation section.
+        normalization_length: Constant to divide the result with.
+        iterator_type: Iterator type that specifies how the two subtrees are traversed.
+
+    Returns:
+        The absolute difference between the downstream path distances of the two children, divided
+        by the normalization length.
+    """
+    _raise_if_not_bifurcation(bif_point)
+    return abs(
+        sf.downstream_pathlength(bif_point.children[0], iterator_type=iterator_type) -
+        sf.downstream_pathlength(bif_point.children[1], iterator_type=iterator_type),
+    ) / normalization_length

neurom/features/neurite.py

@@ -48,7 +48,7 @@
 
 import numpy as np
 from neurom import morphmath
-from neurom.utils import flatten
+from neurom import utils
 from neurom.core.types import NeuriteType
 from neurom.core.morphology import Section, iter_points
 from neurom.core.dataformat import COLS
@@ -69,6 +69,7 @@ def _map_sections(fun, neurite, iterator_type=Section.ipreorder, section_type=Ne
     def homogeneous_filter(section):
         return check_type(section) and Section.is_homogeneous_point(section)
 
+    # forking sections cannot be heterogeneous
     if (
         iterator_type in {Section.ibifurcation_point, Section.iforking_point}
         and section_type != NeuriteType.all
@@ -188,19 +189,14 @@ def section_term_branch_orders(neurite, section_type=NeuriteType.all):
 def section_path_distances(neurite, iterator_type=Section.ipreorder, section_type=NeuriteType.all):
     """Path lengths."""
 
-    def takeuntil(predicate, iterable):
-        """Similar to itertools.takewhile but it returns the last element before stopping."""
-        for x in iterable:
-            yield x
-            if predicate(x):
-                break
-
-    def pl2(node):
+    def path_length(node):
         """Calculate the path length using cached section lengths."""
-        sections = takeuntil(lambda s: s.id == neurite.root_node.id, node.iupstream())
+        sections = utils.takeuntil(lambda s: s.id == neurite.root_node.id, node.iupstream())
         return sum(n.length for n in sections)
 
-    return _map_sections(pl2, neurite, iterator_type=iterator_type, section_type=section_type)
+    return _map_sections(
+        path_length, neurite, iterator_type=iterator_type, section_type=section_type
+    )
 
 
 ################################################################################
@@ -213,7 +209,7 @@ def _map_segments(func, neurite, section_type=NeuriteType.all):
 
     `func` accepts a section and returns list of values corresponding to each segment.
     """
-    return list(flatten(_map_sections(func, neurite, section_type=section_type)))
+    return list(utils.flatten(_map_sections(func, neurite, section_type=section_type)))
 
 
 @feature(shape=(...,))
@@ -293,14 +289,12 @@ def segments_path_length(section):
 @feature(shape=(...,))
 def segment_radial_distances(neurite, origin=None, section_type=NeuriteType.all):
     """Returns the list of distances between all segment mid points and origin."""
-    pos = neurite.root_node.points[0] if origin is None else origin
-
-    def radial_distances(section):
-        """List of distances between the mid point of each segment and pos."""
-        mid_pts = 0.5 * (section.points[:-1, COLS.XYZ] + section.points[1:, COLS.XYZ])
-        return np.linalg.norm(mid_pts - pos[COLS.XYZ], axis=1)
-
-    return _map_segments(radial_distances, neurite, section_type=section_type)
+    origin = neurite.root_node.points[0, COLS.XYZ] if origin is None else origin
+    return _map_segments(
+        func=partial(sf.segment_midpoint_radial_distances, origin=origin),
+        neurite=neurite,
+        section_type=section_type
+    )
 
 
 @feature(shape=(...,))
@@ -337,44 +331,33 @@ def partition_asymmetry(
     :func:`neurom.features.bifurcationfunc.partition_asymmetry`
     """
     if variant not in {'branch-order', 'length'}:
-        raise ValueError('Please provide a valid variant for partition asymmetry,'
-                         f'found {variant}')
+        raise ValueError(
+            "Please provide a valid variant for partition asymmetry. "
+            f"Expected 'branch-order' or 'length', got {variant}."
+        )
     if method not in {'petilla', 'uylings'}:
-        raise ValueError('Please provide a valid method for partition asymmetry,'
-                         'either "petilla" or "uylings"')
-
-    def it_type(section):
-
-        if section_type == NeuriteType.all:
-            return Section.ipreorder(section)
+        raise ValueError(
+            "Please provide a valid method for partition asymmetry. "
+            f"Expected 'petilla' or 'uylings', got {method}."
+        )
 
-        check = is_type(section_type)
-        return (s for s in section.ipreorder() if check(s))
+    # create a downstream iterator that is filtered by the section type
+    it_type = utils.filtered_iterator(is_type(section_type), Section.ipreorder)
 
     if variant == 'branch-order':
-
-        function = partial(
-            bf.partition_asymmetry, uylings=method == 'uylings', iterator_type=it_type
-        )
-
         return _map_sections(
-            function,
+            partial(bf.partition_asymmetry, uylings=method == 'uylings', iterator_type=it_type),
             neurite,
             iterator_type=Section.ibifurcation_point,
             section_type=section_type
         )
 
-    neurite_length = total_length(neurite, section_type=section_type)
-
-    def pathlength_asymmetry_ratio(section):
-        pathlength_diff = abs(
-            sf.downstream_pathlength(section.children[0], iterator_type=it_type) -
-            sf.downstream_pathlength(section.children[1], iterator_type=it_type)
-        )
-        return pathlength_diff / neurite_length
-
     return _map_sections(
-        pathlength_asymmetry_ratio,
+        partial(
+            bf.downstream_pathlength_asymmetry,
+            normalization_length=total_length(neurite, section_type=section_type),
+            iterator_type=it_type,
+        ),
         neurite,
         iterator_type=Section.ibifurcation_point,
         section_type=section_type
@@ -499,7 +482,7 @@ def get_points(section):
 
     # note: duplicate points included but not affect the convex hull calculation
     points = list(
-        flatten(_map_sections(get_points, neurite, section_type=section_type))
+        utils.flatten(_map_sections(get_points, neurite, section_type=section_type))
     )
 
     hull = convex_hull(points)

neurom/features/section.py

@@ -138,6 +138,13 @@ def segment_midpoints(section):
     return np.divide(np.add(pts[:-1], pts[1:]), 2.0).tolist()
 
 
+def segment_midpoint_radial_distances(section, origin=None):
+    """Returns the list of segment midpoint radial distances to the origin."""
+    origin = np.zeros(3, dtype=float) if origin is None else origin
+    midpoints = np.array(segment_midpoints(section))
+    return np.linalg.norm(midpoints - origin, axis=1).tolist()
+
+
 def segment_taper_rates(section):
     """Returns the list of segment taper rates within the section."""
     pts = section.points[:, COLS.XYZR]

neurom/utils.py

@@ -136,3 +136,19 @@ def str_to_plane(plane):
 def flatten(list_of_lists):
     """Flatten one level of nesting."""
     return chain.from_iterable(list_of_lists)
+
+
+def takeuntil(predicate, iterable):
+    """Similar to itertools.takewhile but it returns the last element before stopping."""
+    for x in iterable:
+        yield x
+        if predicate(x):
+            break
+
+
+def filtered_iterator(predicate, iterator_type):
+    """Returns an iterator function that is filtered by the predicate."""
+    @wraps(iterator_type)
+    def composed(*args, **kwargs):
+        return filter(predicate, iterator_type(*args, **kwargs))
+    return composed