Streamlines API

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)