autodmri.estimator

  1import numpy as np
  2
  3try:
  4    from scipy.ndimage import zoom
  5except ImportError:
  6    from scipy.ndimage.interpolation import zoom
  7
  8from scipy.special import gammaincinv
  9
 10from autodmri.gamma import get_noise_distribution
 11from autodmri.blocks import extract_patches
 12
 13from joblib import Parallel, delayed
 14from tqdm import tqdm
 15
 16###########################################
 17# These functions are for over dwis
 18###########################################
 19
 20
 21def estimate_from_dwis(data, axis=-2, return_mask=False, exclude_mask=None, ncores=-1, method='moments', verbose=False, fast_median=False):
 22    '''Given the data, splits over each slice to compute parameters of the gamma distribution
 23
 24    input
 25    ------
 26        data : array, input volume used to identify noise voxels and estimate the noise distribution
 27
 28    optional
 29    --------
 30        axis : (0, 1 or 2) the axis to consider as a slab of uniform noise profile
 31
 32        return_mask : bool, if True returns the identified noise voxels as a mask
 33
 34        exclude_mask : array, mask indicating voxels to remove from all the computations, such as those containing huge artifacts.
 35
 36        ncores : int, number of cores to use for multiprocessing
 37
 38        method='moments' or method='maxlk' : which algorithm to use to estimate sigma and N
 39
 40        verbose : bool, Shows a progress bar for parallel processing
 41
 42        fast_median : Computes the median of medians from each volume.
 43        Useful for large datasets with many volumes (e.g. HCP) since the median requires a full copy of the data and sorting.
 44
 45    output
 46    -------
 47    sigma, N, mask (optional)
 48    '''
 49
 50    # guess a gross upper bound of sigma
 51    # if it's masked, median can be zero, so use the nonzero data
 52    if fast_median:
 53        medians = np.zeros(data.shape[-1])
 54
 55        for idx in range(data.shape[-1]):
 56            chunk = data[..., idx]
 57            median = np.median(chunk)
 58
 59            if median == 0:
 60                median = np.median(chunk[chunk > 0])
 61
 62            medians[idx] = median
 63
 64        median = np.median(medians)
 65    else:
 66        median = np.median(data)
 67
 68        if median == 0:
 69            median = np.median(data[data > 0])
 70
 71    if axis < 0:
 72        axis = data.ndim + axis
 73
 74    ranger = range(data.shape[axis])
 75    swapped_data = data.swapaxes(0, axis)
 76
 77    if exclude_mask is None:
 78        exclude_mask = [None] * swapped_data.shape[0]
 79    else:
 80        exclude_mask = exclude_mask.swapaxes(0, axis)
 81
 82    if verbose:
 83        ranger = tqdm(ranger)
 84
 85    output = Parallel(n_jobs=ncores)(delayed(_inner)(swapped_data[i], median, exclude_mask[i], method) for i in ranger)
 86
 87    # output is each slice we took along axis, so the mask might be reversed
 88    sigma = np.zeros(len(output), dtype=np.float32)
 89    N = np.zeros(len(output), dtype=np.float32)
 90    mask = np.zeros(data.shape[:-1], dtype=np.int16).swapaxes(0, axis)
 91
 92    for i, s in enumerate(output):
 93        sigma[i] = s[0]
 94        N[i] = s[1]
 95        mask[i] = s[2]
 96
 97    if return_mask:
 98        return sigma, N, mask.swapaxes(0, axis)
 99    return sigma, N
100
101
102def _inner(data, median, exclude_mask=None, method='moments', l=50, N_min=1, N_max=12, max_iter=100, eps=1e-3):
103
104    def lambda_cdf(N, alpha_prob):
105        out = gammaincinv(N, alpha_prob)
106        out = np.nan_to_num(out).clip(min=1e-7)
107        return out
108
109    def get_mask(data, N_min, N_max, phi, alpha_prob=0.05):
110        if alpha_prob <= 0 or alpha_prob >= 1:
111            raise ValueError(f'alpha_prob is a parameter in ]0, 1[ with default value 0.05, but has value {alpha_prob}')
112
113        K = np.sum(data > 0, axis=-1)
114        sum_data2 = np.sum(data**2, axis=-1)
115
116        lambda_minus = lambda_cdf(N_min*K, alpha_prob/2)
117        lambda_plus = lambda_cdf(N_max*K, 1 - alpha_prob/2)
118
119        mask_current = np.zeros(data.shape[:-1], dtype=bool)
120        mask_loop = np.zeros(data.shape[:-1], dtype=bool)
121
122        for sigma in phi:
123            s = sum_data2 / (2*sigma**2)
124            mask_current[:] = np.logical_and(lambda_minus < s, s < lambda_plus)
125
126            if mask_current.sum() > mask_loop.sum():
127                mask_loop[:] = mask_current
128
129        return mask_loop
130
131    # Explicitly remove known artifacts
132    if exclude_mask is None:
133        exclude_mask = np.zeros(data.shape[:-1], dtype=bool)
134
135    # we don't know N, so guess parameters iteratively
136    data = data.astype(np.float64) # prevent data**4 overflow
137    data = np.nan_to_num(data).clip(min=0) # cleanup in case we had bad data somewhere to prevent nan propagation
138
139    sigma_prev = -1
140    N_prev = -1
141    sigma_init = median / np.sqrt(2 * lambda_cdf(N_max, 0.5))
142    phi = np.arange(1, l+1) * sigma_init / l
143
144    for _ in range(max_iter):
145
146        mask = get_mask(data, N_min, N_max, phi)
147        mask *= np.logical_not(exclude_mask)
148
149        # empty slice -> mask is zero
150        if mask.sum() == 0:
151            return 0, 0, np.zeros_like(mask)
152
153        datam = data[np.broadcast_to(mask[..., None], data.shape)]
154        sigma, N = get_noise_distribution(datam, method=method)
155
156        if sigma == 0 or N == 0:
157            return 0, 0, np.zeros_like(mask)
158
159        # abs error is small?
160        if (np.abs(N - N_prev) < eps) and (np.abs(sigma - sigma_prev) < eps):
161            break
162
163        # relative error is small?
164        if ((np.abs(N - N_prev) / N) < eps) and ((np.abs(sigma - sigma_prev) / sigma) < eps):
165            break
166
167        N_prev = N
168        sigma_prev = sigma
169
170        N_min = N
171        N_max = N
172
173        phi = np.linspace(.95, 1.05, num=11) * sigma
174
175    return sigma, N, mask
176
177
178###########################################
179# These functions are for over noise maps
180###########################################
181
182
183def estimate_from_nmaps(data, size=5, return_mask=True, method='moments', full=False, ncores=-1, use_rejection=False, verbose=False):
184    '''Given the data, estimates parameters of the gamma distribution in small 3D windows.
185
186    input
187    ------
188        data : noise maps to use for parameter estimation.
189
190    optional
191    --------
192        size : size of the 3D windows (default 5)
193
194        return_mask : bool, if True returns the identified noise voxels as a mask
195
196        method='moments' or method='maxlk' : which algorithm to use to estimate sigma and N
197
198        full : bool, if True estimates are made in overlapping windows
199
200        ncores : int, number of cores to use for multiprocessing
201
202        use_rejection : if True, iterate to reject voxels in each estimated window, but this is much slower than just using all of the data.
203
204        verbose : bool, Shows a progress bar for parallel processing
205
206    output
207    -------
208    sigma, N, mask (optional)
209    '''
210    m_out = np.zeros(data.shape[:-1], dtype=bool)
211    median = np.median(data)
212
213    if median == 0:
214        median = np.median(data[data > 0])
215
216    if full:
217        reshaped_maps = extract_patches(data, (size, size, size, data.shape[-1]), (1, 1, 1, data.shape[-1]), flatten=False)
218
219        sigma = np.zeros(data.shape[:-1], dtype=np.float32)
220        N = np.zeros(data.shape[:-1], dtype=np.float32)
221        count = np.zeros(data.shape[:-1], dtype=np.int32)
222        mask = np.zeros(data.shape[:-1], dtype=np.int32)
223
224        indexer = list(np.ndindex(reshaped_maps.shape[:reshaped_maps.ndim//2 - 1]))
225
226        if verbose:
227            ranger = tqdm(indexer)
228        else:
229            ranger = indexer
230
231        output = Parallel(n_jobs=ncores)(delayed(proc_inner)(reshaped_maps[i], median, size, method, use_rejection) for i in ranger)
232
233        indexer = (np.index_exp[idx[0]:idx[0] + size, idx[1]:idx[1] + size, idx[2]:idx[2] + size] for idx in indexer)
234
235        # We accumulate the value at each voxel then take the average over the overlap
236        for idx, (s, n, m) in zip(indexer, output):
237            sigma[idx] += s
238            N[idx] += n
239
240            mask[idx] = m.sum()
241            count[idx] += 1
242
243        sigma /= count
244        N /= count
245
246        if return_mask:
247            return sigma, N, mask
248        return sigma, N
249
250    else:
251        reshaped_maps = extract_patches(data, (size, size, size, data.shape[-1]), (size, size, size, data.shape[-1]))
252
253        sigma = np.zeros(reshaped_maps.shape[0], dtype=np.float32)
254        N = np.zeros(reshaped_maps.shape[0], dtype=np.float32)
255        mask = np.zeros((reshaped_maps.shape[0], size**3), dtype=bool)
256
257        ranger = range(reshaped_maps.shape[0])
258
259        if verbose:
260            ranger = tqdm(ranger)
261
262        output = Parallel(n_jobs=ncores)(delayed(proc_inner)(reshaped_maps[i], median, size, method, use_rejection) for i in ranger)
263
264        for i, (s, n, m) in enumerate(output):
265            sigma[i] = s
266            N[i] = n
267            mask[i] = np.squeeze(m).sum(axis=-1)
268
269        s_out = sigma.reshape(data.shape[0] // size, data.shape[1] // size, data.shape[2] // size)
270        N_out = N.reshape(data.shape[0] // size, data.shape[1] // size, data.shape[2] // size)
271
272        for n, i in enumerate(np.ndindex(s_out.shape)):
273            i = np.array(i) * size
274            j = i + size
275            m_out[i[0]:j[0], i[1]:j[1], i[2]:j[2]] = mask[n].reshape(size, size, size)
276
277        x, y, z = np.array(s_out.shape) * size
278        interpolated_sigma = np.zeros_like(data[..., 0], dtype=np.float32)
279        interpolated_N = np.zeros_like(data[..., 0], dtype=np.float32)
280
281        interpolated_sigma[:x, :y, :z] = zoom(s_out, size, order=1)
282        interpolated_N[:x, :y, :z] = zoom(N_out, size, order=1)
283
284        if return_mask:
285            return interpolated_sigma, interpolated_N, m_out
286        return interpolated_sigma, interpolated_N
287
288
289def proc_inner(cur_map, median, size, method, use_rejection):
290
291    if use_rejection:
292        cur_map = cur_map.reshape(-1, 1, 1)
293        output = _inner(cur_map, median, method=method)
294    else:
295        cur_map = cur_map.reshape(size**3, 1, -1)
296        sigma, N = get_noise_distribution(cur_map, method=method)
297        mask = np.ones_like(cur_map, dtype=bool)
298
299        if np.isnan(sigma) or np.isnan(N):
300            sigma = 0
301            N = 0
302
303        output = sigma, N, mask
304    return output