1import numpy as np
2from numpy.testing import (assert_equal,
3                           assert_almost_equal,
4                           run_module_suite,
5                           assert_array_equal,
6                           assert_raises)
7from dipy.data import get_fnames, dsi_voxels, default_sphere
8from dipy.reconst.dsi import DiffusionSpectrumModel
9from dipy.reconst.odf import gfa
10from dipy.direction.peaks import peak_directions
11from dipy.sims.voxel import sticks_and_ball
12from dipy.core.sphere import Sphere
13from dipy.core.gradients import gradient_table
14from dipy.core.subdivide_octahedron import create_unit_sphere
15from dipy.core.sphere_stats import angular_similarity
16
17
18def test_dsi():
19    # load repulsion 724 sphere
20    sphere = default_sphere
21
22    # load icosahedron sphere
23    sphere2 = create_unit_sphere(5)
24    btable = np.loadtxt(get_fnames('dsi515btable'))
25    gtab = gradient_table(btable[:, 0], btable[:, 1:])
26    data, golden_directions = sticks_and_ball(gtab, d=0.0015, S0=100,
27                                              angles=[(0, 0), (90, 0)],
28                                              fractions=[50, 50], snr=None)
29
30    ds = DiffusionSpectrumModel(gtab)
31
32    # repulsion724
33    dsfit = ds.fit(data)
34    odf = dsfit.odf(sphere)
35
36    directions, _, _ = peak_directions(odf, sphere)
37    assert_equal(len(directions), 2)
38    assert_almost_equal(angular_similarity(directions, golden_directions),
39                        2, 1)
40
41    # 5 subdivisions
42    dsfit = ds.fit(data)
43    odf2 = dsfit.odf(sphere2)
44    directions, _, _ = peak_directions(odf2, sphere2)
45    assert_equal(len(directions), 2)
46    assert_almost_equal(angular_similarity(directions, golden_directions),
47                        2, 1)
48
49    assert_equal(dsfit.pdf().shape, 3 * (ds.qgrid_size, ))
50    sb_dummies = sticks_and_ball_dummies(gtab)
51    for sbd in sb_dummies:
52        data, golden_directions = sb_dummies[sbd]
53        odf = ds.fit(data).odf(sphere2)
54        directions, _, _ = peak_directions(odf, sphere2)
55        if len(directions) <= 3:
56            assert_equal(len(directions), len(golden_directions))
57        if len(directions) > 3:
58            assert_equal(gfa(odf) < 0.1, True)
59
60    assert_raises(ValueError, DiffusionSpectrumModel, gtab, qgrid_size=16)
61
62
63def test_multivox_dsi():
64    data, gtab = dsi_voxels()
65    DS = DiffusionSpectrumModel(gtab)
66
67    DSfit = DS.fit(data)
68    PDF = DSfit.pdf()
69    assert_equal(data.shape[:-1] + (17, 17, 17), PDF.shape)
70    assert_equal(np.alltrue(np.isreal(PDF)), True)
71
72
73def test_multib0_dsi():
74    data, gtab = dsi_voxels()
75    # Create a new data-set with a b0 measurement:
76    new_data = np.concatenate([data, data[..., 0, None]], -1)
77    new_bvecs = np.concatenate([gtab.bvecs, np.zeros((1, 3))])
78    new_bvals = np.concatenate([gtab.bvals, [0]])
79    new_gtab = gradient_table(new_bvals, new_bvecs)
80    ds = DiffusionSpectrumModel(new_gtab)
81    dsfit = ds.fit(new_data)
82    pdf = dsfit.pdf()
83    dsfit.odf(default_sphere)
84    assert_equal(new_data.shape[:-1] + (17, 17, 17), pdf.shape)
85    assert_equal(np.alltrue(np.isreal(pdf)), True)
86
87    # And again, with one more b0 measurement (two in total):
88    new_data = np.concatenate([data, data[..., 0, None]], -1)
89    new_bvecs = np.concatenate([gtab.bvecs, np.zeros((1, 3))])
90    new_bvals = np.concatenate([gtab.bvals, [0]])
91    new_gtab = gradient_table(new_bvals, new_bvecs)
92    ds = DiffusionSpectrumModel(new_gtab)
93    dsfit = ds.fit(new_data)
94    pdf = dsfit.pdf()
95    dsfit.odf(default_sphere)
96    assert_equal(new_data.shape[:-1] + (17, 17, 17), pdf.shape)
97    assert_equal(np.alltrue(np.isreal(pdf)), True)
98
99
100def sticks_and_ball_dummies(gtab):
101    sb_dummies = {}
102    S, sticks = sticks_and_ball(gtab, d=0.0015, S0=100,
103                                angles=[(0, 0)],
104                                fractions=[100], snr=None)
105    sb_dummies['1fiber'] = (S, sticks)
106    S, sticks = sticks_and_ball(gtab, d=0.0015, S0=100,
107                                angles=[(0, 0), (90, 0)],
108                                fractions=[50, 50], snr=None)
109    sb_dummies['2fiber'] = (S, sticks)
110    S, sticks = sticks_and_ball(gtab, d=0.0015, S0=100,
111                                angles=[(0, 0), (90, 0), (90, 90)],
112                                fractions=[33, 33, 33], snr=None)
113    sb_dummies['3fiber'] = (S, sticks)
114    S, sticks = sticks_and_ball(gtab, d=0.0015, S0=100,
115                                angles=[(0, 0), (90, 0), (90, 90)],
116                                fractions=[0, 0, 0], snr=None)
117    sb_dummies['isotropic'] = (S, sticks)
118    return sb_dummies
119
120
121if __name__ == '__main__':
122    run_module_suite()
123