Explicitly populate feature dictionaries with all available features.

neurom/apps/morph_stats.py

@@ -49,8 +49,9 @@
 from neurom.apps import get_config
 from neurom.core.morphology import Morphology
 from neurom.exceptions import ConfigError
-from neurom.features import _NEURITE_FEATURES, _MORPHOLOGY_FEATURES, _POPULATION_FEATURES, \
-    _get_feature_value_and_func
+from neurom.features import (
+    _NEURITE_FEATURES, _MORPHOLOGY_FEATURES, _POPULATION_FEATURES, get_feature_value_and_func
+)
 from neurom.io.utils import get_files_by_path
 from neurom.utils import flatten, NeuromJSON, warn_deprecated
 
@@ -121,7 +122,7 @@ def _get_feature_stats(feature_name, morphs, modes, kwargs):
     If the feature is 2-dimensional, the feature is flattened on its last axis
     """
     data = {}
-    value, func = _get_feature_value_and_func(feature_name, morphs, **kwargs)
+    value, func = get_feature_value_and_func(feature_name, morphs, **kwargs)
     shape = func.shape
     if len(shape) > 2:
         raise ValueError(f'Len of "{feature_name}" feature shape must be <= 2')  # pragma: no cover

neurom/features/__init__.py

@@ -37,12 +37,16 @@
     >>> ax_sec_len = features.get('section_lengths', m, neurite_type=neurom.AXON)
 """
 import operator
+
+from copy import deepcopy
+import collections.abc
 from enum import Enum
-from functools import reduce
+from functools import reduce, wraps
 
 from neurom.core import Population, Morphology, Neurite
 from neurom.core.morphology import iter_neurites
 from neurom.core.types import NeuriteType, tree_type_checker as is_type
+from neurom.utils import flatten
 from neurom.exceptions import NeuroMError
 
 _NEURITE_FEATURES = {}
@@ -54,86 +58,38 @@ class NameSpace(Enum):
     """The level of morphology abstraction that feature applies to."""
     NEURITE = 'neurite'
     NEURON = 'morphology'
+    MORPHOLOGY = 'morphology'
     POPULATION = 'population'
 
 
-def _flatten_feature(feature_shape, feature_value):
-    """Flattens feature values. Applies for population features for backward compatibility."""
-    if feature_shape == ():
-        return feature_value
-    return reduce(operator.concat, feature_value, [])
-
-
-def _get_neurites_feature_value(feature_, obj, neurite_filter, kwargs):
-    """Collects neurite feature values appropriately to feature's shape."""
-    kwargs.pop('neurite_type', None)  # there is no 'neurite_type' arg in _NEURITE_FEATURES
-    return reduce(operator.add,
-                  (feature_(n, **kwargs) for n in iter_neurites(obj, filt=neurite_filter)),
-                  0 if feature_.shape == () else [])
+_FEATURE_CATEGORIES = {
+    NameSpace.NEURITE: _NEURITE_FEATURES,
+    NameSpace.NEURON: _MORPHOLOGY_FEATURES,
+    NameSpace.MORPHOLOGY: _MORPHOLOGY_FEATURES,
+    NameSpace.POPULATION: _POPULATION_FEATURES,
+}
 
 
-def _get_feature_value_and_func(feature_name, obj, **kwargs):
+def get(feature_name, obj, **kwargs):
     """Obtain a feature from a set of morphology objects.
 
+    Features can be either Neurite, Morphology or Population features. For Neurite features see
+    :mod:`neurom.features.neurite`. For Morphology features see :mod:`neurom.features.morphology`.
+    For Population features see :mod:`neurom.features.population`.
+
     Arguments:
         feature_name(string): feature to extract
-        obj (Neurite|Morphology|Population): neurite, morphology or population
+        obj: a morphology, a morphology population or a neurite tree
         kwargs: parameters to forward to underlying worker functions
 
     Returns:
-        Tuple(List|Number, function): A tuple (feature, func) of the feature value and its function.
-          Feature value can be a list or a number.
+        List|Number: feature value as a list or a single number.
     """
-    # pylint: disable=too-many-branches
-    is_obj_list = isinstance(obj, (list, tuple))
-    if not isinstance(obj, (Neurite, Morphology, Population)) and not is_obj_list:
-        raise NeuroMError('Only Neurite, Morphology, Population or list, tuple of Neurite,'
-                          ' Morphology can be used for feature calculation')
-
-    neurite_filter = is_type(kwargs.get('neurite_type', NeuriteType.all))
-    res, feature_ = None, None
-
-    if isinstance(obj, Neurite) or (is_obj_list and isinstance(obj[0], Neurite)):
-        # input is a neurite or a list of neurites
-        if feature_name in _NEURITE_FEATURES:
-            assert 'neurite_type' not in kwargs, 'Cant apply "neurite_type" arg to a neurite with' \
-                                                 ' a neurite feature'
-            feature_ = _NEURITE_FEATURES[feature_name]
-            if isinstance(obj, Neurite):
-                res = feature_(obj, **kwargs)
-            else:
-                res = [feature_(s, **kwargs) for s in obj]
-    elif isinstance(obj, Morphology):
-        # input is a morphology
-        if feature_name in _MORPHOLOGY_FEATURES:
-            feature_ = _MORPHOLOGY_FEATURES[feature_name]
-            res = feature_(obj, **kwargs)
-        elif feature_name in _NEURITE_FEATURES:
-            feature_ = _NEURITE_FEATURES[feature_name]
-            res = _get_neurites_feature_value(feature_, obj, neurite_filter, kwargs)
-    elif isinstance(obj, Population) or (is_obj_list and isinstance(obj[0], Morphology)):
-        # input is a morphology population or a list of morphs
-        if feature_name in _POPULATION_FEATURES:
-            feature_ = _POPULATION_FEATURES[feature_name]
-            res = feature_(obj, **kwargs)
-        elif feature_name in _MORPHOLOGY_FEATURES:
-            feature_ = _MORPHOLOGY_FEATURES[feature_name]
-            res = _flatten_feature(feature_.shape, [feature_(n, **kwargs) for n in obj])
-        elif feature_name in _NEURITE_FEATURES:
-            feature_ = _NEURITE_FEATURES[feature_name]
-            res = _flatten_feature(
-                feature_.shape,
-                [_get_neurites_feature_value(feature_, n, neurite_filter, kwargs) for n in obj])
-
-    if res is None or feature_ is None:
-        raise NeuroMError(f'Cant apply "{feature_name}" feature. Please check that it exists, '
-                          'and can be applied to your input. See the features documentation page.')
-
-    return res, feature_
+    return get_feature_value_and_func(feature_name, obj, **kwargs)[0]
 
 
-def get(feature_name, obj, **kwargs):
-    """Obtain a feature from a set of morphology objects.
+def get_feature_value_and_func(feature_name, obj, **kwargs):
+    """Obtain a feature's values and corresponding function from a set of morphology objects.
 
     Features can be either Neurite, Morphology or Population features. For Neurite features see
     :mod:`neurom.features.neurite`. For Morphology features see :mod:`neurom.features.morphology`.
@@ -146,8 +102,60 @@ def get(feature_name, obj, **kwargs):
 
     Returns:
         List|Number: feature value as a list or a single number.
+        Callable: feature function used to calculate the value.
+    """
+    try:
+
+        if isinstance(obj, Neurite):
+            feature_function = _NEURITE_FEATURES[feature_name]
+            return feature_function(obj, **kwargs), feature_function
+
+        if isinstance(obj, Morphology):
+            feature_function = _MORPHOLOGY_FEATURES[feature_name]
+            return feature_function(obj, **kwargs), feature_function
+
+        if isinstance(obj, Population):
+            feature_function = _POPULATION_FEATURES[feature_name]
+            return feature_function(obj, **kwargs), feature_function
+
+        if isinstance(obj, collections.abc.Sequence):
+
+            if isinstance(obj[0], Neurite):
+                feature_function = _NEURITE_FEATURES[feature_name]
+                return [feature_function(neu, **kwargs) for neu in obj], feature_function
+
+            if isinstance(obj[0], Morphology):
+                feature_function = _POPULATION_FEATURES[feature_name]
+                return feature_function(obj, **kwargs), feature_function
+
+    except Exception as e:
+
+        raise NeuroMError(
+            f"Cant apply '{feature_name}' feature on {type(obj)}. Please check that it exists, "
+            "and can be applied to your input. See the features documentation page."
+        ) from e
+
+    raise NeuroMError(
+        "Only Neurite, Morphology, Population or list, tuple of Neurite, Morphology can be used for"
+        " feature calculation."
+        f"Got {type(obj)} instead. See the features documentation page."
+    )
+
+
+def feature(shape, namespace: NameSpace, name=None):
+    """Feature decorator to automatically register the feature in the appropriate namespace.
+
+    Arguments:
+        shape(tuple): the expected shape of the feature values
+        namespace(string): a namespace, see :class:`NameSpace`
+        name(string): name of the feature, used to access the feature via `neurom.features.get()`.
     """
-    return _get_feature_value_and_func(feature_name, obj, **kwargs)[0]
+
+    def inner(feature_function):
+        _register_feature(namespace, name or feature_function.__name__, feature_function, shape)
+        return feature_function
+
+    return inner
 
 
 def _register_feature(namespace: NameSpace, name, func, shape):
@@ -162,34 +170,135 @@ def _register_feature(namespace: NameSpace, name, func, shape):
         func(callable): single parameter function of a neurite.
         shape(tuple): the expected shape of the feature values
     """
-    setattr(func, 'shape', shape)
-    _map = {NameSpace.NEURITE: _NEURITE_FEATURES,
-            NameSpace.NEURON: _MORPHOLOGY_FEATURES,
-            NameSpace.POPULATION: _POPULATION_FEATURES}
-    if name in _map[namespace]:
+    if name in _FEATURE_CATEGORIES[namespace]:
         raise NeuroMError(f'A feature is already registered under "{name}"')
-    _map[namespace][name] = func
 
+    setattr(func, "shape", shape)
 
-def feature(shape, namespace: NameSpace, name=None):
-    """Feature decorator to automatically register the feature in the appropriate namespace.
+    _FEATURE_CATEGORIES[namespace][name] = func
 
-    Arguments:
-        shape(tuple): the expected shape of the feature values
-        namespace(string): a namespace, see :class:`NameSpace`
-        name(string): name of the feature, used to access the feature via `neurom.features.get()`.
+
+def _flatten_feature(feature_value, feature_shape):
+    """Flattens feature values. Applies for population features for backward compatibility."""
+    return feature_value if feature_shape == () else reduce(operator.concat, feature_value, [])
+
+
+def _get_neurites_feature_value(feature_, obj, kwargs):
+    """Collects neurite feature values appropriately to feature's shape."""
+    kwargs = deepcopy(kwargs)
+
+    # there is no 'neurite_type' arg in _NEURITE_FEATURES
+    if "neurite_type" in kwargs:
+        neurite_type = kwargs["neurite_type"]
+        del kwargs["neurite_type"]
+    else:
+        neurite_type = NeuriteType.all
+
+    per_neurite_values = (
+        feature_(n, **kwargs) for n in iter_neurites(obj, filt=is_type(neurite_type))
+    )
+
+    return reduce(operator.add, per_neurite_values, 0 if feature_.shape == () else [])
+
+
+def _transform_downstream_features_to_upstream_feature_categories(features):
+    """Adds each feature to all upstream feature categories, adapted for the respective objects.
+
+    If a feature is already defined in the module of an upstream category, it is not overwritten.
+    This allows to achieve both reducible features, which can be defined for instance at the neurite
+    category and then automatically added to the morphology and population categories transformed to
+    work with morphology and population objects respetively.
+
+    However, if a feature is not reducible, which means that an upstream category is not comprised
+    by the accumulation/sum of its components, the feature should be defined on each category
+    module so that the module logic is used instead.
+
+    After the end of this function the _NEURITE_FEATURES, _MORPHOLOGY_FEATURES, _POPULATION_FEATURES
+    are updated so that all features in neurite features are also available in morphology and
+    population dictionaries, and all morphology features are available in the population dictionary.
+
+    Args:
+        features: Dictionary with feature categories.
+
+    Notes:
+        Category     Upstream Categories
+        --------     -----------------
+        morphology   population
+        neurite      morphology, population
     """
+    def apply_neurite_feature_to_population(func):
+        """Transforms a feature in _NEURITE_FEATURES so that it can be applied to a population.
 
-    def inner(func):
-        _register_feature(namespace, name or func.__name__, func, shape)
-        return func
+        Args:
+            func: Feature function.
 
-    return inner
+        Returns:
+            Transformed neurite function to be applied on a population of morphologies.
+        """
+        def apply_to_population(pop, **kwargs):
+
+            per_morphology_values = [
+                _get_neurites_feature_value(func, morph, kwargs) for morph in pop
+            ]
+            return _flatten_feature(per_morphology_values, func.shape)
+
+        return apply_to_population
+
+    def apply_neurite_feature_to_morphology(func):
+        """Transforms a feature in _NEURITE_FEATURES so that it can be applied on neurites.
+
+        Args:
+            func: Feature function.
+
+        Returns:
+            Transformed neurite function to be applied on a morphology.
+        """
+        def apply_to_morphology(morph, **kwargs):
+            return _get_neurites_feature_value(func, morph, kwargs)
+        return apply_to_morphology
+
+    def apply_morphology_feature_to_population(func):
+        """Transforms a feature in _MORPHOLOGY_FEATURES so that it can be applied to a population.
+
+        Args:
+            func: Feature function.
+
+        Returns:
+            Transformed morphology function to be applied on a population of morphologies.
+        """
+        def apply_to_population(pop, **kwargs):
+            per_morphology_values = [func(morph, **kwargs) for morph in pop]
+            return _flatten_feature(per_morphology_values, func.shape)
+        return apply_to_population
+
+    transformations = {
+        (NameSpace.POPULATION, NameSpace.MORPHOLOGY): apply_morphology_feature_to_population,
+        (NameSpace.POPULATION, NameSpace.NEURITE): apply_neurite_feature_to_population,
+        (NameSpace.MORPHOLOGY, NameSpace.NEURITE): apply_neurite_feature_to_morphology,
+    }
+
+    for (upstream_category, category), transformation in transformations.items():
+
+        features = _FEATURE_CATEGORIES[category]
+        upstream_features = _FEATURE_CATEGORIES[upstream_category]
+
+        for feature_name, feature_function in features.items():
+
+            if feature_name in upstream_features:
+                continue
+
+            upstream_features[feature_name] = transformation(feature_function)
+            setattr(upstream_features[feature_name], "shape", feature_function.shape)
 
 
 # These imports are necessary in order to register the features
-from neurom.features import neurite, morphology, \
-    population  # noqa, pylint: disable=wrong-import-position
+from neurom.features import neurite, morphology, population  # noqa, pylint: disable=wrong-import-position
+
+
+# Update the feature dictionaries so that features from lower categories are transformed and usable
+# by upstream categories. For example, a neurite feature will be added to morphology and population
+# feature dictionaries, transformed so that it works with the respective objects.
+_transform_downstream_features_to_upstream_feature_categories(_FEATURE_CATEGORIES)
 
 
 def _features_catalogue():

tests/features/test_get_features.py

@@ -65,10 +65,9 @@ def _stats(seq):
 
 
 def test_get_raises():
-    with pytest.raises(NeuroMError,
-                       match='Only Neurite, Morphology, Population or list, tuple of Neurite, Morphology'):
+    with pytest.raises(NeuroMError, match="Only Neurite, Morphology, Population or list, tuple of Neurite, Morphology"):
         features.get('soma_radius', (n for n in POP))
-    with pytest.raises(NeuroMError, match='Cant apply "invalid" feature'):
+    with pytest.raises(NeuroMError, match="Cant apply 'invalid' feature"):
         features.get('invalid', NRN)