Merge pull request #391 from MarcCote/streamlines_api

MRG: Streamlines API

Provides a general API to manage different streamlines file format (the code can all be found in `nibabel.streamlines.`). This is related to this somewhat outdated [BIAP5](https://github.com/nipy/nibabel/wiki/BIAP5).

This new streamline API is divided in two main parts:

1. Representation of streamlines in memory (`Tractogram` and `ArraySequence` classes)
2. Loading and saving streamlines from/to different file format (`TractogramFile` class)

*We use the term 'Tractogram' to refer to a set of streamlines, usually the output of a tractography algorithm.*

## ArraySequence

The `ArraySequence` class needs to be explained first as it is at the heart of the streamlines representation. I guess it could be seen as a sort of jagged/ragged array to store a list of 2D arrays that have different length but a same number of dimensions (e.g. a list of list of points). For efficiency, we store a list of streamlines contiguously in memory using what we call a `ArraySequence`. This was needed as @Garyfallidis and I realized that a list of numpy arrays has too much overhead in term of memory when handling >1M streamlines (loading a tractogram of 2.5Gb could take up to 5Gb in RAM!).

An `ArraySequence` behaves exactly like a `list` of `2D ndarray`. That is, it supports appending, extending, indexing (supports slicing too), iterating, etc. as a `list` object. For example, indexing a `ArraySequence` object will return a 2D `ndarray`, as one would expect, representing the 3D points of the selected streamline. Note that, while `ArraySequence` does not support deletion (i.e. no `del` nor `remove`), one can efficiently use slicing to achieve the same goal.

An `ArraySequence` object can be instantiated using any iterable yielding `2D array`:

```
import numpy as np
import nibabel as nib
data = [np.random.rand(10, 3), np.random.rand(3, 3)]
seq = nib.streamlines.ArraySequence(data)
```

Slicing a `ArraySequence` object returns a new `ArraySequence` object acting as a view (*no memory duplication*). Of course, one can use the `copy` method to "force" the memory duplication.

## Tractogram

This part handles streamlines in memory using the Python class `Tractogram`. This class provides a mechanism to keep data associated to points of every streamline (`data_per_point` property). It can also keep data associated to each streamline (`data_per_streamline` property). 

Regarding the reference space we used for the streamlines of a tractogram live in, we decided to follow the NiBabel convention: the streamlines are assumed to be in **RAS+**, points are expressed in **millimeter** and the coordinate of a voxel (eg. (0,0,0) or (1,2,3)) refers to the **center of that voxel**.

A `Tractogram` objects supports indexing, slicing and iterating. For example, indexing a `Tractogram` object will return a `TractogramItem` object, representing a single streamline with its associated data. Note that, while `Tractogram` does not support deletion (i.e. no `del` nor `remove`), one can efficiently use slicing to achieve the same goal. Again, slicing a `Tractogram` object returns a new `Tractogram` object acting as a view (*no memory duplication*).

Here a small example of instantiating a `Tractogram` object:

```
import numpy as np
import nibabel as nib
data = [np.random.rand(10, 3), np.random.rand(3, 3)]
tractogram = nib.streamlines.Tractogram(streamlines=data)
colors = [(1, 0, 0), (0, 1, 0)]
tractogram.data_per_streamlines['colors'] = colors
```

## Tractogram File

This part handles the saving and loading of `Tractogram` objects from/to the disk. Right now, as a proof of concept, only the TRK file format is supported (support for TCK is already in progress and will be added in a following PR).

The `TractogramFile` class is abstract. Each file format should have its own class inheriting `TractogramFile` (e.g., `TrkFile`).

## Realistic example

Here a small example of loading a TRK file, computing the length of each streamlines, adding that information to the tractogram and saving it back:

```
import nibabel as nib
trk = nib.streamlines.load('my_streamlines.trk')

from dipy.tracking.streamline import length
lengths = length(trk.tractogram.streamlines)

trk.tractogram.data_per_streamline['lengths'] = lengths

new_trk = nib.streamlines.TrkFile(trk.tractogram, trk.header)
nib.streamlines.save(new_trk, "my_streamlines.trk")
```

Author: matthew-brett

Changelog

@@ -36,6 +36,9 @@ References like "pr/298" refer to github pull request numbers.
     are raising a DataError if the track is truncated when ``strict=True``
     (the default), rather than a TypeError when trying to create the points
     array.
+  * New API for managing streamlines and their different file formats. This
+    adds a new module ``nibabel.streamlines`` that will eventually deprecate
+    the current trackvis reader found in ``nibabel.trackvis``.
 
 * 2.0.2 (Monday 23 November 2015)
 
@@ -251,7 +254,7 @@ References like "pr/298" refer to github pull request numbers.
     the ability to transform to the image with data closest to the cononical
     image orientation (first axis left-to-right, second back-to-front, third
     down-to-up) (MB, Jonathan Taylor)
-  * Gifti format read and write support (preliminary) (Stephen Gerhard) 
+  * Gifti format read and write support (preliminary) (Stephen Gerhard)
   * Added utilities to use nipy-style data packages, by rip then edit of nipy
     data package code (MB)
   * Some improvements to release support (Jarrod Millman, MB, Fernando Perez)
@@ -469,7 +472,7 @@ visiting the URL::
 
   * Removed functionality for "NiftiImage.save() raises an IOError
     exception when writing the image file fails." (Yaroslav Halchenko)
-  * Added ability to force a filetype when setting the filename or saving 
+  * Added ability to force a filetype when setting the filename or saving
     a file.
   * Reverse the order of the 'header' and 'load' argument in the NiftiImage
     constructor. 'header' is now first as it seems to be used more often.
@@ -481,7 +484,7 @@ visiting the URL::
 
 * 0.20070301.2 (Thu, 1 Mar 2007)
 
-  * Fixed wrong link to the source tarball in README.html. 
+  * Fixed wrong link to the source tarball in README.html.
 
 
 * 0.20070301.1 (Thu, 1 Mar 2007)

nibabel/__init__.py

@@ -64,6 +64,7 @@
 from .imageclasses import class_map, ext_map, all_image_classes
 from . import trackvis
 from . import mriutils
+from . import streamlines
 from . import viewers
 
 # be friendly on systems with ancient numpy -- no tests, but at least

nibabel/benchmarks/bench_streamlines.py

@@ -0,0 +1,95 @@
+""" Benchmarks for load and save of streamlines
+
+Run benchmarks with::
+
+    import nibabel as nib
+    nib.bench()
+
+If you have doctests enabled by default in nose (with a noserc file or
+environment variable), and you have a numpy version <= 1.6.1, this will also run
+the doctests, let's hope they pass.
+
+Run this benchmark with:
+
+    nosetests -s --match '(?:^|[\\b_\\.//-])[Bb]ench' /path/to/bench_streamlines.py
+"""
+from __future__ import division, print_function
+
+import numpy as np
+
+from nibabel.externals.six.moves import zip
+from nibabel.tmpdirs import InTemporaryDirectory
+
+from numpy.testing import assert_array_equal
+from nibabel.streamlines import Tractogram
+from nibabel.streamlines import TrkFile
+
+import nibabel as nib
+import nibabel.trackvis as tv
+
+from numpy.testing import measure
+
+
+def bench_load_trk():
+    rng = np.random.RandomState(42)
+    dtype = 'float32'
+    NB_STREAMLINES = 5000
+    NB_POINTS = 1000
+    points = [rng.rand(NB_POINTS, 3).astype(dtype)
+              for i in range(NB_STREAMLINES)]
+    scalars = [rng.rand(NB_POINTS, 10).astype(dtype)
+               for i in range(NB_STREAMLINES)]
+
+    repeat = 10
+
+    with InTemporaryDirectory():
+        trk_file = "tmp.trk"
+        tractogram = Tractogram(points, affine_to_rasmm=np.eye(4))
+        TrkFile(tractogram).save(trk_file)
+
+        streamlines_old = [d[0] - 0.5
+                           for d in tv.read(trk_file, points_space="rasmm")[0]]
+        mtime_old = measure('tv.read(trk_file, points_space="rasmm")', repeat)
+        print("Old: Loaded {:,} streamlines in {:6.2f}".format(NB_STREAMLINES,
+                                                               mtime_old))
+
+        trk = nib.streamlines.load(trk_file, lazy_load=False)
+        streamlines_new = trk.streamlines
+        mtime_new = measure('nib.streamlines.load(trk_file, lazy_load=False)',
+                            repeat)
+        print("\nNew: Loaded {:,} streamlines in {:6.2}".format(NB_STREAMLINES,
+                                                                mtime_new))
+        print("Speedup of {:.2f}".format(mtime_old / mtime_new))
+        for s1, s2 in zip(streamlines_new, streamlines_old):
+            assert_array_equal(s1, s2)
+
+    # Points and scalars
+    with InTemporaryDirectory():
+
+        trk_file = "tmp.trk"
+        tractogram = Tractogram(points,
+                                data_per_point={'scalars': scalars},
+                                affine_to_rasmm=np.eye(4))
+        TrkFile(tractogram).save(trk_file)
+
+        streamlines_old = [d[0] - 0.5
+                           for d in tv.read(trk_file, points_space="rasmm")[0]]
+
+        scalars_old = [d[1]
+                       for d in tv.read(trk_file, points_space="rasmm")[0]]
+        mtime_old = measure('tv.read(trk_file, points_space="rasmm")', repeat)
+        msg = "Old: Loaded {:,} streamlines with scalars in {:6.2f}"
+        print(msg.format(NB_STREAMLINES, mtime_old))
+
+        trk = nib.streamlines.load(trk_file, lazy_load=False)
+        scalars_new = trk.tractogram.data_per_point['scalars']
+        mtime_new = measure('nib.streamlines.load(trk_file, lazy_load=False)',
+                            repeat)
+        msg = "New: Loaded {:,} streamlines with scalars in {:6.2f}"
+        print(msg.format(NB_STREAMLINES, mtime_new))
+        print("Speedup of {:2f}".format(mtime_old / mtime_new))
+        for s1, s2 in zip(scalars_new, scalars_old):
+            assert_array_equal(s1, s2)
+
+if __name__ == '__main__':
+    bench_load_trk()

nibabel/externals/six.py

@@ -143,6 +143,7 @@ class _MovedItems(types.ModuleType):
     MovedAttribute("StringIO", "StringIO", "io"),
     MovedAttribute("xrange", "__builtin__", "builtins", "xrange", "range"),
     MovedAttribute("zip", "itertools", "builtins", "izip", "zip"),
+    MovedAttribute("zip_longest", "itertools", "itertools", "izip_longest", "zip_longest"),
 
     MovedModule("builtins", "__builtin__"),
     MovedModule("configparser", "ConfigParser"),

nibabel/streamlines/__init__.py

@@ -0,0 +1,133 @@
+import os
+import warnings
+from ..externals.six import string_types
+
+from .header import Field
+from .array_sequence import ArraySequence
+from .tractogram import Tractogram, LazyTractogram
+from .tractogram_file import ExtensionWarning
+
+from .trk import TrkFile
+
+# List of all supported formats
+FORMATS = {".trk": TrkFile}
+
+
+def is_supported(fileobj):
+    """ Checks if the file-like object if supported by NiBabel.
+
+    Parameters
+    ----------
+    fileobj : string or file-like object
+        If string, a filename; otherwise an open file-like object pointing
+        to a streamlines file (and ready to read from the beginning of the
+        header)
+
+    Returns
+    -------
+    is_supported : boolean
+    """
+    return detect_format(fileobj) is not None
+
+
+def detect_format(fileobj):
+    """ Returns the StreamlinesFile object guessed from the file-like object.
+
+    Parameters
+    ----------
+    fileobj : string or file-like object
+        If string, a filename; otherwise an open file-like object pointing
+        to a tractogram file (and ready to read from the beginning of the
+        header)
+
+    Returns
+    -------
+    tractogram_file : :class:`TractogramFile` class
+        The class type guessed from the content of `fileobj`.
+    """
+    for format in FORMATS.values():
+        try:
+            if format.is_correct_format(fileobj):
+                return format
+        except IOError:
+            pass
+
+    if isinstance(fileobj, string_types):
+        _, ext = os.path.splitext(fileobj)
+        return FORMATS.get(ext.lower())
+
+    return None
+
+
+def load(fileobj, lazy_load=False):
+    """ Loads streamlines in *RAS+* and *mm* space from a file-like object.
+
+    Parameters
+    ----------
+    fileobj : string or file-like object
+        If string, a filename; otherwise an open file-like object
+        pointing to a streamlines file (and ready to read from the beginning
+        of the streamlines file's header).
+    lazy_load : {False, True}, optional
+        If True, load streamlines in a lazy manner i.e. they will not be kept
+        in memory and only be loaded when needed.
+        Otherwise, load all streamlines in memory.
+
+    Returns
+    -------
+    tractogram_file : :class:`TractogramFile` object
+        Returns an instance of a :class:`TractogramFile` containing data and
+        metadata of the tractogram loaded from `fileobj`.
+
+    Notes
+    -----
+    The streamline coordinate (0,0,0) refers to the center of the voxel.
+    """
+    tractogram_file = detect_format(fileobj)
+
+    if tractogram_file is None:
+        raise ValueError("Unknown format for 'fileobj': {}".format(fileobj))
+
+    return tractogram_file.load(fileobj, lazy_load=lazy_load)
+
+
+def save(tractogram, filename, **kwargs):
+    """ Saves a tractogram to a file.
+
+    Parameters
+    ----------
+    tractogram : :class:`Tractogram` object or :class:`TractogramFile` object
+        If :class:`Tractogram` object, the file format will be guessed from
+        `filename` and a :class:`TractogramFile` object will be created using
+        provided keyword arguments.
+        If :class:`TractogramFile` object, the file format is known and will
+        be used to save its content to `filename`.
+    filename : str
+        Name of the file where the tractogram will be saved.
+    \*\*kwargs : keyword arguments
+        Keyword arguments passed to :class:`TractogramFile` constructor.
+        Should not be specified if `tractogram` is already an instance of
+        :class:`TractogramFile`.
+    """
+    tractogram_file_class = detect_format(filename)
+    if isinstance(tractogram, Tractogram):
+        if tractogram_file_class is None:
+            msg = "Unknown tractogram file format: '{}'".format(filename)
+            raise ValueError(msg)
+
+        tractogram_file = tractogram_file_class(tractogram, **kwargs)
+
+    else:  # Assume it's a TractogramFile object.
+        tractogram_file = tractogram
+        if (tractogram_file_class is None or
+                not isinstance(tractogram_file, tractogram_file_class)):
+            msg = ("The extension you specified is unusual for the provided"
+                   " 'TractogramFile' object.")
+            warnings.warn(msg, ExtensionWarning)
+
+        if len(kwargs) > 0:
+            msg = ("A 'TractogramFile' object was provided, no need for"
+                   " keyword arguments.")
+            raise ValueError(msg)
+
+    tractogram_file.save(filename)