dtifit.py

#!/usr/bin/env python
#
# dtifit.py - The DTIFitTensor class, and some related utility functions.
#
# Author: Paul McCarthy <pauldmccarthy@gmail.com>
#
"""This module provides the :class:`.DTIFitTensor` class, which encapsulates
the diffusion tensor data generated by the FSL ``dtifit`` tool.

The following utility functions are also defined:

  .. autosummary::
     :nosignatures:

     getDTIFitDataPrefix
     isDTIFitPath
     looksLikeTensorImage
     decomposeTensorMatrix
"""


import                 logging
import                 re
import                 glob
import os.path      as op

import numpy        as np
import numpy.linalg as npla

from . import image as fslimage


log = logging.getLogger(__name__)


def getDTIFitDataPrefix(path):
    """Returns the prefix (a.k,a, base name) used for the ``dtifit`` file
    names in the given directory, or ``None`` if the ``dtifit`` files could
    not be identified.
    """

    v1s   = glob.glob(op.join(path, '*_V1.*'))
    v2s   = glob.glob(op.join(path, '*_V2.*'))
    v3s   = glob.glob(op.join(path, '*_V3.*'))
    l1s   = glob.glob(op.join(path, '*_L1.*'))
    l2s   = glob.glob(op.join(path, '*_L2.*'))
    l3s   = glob.glob(op.join(path, '*_L3.*'))
    files = [v1s, v2s, v3s, l1s, l2s, l3s]

    # Gather all of the existing file
    # prefixes into a dictionary of
    # prefix : [file list] mappings.
    pattern  = '^(.*)_(?:V1|V2|V3|L1|L2|L3).*$'
    prefixes = {}

    for f in [f for flist in files for f in flist]:
        prefix = re.findall(pattern, f)[0]

        if prefix not in prefixes: prefixes[prefix] = [f]
        else:                      prefixes[prefix].append(f)

    # Discard any prefixes which are
    # not present for every file type.
    for prefix, files in list(prefixes.items()):
        if len(files) != 6:
            prefixes.pop(prefix)

    # Discard any prefixes which
    # match any files that do
    # not look like image files
    for prefix, files in list(prefixes.items()):
        if not all([fslimage.looksLikeImage(f) for f in files]):
            prefixes.pop(prefix)

    prefixes = list(prefixes.keys())

    # No more prefixes remaining -
    # this is probably not a dtifit
    # directory
    if len(prefixes) == 0:
        return None

    # If there's more than one remaining
    # prefix, I don't know what to do -
    # just return the first one.
    if len(prefixes) > 1:
        log.warning('Multiple dtifit prefixes detected: {}'.format(prefixes))

    return op.basename(sorted(prefixes)[0])


def isDTIFitPath(path):
    """Returns ``True`` if the given directory path looks like it contains
    ``dtifit`` data, ``False`` otherwise.
    """

    return getDTIFitDataPrefix(path) is not None


def looksLikeTensorImage(image):
    """Returns ``True`` if the given :class:`.Image` looks like it could
    contain tensor matrix data, ``False`` otherwise.
    """

    return len(image.shape) == 4 and image.shape[3] == 6


def decomposeTensorMatrix(data):
    """Decomposes the given ``numpy`` array into six separate arrays,
    containing the eigenvectors and eigenvalues of the tensor matrix
    decompositions.

    :arg image: A 4D ``numpy`` array with 6 volumes, which contains
                the unique elements of diffusion tensor matrices at
                every voxel.

    :returns:   A tuple containing the principal eigenvectors and
                eigenvalues of the tensor matrix.
    """

    # The image contains 6 volumes, corresponding
    # to the Dxx, Dxy, Dxz, Dyy, Dyz, Dzz elements
    # of the tensor matrix, at each voxel.
    #
    # We need to re-organise this into a series of
    # complete 3x3 tensor matrices, one for each
    # voxel.

    shape    = data.shape[:3]
    nvoxels  = np.prod(shape)
    matrices = np.zeros((nvoxels, 3, 3), dtype=np.float32)

    # Copy the tensor matrix elements
    # into their respective locations
    matrices[:, 0, 0] = data[..., 0].flat
    matrices[:, 0, 1] = data[..., 1].flat
    matrices[:, 1, 0] = data[..., 1].flat
    matrices[:, 0, 2] = data[..., 2].flat
    matrices[:, 2, 0] = data[..., 2].flat
    matrices[:, 1, 1] = data[..., 3].flat
    matrices[:, 1, 2] = data[..., 4].flat
    matrices[:, 2, 1] = data[..., 4].flat
    matrices[:, 2, 2] = data[..., 5].flat

    # Calculate the eigenvectors and
    # values on all of those matrices
    vals, vecs = npla.eig(matrices)
    vecShape   = list(shape) + [3]

    # Grr, np.linalg.eig does not
    # sort the eigenvalues/vectors,
    # so we have to do it ourselves.
    order = vals.argsort(axis=1)
    i     = np.arange(nvoxels)[:, np.newaxis]
    vecs  = vecs.transpose(0, 2, 1)
    vals  = vals[i, order]
    vecs  = vecs[i, order, :]

    l1 = vals[:, 2]   .reshape(shape)
    l2 = vals[:, 1]   .reshape(shape)
    l3 = vals[:, 0]   .reshape(shape)
    v1 = vecs[:, 2, :].reshape(vecShape)
    v2 = vecs[:, 1, :].reshape(vecShape)
    v3 = vecs[:, 0, :].reshape(vecShape)

    return v1, v2, v3, l1, l2, l3


class DTIFitTensor(fslimage.Nifti):
    """The ``DTIFitTensor`` class is able to load and encapsulate the diffusion
    tensor data generated by the FSL ``dtifit`` tool.  The ``DtiFitTensor``
    class supports tensor model data generated by ``dtifit``, where the
    eigenvectors and eigenvalues of the tensor matrices have been saved as six
    separate NIFTI images.
    """


    def __init__(self, path):
        """Create a ``DTIFitTensor``.

        :arg path: A path to a ``dtifit`` directory.
        """

        prefix      = getDTIFitDataPrefix(path)
        isDTIfitDir = prefix is not None

        if not isDTIfitDir:
            raise ValueError('{} does not look like a dtifit '
                             'output directory!'.format(path))

        # DTIFit output directory with separate
        # eigenvector/eigenvalue images

        self.__v1 = fslimage.Image(op.join(path, '{}_V1'.format(prefix)))
        self.__v2 = fslimage.Image(op.join(path, '{}_V2'.format(prefix)))
        self.__v3 = fslimage.Image(op.join(path, '{}_V3'.format(prefix)))
        self.__l1 = fslimage.Image(op.join(path, '{}_L1'.format(prefix)))
        self.__l2 = fslimage.Image(op.join(path, '{}_L2'.format(prefix)))
        self.__l3 = fslimage.Image(op.join(path, '{}_L3'.format(prefix)))

        fslimage.Nifti.__init__(self, self.__l1.header)

        self.dataSource = op.abspath(path)
        self.name       = '{}'.format(op.basename(path))

    def V1(self):
        return self.__v1
    def V2(self):
        return self.__v2
    def V3(self):
        return self.__v3
    def L1(self):
        return self.__l1
    def L2(self):
        return self.__l2
    def L3(self):
        return self.__l3