Can pass List[Population] args to neurom.app.morph_stats.extract_data…

…frame (#1076)

neurom/apps/morph_stats.py

@@ -47,6 +47,7 @@
 import neurom as nm
 from neurom.apps import get_config
 from neurom.core.morphology import Morphology, Neurite
+from neurom.core.population import Population
 from neurom.exceptions import ConfigError
 from neurom.features import _NEURITE_FEATURES, _MORPHOLOGY_FEATURES, _POPULATION_FEATURES, \
     _get_feature_value_and_func
@@ -61,16 +62,17 @@
 
 def _run_extract_stats(morph, config):
     """The function to be called by multiprocessing.Pool.imap_unordered."""
-    if not isinstance(morph, Morphology):
-        morph = nm.load_morphology(morph)
+    if not isinstance(morph, (Morphology, Population)):
+        morph = nm.load_morphologies(morph)
     return morph.name, extract_stats(morph, config)
 
 
 def extract_dataframe(morphs, config, n_workers=1):
     """Extract stats grouped by neurite type from morphs.
 
     Arguments:
-        morphs: a morphology, population, neurite tree or list of morphology paths
+        morphs: a morphology, population, neurite tree, list of populations or list of morphology
+            paths
         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

tests/apps/test_morph_stats.py

@@ -34,6 +34,7 @@
 import neurom as nm
 import pandas as pd
 from neurom.apps import morph_stats as ms
+from neurom.core.population import Population
 from neurom.exceptions import ConfigError
 from neurom.features import _NEURITE_FEATURES, _MORPHOLOGY_FEATURES, _POPULATION_FEATURES
 
@@ -337,6 +338,30 @@ def test_extract_dataframe():
     assert_frame_equal(actual, expected.iloc[[0]], check_dtype=False)
     assert REF_CONFIG_NEW == initial_config
 
+    # Test with a List[Population] argument
+    pop1 = nm.load_morphologies([SWC_PATH / 'Neuron.swc', SWC_PATH / 'simple.swc'], name="Pop1")
+    pop2 = Population(pop1)
+    pop2.name = "Pop2"
+    actual = ms.extract_dataframe([pop1, pop2], REF_CONFIG_NEW)
+    actual = actual.drop(columns='raw_section_branch_orders', level=1)
+    aggregated_expected = pd.concat(
+        [
+            expected[[col for col in expected.columns if col[1].startswith("mean_")]].mean().to_frame().T,
+            expected[[col for col in expected.columns if col[1].startswith("max_")]].max().to_frame().T,
+            expected[[col for col in expected.columns if col[1].startswith("min_")]].min().to_frame().T,
+            expected[[col for col in expected.columns if col[1].startswith("sum_")]].sum().to_frame().T,
+        ],
+        axis=1,
+    )
+    aggregated_expected.loc[1] = aggregated_expected.loc[0]
+    aggregated_expected[("property", "name")] = ["Pop1", "Pop2"]
+    assert (actual.columns.sort_values() == aggregated_expected.columns.sort_values()).all()
+
+    actual = actual[actual.columns.sort_values()]
+    aggregated_expected = aggregated_expected[aggregated_expected.columns.sort_values()]
+    assert_frame_equal(actual, aggregated_expected, check_dtype=False)
+    assert REF_CONFIG_NEW == initial_config
+
     # Test with a config without the 'morphology' key
     morphs = nm.load_morphologies([Path(SWC_PATH, name)
                                  for name in ['Neuron.swc', 'simple.swc']])