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
