def __init__(self, ksp, calib_width=24, thresh=0.01, kernel_width=6, crop=0.8, max_iter=100, device=sp.cpu_device, output_eigenvalue=False, show_pbar=True): self.device = sp.Device(device) self.output_eigenvalue = output_eigenvalue self.crop = crop img_ndim = ksp.ndim - 1 num_coils = len(ksp) with sp.get_device(ksp): # Get calibration region calib_shape = [num_coils] + [calib_width] * img_ndim calib = sp.resize(ksp, calib_shape) calib = sp.to_device(calib, device) xp = self.device.xp with self.device: # Get calibration matrix kernel_shape = [num_coils] + [kernel_width] * img_ndim kernel_strides = [1] * (img_ndim + 1) mat = sp.array_to_blocks(calib, kernel_shape, kernel_strides) mat = mat.reshape([-1, sp.prod(kernel_shape)]) # Perform SVD on calibration matrix _, S, VH = xp.linalg.svd(mat, full_matrices=False) VH = VH[S > thresh * S.max(), :] # Get kernels num_kernels = len(VH) kernels = VH.reshape([num_kernels] + kernel_shape) img_shape = ksp.shape[1:] # Get covariance matrix in image domain AHA = xp.zeros(img_shape[::-1] + (num_coils, num_coils), dtype=ksp.dtype) for kernel in kernels: img_kernel = sp.ifft(sp.resize(kernel, ksp.shape), axes=range(-img_ndim, 0)) aH = xp.expand_dims(img_kernel.T, axis=-1) a = xp.conj(aH.swapaxes(-1, -2)) AHA += aH @ a AHA *= (sp.prod(img_shape) / kernel_width**img_ndim) self.mps = xp.ones(ksp.shape[::-1] + (1, ), dtype=ksp.dtype) alg = sp.alg.PowerMethod( lambda x: AHA @ x, self.mps, norm_func=lambda x: xp.sum( xp.abs(x)**2, axis=-2, keepdims=True)**0.5, max_iter=max_iter) super().__init__(alg, show_pbar=show_pbar)
def _get_A(self): input_t_size = sp.prod(self.input_t.shape[1:]) output_t_size = sp.prod(self.output_t.shape[1:]) Ri = sp.linop.Reshape([input_t_size, output_t_size], self.mat.shape) M = sp.linop.MatMul([input_t_size, output_t_size], self.input_t.reshape([self.batch_size, -1])) Ro = sp.linop.Reshape(self.output_t.shape, [self.batch_size, output_t_size]) self.A = Ro * M * Ri
def _get_G(self, j): b_j = [min(i, self.blk_widths[j]) for i in self.img_shape] s_j = [(b + 1) // 2 for b in b_j] i_j = [ ceil((i - b + s) / s) * s + b - s for i, b, s in zip(self.img_shape, b_j, s_j) ] n_j = [(i - b + s) // s for i, b, s in zip(i_j, b_j, s_j)] M_j = sp.prod(b_j) P_j = sp.prod(n_j) return M_j**0.5 + self.T**0.5 + (2 * np.log(P_j))**0.5
def array_to_image(arr, color=False): """ Flattens all dimensions except the last two """ if color: arr = np.divide(arr, np.abs(arr).max(), out=np.zeros_like(arr), where=arr != 0) if arr.ndim == 2: return arr elif color and arr.ndim == 3: return arr if color: ndim = 3 else: ndim = 2 arr = sp.resize(arr, arr.shape[:-2] + (arr.shape[-2] + 2, arr.shape[-1] + 2)) shape = arr.shape batch = sp.prod(shape[:-ndim]) mshape = mosaic_shape(batch) if sp.prod(mshape) == batch: img = arr.reshape((batch, ) + shape[-ndim:]) else: img = np.zeros((sp.prod(mshape), ) + shape[-ndim:], dtype=arr.dtype) img[:batch, ...] = arr.reshape((batch, ) + shape[-ndim:]) img = img.reshape(mshape + shape[-ndim:]) if color: img = np.transpose(img, (0, 2, 1, 3, 4)) img = img.reshape( (shape[-3] * mshape[-2], shape[-2] * mshape[-1], shape[-1])) else: img = np.transpose(img, (0, 2, 1, 3)) img = img.reshape((shape[-2] * mshape[-2], shape[-1] * mshape[-1])) return img
def array_to_image(arr, color=False): """ Flattens all dimensions except the last two Args: arr (array): shape [z, x, y, c] if color, else [z, x, y] """ if color and not (arr.max() == 0 and arr.min() == 0): arr = arr / np.abs(arr).max() if arr.ndim == 2: return arr elif color and arr.ndim == 3: return arr if color: img_shape = arr.shape[-3:] batch = sp.prod(arr.shape[:-3]) mshape = mosaic_shape(batch) else: img_shape = arr.shape[-2:] batch = sp.prod(arr.shape[:-2]) mshape = mosaic_shape(batch) if sp.prod(mshape) == batch: img = arr.reshape((batch, ) + img_shape) else: img = np.zeros((sp.prod(mshape), ) + img_shape, dtype=arr.dtype) img[:batch, ...] = arr.reshape((batch, ) + img_shape) img = img.reshape(mshape + img_shape) if color: img = np.transpose(img, (0, 2, 1, 3, 4)) img = img.reshape((img_shape[0] * mshape[0], img_shape[1] * mshape[1], 3)) else: img = np.transpose(img, (0, 2, 1, 3)) img = img.reshape((img_shape[0] * mshape[0], img_shape[1] * mshape[1])) return img
def mosaic_shape(batch): mshape = [int(batch**0.5), batch // int(batch**0.5)] while (sp.prod(mshape) < batch): mshape[1] += 1 if (mshape[0] - 1) * (mshape[1] + 1) == batch: mshape[0] -= 1 mshape[1] += 1 return tuple(mshape)
def __init__(self, shape, L, R, res=None): self.shape = tuple(shape) self.img_shape = self.shape[1:] self.T = self.shape[0] self.size = sp.prod(self.shape) self.ndim = len(self.shape) self.dtype = L[0].dtype self.J = len(L) self.D = self.ndim - 1 self.blk_widths = [max(L[j].shape[-self.D:]) for j in range(self.J)] self.L = L self.R = R self.device = sp.cpu_device if res is None: self.res = (1, ) * self.D
# PDHG with dcf # Compute preconditioner precond_dcf = mr.pipe_menon_dcf(coord, device=device) print(f'DCF shape: {precond_dcf.shape}') print(f'DCF dtype: {precond_dcf.dtype}') f, ax = plt.subplots(1, 1) ax.imshow(precond_dcf) precond_dcf = xp.tile(precond_dcf, [len(mps)] + [1] * (mps.ndim - 1)) img_shape = mps.shape[1:] G = sp.linop.FiniteDifference(img_shape) max_eig_G = sp.app.MaxEig(G.H * G).run() sigma2 = xp.ones([sp.prod(img_shape) * len(img_shape)], dtype=ksp.dtype) / max_eig_G sigma = xp.concatenate([precond_dcf.ravel(), sigma2.ravel()]) pdhg_dcf_app = mr.app.TotalVariationRecon(ksp, mps, lamda=lamda, coord=coord, sigma=sigma, max_iter=max_iter, device=device, save_objective_values=True) print(f'Name of solver: {pdhg_dcf_app.alg_name}') pdhg_dcf_img = pdhg_dcf_app.run() pl.ImagePlot(pdhg_dcf_img)
def circulant_precond(mps, weights=None, coord=None, lamda=0, device=sp.cpu_device): r"""Compute circulant preconditioner. Considers the optimization problem: .. math:: \min_P \| A^H A - F P F^H \|_2^2 where A is the Sense operator, and F is a unitary Fourier transform operator. Args: mps (array): sensitivity maps of shape [num_coils] + image shape. weights (array): k-space weights. coord (array): k-space coordinates of shape [...] + [ndim]. lamda (float): regularization. Returns: array: circulant preconditioner of image shape. """ if coord is not None: coord = sp.to_device(coord, device) if weights is not None: weights = sp.to_device(weights, device) dtype = mps.dtype device = sp.Device(device) xp = device.xp mps_shape = list(mps.shape) img_shape = mps_shape[1:] img2_shape = [i * 2 for i in img_shape] ndim = len(img_shape) scale = sp.prod(img2_shape)**1.5 / sp.prod(img_shape)**2 with device: idx = (slice(None, None, 2), ) * ndim if coord is None: ones = xp.zeros(img2_shape, dtype=dtype) if weights is None: ones[idx] = 1 else: ones[idx] = weights**0.5 psf = sp.ifft(ones) else: coord2 = coord * 2 ones = xp.ones(coord.shape[:-1], dtype=dtype) if weights is not None: ones *= weights**0.5 psf = sp.nufft_adjoint(ones, coord2, img2_shape) p_inv = 0 for mps_i in mps: mps_i = sp.to_device(mps_i, device) xcorr_fourier = xp.abs(sp.fft(xp.conj(mps_i), img2_shape))**2 xcorr = sp.ifft(xcorr_fourier) xcorr *= psf p_inv_i = sp.fft(xcorr) p_inv_i = p_inv_i[idx] p_inv_i *= scale if weights is not None: p_inv_i *= weights**0.5 p_inv += p_inv_i p_inv += lamda p_inv[p_inv == 0] = 1 p = 1 / p_inv return p.astype(dtype)
def kspace_precond(mps, weights=None, coord=None, lamda=0, device=sp.cpu_device, oversamp=1.25): r"""Compute a diagonal preconditioner in k-space. Considers the optimization problem: .. math:: \min_P \| P A A^H - I \|_F^2 where A is the Sense operator. Args: mps (array): sensitivity maps of shape [num_coils] + image shape. weights (array): k-space weights. coord (array): k-space coordinates of shape [...] + [ndim]. lamda (float): regularization. Returns: array: k-space preconditioner of same shape as k-space. """ dtype = mps.dtype if weights is not None: weights = sp.to_device(weights, device) device = sp.Device(device) xp = device.xp mps_shape = list(mps.shape) img_shape = mps_shape[1:] img2_shape = [i * 2 for i in img_shape] ndim = len(img_shape) scale = sp.prod(img2_shape)**1.5 / sp.prod(img_shape) with device: if coord is None: idx = (slice(None, None, 2), ) * ndim ones = xp.zeros(img2_shape, dtype=dtype) if weights is None: ones[idx] = 1 else: ones[idx] = weights**0.5 psf = sp.ifft(ones) else: coord2 = coord * 2 ones = xp.ones(coord.shape[:-1], dtype=dtype) if weights is not None: ones *= weights**0.5 psf = sp.nufft_adjoint(ones, coord2, img2_shape, oversamp=oversamp) p_inv = [] for mps_i in mps: mps_i = sp.to_device(mps_i, device) mps_i_norm2 = xp.linalg.norm(mps_i)**2 xcorr_fourier = 0 for mps_j in mps: mps_j = sp.to_device(mps_j, device) xcorr_fourier += xp.abs( sp.fft(mps_i * xp.conj(mps_j), img2_shape))**2 xcorr = sp.ifft(xcorr_fourier) xcorr *= psf if coord is None: p_inv_i = sp.fft(xcorr)[idx] else: p_inv_i = sp.nufft(xcorr, coord2, oversamp=oversamp) if weights is not None: p_inv_i *= weights**0.5 p_inv.append(p_inv_i * scale / mps_i_norm2) p_inv = (xp.abs(xp.stack(p_inv, axis=0)) + lamda) / (1 + lamda) p_inv[p_inv == 0] = 1 p = 1 / p_inv return p.astype(dtype)
def espirit_maps(ksp, calib_width=24, thresh=0.001, kernel_width=6, crop=0.8, max_power_iter=30, device=sp.cpu_device, output_eigenvalue=False): """Generate ESPIRiT maps from k-space. Currently only supports outputting one set of maps. Args: ksp (array): k-space array of shape [num_coils, n_ndim, ..., n_1] calib (tuple of ints): length-2 image shape. thresh (float): threshold for the calibration matrix. kernel_width (int): kernel width for the calibration matrix. max_power_iter (int): maximum number of power iterations. device (Device): computing device. crop (int): cropping threshold. Returns: array: ESPIRiT maps of the same shape as ksp. References: Martin Uecker, Peng Lai, Mark J. Murphy, Patrick Virtue, Michael Elad, John M. Pauly, Shreyas S. Vasanawala, and Michael Lustig ESPIRIT - An Eigenvalue Approach to Autocalibrating Parallel MRI: Where SENSE meets GRAPPA. Magnetic Resonance in Medicine, 71:990-1001 (2014) """ img_ndim = ksp.ndim - 1 num_coils = len(ksp) with sp.get_device(ksp): # Get calibration region calib_shape = [num_coils] + [calib_width] * img_ndim calib = sp.resize(ksp, calib_shape) calib = sp.to_device(calib, device) device = sp.Device(device) xp = device.xp with device: # Get calibration matrix kernel_shape = [num_coils] + [kernel_width] * img_ndim kernel_strides = [1] * (img_ndim + 1) mat = sp.array_to_blocks(calib, kernel_shape, kernel_strides) mat = mat.reshape([-1, sp.prod(kernel_shape)]) # Perform SVD on calibration matrix _, S, VH = xp.linalg.svd(mat, full_matrices=False) VH = VH[S > thresh * S.max(), :] # Get kernels num_kernels = len(VH) kernels = VH.reshape([num_kernels] + kernel_shape) img_shape = ksp.shape[1:] # Get covariance matrix in image domain AHA = xp.zeros(img_shape[::-1] + (num_coils, num_coils), dtype=ksp.dtype) for kernel in kernels: img_kernel = sp.ifft(sp.resize(kernel, ksp.shape), axes=range(-img_ndim, 0)) aH = xp.expand_dims(img_kernel.T, axis=-1) a = xp.conj(aH.swapaxes(-1, -2)) AHA += aH @ a AHA *= (sp.prod(img_shape) / kernel_width**img_ndim) # Power Iteration to compute top eigenvector mps = xp.ones(ksp.shape[::-1] + (1, ), dtype=ksp.dtype) for _ in range(max_power_iter): sp.copyto(mps, AHA @ mps) eig_value = xp.sum(xp.abs(mps)**2, axis=-2, keepdims=True)**0.5 mps /= eig_value # Normalize phase with respect to first channel mps = mps.T[0] mps *= xp.conj(mps[0] / xp.abs(mps[0])) # Crop maps by thresholding eigenvalue eig_value = eig_value.T[0] mps *= eig_value > crop if output_eigenvalue: return mps, eig_value else: return mps