1import numpy as np
2from numpy.testing import (assert_equal, run_module_suite,
3                           assert_array_almost_equal)
4
5from dipy.align.streamlinear import (compose_matrix44, decompose_matrix44,
6                                     transform_streamlines, whole_brain_slr,
7                                     slr_with_qbx)
8from dipy.io.streamline import load_tractogram
9from dipy.data import get_fnames
10from dipy.tracking.streamline import Streamlines
11from dipy.tracking.distances import bundles_distances_mam
12
13
14def test_whole_brain_slr():
15    fname = get_fnames('fornix')
16
17    fornix = load_tractogram(fname, 'same',
18                             bbox_valid_check=False).streamlines
19
20    f = Streamlines(fornix)
21    f1 = f.copy()
22    f2 = f.copy()
23
24    # check translation
25    f2._data += np.array([50, 0, 0])
26
27    moved, transform, qb_centroids1, qb_centroids2 = whole_brain_slr(
28            f1, f2, x0='affine', verbose=True, rm_small_clusters=2,
29            greater_than=0, less_than=np.inf,
30            qbx_thr=[5, 2, 1], progressive=False)
31
32    # we can check the quality of registration by comparing the matrices
33    # MAM streamline distances before and after SLR
34    D12 = bundles_distances_mam(f1, f2)
35    D1M = bundles_distances_mam(f1, moved)
36
37    d12_minsum = np.sum(np.min(D12, axis=0))
38    d1m_minsum = np.sum(np.min(D1M, axis=0))
39
40    print("distances= ", d12_minsum, " ", d1m_minsum)
41
42    assert_equal(d1m_minsum < d12_minsum, True)
43
44    assert_array_almost_equal(transform[:3, 3], [-50, -0, -0], 2)
45
46    # check rotation
47
48    mat = compose_matrix44([0, 0, 0, 15, 0, 0])
49
50    f3 = f.copy()
51    f3 = transform_streamlines(f3, mat)
52
53    moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
54            f1, f3, verbose=False, rm_small_clusters=1, greater_than=20,
55            less_than=np.inf, qbx_thr=[2],
56            progressive=True)
57
58    # we can also check the quality by looking at the decomposed transform
59
60    assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
61
62    moved, transform, qb_centroids1, qb_centroids2 = slr_with_qbx(
63            f1, f3, verbose=False, rm_small_clusters=1, select_random=400,
64            greater_than=20, less_than=np.inf, qbx_thr=[2],
65            progressive=True)
66
67    # we can also check the quality by looking at the decomposed transform
68
69    assert_array_almost_equal(decompose_matrix44(transform)[3], -15, 2)
70
71if __name__ == '__main__':
72    # run_module_suite()
73    test_whole_brain_slr()
74