Repository: BlueBrain/NeuroM

Issue:
Is there a way to define optional arguments from morph_stats application?
I would like to use `extract_dataframe` from `morph_stats` with the shape features, such as `aspect_ratio`. 

This function works with a config as follows:

``` python
from neurom.apps import morph_stats
full_cfg = morph_stats.full_config()
df = morph_stats.extract_dataframe(neuron, config=full_cfg)
```
How can I select the options such as `projection_plane` from the input `config`? 

I am not sure if this is possible with the current code, if not would it be possible to add this? 

Discussion:
codecov-commenter: # [Codecov](https://codecov.io/gh/BlueBrain/NeuroM/pull/1021?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=BlueBrain) Report
> Merging [#1021](https://codecov.io/gh/BlueBrain/NeuroM/pull/1021?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=BlueBrain) (147c05f) into [master](https://codecov.io/gh/BlueBrain/NeuroM/commit/f4cb0398ae72d8eec2f9e0d0102203b3a2c6e0ba?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=BlueBrain) (f4cb039) will **not change** coverage.
> The diff coverage is `100.00%`.



```diff
@@            Coverage Diff            @@
##            master     #1021   +/-   ##
=========================================
  Coverage   100.00%   100.00%           
=========================================
  Files           36        36           
  Lines         2401      2395    -6     
=========================================
- Hits          2401      2395    -6     
```

lidakanari: Yep, that works well for me, thanks! 

Patch:
diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index 8606975..7a3dfef 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -4,6 +4,7 @@ Changelog
 Version 3.2.0
 -------------
 
+- List of multiple kwargs configurations are now allowed in``neurom.apps.morph_stats``.
 - ``neurom.features.neurite.principal_direction_extents`` directions correspond to extents
   ordered in a descending order.
 - Add features ``neurom.features.morphology.(aspect_ration, circularity, shape_factor)```
diff --git a/doc/source/morph_stats.rst b/doc/source/morph_stats.rst
index 21201a3..47917a4 100644
--- a/doc/source/morph_stats.rst
+++ b/doc/source/morph_stats.rst
@@ -58,7 +58,7 @@ Short format (prior version 3.0.0)
 An example config:
 
 .. code-block:: yaml
-    
+
     neurite:
         section_lengths:
             - max
@@ -67,13 +67,13 @@ An example config:
             - total
         section_branch_orders:
             - max
-    
+
     neurite_type:
         - AXON
         - APICAL_DENDRITE
         - BASAL_DENDRITE
         - ALL
-    
+
     neuron:
         soma_radius:
             - mean
@@ -93,7 +93,7 @@ function, e.g.
 * ``raw``: array of raw values
 * ``max``, ``min``, ``mean``, ``median``, ``std``: self-explanatory.
 * ``total``: sum of the raw values
-  
+
 An additional field ``neurite_type`` specifies the neurite types into which the morphometrics
 are to be split. It applies only to ``neurite`` features. A sample output using the above
 configuration:
@@ -198,6 +198,65 @@ So the example config from `Short format (prior version 3.0.0)`_ looks:
                - mean
 
 
+List of features format (starting version 3.2.0)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The kwargs entry was converted into a list to allow running the same feature with different arguments. The ``partition_asymmetry`` feature in the example above can be specified multiple times with different arguments as follows:
+
+.. code-block:: yaml
+
+    neurite:
+      partition_asymmetry:
+        kwargs:
+        - method: petilla
+          variant: length
+        - method: uylings
+          variant: branch-order
+        modes:
+        - max
+        - sum
+
+To allow differentiation between the feature multiples, the keys and values of the kwargs are appended at the end of the feature name:
+
+.. code-block::
+
+    partition_asymmetry__variant:length__method:petilla
+    partition_asymmetry__variant:branch-order__method:uylings
+
+The example config from `Short format (prior version 3.0.0)`_ becomes:
+
+.. code-block:: yaml
+
+    neurite:
+      section_branch_orders:
+        kwargs:
+        - {}
+        modes:
+        - max
+      section_lengths:
+        kwargs:
+        - {}
+        modes:
+        - max
+        - sum
+      section_volumes:
+        kwargs:
+        - {}
+        modes:
+        - sum
+    morphology:
+      soma_radius:
+        kwargs:
+        - {}
+        modes:
+        - mean
+    neurite_type:
+    - AXON
+    - APICAL_DENDRITE
+    - BASAL_DENDRITE
+    - ALL
+
+
 Features
 --------
 
diff --git a/neurom/apps/morph_stats.py b/neurom/apps/morph_stats.py
index 9bcb948..b1d90d9 100644
--- a/neurom/apps/morph_stats.py
+++ b/neurom/apps/morph_stats.py
@@ -38,10 +38,8 @@ from collections import defaultdict
 from collections.abc import Sized
 from copy import deepcopy
 from functools import partial
-from itertools import chain, product
 from pathlib import Path
 import pkg_resources
-from tqdm import tqdm
 import numpy as np
 import pandas as pd
 from morphio import SomaError
@@ -53,7 +51,7 @@ from neurom.exceptions import ConfigError
 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
+from neurom.utils import flatten, NeuromJSON
 
 L = logging.getLogger(__name__)
 
@@ -76,9 +74,11 @@ def extract_dataframe(morphs, config, n_workers=1):
         config (dict): configuration dict. The keys are:
             - neurite_type: a list of neurite types for which features are extracted
               If not provided, all neurite_type will be used
-            - neurite: a dictionary {{neurite_feature: mode}} where:
-                - neurite_feature is a string from NEURITEFEATURES or NEURONFEATURES
-                - mode is an aggregation operation provided as a string such as:
+            - neurite:
+                Either a list of features: [feature_name, {kwargs: {}, modes: []}] or
+                a dictionary of features {feature_name: {kwargs: {}, modes: []}}.
+                - kwargs is an optional entry allowing to pass kwargs to the feature function
+                - modes is an aggregation operation provided as a string such as:
                   ['min', 'max', 'median', 'mean', 'std', 'raw', 'sum']
             - morphology: same as neurite entry, but it will not be run on each neurite_type,
               but only once on the whole morphology.
@@ -90,7 +90,6 @@ def extract_dataframe(morphs, config, n_workers=1):
     Note:
         An example config can be found at:
 
-    {config_path}
     """
     if isinstance(morphs, Morphology):
         morphs = [morphs]
@@ -112,7 +111,7 @@ def extract_dataframe(morphs, config, n_workers=1):
     return pd.DataFrame(columns=pd.MultiIndex.from_tuples(columns), data=rows)
 
 
-extract_dataframe.__doc__ = extract_dataframe.__doc__.format(config_path=EXAMPLE_CONFIG)
+extract_dataframe.__doc__ += str(EXAMPLE_CONFIG)
 
 
 def _get_feature_stats(feature_name, morphs, modes, kwargs):
@@ -120,6 +119,21 @@ def _get_feature_stats(feature_name, morphs, modes, kwargs):
 
     If the feature is 2-dimensional, the feature is flattened on its last axis
     """
+    def stat_name_format(mode, feature_name, kwargs):
+        """Returns the key name for the data dictionary.
+
+        The key is a combination of the mode, feature_name and an optional suffix of all the extra
+        kwargs that are passed in the feature function (apart from neurite_type).
+        """
+        suffix = "__".join(
+            [f"{key}:{value}" for key, value in kwargs.items() if key != "neurite_type"]
+        )
+
+        if suffix:
+            return f"{mode}_{feature_name}__{suffix}"
+
+        return f"{mode}_{feature_name}"
+
     data = {}
     value, func = _get_feature_value_and_func(feature_name, morphs, **kwargs)
     shape = func.shape
@@ -127,7 +141,8 @@ def _get_feature_stats(feature_name, morphs, modes, kwargs):
         raise ValueError(f'Len of "{feature_name}" feature shape must be <= 2')  # pragma: no cover
 
     for mode in modes:
-        stat_name = f'{mode}_{feature_name}'
+
+        stat_name = stat_name_format(mode, feature_name, kwargs)
 
         stat = value
         if isinstance(value, Sized):
@@ -154,7 +169,12 @@ def extract_stats(morphs, config):
         config (dict): configuration dict. The keys are:
             - neurite_type: a list of neurite types for which features are extracted
               If not provided, all neurite_type will be used.
-            - neurite: a dictionary {{neurite_feature: mode}} where:
+            - neurite:
+                Either a list of features: [feature_name, {kwargs: {}, modes: []}] or
+                a dictionary of features {feature_name: {kwargs: {}, modes: []}}.
+                - kwargs is an optional entry allowing to pass kwargs to the feature function
+                - modes is an aggregation operation provided as a string such as:
+                  ['min', 'max', 'median', 'mean', 'std', 'raw', 'sum']
                 - neurite_feature is a string from NEURITEFEATURES or NEURONFEATURES
                 - mode is an aggregation operation provided as a string such as:
                   ['min', 'max', 'median', 'mean', 'std', 'raw', 'sum']
@@ -167,56 +187,60 @@ def extract_stats(morphs, config):
     Note:
         An example config can be found at:
 
-    {config_path}
     """
-    stats = defaultdict(dict)
-    neurite_features = product(['neurite'], config.get('neurite', {}).items())
-    if 'neuron' in config:    # pragma: no cover
-        warn_deprecated('Usage of "neuron" is deprecated in configs of `morph_stats` package. '
-                        'Use "morphology" instead.')
-        config['morphology'] = config['neuron']
-        del config['neuron']
-    morph_features = product(['morphology'], config.get('morphology', {}).items())
-    population_features = product(['population'], config.get('population', {}).items())
+    config = _sanitize_config(config)
+
     neurite_types = [_NEURITE_MAP[t] for t in config.get('neurite_type', _NEURITE_MAP.keys())]
 
-    for namespace, (feature_name, opts) in chain(neurite_features, morph_features,
-                                                 population_features):
-        if isinstance(opts, dict):
-            kwargs = deepcopy(opts.get('kwargs', {}))
-            modes = opts.get('modes', [])
-        else:
-            kwargs = {}
-            modes = opts
-        if namespace == 'neurite':
-            if 'neurite_type' not in kwargs and neurite_types:
-                for t in neurite_types:
-                    kwargs['neurite_type'] = t
-                    stats[t.name].update(_get_feature_stats(feature_name, morphs, modes, kwargs))
-            else:
-                t = _NEURITE_MAP[kwargs.get('neurite_type', 'ALL')]
-                kwargs['neurite_type'] = t
-                stats[t.name].update(_get_feature_stats(feature_name, morphs, modes, kwargs))
-        else:
-            stats[namespace].update(_get_feature_stats(feature_name, morphs, modes, kwargs))
+    stats = defaultdict(dict)
+    for category in ("neurite", "morphology", "population"):
+        for feature_name, opts in config[category].items():
+
+            list_of_kwargs = opts["kwargs"]
+            modes = opts["modes"]
+
+            for feature_kwargs in list_of_kwargs:
+
+                if category == 'neurite':
+
+                    # mutated below, need a copy
+                    feature_kwargs = deepcopy(feature_kwargs)
+
+                    types = (
+                        neurite_types
+                        if 'neurite_type' not in feature_kwargs and neurite_types
+                        else [_NEURITE_MAP[feature_kwargs.get('neurite_type', 'ALL')]]
+                    )
+
+                    for neurite_type in types:
+                        feature_kwargs["neurite_type"] = neurite_type
+                        stats[neurite_type.name].update(
+                            _get_feature_stats(feature_name, morphs, modes, feature_kwargs)
+                        )
+
+                else:
+                    stats[category].update(
+                        _get_feature_stats(feature_name, morphs, modes, feature_kwargs)
+                    )
 
     return dict(stats)
 
 
-extract_stats.__doc__ = extract_stats.__doc__.format(config_path=EXAMPLE_CONFIG)
+extract_stats.__doc__ += str(EXAMPLE_CONFIG)
 
 
-def get_header(results):
+def _get_header(results):
     """Extracts the headers, using the first value in the dict as the template."""
-    ret = ['name', ]
     values = next(iter(results.values()))
-    for k, v in values.items():
-        for metric in v.keys():
-            ret.append('%s:%s' % (k, metric))
-    return ret
+
+    return ['name'] + [
+        f'{k}:{metric}'
+        for k, v in values.items()
+        for metric in v.keys()
+    ]
 
 
-def generate_flattened_dict(headers, results):
+def _generate_flattened_dict(headers, results):
     """Extract from results the fields in the headers list."""
     for name, values in results.items():
         row = []
@@ -224,7 +248,9 @@ def generate_flattened_dict(headers, results):
             if header == 'name':
                 row.append(name)
             else:
-                neurite_type, metric = header.split(':')
+                # split on first occurence of `:` because feature kwargs may
+                # use a colon for separating key and value.
+                neurite_type, metric = header.split(':', 1)
                 row.append(values[neurite_type][metric])
         yield row
 
@@ -240,24 +266,83 @@ _NEURITE_MAP = {
 def full_config():
     """Returns a config with all features, all modes, all neurite types."""
     modes = ['min', 'max', 'median', 'mean', 'std']
-    return {
-        'neurite': {feature: modes for feature in _NEURITE_FEATURES},
-        'morphology': {feature: modes for feature in _MORPHOLOGY_FEATURES},
-        'population': {feature: modes for feature in _POPULATION_FEATURES},
-        'neurite_type': list(_NEURITE_MAP.keys()),
+
+    categories = {
+        "neurite": _NEURITE_FEATURES,
+        "morphology": _MORPHOLOGY_FEATURES,
+        "population": _POPULATION_FEATURES
     }
 
+    config = {
+        category: {name: {"kwargs": [{}], "modes": modes} for name in features}
+        for category, features in categories.items()
+    }
+
+    config["neurite_type"] = list(_NEURITE_MAP.keys())
+
+    return config
+
+
+def _standardize_layout(category_features):
+    """Standardizes the dictionary of features to a single format.
+
+    Args:
+        category_features: A dictionary the keys of which are features names and its values are
+        either a list of modes ([]), a dictionary of kwargs and modes {kwargs: {}, modes: []}, or
+        the standardized layout where the kwargs take a list of dicts {kwargs: [{}], modes: []}.
+
+    Returns:
+        The standardized features layout {feature: {kwargs: [{}], modes: []}}
+
+    Notes:
+        And example of the final layout is:
+
+            - feature1:
+                kwargs:
+                  - kwargs1
+                  - kwargs2
+                modes:
+                  - mode1
+                  - mode2
+            - feature2:
+                kwargs:
+                  - kwargs1
+                  - kwargs2
+                modes:
+                  - mode1
+                  - mode2
+    """
+    def standardize_options(options):
+        """Returns options as a dict with two keys: 'kwargs' and 'modes'."""
+        # convert short format
+        if isinstance(options, list):
+            return {"kwargs": [{}], "modes": options}
+
+        modes = options.get("modes", [])
 
-def sanitize_config(config):
+        if "kwargs" not in options:
+            return {"kwargs": [{}], "modes": modes}
+
+        kwargs = options["kwargs"]
+
+        # previous format where kwargs were a single entry
+        if isinstance(kwargs, dict):
+            return {"kwargs": [kwargs], "modes": modes}
+
+        return {"kwargs": kwargs, "modes": modes}
+
+    return {name: standardize_options(options) for name, options in category_features.items()}
+
+
+def _sanitize_config(config):
     """Check that the config has the correct keys, add missing keys if necessary."""
-    if 'neurite' in config:
-        if 'neurite_type' not in config:
-            raise ConfigError('"neurite_type" missing from config, but "neurite" set')
-    else:
-        config['neurite'] = {}
+    config = deepcopy(config)
+
+    if "neuron" in config:
+        config["morphology"] = config.pop("neuron")
 
-    if 'morphology' not in config:
-        config['morphology'] = {}
+    for category in ("neurite", "morphology", "population"):
+        config[category] = _standardize_layout(config[category]) if category in config else {}
 
     return config
 
@@ -273,28 +358,25 @@ def main(datapath, config, output_file, is_full_config, as_population, ignored_e
         as_population (bool): treat ``datapath`` as directory of morphologies population
         ignored_exceptions (list|tuple|None): exceptions to ignore when loading a morphology
     """
-    if is_full_config:
-        config = full_config()
-    else:
-        try:
-            config = get_config(config, EXAMPLE_CONFIG)
-            config = sanitize_config(config)
-        except ConfigError as e:
-            L.error(e)
-            raise
+    config = full_config() if is_full_config else get_config(config, EXAMPLE_CONFIG)
+
+    if 'neurite' in config and 'neurite_type' not in config:
+        error = ConfigError('"neurite_type" missing from config, but "neurite" set')
+        L.error(error)
+        raise error
 
     if ignored_exceptions is None:
         ignored_exceptions = ()
-    ignored_exceptions = tuple(IGNORABLE_EXCEPTIONS[k] for k in ignored_exceptions)
-    morphs = nm.load_morphologies(get_files_by_path(datapath),
-                                  ignored_exceptions=ignored_exceptions)
 
-    results = {}
+    morphs = nm.load_morphologies(
+        get_files_by_path(datapath),
+        ignored_exceptions=tuple(IGNORABLE_EXCEPTIONS[k] for k in ignored_exceptions)
+    )
+
     if as_population:
-        results[datapath] = extract_stats(morphs, config)
+        results = {datapath: extract_stats(morphs, config)}
     else:
-        for m in tqdm(morphs):
-            results[m.name] = extract_stats(m, config)
+        results = {m.name: extract_stats(m, config) for m in morphs}
 
     if not output_file:
         print(json.dumps(results, indent=2, separators=(',', ':'), cls=NeuromJSON))
@@ -304,7 +386,7 @@ def main(datapath, config, output_file, is_full_config, as_population, ignored_e
     else:
         with open(output_file, 'w') as f:
             csvwriter = csv.writer(f)
-            header = get_header(results)
+            header = _get_header(results)
             csvwriter.writerow(header)
-            for line in generate_flattened_dict(header, dict(results)):
+            for line in _generate_flattened_dict(header, dict(results)):
                 csvwriter.writerow(line)


Test patch:
diff --git a/tests/apps/test_morph_stats.py b/tests/apps/test_morph_stats.py
index 6c15ef5..1b1072a 100644
--- a/tests/apps/test_morph_stats.py
+++ b/tests/apps/test_morph_stats.py
@@ -43,6 +43,7 @@ from pandas.testing import assert_frame_equal
 
 DATA_PATH = Path(__file__).parent.parent / 'data'
 SWC_PATH = DATA_PATH / 'swc'
+
 REF_CONFIG = {
     'neurite': {
         'section_lengths': ['max', 'sum'],
@@ -73,6 +74,8 @@ REF_CONFIG_NEW = {
     }
 }
 
+
+
 REF_OUT = {
     'morphology': {
         'mean_soma_radius': 0.13065629648763766,
@@ -181,6 +184,122 @@ def test_extract_stats_scalar_feature():
                    'morphology': {'sum_soma_volume': 1424.4383771584492}}
 
 
+
+def test_extract_stats__kwarg_modes_multiple_features():
+
+    m = nm.load_morphology(SWC_PATH / 'Neuron.swc')
+    config = {
+        'neurite': {
+            'principal_direction_extents': {
+                'kwargs': [
+                    {"direction": 2},
+                    {"direction": 1},
+                    {"direction": 0},
+                ],
+                'modes': ['sum', "min"]
+            },
+        },
+        'neurite_type': ['AXON', 'APICAL_DENDRITE', 'BASAL_DENDRITE', 'ALL'],
+        'morphology': {
+            'soma_radius': {'modes': ['mean']},
+            'partition_asymmetry': {
+                'kwargs': [
+                    {'variant': 'branch-order', 'method': 'petilla'},
+                    {'variant': 'length', 'method': 'uylings'},
+                ],
+                'modes': ['min', 'max'],
+            },
+        }
+    }
+
+    res = ms.extract_stats(m, config)
+
+    assert set(res.keys()) == {"axon", "basal_dendrite", "apical_dendrite", "all", "morphology"}
+
+    for key in ("axon", "basal_dendrite", "apical_dendrite", "all"):
+
+        assert set(res[key].keys()) == {
+            "sum_principal_direction_extents__direction:2",
+            "min_principal_direction_extents__direction:2",
+            "sum_principal_direction_extents__direction:1",
+            "min_principal_direction_extents__direction:1",
+            "sum_principal_direction_extents__direction:0",
+            "min_principal_direction_extents__direction:0",
+        }
+
+    assert set(res["morphology"].keys()) == {
+        "mean_soma_radius",
+        "min_partition_asymmetry__variant:branch-order__method:petilla",
+        "max_partition_asymmetry__variant:branch-order__method:petilla",
+        "min_partition_asymmetry__variant:length__method:uylings",
+        "max_partition_asymmetry__variant:length__method:uylings",
+    }
+
+
+def test_extract_dataframe__kwarg_modes_multiple_features():
+    m = nm.load_morphology(SWC_PATH / 'Neuron.swc')
+    config = {
+        'neurite': {
+            'principal_direction_extents': {
+                'kwargs': [
+                    {"direction": 2},
+                    {"direction": 1},
+                    {"direction": 0},
+                ],
+                'modes': ['sum', "min"],
+            },
+        },
+        'neurite_type': ['AXON', 'APICAL_DENDRITE', 'BASAL_DENDRITE', 'ALL'],
+        'morphology': {
+            'soma_radius': {'modes': ['mean']},
+            'partition_asymmetry': {
+                'kwargs': [
+                    {'variant': 'branch-order', 'method': 'petilla'},
+                    {'variant': 'length', 'method': 'uylings'},
+                ],
+                'modes': ['min', 'max'],
+            },
+        },
+    }
+
+    res = ms.extract_dataframe(m, config)
+
+    expected_columns = pd.MultiIndex.from_tuples([
+        ('property', 'name'),
+        ('axon', 'sum_principal_direction_extents__direction:2'),
+        ('axon', 'min_principal_direction_extents__direction:2'),
+        ('axon', 'sum_principal_direction_extents__direction:1'),
+        ('axon', 'min_principal_direction_extents__direction:1'),
+        ('axon', 'sum_principal_direction_extents__direction:0'),
+        ('axon', 'min_principal_direction_extents__direction:0'),
+        ('apical_dendrite', 'sum_principal_direction_extents__direction:2'),
+        ('apical_dendrite', 'min_principal_direction_extents__direction:2'),
+        ('apical_dendrite', 'sum_principal_direction_extents__direction:1'),
+        ('apical_dendrite', 'min_principal_direction_extents__direction:1'),
+        ('apical_dendrite', 'sum_principal_direction_extents__direction:0'),
+        ('apical_dendrite', 'min_principal_direction_extents__direction:0'),
+        ('basal_dendrite', 'sum_principal_direction_extents__direction:2'),
+        ('basal_dendrite', 'min_principal_direction_extents__direction:2'),
+        ('basal_dendrite', 'sum_principal_direction_extents__direction:1'),
+        ('basal_dendrite', 'min_principal_direction_extents__direction:1'),
+        ('basal_dendrite', 'sum_principal_direction_extents__direction:0'),
+        ('basal_dendrite', 'min_principal_direction_extents__direction:0'),
+        ('all', 'sum_principal_direction_extents__direction:2'),
+        ('all', 'min_principal_direction_extents__direction:2'),
+        ('all', 'sum_principal_direction_extents__direction:1'),
+        ('all', 'min_principal_direction_extents__direction:1'),
+        ('all', 'sum_principal_direction_extents__direction:0'),
+        ('all', 'min_principal_direction_extents__direction:0'),
+        ('morphology', 'mean_soma_radius'),
+        ('morphology', 'min_partition_asymmetry__variant:branch-order__method:petilla'),
+        ('morphology', 'max_partition_asymmetry__variant:branch-order__method:petilla'),
+        ('morphology', 'min_partition_asymmetry__variant:length__method:uylings'),
+        ('morphology', 'max_partition_asymmetry__variant:length__method:uylings'),
+    ])
+
+    pd.testing.assert_index_equal(res.columns, expected_columns)
+
+
 def test_extract_dataframe():
     # Vanilla test
     initial_config = deepcopy(REF_CONFIG_NEW)
@@ -279,6 +398,8 @@ def test_extract_dataframe_with_kwargs():
     assert_frame_equal(actual, expected, check_dtype=False)
 
 
+
+
 def test_extract_dataframe_multiproc():
     morphs = [Path(SWC_PATH, name)
             for name in ['Neuron.swc', 'simple.swc']]
@@ -303,23 +424,152 @@ def test_get_header():
                     'fake_name1': REF_OUT,
                     'fake_name2': REF_OUT,
                     }
-    header = ms.get_header(fake_results)
+    header = ms._get_header(fake_results)
+
     assert 1 + 2 + 4 * (4 + 5) == len(header)  # name + everything in REF_OUT
     assert 'name' in header
     assert 'morphology:mean_soma_radius' in header
 
 
+def test_get_header__with_kwargs():
+
+    fake_results = {
+        "fake_name0": {
+            'axon': {
+                'sum_principal_direction_extents__direction:2': 4.236138323156951,
+                'min_principal_direction_extents__direction:2': 4.236138323156951,
+                'sum_principal_direction_extents__direction:1': 8.070668782620396,
+                'max_principal_direction_extents__direction:1': 8.070668782620396,
+                'mean_principal_direction_extents__direction:0': 82.38543140446015
+            },
+            'apical_dendrite': {
+                'sum_principal_direction_extents__direction:2': 3.6493184467335213,
+                'min_principal_direction_extents__direction:2': 3.6493184467335213,
+                'sum_principal_direction_extents__direction:1': 5.5082642304864695,
+                'max_principal_direction_extents__direction:1': 5.5082642304864695,
+                'mean_principal_direction_extents__direction:0': 99.57940514500457
+            },
+            'basal_dendrite': {
+                'sum_principal_direction_extents__direction:2': 7.32638745131256,
+                'min_principal_direction_extents__direction:2': 3.10141343122575,
+                'sum_principal_direction_extents__direction:1': 11.685447149154676,
+                'max_principal_direction_extents__direction:1': 6.410958014733595,
+                'mean_principal_direction_extents__direction:0': 87.2112016874677
+            },
+            'all': {
+                'sum_principal_direction_extents__direction:2': 15.211844221203034,
+                'min_principal_direction_extents__direction:2': 3.10141343122575,
+                'sum_principal_direction_extents__direction:1': 25.26438016226154,
+                'max_principal_direction_extents__direction:1': 8.070668782620396,
+                'mean_principal_direction_extents__direction:0': 89.09680998110002
+            },
+            'morphology': {
+                'mean_soma_radius': 0.13065629977308288,
+                'min_partition_asymmetry__variant:branch-order__method:petilla': 0.0,
+                'max_partition_asymmetry__variant:branch-order__method:petilla': 0.9,
+                'min_partition_asymmetry__variant:length__method:uylings': 0.00030289197373727377,
+                'max_partition_asymmetry__variant:length__method:uylings': 0.8795344229855895}
+            }
+    }
+
+    assert ms._get_header(fake_results) == [
+        'name',
+        'axon:sum_principal_direction_extents__direction:2',
+        'axon:min_principal_direction_extents__direction:2',
+        'axon:sum_principal_direction_extents__direction:1',
+        'axon:max_principal_direction_extents__direction:1',
+        'axon:mean_principal_direction_extents__direction:0',
+        'apical_dendrite:sum_principal_direction_extents__direction:2',
+        'apical_dendrite:min_principal_direction_extents__direction:2',
+        'apical_dendrite:sum_principal_direction_extents__direction:1',
+        'apical_dendrite:max_principal_direction_extents__direction:1',
+        'apical_dendrite:mean_principal_direction_extents__direction:0',
+        'basal_dendrite:sum_principal_direction_extents__direction:2',
+        'basal_dendrite:min_principal_direction_extents__direction:2',
+        'basal_dendrite:sum_principal_direction_extents__direction:1',
+        'basal_dendrite:max_principal_direction_extents__direction:1',
+        'basal_dendrite:mean_principal_direction_extents__direction:0',
+        'all:sum_principal_direction_extents__direction:2',
+        'all:min_principal_direction_extents__direction:2',
+        'all:sum_principal_direction_extents__direction:1',
+        'all:max_principal_direction_extents__direction:1',
+        'all:mean_principal_direction_extents__direction:0',
+        'morphology:mean_soma_radius',
+        'morphology:min_partition_asymmetry__variant:branch-order__method:petilla',
+        'morphology:max_partition_asymmetry__variant:branch-order__method:petilla',
+        'morphology:min_partition_asymmetry__variant:length__method:uylings',
+        'morphology:max_partition_asymmetry__variant:length__method:uylings'
+    ]
+
+
 def test_generate_flattened_dict():
     fake_results = {'fake_name0': REF_OUT,
                     'fake_name1': REF_OUT,
                     'fake_name2': REF_OUT,
                     }
-    header = ms.get_header(fake_results)
-    rows = list(ms.generate_flattened_dict(header, fake_results))
+    header = ms._get_header(fake_results)
+    rows = list(ms._generate_flattened_dict(header, fake_results))
     assert 3 == len(rows)  # one for fake_name[0-2]
     assert 1 + 2 + 4 * (4 + 5) == len(rows[0])  # name + everything in REF_OUT
 
 
+def test_generate_flattened_dict__with_kwargs():
+
+    results = {
+        'axon': {
+            'sum_principal_direction_extents__direction:2': 0.0,
+            'min_principal_direction_extents__direction:2': 1.0,
+            'sum_principal_direction_extents__direction:1': 2.0,
+            'max_principal_direction_extents__direction:1': 3.0,
+            'mean_principal_direction_extents__direction:0': 4.0,
+        },
+        'apical_dendrite': {
+            'sum_principal_direction_extents__direction:2': 5.0,
+            'min_principal_direction_extents__direction:2': 6.0,
+            'sum_principal_direction_extents__direction:1': 7.0,
+            'max_principal_direction_extents__direction:1': 8.0,
+            'mean_principal_direction_extents__direction:0': 9.0,
+        },
+        'basal_dendrite': {
+            'sum_principal_direction_extents__direction:2': 1.0,
+            'min_principal_direction_extents__direction:2': 2.0,
+            'sum_principal_direction_extents__direction:1': 3.0,
+            'max_principal_direction_extents__direction:1': 4.0,
+            'mean_principal_direction_extents__direction:0': 5.0,
+        },
+        'all': {
+            'sum_principal_direction_extents__direction:2': 6.0,
+            'min_principal_direction_extents__direction:2': 7.0,
+            'sum_principal_direction_extents__direction:1': 8.0,
+            'max_principal_direction_extents__direction:1': 9.0,
+            'mean_principal_direction_extents__direction:0': 1.0,
+        },
+        'morphology': {
+            'mean_soma_radius': 2.0,
+            'min_partition_asymmetry__variant:branch-order__method:petilla': 3.0,
+            'max_partition_asymmetry__variant:branch-order__method:petilla': 4.0,
+            'min_partition_asymmetry__variant:length__method:uylings': 5.0,
+            'max_partition_asymmetry__variant:length__method:uylings': 6.0,
+        }
+    }
+
+    fake_results = {
+        "fake_name0": results,
+        "fake_name1": results,
+    }
+
+    header = ms._get_header(fake_results)
+
+    assert list(ms._generate_flattened_dict(header, fake_results)) == [
+        [
+            'fake_name0', 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, 3.0, 4.0,
+            5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0],
+        [
+            'fake_name1', 0.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, 3.0, 4.0,
+            5.0, 6.0, 7.0, 8.0, 9.0, 1.0, 2.0, 3.0, 4.0, 5.0, 6.0]
+    ]
+
+
 def test_full_config():
     config = ms.full_config()
     assert set(config.keys()) == {'neurite', 'population', 'morphology', 'neurite_type'}
@@ -329,13 +579,51 @@ def test_full_config():
     assert set(config['population'].keys()) == set(_POPULATION_FEATURES.keys())
 
 
-def test_sanitize_config():
+def test_standardize_layout():
+    """Converts the config category entries (e.g. neurite, morphology, population) to using
+    the kwarg and modes layout.
+    """
+    # from short format
+    entry = {"f1": ["min", "max"], "f2": ["min"], "f3": []}
+    assert ms._standardize_layout(entry) == {
+        "f1": {"kwargs": [{}], "modes": ["min", "max"]},
+        "f2": {"kwargs": [{}], "modes": ["min"]},
+        "f3": {"kwargs": [{}], "modes": []},
+    }
+
+    # from kwarg/modes with missing options
+    entry = {
+        "f1": {"kwargs": {"a1": 1, "a2": 2}, "modes": ["min", "max"]},
+        "f2": {"modes": ["min", "median"]},
+        "f3": {"kwargs": {"a1": 1, "a2": 2}},
+        "f4": {},
+    }
+    assert ms._standardize_layout(entry) == {
+        "f1": {"kwargs": [{"a1": 1, "a2": 2}], "modes": ["min", "max"]},
+        "f2": {"kwargs": [{}], "modes": ["min", "median"]},
+        "f3": {"kwargs": [{"a1": 1, "a2": 2}], "modes": []},
+        "f4": {"kwargs": [{}], "modes": []},
+    }
+
+    # from list of kwargs format
+    entry = {
+        "f1": {"kwargs": [{"a1": 1, "a2": 2}], "modes": ["min", "max"]},
+        "f2": {"modes": ["min", "median"]},
+        "f3": {"kwargs": [{"a1": 1, "a2": 2}]},
+        "f4": {},
+    }
+    assert ms._standardize_layout(entry) == {
+        "f1": {"kwargs": [{"a1": 1, "a2": 2}], "modes": ["min", "max"]},
+        "f2": {"kwargs": [{}], "modes": ["min", "median"]},
+        "f3": {"kwargs": [{"a1": 1, "a2": 2}], "modes": []},
+        "f4": {"kwargs": [{}], "modes": []},
+    }
+
 
-    with pytest.raises(ConfigError):
-        ms.sanitize_config({'neurite': []})
+def test_sanitize_config():
 
-    new_config = ms.sanitize_config({})  # empty
-    assert 2 == len(new_config)  # neurite & morphology created
+    new_config = ms._sanitize_config({})  # empty
+    assert 3 == len(new_config)  # neurite & morphology & population created
 
     full_config = {
         'neurite': {
@@ -348,8 +636,28 @@ def test_sanitize_config():
             'soma_radius': ['mean']
         }
     }
-    new_config = ms.sanitize_config(full_config)
-    assert 3 == len(new_config)  # neurite, neurite_type & morphology
+    new_config = ms._sanitize_config(full_config)
+
+    expected_config = {
+        'neurite': {
+            'section_lengths': {"kwargs": [{}], "modes": ['max', 'sum']},
+            'section_volumes': {"kwargs": [{}], "modes": ['sum']},
+            'section_branch_orders': {"kwargs": [{}], "modes": ['max']},
+        },
+        'neurite_type': ['AXON', 'APICAL_DENDRITE', 'BASAL_DENDRITE', 'ALL'],
+        'morphology': {
+            'soma_radius': {"kwargs": [{}], "modes": ["mean"]},
+        },
+        "population": {},
+    }
+    assert new_config == expected_config
+
+    # check that legacy neuron entries are converted to morphology ones
+    full_config["neuron"] = full_config.pop("morphology")
+    assert ms._sanitize_config(full_config) == expected_config
+
+    # check that all formats are converted to the same sanitized config:
+    assert ms._sanitize_config(REF_CONFIG) == ms._sanitize_config(REF_CONFIG_NEW)
 
 
 def test_multidimensional_features():