Repository: yeatmanlab/pyAFQ
Pull request #125 (2019-07-01 20:20:37)

Reduce Tensor ODF Memory Usage


Patch:
diff --git a/AFQ/_fixes.py b/AFQ/_fixes.py
index 10b44bbf..993b02c5 100644
--- a/AFQ/_fixes.py
+++ b/AFQ/_fixes.py
@@ -13,3 +13,95 @@ def spherical_harmonics(m, n, theta, phi):
     val *= np.exp(0.5 * (gammaln(n - m + 1) - gammaln(n + m + 1)))
     val = val * np.exp(1j * m * theta)
     return val
+
+
+def in_place_norm(vec, axis=-1, keepdims=False, delvec=True):
+    """ Return Vectors with Euclidean (L2) norm
+
+    See :term:`unit vector` and :term:`Euclidean norm`
+
+    Parameters
+    -------------
+    vec : array_like
+        Vectors to norm. Squared in the process of calculating the norm.
+    axis : int, optional
+        Axis over which to norm. By default norm over last axis. If `axis` is
+        None, `vec` is flattened then normed. Default is -1.
+    keepdims : bool, optional
+        If True, the output will have the same number of dimensions as `vec`,
+        with shape 1 on `axis`. Default is False.
+    delvec : bool, optional
+        If True, vec is deleted as soon as possible.
+        If False, vec is not deleted, but still squared. Default is True.
+
+    Returns
+    ---------
+    norm : array
+        Euclidean norms of vectors.
+
+    Examples
+    --------
+    >>> vec = [[8, 15, 0], [0, 36, 77]]
+    >>> in_place_norm(vec)
+    array([ 17.,  85.])
+    >>> vec = [[8, 15, 0], [0, 36, 77]]
+    >>> in_place_norm(vec, keepdims=True)
+    array([[ 17.],
+           [ 85.]])
+    >>> vec = [[8, 15, 0], [0, 36, 77]]
+    >>> in_place_norm(vec, axis=0)
+    array([  8.,  39.,  77.])
+    """
+    vec = np.asarray(vec)
+
+    if keepdims:
+        ndim = vec.ndim
+        shape = vec.shape
+
+    np.square(vec, out=vec)
+    vec_norm = vec.sum(axis)
+    if delvec:
+        del vec
+    try:
+        np.sqrt(vec_norm, out=vec_norm)
+    except TypeError:
+        vec_norm = vec_norm.astype(float)
+        np.sqrt(vec_norm, out=vec_norm)
+
+    if keepdims:
+        if axis is None:
+            shape = [1] * ndim
+        else:
+            shape = list(shape)
+            shape[axis] = 1
+        vec_norm = vec_norm.reshape(shape)
+
+    return vec_norm
+
+
+def tensor_odf(evals, evecs, sphere):
+    """
+    Calculate the tensor Orientation Distribution Function
+
+    Parameters
+    ----------
+    evals : array (4D)
+        Eigenvalues of a tensor. Shape (x, y, z, 3).
+    evecs : array (5D)
+        Eigenvectors of a tensor. Shape (x, y, z, 3, 3)
+    sphere : sphere object
+        The ODF will be calculated in each vertex of this sphere.
+    """
+    mask = np.where((evals[..., 0] > 0)
+                    & (evals[..., 1] > 0)
+                    & (evals[..., 2] > 0))
+
+    projection = np.dot(sphere.vertices, evecs[mask])
+    projection /= np.sqrt(evals[mask])
+    proj_norm = in_place_norm(projection)
+    proj_norm **= -3
+    proj_norm /= 4 * np.pi * np.sqrt(np.prod(evals[mask], -1))
+
+    odf = np.zeros((evals.shape[:3] + (sphere.vertices.shape[0],)))
+    odf[mask] = proj_norm.T
+    return odf
diff --git a/AFQ/dti.py b/AFQ/dti.py
index cf5d9b78..31ad68f9 100644
--- a/AFQ/dti.py
+++ b/AFQ/dti.py
@@ -9,7 +9,7 @@ from dipy.reconst import dti
 
 import AFQ.utils.models as ut
 
-__all__ = ["fit_dti", "predict", "tensor_odf"]
+__all__ = ["fit_dti", "predict"]
 
 
 def _fit(gtab, data, mask=None):
@@ -115,28 +115,3 @@ def predict(params_file, gtab, S0_file=None, out_dir=None):
     nib.save(nib.Nifti1Image(pred, img.affine), fname)
 
     return fname
-
-
-def tensor_odf(evals, evecs, sphere):
-    """
-    Calculate the tensor Orientation Distribution Function
-
-    Parameters
-    ----------
-    evals : array (4D)
-        Eigenvalues of a tensor. Shape (x, y, z, 3).
-    evecs : array (5D)
-        Eigenvectors of a tensor. Shape (x, y, z, 3, 3)
-    sphere : sphere object
-        The ODF will be calculated in each vertex of this sphere.
-    """
-    odf = np.zeros((evals.shape[:3] + (sphere.vertices.shape[0],)))
-    mask = np.where((evals[..., 0] > 0)
-                    & (evals[..., 1] > 0)
-                    & (evals[..., 2] > 0))
-
-    lower = 4 * np.pi * np.sqrt(np.prod(evals[mask], -1))
-    projection = np.dot(sphere.vertices, evecs[mask])
-    projection /= np.sqrt(evals[mask])
-    odf[mask] = ((vector_norm(projection) ** -3) / lower).T
-    return odf
diff --git a/AFQ/segmentation.py b/AFQ/segmentation.py
index 6d9fea1b..3514dddb 100644
--- a/AFQ/segmentation.py
+++ b/AFQ/segmentation.py
@@ -284,7 +284,7 @@ def segment(fdata, fbval, fbvec, streamlines, bundle_dict, mapping,
                 roi = roi.get_data()
             warped_roi = auv.patch_up_roi(
                 (mapping.transform_inverse(
-                    roi.astype(np.float32),
+                    roi,
                     interpolation='linear')) > 0)
 
             if rule:
diff --git a/AFQ/tractography.py b/AFQ/tractography.py
index a8ccfcdb..d85c0b44 100644
--- a/AFQ/tractography.py
+++ b/AFQ/tractography.py
@@ -9,7 +9,7 @@ import dipy.tracking.utils as dtu
 import dipy.tracking.streamline as dts
 from dipy.tracking.local import ThresholdTissueClassifier, LocalTracking
 
-from AFQ.dti import tensor_odf
+from AFQ._fixes import tensor_odf
 
 
 def track(params_file, directions="det", max_angle=30., sphere=None,
diff --git a/examples/plot_tract_profile.py b/examples/plot_tract_profile.py
index 9e097dea..03140330 100644
--- a/examples/plot_tract_profile.py
+++ b/examples/plot_tract_profile.py
@@ -83,7 +83,9 @@ fiber_groups = seg.segment(hardi_fdata,
                            streamlines,
                            bundles,
                            reg_template=MNI_T2_img,
-                           mapping=mapping)
+                           mapping=mapping,
+                           as_generator=False,
+                           affine=img.affine)
 
 
 print("Cleaning fiber groups...")
diff --git a/requirements-dev.txt b/requirements-dev.txt
index c847cfe9..9ba1ec9b 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -5,5 +5,4 @@ flake8
 sphinx_gallery
 sphinx_rtd_theme
 numpydoc
-sphinx-autoapi
-six==1.12.0
+sphinx-autoapi
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index a5be7eea..7397e77c 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -10,5 +10,4 @@ palettable==3.1.1
 cloudpickle==0.5.5
 toolz==0.9.0
 boto3==1.9.2
-six==1.12.0
 s3fs==0.2.0
\ No newline at end of file


Test patch:
diff --git a/AFQ/tests/test_dti.py b/AFQ/tests/test_dti.py
index bbaebf7d..7ccb83ee 100644
--- a/AFQ/tests/test_dti.py
+++ b/AFQ/tests/test_dti.py
@@ -8,10 +8,12 @@ import numpy.testing as npt
 import nibabel as nib
 import nibabel.tmpdirs as nbtmp
 
+from dipy.core.geometry import vector_norm
 import dipy.core.gradients as dpg
 import dipy.data as dpd
 
 from AFQ import dti
+from AFQ._fixes import in_place_norm
 from AFQ.utils.testing import make_dti_data
 
 
@@ -70,4 +72,26 @@ def test_cli():
              "-c", op.join(tmpdir, "dti.bvec")])
         out = os.system(cmd)
         assert out ==  0
-        assert op.exists(op.join(tmpdir, 'dti', 'dti_params.nii.gz'))
\ No newline at end of file
+        assert op.exists(op.join(tmpdir, 'dti', 'dti_params.nii.gz'))
+
+
+def test_inplace_norm():
+    vec = [[8, 15, 0], [0, 36, 77]]
+    norm1 = vector_norm(vec)
+    norm2 = in_place_norm(vec)
+    npt.assert_equal(norm1, norm2)
+
+    vec = [[8.0, 15.0, 0.0], [0.0, 36.0, 77.0]]
+    norm1 = vector_norm(vec)
+    norm2 = in_place_norm(vec)
+    npt.assert_equal(norm1, norm2)
+
+    vec = [[8, 15, 0], [0, 36, 77]]
+    norm1 = vector_norm(vec, keepdims=True)
+    norm2 = in_place_norm(vec, keepdims=True)
+    npt.assert_equal(norm1, norm2)
+
+    vec = [[8, 15, 0], [0, 36, 77]]
+    norm1 = vector_norm(vec, axis=0)
+    norm2 = in_place_norm(vec, axis=0)
+    npt.assert_equal(norm1, norm2)