ENH: test kurtosis with 3 components

unit_tests/fdt/dtifit/feedsRun

@@ -38,7 +38,7 @@ def gen_data():
     # eigen-values can not be the same or eigen-vectors will be ill-defined
     eigen_values = np.array([
         [1.2, 1., 0.4],
-        [1., 0.9, 0.],
+        [1., 0.9, 0.6],
         [0.8, .6, .3],
     ])[None, :, None, :, None] * 1e-3
 
@@ -68,9 +68,11 @@ def gen_data():
             testing.assert_allclose((bvecs ** 2).sum(-1), 1.)
 
             if kurt == 2:
-                beig = bvals[:, None] * eigen_values[..., None, :, 0]
-                beig[..., 0] -= 0.1 * (eigen_values[..., None, 0, 0] * bvals) ** 2 / 6
-                beig[..., 1:] -= 0.05 * (np.mean(eigen_values[..., None, 1:, 0], -1) * bvals)[..., None] ** 2 / 6
+                kvec = np.array([0.1, 0.05, 0.2])
+                beig = (
+                               bvals[:, None] * eigen_values[..., None, :, 0] -
+                               kvec * (eigen_values[..., None, :, 0] * bvals[:, None]) ** 2 / 6
+                )
                 data = S0[..., None] * np.exp(np.sum(
                     -beig *
                     np.sum(bvecs[:, None, :] * eigen_vectors[..., None, :, :], -1) ** 2, axis=-1
@@ -108,7 +110,6 @@ def gen_data():
                 MD = eigen_values[..., 0].mean(-1)
                 FA = np.sqrt(1.5 * ((eigen_values[..., 0] - MD[..., None]) ** 2).sum(-1) / (eigen_values[..., 0] ** 2).sum(-1))
                 mode = 3 * np.sqrt(6) * np.prod(eigen_values[..., 0] - MD[..., None], -1) / ((eigen_values[..., 0] - MD[..., None]) ** 2).sum(-1) ** 1.5
-                print('mode', mode[0, :, 0])
                 nib.Nifti1Image(MD, affine=affine).to_filename(f'{fdir}/ref_MD.nii.gz')
                 nib.Nifti1Image(FA, affine=affine).to_filename(f'{fdir}/ref_FA.nii.gz')
                 nib.Nifti1Image(mode, affine=affine).to_filename(f'{fdir}/ref_MO.nii.gz')
@@ -152,7 +153,7 @@ def fit_data(directory):
 
 for directory, multi_shell, kurt in gen_data():
     for base_output, fkurt, fkurtdir in fit_data(directory):
-        print(base_output)
+        print('testing', base_output)
         if (
             (kurt == 0 and (multi_shell or not (fkurt or fkurtdir))) or
             (kurt == 1 and fkurt and not fkurtdir) or
@@ -162,7 +163,7 @@ for directory, multi_shell, kurt in gen_data():
                 ref = nib.load(f'{directory}/ref_{name}.nii.gz').get_fdata()
                 fit = nib.load(f'{base_output}_{name}.nii.gz').get_fdata()
                 assert ref.shape == fit.shape, f'incorrect NIFTI image shape for {name}'
-                print(name, ref[0, 0, 0], fit[0, 0, 0])
+                #print(name, ref[0, 0, 0], fit[0, 0, 0])
                 testing.assert_allclose(ref, fit, atol=1e-3 if fkurtdir else 1e-6, rtol=0.1 if fkurtdir else 1e-3,
                                         err_msg=f'mismatch in {name}')
 
@@ -193,21 +194,29 @@ for directory, multi_shell, kurt in gen_data():
             else:
                 assert not os.path.isfile(f'{base_output}_kurt.nii.gz')
             if fkurtdir:
-                kurt_para = nib.load(f'{base_output}_kurt_para.nii.gz').get_fdata()
-                kurt_perp = nib.load(f'{base_output}_kurt_perp.nii.gz').get_fdata()
+                for idx, ref in zip(range(1, 4), [0.1, 0.05, 0.2]):
+                    fit = nib.load(f'{base_output}_kurt{idx}.nii.gz').get_fdata()
+                    #print('kurtosis', idx, ref, fit[0, 0, 0])
+                    if kurt == 0:
+                        testing.assert_allclose(fit, 0., atol=1e-5)
+                    elif kurt == 1:
+                        assert not np.allclose(fit, 0.1, rtol=0.02)
+                    elif kurt == 2:
+                        testing.assert_allclose(fit, ref, rtol=0.1)
+
+                MK = nib.load(f'{base_output}_MK.nii.gz').get_fdata()
                 if kurt == 0:
-                    testing.assert_allclose(kurt_para, 0., atol=1e-5)
-                    testing.assert_allclose(kurt_perp, 0., atol=1e-5)
+                    testing.assert_allclose(kurt, 0., atol=1e-5)
                 elif kurt == 1:
                     # k_para and k_perp = 1 is different from MK=1...
-                    assert not np.allclose(kurt, 0.1, rtol=0.01)
-                    assert not np.allclose(kurt, 0.1, rtol=0.01)
+                    assert not np.allclose(MK, 0.1, rtol=0.01)
                 elif kurt == 2:
-                    testing.assert_allclose(kurt_para, 0.1, rtol=0.1)
-                    testing.assert_allclose(kurt_perp, 0.05, rtol=0.1)
+                    testing.assert_allclose(MK, 0.35 / 3, rtol=0.1)
             else:
-                assert not os.path.isfile(f'{base_output}_kurt_para.nii.gz')
-                assert not os.path.isfile(f'{base_output}_kurt_perp.nii.gz')
+                assert not os.path.isfile(f'{base_output}_kurt1.nii.gz')
+                assert not os.path.isfile(f'{base_output}_kurt2.nii.gz')
+                assert not os.path.isfile(f'{base_output}_kurt3.nii.gz')
+                assert not os.path.isfile(f'{base_output}_MK.nii.gz')
         else:
             if fkurtdir and (not multi_shell or kurt == 1):
                 continue  # unpredictable whether this works or not