Centroid based metric

.gitignore

@@ -26,4 +26,5 @@ nibabel*egg/
 pyx-stamps
 __config__.py
 .DS_Store
-.eggs
+.eggs
+*.idea*

dipy/align/bundlemin.pyx

@@ -206,6 +206,102 @@ def _bundle_minimum_distance(double [:, ::1] stat,
 
     return dist
 
+def _bundle_centroids_minimum_distance(double [:, ::1] stat,
+                             double [:, ::1] mov,
+                             int [:] stat_clus_size,
+                             int [:] mov_clus_size,
+                             cnp.npy_intp static_size,
+                             cnp.npy_intp moving_size,
+                             cnp.npy_intp rows):
+    """ MDF-based pairwise distance optimization function
+
+    We minimize the distance between moving centroids of the same number of
+    points as they align with the static centroids.
+
+    Parameters
+    -----------
+    stat : array
+        Static centroids
+    mov : array
+        Moving centroids
+    stat_clus_size : array
+        Size of clusters
+    mov_clus_size : array
+        Size of clusters
+    static_size : int
+        Number of static centroids
+    moving_size : int
+        Number of moving centroids
+    rows : int
+        Number of points per centroids
+
+    Returns
+    -------
+    cost : double
+
+    Notes
+    -----
+    The difference with ``_bundle_minimum_distance_matrix`` is that it does not
+    save the full distance matrix and therefore needs much less memory.
+    """
+
+    cdef:
+        cnp.npy_intp i=0, j=0
+        double sum_i=0, sum_j=0, tmp=0
+        double inf = np.finfo('f8').max
+        double dist=0
+        double * min_j
+        double * min_i
+        openmp.omp_lock_t lock
+
+    with nogil:
+
+        if have_openmp:
+            openmp.omp_init_lock(&lock)
+
+        min_j = <double *> malloc(static_size * sizeof(double))
+        min_i = <double *> malloc(moving_size * sizeof(double))
+
+        for i in range(static_size):
+            min_j[i] = inf
+
+        for j in range(moving_size):
+            min_i[j] = inf
+
+        for i in prange(static_size):
+
+            for j in range(moving_size):
+
+                tmp = min_direct_flip_dist(&stat[i * rows, 0],
+                                       &mov[j * rows, 0], rows)
+
+                if have_openmp:
+                    openmp.omp_set_lock(&lock)
+                if tmp < min_j[i]:
+                    min_j[i] = tmp
+
+                if tmp < min_i[j]:
+                    min_i[j] = tmp
+                if have_openmp:
+                    openmp.omp_unset_lock(&lock)
+
+        if have_openmp:
+            openmp.omp_destroy_lock(&lock)
+
+        for i in range(static_size):
+            sum_i += stat_clus_size[i]*min_j[i]
+
+        for j in range(moving_size):
+            sum_j += mov_clus_size[j]*min_i[j]
+
+        free(min_j)
+        free(min_i)
+
+        dist = (sum_i / <double>static_size + sum_j / <double>moving_size)
+
+        dist = 0.25 * dist * dist
+
+    return dist
 
 def _bundle_minimum_distance_static(double [:, ::1] stat,
                                     double [:, ::1] mov,

dipy/align/streamlinear.py

@@ -3,6 +3,7 @@
 from dipy.utils.six import with_metaclass
 from dipy.core.optimize import Optimizer
 from dipy.align.bundlemin import (_bundle_minimum_distance,
+                                  _bundle_centroids_minimum_distance,
                                   _bundle_minimum_distance_static,
                                   distance_matrix_mdf)
 from dipy.tracking.streamline import (transform_streamlines,
@@ -11,7 +12,7 @@
                                       set_number_of_points,
                                       select_random_set_of_streamlines,
                                       length)
-from dipy.segment.clustering import QuickBundles
+from dipy.segment.clustering import QuickBundles, qbx_with_merge
 from dipy.core.geometry import (compose_transformations,
                                 compose_matrix,
                                 decompose_matrix)
@@ -102,6 +103,70 @@ def distance(self, xopt):
                                         self.moving_centered_pts,
                                         self.block_size)
 
+class BundleCentroidsMinDistanceMetric(StreamlineDistanceMetric):
+    """ Bundle-based Minimum Distance aka BMD
+        Using centroids and their size as weight
+
+    This is the cost function used by the StreamlineLinearRegistration
+
+    Methods
+    -------
+    setup(static, moving)
+    distance(xopt)
+
+    References
+    ----------
+    .. [Garyfallidis14] Garyfallidis et al., "Direct native-space fiber
+                        bundle alignment for group comparisons", ISMRM,
+                        2014.
+    """
+
+    def setup(self, static, moving):
+        """ Setup static and moving sets of streamlines
+
+        Parameters
+        ----------
+        static : tuple - (centroids, cluster sizes)
+            Fixed or reference set of streamlines.
+        moving : tuple - (centroids, cluster sizes)
+            Moving streamlines.
+
+        Notes
+        -----
+        Call this after the object is initiated and before distance.
+        """
+
+        self._set_static_centroids(static)
+        self._set_moving_centroids(moving)
+
+    def _set_static_centroids(self, static):
+        static_centered_pts, st_idx = unlist_streamlines(static[0])
+        self.static_centered_pts = np.ascontiguousarray(static_centered_pts,
+                                                        dtype=np.float64)
+        self.block_size = st_idx[0]
+        self.static_clusters_size = static[1]
+
+    def _set_moving_centroids(self, moving):
+        self.moving_centered_pts, _ = unlist_streamlines(moving[0])
+        self.moving_clusters_size = moving[1]
+
+
+    def distance(self, xopt):
+        """ Distance calculated from this Metric
+
+        Parameters
+        ----------
+        xopt : sequence
+            List of affine parameters as an 1D vector,
+
+        """
+        return bundle_min_distance_centroids_fast(xopt,
+                                        self.static_centered_pts,
+                                        self.moving_centered_pts,
+                                        self.static_clusters_size,
+                                        self.moving_clusters_size,
+                                        self.block_size)
+
 class BundleMinDistanceStaticMetric(BundleMinDistanceMetric):
 
     def distance(self, xopt):
@@ -283,9 +348,9 @@ def optimize(self, static, moving, mat=None):
         Parameters
         ----------
 
-        static : streamlines
+        static : streamlines OR tuple - (centroids, weight)
             Reference or fixed set of streamlines.
-        moving : streamlines
+        moving : streamlines OR tuple - (centroids, weight)
             Moving set of streamlines.
         mat : array
             Transformation (4, 4) matrix to start the registration. ``mat``
@@ -298,34 +363,48 @@ def optimize(self, static, moving, mat=None):
         map : StreamlineRegistrationMap
 
         """
+        #If we are using centroids and their weight, static/moving are tuple
+        #The function only need the centroids
+        if type(static) is tuple and type(moving) is tuple:
+            static_centroids = static[0]
+            static_cluster_sizes = static[1]
+            moving_centroids = moving[0]
+            moving_cluster_sizes = moving[1]
+        else:
+            static_centroids = static
+            moving_centroids = moving
 
         msg = 'need to have the same number of points. Use '
         msg += 'set_number_of_points from dipy.tracking.streamline'
 
-        if not np.all(np.array(list(map(len, static))) == static[0].shape[0]):
+        if not np.all(np.array(list(map(len, static_centroids))) == static_centroids[0].shape[0]):
             raise ValueError('Static streamlines ' + msg)
 
-        if not np.all(np.array(list(map(len, moving))) == moving[0].shape[0]):
+        if not np.all(np.array(list(map(len, moving_centroids))) == moving_centroids[0].shape[0]):
             raise ValueError('Moving streamlines ' + msg)
 
-        if not np.all(np.array(list(map(len, moving))) == static[0].shape[0]):
+        if not np.all(np.array(list(map(len, moving_centroids))) == static_centroids[0].shape[0]):
             raise ValueError('Static and moving streamlines ' + msg)
 
         if mat is None:
-            static_centered, static_shift = center_streamlines(static)
-            moving_centered, moving_shift = center_streamlines(moving)
+            static_centered, static_shift = center_streamlines(static_centroids)
+            moving_centered, moving_shift = center_streamlines(moving_centroids)
+
             static_mat = compose_matrix44([static_shift[0], static_shift[1],
                                            static_shift[2], 0, 0, 0])
 
             moving_mat = compose_matrix44([-moving_shift[0], -moving_shift[1],
                                            -moving_shift[2], 0, 0, 0])
         else:
-            static_centered = static
-            moving_centered = transform_streamlines(moving, mat)
+            static_centered = static_centroids
+            moving_centered = transform_streamlines(moving_centroids, mat)
             static_mat = np.eye(4)
             moving_mat = mat
 
-        self.metric.setup(static_centered, moving_centered)
+        if type(static) is tuple and type(moving) is tuple:
+            self.metric.setup((static_centered, static_cluster_sizes), (moving_centered, moving_cluster_sizes))
+        else:
+            self.metric.setup(static_centered, moving_centered)
 
         distance = self.metric.distance
 
@@ -338,7 +417,7 @@ def optimize(self, static, moving, mat=None):
                             method=self.method, options=self.options,
                             evolution=self.evolution)
 
-        if self.method == 'L-BFGS-B':
+        if self.method == 'L-BFGS-B' or self.method == 'centroids':
 
             if self.options is None:
                 self.options = {'maxcor': 10, 'ftol': 1e-7,
@@ -613,6 +692,77 @@ def bundle_min_distance_fast(t, static, moving, block_size):
                                     block_size)
 
 
+def bundle_min_distance_centroids_fast(t, static, moving,
+                                       static_clusters_size, moving_clusters_size,
+                                       block_size):
+    """ MDF-based pairwise distance optimization function (MIN)
+
+    We minimize the distance between moving streamlines as they align
+    with the static streamlines.
+
+    Parameters
+    -----------
+    t : array
+        1D array. t is a vector of of affine transformation parameters with
+        size at least 6.
+        If size is 6, t is interpreted as translation + rotation.
+        If size is 7, t is interpreted as translation + rotation +
+        isotropic scaling.
+        If size is 12, t is interpreted as translation + rotation +
+        scaling + shearing.
+
+    static : array
+        N*M x 3 array. All the points of the static centroids. With order of
+        centroids intact. Where N is the number of centroids and M
+        is the number of points per centroids.
+
+    moving : array
+        K*M x 3 array. All the points of the moving centroids. With order of
+        centroids intact. Where K is the number of centroids and M
+        is the number of points per centroids.
+
+
+    static_clusters_size : array
+        An array of size K. Int representing the size of the cluster associated
+        with their respective centroid (static)
+
+    moving_clusters_size : array
+        An array of size K. Int representing the size of the cluster associated
+        with their respective centroid (moving)
+
+    block_size : int
+        Number of points per streamline. All streamlines in static and moving
+        should have the same number of points M.
+
+    Returns
+    -------
+    cost: float
+
+    Notes
+    -----
+    This is a faster implementation of ``bundle_min_distance``, which requires
+    that all the points of each streamline are allocated into an ndarray
+    (of shape N*M by 3, with N the number of points per streamline and M the
+    number of streamlines). This can be done by calling
+    `dipy.tracking.streamlines.unlist_streamlines`.
+
+    """
+
+    aff = compose_matrix44(t)
+    moving = np.dot(aff[:3, :3], moving.T).T + aff[:3, 3]
+    moving = np.ascontiguousarray(moving, dtype=np.float64)
+
+    rows = static.shape[0] / block_size
+    cols = moving.shape[0] / block_size
+
+    static_clusters_size = np.asarray(static_clusters_size, dtype=np.int32)
+    moving_clusters_size = np.asarray(moving_clusters_size, dtype=np.int32)
+    return _bundle_centroids_minimum_distance(static, moving,
+                                    static_clusters_size,
+                                    moving_clusters_size,
+                                    rows, cols,
+                                    block_size)
+
 def bundle_min_distance_static_fast(t, static, moving, block_size):
     """ MDF-based pairwise distance optimization function (MIN)
 
@@ -803,10 +953,9 @@ def check_range(streamline, gt=greater_than, lt=less_than):
     else:
         rstreamlines1 = streamlines1
 
-    rstreamlines1 = set_number_of_points(rstreamlines1, nb_pts)
-    qb1 = QuickBundles(threshold=qb_thr)
-    rstreamlines1 = [s.astype('f4') for s in rstreamlines1]
-    cluster_map1 = qb1.cluster(rstreamlines1)
+    thresholds = [40, 30, 20, qb_thr]
+
+    cluster_map1 = qbx_with_merge(rstreamlines1, thresholds, nb_pts=nb_pts)
     clusters1 = remove_clusters_by_size(cluster_map1, rm_small_clusters)
     qb_centroids1 = [cluster.centroid for cluster in clusters1]
 
@@ -816,10 +965,7 @@ def check_range(streamline, gt=greater_than, lt=less_than):
     else:
         rstreamlines2 = streamlines2
 
-    rstreamlines2 = set_number_of_points(rstreamlines2, nb_pts)
-    qb2 = QuickBundles(threshold=qb_thr)
-    rstreamlines2 = [s.astype('f4') for s in rstreamlines2]
-    cluster_map2 = qb2.cluster(rstreamlines2)
+    cluster_map2 = qbx_with_merge(rstreamlines2, thresholds, nb_pts=nb_pts)
     clusters2 = remove_clusters_by_size(cluster_map2, rm_small_clusters)
     qb_centroids2 = [cluster.centroid for cluster in clusters2]
 
@@ -836,7 +982,7 @@ def check_range(streamline, gt=greater_than, lt=less_than):
                   (-10, 10), (-10, 10), (-10, 10)]
         slm = progressive_slr(qb_centroids1, qb_centroids2,
                               x0=x0, metric=None,
-                              bounds=bounds)
+                              bounds=bounds,verbose=verbose)
 
     if verbose:
         print('QB static centroids size %d' % len(qb_centroids1,))