Merge branch 'enh_convert_tensors' into 'master'

Convert diffusion tensors between 3 formats

See merge request fsl/fslpy!114

fsl/data/dtifit.py

@@ -7,7 +7,14 @@
 """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:
+There are also conversion tools between the diffusion tensors defined in 3 formats:
+
+* (..., 3, 3) array with the full diffusion tensor
+* (..., 6) array with the unique components (Dxx, Dxy, Dxz, Dyy, Dyz, Dzz)
+* Tuple with the eigenvectors and eigenvalues (V1, V2, V3, L1, L2, L3)
+
+
+Finally the following utility functions are also defined:
 
   .. autosummary::
      :nosignatures:
@@ -33,6 +40,123 @@ from . import image as fslimage
 log = logging.getLogger(__name__)
 
 
+def eigendecompositionToTensor(V1, V2, V3, L1, L2, L3):
+    """
+    Converts the eigenvalues/eigenvectors into a 3x3 diffusion tensor
+
+    :param V1: (..., 3) shaped array with the first eigenvector
+    :param V2: (..., 3) shaped array with the second eigenvector
+    :param V3: (..., 3) shaped array with the third eigenvector
+    :param L1: (..., ) shaped array with the first eigenvalue
+    :param L2: (..., ) shaped array with the second eigenvalue
+    :param L3: (..., ) shaped array with the third eigenvalue
+    :return: (..., 3, 3) array with the diffusion tensor
+    """
+    check_shape = L1.shape
+    for eigen_value in (L2, L3):
+        if eigen_value.shape != check_shape:
+            raise ValueError("Not all eigenvalues have the same shape")
+    for eigen_vector in (V1, V2, V3):
+        if eigen_vector.shape != eigen_value.shape + (3, ):
+            raise ValueError("Not all eigenvectors have the same shape as the eigenvalues")
+    return (
+        L1[..., None, None] * V1[..., None, :] * V1[..., :, None] +
+        L2[..., None, None] * V2[..., None, :] * V2[..., :, None] +
+        L3[..., None, None] * V3[..., None, :] * V3[..., :, None]
+    )
+
+
+def tensorToEigendecomposition(matrices):
+    """
+    Decomposes the 3x3 diffusion tensor into eigenvalues and eigenvectors
+
+    :param matrices: (..., 3, 3) array-like with diffusion tensor
+    :return: Tuple containing the eigenvectors and eigenvalues (V1, V2, V3, L1, L2, L3)
+    """
+    matrices = np.asanyarray(matrices)
+    if matrices.shape[-2:] != (3, 3):
+        raise ValueError("Expected 3x3 diffusion tensors")
+
+    shape = matrices.shape[:-2]
+    nvoxels = np.prod(shape)
+
+    # Calculate the eigenvectors and
+    # values on all of those matrices
+    flat_matrices = matrices.reshape((-1, 3, 3))
+    vals, vecs = npla.eigh(flat_matrices)
+    vecShape   = shape + (3, )
+
+    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
+
+
+def tensorToComponents(matrices):
+    """
+    Extracts the 6 unique components from a 3x3 diffusion tensor
+
+    :param matrices: (..., 3, 3) array-like with diffusion tensors
+    :return: (..., 6) array with the unique components sorted like Dxx, Dxy, Dxz, Dyy, Dyz, Dzz
+    """
+    matrices = np.asanyarray(matrices)
+    if matrices.shape[-2:] != (3, 3):
+        raise ValueError("Expected 3x3 diffusion tensors")
+    return np.stack([
+        matrices[..., 0, 0],
+        matrices[..., 0, 1],
+        matrices[..., 0, 2],
+        matrices[..., 1, 1],
+        matrices[..., 1, 2],
+        matrices[..., 2, 2],
+    ], -1)
+
+
+def componentsToTensor(components):
+    """
+    Creates 3x3 diffusion tensors from the 6 unique components
+
+    :param components: (..., 6) array-like with Dxx, Dxy, Dxz, Dyy, Dyz, Dzz
+    :return: (..., 3, 3) array with the diffusion tensors
+    """
+    components = np.asanyarray(components)
+    if components.shape[-1] != 6:
+        raise ValueError("Expected 6 unique components of diffusion tensor")
+    first = np.stack([components[..., index] for index in (0, 1, 2)], -1)
+    second = np.stack([components[..., index] for index in (1, 3, 4)], -1)
+    third = np.stack([components[..., index] for index in (2, 4, 5)], -1)
+    return np.stack([first, second, third], -1)
+
+
+def eigendecompositionToComponents(V1, V2, V3, L1, L2, L3):
+    """
+    Converts the eigenvalues/eigenvectors into the 6 unique components of the diffusion tensor
+
+    :param V1: (..., 3) shaped array with the first eigenvector
+    :param V2: (..., 3) shaped array with the second eigenvector
+    :param V3: (..., 3) shaped array with the third eigenvector
+    :param L1: (..., ) shaped array with the first eigenvalue
+    :param L2: (..., ) shaped array with the second eigenvalue
+    :param L3: (..., ) shaped array with the third eigenvalue
+    :return: (..., 6) array with the unique components sorted like Dxx, Dxy, Dxz, Dyy, Dyz, Dzz
+    """
+    return tensorToComponents(eigendecompositionToTensor(V1, V2, V3, L1, L2, L3))
+
+
+def componentsToEigendecomposition(components):
+    """
+    Decomposes diffusion tensor defined by its 6 unique components
+
+    :param components: (..., 6) array-like with Dxx, Dxy, Dxz, Dyy, Dyz, Dzz
+    :return: Tuple containing the eigenvectors and eigenvalues (V1, V2, V3, L1, L2, L3)
+    """
+    return tensorToEigendecomposition(componentsToTensor(components))
+
+
+
 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
@@ -117,53 +241,7 @@ def decomposeTensorMatrix(data):
     :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
+    return componentsToEigendecomposition(data)
 
 
 class DTIFitTensor(fslimage.Nifti):

tests/test_dtifit.py

@@ -137,6 +137,23 @@ def test_decomposeTensorMatrix():
             assert np.all(np.isclose(resvec,  expvec)) or \
                    np.all(np.isclose(resvec, -expvec))
 
+    assert np.allclose(
+            dtifit.eigendecompositionToComponents(expV1, expV2, expV3, expL1, expL2, expL3),
+            tensorMatrices
+    )
+
+    random_tensor = np.random.randn(6)
+    assert np.allclose(
+            random_tensor,
+            dtifit.eigendecompositionToComponents(*dtifit.componentsToEigendecomposition(random_tensor))
+    )
+
+    random_tensor = np.random.randn(2, 2, 3, 6)
+    assert np.allclose(
+            random_tensor,
+            dtifit.eigendecompositionToComponents(*dtifit.componentsToEigendecomposition(random_tensor))
+    )
+
 
 def test_DTIFitTensor():