[TEST] Fixes tests

Author: rmarkello

netneurotools/stats.py

@@ -580,9 +580,9 @@ def gen_spinsamples(coords, hemiid, n_rotate=1000, check_duplicates=True,
         >>> nnstats.gen_spinsamples(coords, hemi, n_rotate=1, seed=1,
         ...                         method='original', check_duplicates=False)
         array([[0],
-               [1],
+               [0],
                [2],
-               [2]], dtype=int32)
+               [3]], dtype=int32)
 
     While this is reasonable in most circumstances, if you feel incredibly
     strongly about having a perfect "permutation" (i.e., all indices appear
@@ -675,12 +675,6 @@ 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):
@@ -731,7 +725,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
                 elif method == 'original':
-                    dist, col = kdtrees[h].query(coor @ rot, 1)
+                    dist, col = spatial.cKDTree(coor @ rot).query(coor, 1)
                     cost[hinds, n] = dist
 
                 resampled[hinds] = inds[hinds][col]

netneurotools/tests/test_stats.py

@@ -135,17 +135,21 @@ def test_gen_spinsamples():
     hemi = np.hstack([np.zeros(len(coords) // 2), np.ones(len(coords) // 2)])
 
     # generate "normal" test spins
-    spins, cost = stats.gen_spinsamples(coords, hemi, n_rotate=10, seed=1234)
+    spins, cost = stats.gen_spinsamples(coords, hemi, n_rotate=10, seed=1234,
+                                        return_cost=True)
     assert spins.shape == (len(coords), 10)
     assert cost.shape == (len(coords), 10)
 
     # confirm that `exact` parameter functions as desired
-    spin_exact, cost_exact = stats.gen_spinsamples(coords, hemi, n_rotate=10,
-                                                   exact=True, seed=1234)
-    assert spin_exact.shape == (len(coords), 10)
-    assert cost.shape == (len(coords), 10)
-    for s in spin_exact.T:
-        assert len(np.unique(s)) == len(s)
+    for method in ['vasa', 'hungarian']:
+        spin_exact, cost_exact = stats.gen_spinsamples(coords, hemi,
+                                                       n_rotate=10, seed=1234,
+                                                       method=method,
+                                                       return_cost=True)
+        assert spin_exact.shape == (len(coords), 10)
+        assert cost.shape == (len(coords), 10)
+        for s in spin_exact.T:
+            assert len(np.unique(s)) == len(s)
 
     # confirm that check_duplicates will raise warnings
     # since spins aren't exact permutations we need to use 4C4 with repeats