[FIX] Fixes ordering of spins for spintest

Author: rmarkello

netneurotools/stats.py

@@ -646,6 +646,12 @@ def gen_spinsamples(coords, hemiid, n_rotate=1000, check_duplicates=True,
     cost = np.zeros((len(coords), n_rotate))
     inds = np.arange(len(coords), dtype='int32')
 
+    # precompute kd trees
+    kdtrees = (
+        spatial.cKDTree(coords[hemiid == 0]),
+        spatial.cKDTree(coords[hemiid == 1])
+    )
+
     # generate rotations and resampling array!
     msg, warned = '', False
     for n in range(n_rotate):
@@ -697,7 +703,7 @@ def gen_spinsamples(coords, hemiid, n_rotate=1000, check_duplicates=True,
                 # huge thanks to https://stackoverflow.com/a/47779290 for this
                 # memory-efficient method
                 else:
-                    dist, col = spatial.cKDTree(coor @ rot).query(coor, 1)
+                    dist, col = kdtrees[h].query(coor @ rot, 1)
                     cost[hinds, n] = dist
 
                 resampled[hinds] = inds[hinds][col]
@@ -722,4 +728,7 @@ def gen_spinsamples(coords, hemiid, n_rotate=1000, check_duplicates=True,
     if verbose:
         print(' ' * len(msg) + '\b' * len(msg), end='', flush=True)
 
-    return spinsamples, cost
+    if return_cost:
+        return spinsamples, cost
+
+    return spinsamples