Don't force loading of neurons into memory (#922)

Don't force loading of neurons into memory for `Population` and `load_neurons`

CHANGELOG.rst

@@ -3,12 +3,15 @@ Changelog
 
 Version 2.2.0
 -------------
+- Don't force loading of neurons into memory for Population (#922). See new API of
+  :class:`Population<neurom.core.population.Population>` and
+  :func:`load_neurons<neurom.io.utils.load_neurons>`
 - Move ``total_length`` feature to from ``neuritefunc`` to ``neuronfunc``. Use ``neurite_lengths``
   feature for neurites
 - Include morphology filename extension into Neuron's name
 - Extend ``tree_type_checker`` to accept a single tuple as an argument. Additionally validate
-  function's arguments (#909, #912)
-- Optimize Sholl analysis code (#905)
+  function's arguments (#912, #914)
+- Optimize Sholl analysis code (#905, #919)
 
 Version 2.1.2
 -------------

doc/source/api.rst

@@ -47,6 +47,7 @@ API Documentation
    neurom.check.neuron_checks
    neurom.core.types
    neurom.core.neuron
+   neurom.core.population
    neurom.core.soma
    neurom.core.dataformat
    neurom.io.utils

doc/source/conf.py

@@ -141,6 +141,7 @@
 suppress_warnings = ["ref.python"]
 autosummary_generate = True
 autosummary_imported_members = True
+autoclass_content = 'both'
 autodoc_default_options = {
     'members': True,
     'imported-members': True,

neurom/core/population.py

@@ -27,42 +27,90 @@
 # SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
 """Neuron Population Classes and Functions."""
+import logging
 
-from itertools import chain
+from morphio import MorphioError
+import neurom
+from neurom.exceptions import NeuroMError
+
+
+L = logging.getLogger(__name__)
 
 
 class Population:
     """Neuron Population Class.
 
-    Features:
-        - flattened collection of neurites.
-        - collection of somas, neurons.
-        - iterable-like iteration over neurons.
+    Offers an iterator over neurons within population, neurites of neurons, somas of neurons.
+    It does not store the loaded neuron in memory unless the neuron has been already passed
+    as loaded (instance of ``Neuron``).
     """
-    def __init__(self, neurons, name='Population'):
+    def __init__(self, files, name='Population', ignored_exceptions=(), cache=False):
         """Construct a neuron population.
 
         Arguments:
-            neurons: iterable of neuron objects.
-            name: Optional name for this Population.
+            files (collections.abc.Sequence[str|Path|Neuron]): collection of neuron files or
+                paths to them
+            name (str): Optional name for this Population
+            ignored_exceptions (tuple): NeuroM and MorphIO exceptions that you want to ignore when
+                loading neurons.
+            cache (bool): whether to cache the loaded neurons in memory. If false then a neuron
+                will be loaded everytime it is accessed within the population. Which is good when
+                population is big. If true then all neurons will be loaded upon the construction
+                and kept in memory.
         """
-        self.neurons = tuple(neurons)
-        self.somata = tuple(neu.soma for neu in neurons)
-        self.neurites = tuple(chain.from_iterable(neu.neurites for neu in neurons))
+        self._ignored_exceptions = ignored_exceptions
         self.name = name
+        if cache:
+            self._files = [self._load_file(f) for f in files if f is not None]
+        else:
+            self._files = files
+
+    @property
+    def neurons(self):
+        """Iterator to populations's somas."""
+        return (n for n in self)
+
+    @property
+    def somata(self):
+        """Iterator to populations's somas."""
+        return (n.soma for n in self)
+
+    @property
+    def neurites(self):
+        """Iterator to populations's neurites."""
+        return (neurite for n in self for neurite in n.neurites)
+
+    def _load_file(self, f):
+        if isinstance(f, neurom.core.neuron.Neuron):
+            return f
+        try:
+            return neurom.load_neuron(f)
+        except (NeuroMError, MorphioError) as e:
+            if isinstance(e, self._ignored_exceptions):
+                L.info('Ignoring exception "%s" for file %s', e, f.name)
+            else:
+                raise NeuroMError('`load_neurons` failed') from e
+        return None
 
     def __iter__(self):
         """Iterator to populations's neurons."""
-        return iter(self.neurons)
+        for f in self._files:
+            nrn = self._load_file(f)
+            if nrn is None:
+                continue
+            yield nrn
 
     def __len__(self):
         """Length of neuron collection."""
-        return len(self.neurons)
+        return len(self._files)
 
     def __getitem__(self, idx):
         """Get neuron at index idx."""
-        return self.neurons[idx]
+        if idx > len(self):
+            raise ValueError(
+                f'no {idx} index in "{self.name}" population, max possible index is {len(self)}')
+        return self._load_file(self._files[idx])
 
     def __str__(self):
         """Return a string representation."""
-        return 'Population <name: %s, nneurons: %d>' % (self.name, len(self.neurons))
+        return 'Population <name: %s, nneurons: %d>' % (self.name, len(self))

neurom/io/utils.py

@@ -38,7 +38,6 @@
 from pathlib import Path
 
 import morphio
-from morphio import MorphioError
 from neurom.core.neuron import Neuron
 from neurom.core.population import Population
 from neurom.exceptions import NeuroMError
@@ -168,47 +167,29 @@ def load_neuron(neuron, reader=None):
 
 
 def load_neurons(neurons,
-                 neuron_loader=load_neuron,
                  name=None,
-                 population_class=Population,
-                 ignored_exceptions=()):
+                 ignored_exceptions=(),
+                 cache=False):
     """Create a population object.
 
     From all morphologies in a directory of from morphologies in a list of file names.
 
     Arguments:
-        neurons: directory path or list of neuron file paths
-        neuron_loader: function taking a filename and returning a neuron
-        population_class: class representing populations
+        neurons(str|Path|Iterable[Path]): path to a folder or list of paths to neuron files
         name (str): optional name of population. By default 'Population' or\
             filepath basename depending on whether neurons is list or\
             directory path respectively.
         ignored_exceptions (tuple): NeuroM and MorphIO exceptions that you want to ignore when
-            loading neurons.
+            loading neurons
+        cache (bool): whether to cache the loaded neurons in memory
 
     Returns:
-        neuron population object
+        Population: population object
     """
-    if isinstance(neurons, str):
-        neurons = Path(neurons)
-
-    if isinstance(neurons, Path):
+    if isinstance(neurons, (str, Path)):
         files = get_files_by_path(neurons)
-        name = name or neurons.name
+        name = name or Path(neurons).name
     else:
         files = neurons
         name = name or 'Population'
-
-    ignored_exceptions = tuple(ignored_exceptions)
-    pop = []
-    for f in files:
-        try:
-            pop.append(neuron_loader(f))
-        except (NeuroMError, MorphioError) as e:
-            if isinstance(e, ignored_exceptions):
-                L.info('Ignoring exception "%s" for file %s',
-                       e, f.name)
-                continue
-            raise NeuroMError('`load_neurons` failed') from e
-
-    return population_class(pop, name=name)
+    return Population(files, name, ignored_exceptions, cache)

tests/core/test_population.py

@@ -29,50 +29,74 @@
 from pathlib import Path
 
 from neurom.core.population import Population
+from neurom.core.neuron import Neuron
 from neurom import load_neuron
 
+import pytest
+
 DATA_PATH = Path(__file__).parent.parent / 'data'
 
-NRN1 = load_neuron(DATA_PATH / 'swc/Neuron.swc')
-NRN2 = load_neuron(DATA_PATH / 'swc/Single_basal.swc')
-NRN3 = load_neuron(DATA_PATH / 'swc/Neuron_small_radius.swc')
+FILES = [DATA_PATH / 'swc/Neuron.swc',
+         DATA_PATH / 'swc/Single_basal.swc',
+         DATA_PATH / 'swc/Neuron_small_radius.swc']
 
-NEURONS = [NRN1, NRN2, NRN3]
+NEURONS = [load_neuron(f) for f in FILES]
 TOT_NEURITES = sum(len(N.neurites) for N in NEURONS)
-POP = Population(NEURONS, name='foo')
+populations = [Population(NEURONS, name='foo'),
+               Population(FILES, name='foo', cache=True)]
+
 
+@pytest.mark.parametrize('pop', populations)
+def test_names(pop):
+    assert pop[0].name, 'Neuron'
+    assert pop[1].name, 'Single_basal'
+    assert pop[2].name, 'Neuron_small_radius'
+    assert pop.name == 'foo'
 
-def test_population():
-    assert len(POP.neurons) == 3
-    assert POP.neurons[0].name, 'Neuron'
-    assert POP.neurons[1].name, 'Single_basal'
-    assert POP.neurons[2].name, 'Neuron_small_radius'
 
-    assert len(POP.somata) == 3
+def test_indexing():
+    pop = populations[0]
+    for i, n in enumerate(NEURONS):
+        assert n is pop[i]
+    with pytest.raises(ValueError, match='no 10 index'):
+        pop[10]
 
-    assert len(POP.neurites) == TOT_NEURITES
 
-    assert POP.name == 'foo'
+def test_cache():
+    pop = populations[1]
+    for n in pop._files:
+        assert isinstance(n, Neuron)
 
 
-def test_neurons():
+def test_double_indexing():
+    pop = populations[0]
     for i, n in enumerate(NEURONS):
-        assert n is POP.neurons[i]
+        assert n is pop[i]
+    # second time to assure that generator is available again
+    for i, n in enumerate(NEURONS):
+        assert n is pop[i]
 
 
-def test_iterate_neurons():
-    for a, b in zip(NEURONS, POP):
+def test_iterating():
+    pop = populations[0]
+    for a, b in zip(NEURONS, pop):
         assert a is b
 
+    for a, b in zip(NEURONS, pop.somata):
+        assert a.soma is b
+
 
-def test_len():
-    assert len(POP) == len(NEURONS)
+@pytest.mark.parametrize('pop', populations)
+def test_len(pop):
+    assert len(pop) == len(NEURONS)
 
 
 def test_getitem():
+    pop = populations[0]
     for i in range(len(NEURONS)):
-        assert POP[i] is NEURONS[i]
+        assert pop[i] is NEURONS[i]
 
 
-def test_str():
-    assert 'Population' in str(POP)
+@pytest.mark.parametrize('pop', populations)
+def test_str(pop):
+    assert 'Population' in str(pop)