Streamlines API

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
 
 # be friendly on systems with ancient numpy -- no tests, but at least
 # importable

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,131 @@
+import os
+from ..externals.six import string_types
+
+from .header import TractogramHeader
+from .compact_list import CompactList
+from .tractogram import Tractogram, LazyTractogram
+
+from nibabel.streamlines.trk import TrkFile
+#from nibabel.streamlines.tck import TckFile
+#from nibabel.streamlines.vtk import VtkFile
+
+# List of all supported formats
+FORMATS = {".trk": TrkFile,
+           #".tck": TckFile,
+           #".vtk": VtkFile,
+           }
+
+
+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 : ``TractogramFile`` class
+        Returns an instance of a `TractogramFile` class containing data and
+        metadata of the tractogram contained from `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, ref=None):
+    ''' Loads streamlines from a file-like object in voxel space.
+
+    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 : boolean (optional)
+        Load streamlines in a lazy manner i.e. they will not be kept
+        in memory.
+
+    ref : filename | `Nifti1Image` object | 2D array (4,4) (optional)
+        Reference space where streamlines will live in `fileobj`.
+
+    Returns
+    -------
+    tractogram_file : ``TractogramFile``
+        Returns an instance of a `TractogramFile` class containing data and
+        metadata of the tractogram loaded from `fileobj`.
+    '''
+    tractogram_file = detect_format(fileobj)
+
+    if tractogram_file is None:
+        raise TypeError("Unknown format for 'fileobj': {}".format(fileobj))
+
+    return tractogram_file.load(fileobj, lazy_load=lazy_load)
+
+
+def save(tractogram_file, filename):
+    ''' Saves a tractogram to a file.
+
+    Parameters
+    ----------
+    tractogram_file : ``TractogramFile`` object
+        Tractogram to be saved on disk.
+
+    filename : str
+        Name of the file where the tractogram will be saved. The format will
+        be guessed from `filename`.
+    '''
+    tractogram_file.save(filename)
+
+
+def save_tractogram(tractogram, filename, **kwargs):
+    ''' Saves a tractogram to a file.
+
+    Parameters
+    ----------
+    tractogram : ``Tractogram`` object
+        Tractogram to be saved.
+
+    filename : str
+        Name of the file where the tractogram will be saved. The format will
+        be guessed from `filename`.
+    '''
+    tractogram_file_class = detect_format(filename)
+
+    if tractogram_file_class is None:
+        raise TypeError("Unknown tractogram file format: '{}'".format(filename))
+
+    tractogram_file = tractogram_file_class(tractogram, **kwargs)
+    tractogram_file.save(filename)

nibabel/streamlines/compact_list.py

@@ -0,0 +1,221 @@
+import numpy as np
+
+
+class CompactList(object):
+    """ Class for compacting list of ndarrays with matching shape except for
+    the first dimension.
+    """
+
+    BUFFER_SIZE = 10000000  # About 128 Mb if item shape is 3.
+
+    def __init__(self, iterable=None):
+        """
+        Parameters
+        ----------
+        iterable : iterable (optional)
+            If specified, create a ``CompactList`` object initialized from
+            iterable's items. Otherwise, create an empty ``CompactList``.
+
+        Notes
+        -----
+        If `iterable` is a ``CompactList`` object, a view is returned and no
+        memory is allocated. For an actual copy use the `.copy()` method.
+        """
+        # Create new empty `CompactList` object.
+        self._data = None
+        self._offsets = []
+        self._lengths = []
+
+        if isinstance(iterable, CompactList):
+            # Create a view.
+            self._data = iterable._data
+            self._offsets = iterable._offsets
+            self._lengths = iterable._lengths
+
+        elif iterable is not None:
+            # Initialize the `CompactList` object from iterable's item.
+            offset = 0
+            for i, e in enumerate(iterable):
+                e = np.asarray(e)
+                if i == 0:
+                    new_shape = (CompactList.BUFFER_SIZE,) + e.shape[1:]
+                    self._data = np.empty(new_shape, dtype=e.dtype)
+
+                end = offset + len(e)
+                if end >= len(self._data):
+                    # Resize needed, adding `len(e)` new items plus some buffer.
+                    nb_points = len(self._data)
+                    nb_points += len(e) + CompactList.BUFFER_SIZE
+                    self._data.resize((nb_points,) + self.shape)
+
+                self._offsets.append(offset)
+                self._lengths.append(len(e))
+                self._data[offset:offset+len(e)] = e
+                offset += len(e)
+
+            # Clear unused memory.
+            if self._data is not None:
+                self._data.resize((offset,) + self.shape)
+
+    @property
+    def shape(self):
+        """ Returns the matching shape of the elements in this compact list. """
+        if self._data is None:
+            return None
+
+        return self._data.shape[1:]
+
+    def append(self, element):
+        """ Appends `element` to this compact list.
+
+        Parameters
+        ----------
+        element : ndarray
+            Element to append. The shape must match already inserted elements
+            shape except for the first dimension.
+
+        Notes
+        -----
+        If you need to add multiple elements you should consider
+        `CompactList.extend`.
+        """
+        if self._data is None:
+            self._data = np.asarray(element).copy()
+            self._offsets.append(0)
+            self._lengths.append(len(element))
+            return
+
+        if element.shape[1:] != self.shape:
+            raise ValueError("All dimensions, except the first one,"
+                             " must match exactly")
+
+        self._offsets.append(len(self._data))
+        self._lengths.append(len(element))
+        self._data = np.append(self._data, element, axis=0)
+
+    def extend(self, elements):
+        """ Appends all `elements` to this compact list.
+
+        Parameters
+        ----------
+        elements : list of ndarrays, ``CompactList`` object
+            Elements to append. The shape must match already inserted elements
+            shape except for the first dimension.
+
+        """
+        if self._data is None:
+            elem = np.asarray(elements[0])
+            self._data = np.zeros((0, elem.shape[1]), dtype=elem.dtype)
+
+        next_offset = self._data.shape[0]
+
+        if isinstance(elements, CompactList):
+            self._data.resize((self._data.shape[0]+sum(elements._lengths),
+                               self._data.shape[1]))
+
+            for offset, length in zip(elements._offsets, elements._lengths):
+                self._offsets.append(next_offset)
+                self._lengths.append(length)
+                self._data[next_offset:next_offset+length] = elements._data[offset:offset+length]
+                next_offset += length
+
+        else:
+            self._data = np.concatenate([self._data] + list(elements), axis=0)
+            lengths = list(map(len, elements))
+            self._lengths.extend(lengths)
+            self._offsets.extend(np.cumsum([next_offset] + lengths).tolist()[:-1])
+
+    def copy(self):
+        """ Creates a copy of this ``CompactList`` object. """
+        # We do not simply deepcopy this object since we might have a chance
+        # to use less memory. For example, if the compact list being copied
+        # is the result of a slicing operation on a compact list.
+        clist = CompactList()
+        total_lengths = np.sum(self._lengths)
+        clist._data = np.empty((total_lengths,) + self._data.shape[1:],
+                               dtype=self._data.dtype)
+
+        idx = 0
+        for offset, length in zip(self._offsets, self._lengths):
+            clist._offsets.append(idx)
+            clist._lengths.append(length)
+            clist._data[idx:idx+length] = self._data[offset:offset+length]
+            idx += length
+
+        return clist
+
+    def __getitem__(self, idx):
+        """ Gets element(s) through indexing.
+
+        Parameters
+        ----------
+        idx : int, slice or list
+            Index of the element(s) to get.
+
+        Returns
+        -------
+        ndarray object(s)
+            When `idx` is an int, returns a single ndarray.
+            When `idx` is either a slice or a list, returns a list of ndarrays.
+        """
+        if isinstance(idx, int) or isinstance(idx, np.integer):
+            start = self._offsets[idx]
+            return self._data[start:start+self._lengths[idx]]
+
+        elif isinstance(idx, slice):
+            clist = CompactList()
+            clist._data = self._data
+            clist._offsets = self._offsets[idx]
+            clist._lengths = self._lengths[idx]
+            return clist
+
+        elif isinstance(idx, list):
+            clist = CompactList()
+            clist._data = self._data
+            clist._offsets = [self._offsets[i] for i in idx]
+            clist._lengths = [self._lengths[i] for i in idx]
+            return clist
+
+        elif isinstance(idx, np.ndarray) and idx.dtype == np.bool:
+            clist = CompactList()
+            clist._data = self._data
+            clist._offsets = [self._offsets[i]
+                              for i, take_it in enumerate(idx) if take_it]
+            clist._lengths = [self._lengths[i]
+                              for i, take_it in enumerate(idx) if take_it]
+            return clist
+
+        raise TypeError("Index must be either an int, a slice, a list of int"
+                        " or a ndarray of bool! Not " + str(type(idx)))
+
+    def __iter__(self):
+        if len(self._lengths) != len(self._offsets):
+            raise ValueError("CompactList object corrupted:"
+                             " len(self._lengths) != len(self._offsets)")
+
+        for offset, lengths in zip(self._offsets, self._lengths):
+            yield self._data[offset: offset+lengths]
+
+    def __len__(self):
+        return len(self._offsets)
+
+    def __repr__(self):
+        return repr(list(self))
+
+
+def save_compact_list(filename, clist):
+    """ Saves a `CompactList` object to a .npz file. """
+    np.savez(filename,
+             data=clist._data,
+             offsets=clist._offsets,
+             lengths=clist._lengths)
+
+
+def load_compact_list(filename):
+    """ Loads a `CompactList` object from a .npz file. """
+    content = np.load(filename)
+    clist = CompactList()
+    clist._data = content["data"]
+    clist._offsets = content["offsets"].tolist()
+    clist._lengths = content["lengths"].tolist()
+    return clist