1import numpy as np
2
3
4def dki_design_matrix(gtab):
5    r"""Construct B design matrix for DKI.
6
7    Parameters
8    ----------
9    gtab : GradientTable
10        Measurement directions.
11
12    Returns
13    -------
14    B : array (N, 22)
15        Design matrix or B matrix for the DKI model
16        B[j, :] = (Bxx, Bxy, Bzz, Bxz, Byz, Bzz,
17                   Bxxxx, Byyyy, Bzzzz, Bxxxy, Bxxxz,
18                   Bxyyy, Byyyz, Bxzzz, Byzzz, Bxxyy,
19                   Bxxzz, Byyzz, Bxxyz, Bxyyz, Bxyzz,
20                   BlogS0)
21
22    """
23    b = gtab.bvals
24    bvec = gtab.bvecs
25
26    B = np.zeros((len(b), 22))
27    B[:, 0] = -b * bvec[:, 0] * bvec[:, 0]
28    B[:, 1] = -2 * b * bvec[:, 0] * bvec[:, 1]
29    B[:, 2] = -b * bvec[:, 1] * bvec[:, 1]
30    B[:, 3] = -2 * b * bvec[:, 0] * bvec[:, 2]
31    B[:, 4] = -2 * b * bvec[:, 1] * bvec[:, 2]
32    B[:, 5] = -b * bvec[:, 2] * bvec[:, 2]
33    B[:, 6] = b * b * bvec[:, 0]**4 / 6
34    B[:, 7] = b * b * bvec[:, 1]**4 / 6
35    B[:, 8] = b * b * bvec[:, 2]**4 / 6
36    B[:, 9] = 4 * b * b * bvec[:, 0]**3 * bvec[:, 1] / 6
37    B[:, 10] = 4 * b * b * bvec[:, 0]**3 * bvec[:, 2] / 6
38    B[:, 11] = 4 * b * b * bvec[:, 1]**3 * bvec[:, 0] / 6
39    B[:, 12] = 4 * b * b * bvec[:, 1]**3 * bvec[:, 2] / 6
40    B[:, 13] = 4 * b * b * bvec[:, 2]**3 * bvec[:, 0] / 6
41    B[:, 14] = 4 * b * b * bvec[:, 2]**3 * bvec[:, 1] / 6
42    B[:, 15] = b * b * bvec[:, 0]**2 * bvec[:, 1]**2
43    B[:, 16] = b * b * bvec[:, 0]**2 * bvec[:, 2]**2
44    B[:, 17] = b * b * bvec[:, 1]**2 * bvec[:, 2]**2
45    B[:, 18] = 2 * b * b * bvec[:, 0]**2 * bvec[:, 1] * bvec[:, 2]
46    B[:, 19] = 2 * b * b * bvec[:, 1]**2 * bvec[:, 0] * bvec[:, 2]
47    B[:, 20] = 2 * b * b * bvec[:, 2]**2 * bvec[:, 0] * bvec[:, 1]
48    B[:, 21] = -np.ones(len(b))
49
50    return B
51
52
53def _roi_in_volume(data_shape, roi_center, roi_radii):
54    """Ensures that a cuboid ROI is in a data volume.
55
56    Parameters
57    ----------
58    data_shape : ndarray
59        Shape of the data
60    roi_center : ndarray, (3,)
61        Center of ROI in data.
62    roi_radii : ndarray, (3,)
63        Radii of cuboid ROI
64
65    Returns
66    -------
67    roi_radii : ndarray, (3,)
68        Truncated radii of cuboid ROI. It remains unchanged if
69        the ROI was already contained inside the data volume.
70    """
71
72    for i in range(len(roi_center)):
73        inf_lim = int(roi_center[i] - roi_radii[i])
74        sup_lim = int(roi_center[i] + roi_radii[i])
75        if inf_lim < 0 or sup_lim >= data_shape[i]:
76            roi_radii[i] = min(int(roi_center[i]),
77                               int(data_shape[i] - roi_center[i]))
78    return roi_radii
79
80
81def _mask_from_roi(data_shape, roi_center, roi_radii):
82    """Produces a mask from a cuboid ROI defined by center and radii.
83
84    Parameters
85    ----------
86    data_shape : array-like, (3,)
87        Shape of the data from which the ROI is taken.
88    roi_center : array-like, (3,)
89        Center of ROI in data.
90    roi_radii : array-like, (3,)
91        Radii of cuboid ROI.
92
93    Returns
94    -------
95    mask : ndarray
96        Mask of the cuboid ROI.
97    """
98
99    ci, cj, ck = roi_center
100    wi, wj, wk = roi_radii
101    interval_i = slice(int(ci - wi), int(ci + wi) + 1)
102    interval_j = slice(int(cj - wj), int(cj + wj) + 1)
103    interval_k = slice(int(ck - wk), int(ck + wk) + 1)
104
105    if wi == 0:
106        interval_i = ci
107    elif wj == 0:
108        interval_j = cj
109    elif wk == 0:
110        interval_k = ck
111
112    mask = np.zeros(data_shape, dtype=np.int64)
113    mask[interval_i, interval_j, interval_k] = 1
114
115    return mask
116