From 3599701863358edecec1179dc9f9534c2b4c1e94 Mon Sep 17 00:00:00 2001
From: asanin-epfl <53935643+asanin-epfl@users.noreply.github.com>
Date: Mon, 7 Jun 2021 10:29:54 +0200
Subject: [PATCH] Fix 'section_path_lengths' feature for Population (#927)

---
 CHANGELOG.rst                       | 4 ++++
 neurom/features/neuritefunc.py      | 8 +-------
 tests/features/test_get_features.py | 6 ++++++
 3 files changed, 11 insertions(+), 7 deletions(-)

diff --git a/CHANGELOG.rst b/CHANGELOG.rst
index afe65b83..835b80c1 100644
--- a/CHANGELOG.rst
+++ b/CHANGELOG.rst
@@ -1,6 +1,10 @@
 Changelog
 =========
 
+Version 2.2.1
+-------------
+- Fix 'section_path_lengths' feature for Population
+
 Version 2.2.0
 -------------
 - Don't force loading of neurons into memory for Population (#922). See new API of
diff --git a/neurom/features/neuritefunc.py b/neurom/features/neuritefunc.py
index 8da8f580..27cf27e0 100644
--- a/neurom/features/neuritefunc.py
+++ b/neurom/features/neuritefunc.py
@@ -162,15 +162,9 @@ def section_path_lengths(neurites, neurite_type=NeuriteType.all):
     # then queries the lengths in the path length iterations.
     # This avoids repeatedly calculating the lengths of the
     # same sections.
-    dist = {}
-    neurite_filter = is_type(neurite_type)
-
-    for s in iter_sections(neurites, neurite_filter=neurite_filter):
-        dist[s] = s.length
-
     def pl2(node):
         """Calculate the path length using cached section lengths."""
-        return sum(dist[n] for n in node.iupstream())
+        return sum(n.length for n in node.iupstream())
 
     return _map_sections(pl2, neurites, neurite_type=neurite_type)
 
diff --git a/tests/features/test_get_features.py b/tests/features/test_get_features.py
index ab719f43..53dfde73 100644
--- a/tests/features/test_get_features.py
+++ b/tests/features/test_get_features.py
@@ -639,6 +639,12 @@ def test_section_path_distances_axon():
     assert len(path_lengths) == 21
 
 
+def test_section_path_distances_pop():
+    path_distances = get_feature('section_path_distances', POP)
+    assert len(path_distances) == 328
+    assert sum(len(get_feature('section_path_distances', nrn)) for nrn in POP) == 328
+
+
 def test_segment_lengths():
     ref_seglen = nf.segment_lengths(NEURON)
     seglen = get_feature('segment_lengths', NEURON)