def nufft_adjoint(input, coord, oshape=None, oversamp=1.25, width=4):
"""Adjoint non-uniform Fast Fourier Transform.
Args:
input (array): input Fourier domain array of shape
(...) + coord.shape[:-1]. That is, the last dimensions
of input must match the first dimensions of coord.
The nufft_adjoint is applied on the last coord.ndim - 1 axes,
and looped over the remaining axes.
coord (array): Fourier domain coordinate array of shape (..., ndim).
ndim determines the number of dimension to apply nufft adjoint.
coord[..., i] should be scaled to have its range between
-n_i // 2, and n_i // 2.
oshape (tuple of ints): output shape of the form
(..., n_{ndim - 1}, ..., n_1, n_0).
oversamp (float): oversampling factor.
width (float): interpolation kernel full-width in terms of
oversampled grid.
n (int): number of sampling points of the interpolation kernel.
Returns:
array: signal domain array with shape specified by oshape.
See Also:
:func:`sigpy.nufft.nufft`
"""
ndim = coord.shape[-1]
beta = np.pi * (((width / oversamp) * (oversamp - 0.5))**2 - 0.8)**0.5
if oshape is None:
oshape = list(input.shape[:-coord.ndim + 1]) + estimate_shape(coord)
else:
oshape = list(oshape)
os_shape = _get_oversamp_shape(oshape, ndim, oversamp)
# Gridding
coord = _scale_coord(coord, oshape, oversamp)
output = interp.gridding(input,
coord,
os_shape,
kernel='kaiser_bessel',
width=width,
param=beta)
output /= width**ndim
# IFFT
output = ifft(output, axes=range(-ndim, 0), norm=None)
# Crop
output = util.resize(output, oshape)
output *= util.prod(os_shape[-ndim:]) / util.prod(oshape[-ndim:])**0.5
# Apodize
_apodize(output, ndim, oversamp, width, beta)
return output
def _get_convolve_adjoint_input_params(W, y, input_multi_channel,
output_multi_channel):
ndim = W.ndim - input_multi_channel - output_multi_channel
output_shape = y.shape[-ndim:]
filter_shape = W.shape[-ndim:]
batch_shape = y.shape[:-ndim - output_multi_channel]
batch_size = util.prod(batch_shape)
if y.dtype != W.dtype:
raise TypeError(
'y and W must have the same dtype, got {} and {}.'.format(
y.dtype, W.dtype))
if backend.get_device(y) != backend.get_device(W):
raise TypeError(
'y and W must be on the same device, got {} and {}.'.format(
backend.get_device(y), backend.get_device(W)))
if input_multi_channel:
input_channel = W.shape[-ndim - 1]
else:
input_channel = 1
if output_multi_channel:
output_channel = y.shape[-ndim - 1]
else:
output_channel = 1
return ndim, output_shape, filter_shape, batch_shape, \
batch_size, input_channel, output_channel
def _get_convolve_adjoint_filter_params(x, y, ndim, input_multi_channel,
output_multi_channel):
output_shape = y.shape[-ndim:]
input_shape = x.shape[-ndim:]
batch_shape = y.shape[:-ndim - output_multi_channel]
batch_size = util.prod(batch_shape)
if x.dtype != y.dtype:
raise TypeError(
'x and y must have the same dtype, got {} and {}.'.format(
x.dtype, y.dtype))
if backend.get_device(y) != backend.get_device(x):
raise TypeError(
'y and x must be on the same device, got {} and {}.'.format(
backend.get_device(y), backend.get_device(x)))
if input_multi_channel:
input_channel = x.shape[-ndim - 1]
else:
input_channel = 1
if output_multi_channel:
output_channel = y.shape[-ndim - 1]
else:
output_channel = 1
return input_shape, output_shape, batch_shape, \
batch_size, input_channel, output_channel
def nufft(input, coord, oversamp=1.25, width=4):
"""Non-uniform Fast Fourier Transform.
Args:
input (array): input signal domain array of shape
(..., n_{ndim - 1}, ..., n_1, n_0),
where ndim is specified by coord.shape[-1]. The nufft
is applied on the last ndim axes, and looped over
the remaining axes.
coord (array): Fourier domain coordinate array of shape (..., ndim).
ndim determines the number of dimensions to apply the nufft.
coord[..., i] should be scaled to have its range between
-n_i // 2, and n_i // 2.
oversamp (float): oversampling factor.
width (float): interpolation kernel full-width in terms of
oversampled grid.
n (int): number of sampling points of the interpolation kernel.
Returns:
array: Fourier domain data of shape
input.shape[:-ndim] + coord.shape[:-1].
References:
Fessler, J. A., & Sutton, B. P. (2003).
Nonuniform fast Fourier transforms using min-max interpolation
IEEE Transactions on Signal Processing, 51(2), 560-574.
Beatty, P. J., Nishimura, D. G., & Pauly, J. M. (2005).
Rapid gridding reconstruction with a minimal oversampling ratio.
IEEE transactions on medical imaging, 24(6), 799-808.
"""
ndim = coord.shape[-1]
beta = np.pi * (((width / oversamp) * (oversamp - 0.5))**2 - 0.8)**0.5
os_shape = _get_oversamp_shape(input.shape, ndim, oversamp)
output = input.copy()
# Apodize
_apodize(output, ndim, oversamp, width, beta)
# Zero-pad
output /= util.prod(input.shape[-ndim:])**0.5
output = util.resize(output, os_shape)
# FFT
output = fft(output, axes=range(-ndim, 0), norm=None)
# Interpolate
coord = _scale_coord(coord, input.shape, oversamp)
output = interp.interpolate(output,
coord,
kernel='kaiser_bessel',
width=width,
param=beta)
output /= width**ndim
return output
def __init__(self, proxs):
self.nops = len(proxs)
assert(self.nops > 0)
self.proxs = proxs
self.shapes = [prox.shape for prox in proxs]
shape = [sum(util.prod(prox.shape) for prox in proxs)]
super().__init__(shape)
def _fft_convolve_adjoint_filter(x, y, mode='full'):
ndim = x.ndim - 2
batch_size = len(x)
output_channel = y.shape[1]
input_channel = x.shape[1]
output_shape = y.shape[-ndim:]
input_shape = x.shape[-ndim:]
if mode == 'full':
filter_shape = tuple(p - m + 1
for m, p in zip(input_shape, output_shape))
pad_shape = output_shape
elif mode == 'valid':
filter_shape = tuple(m - p + 1
for m, p in zip(input_shape, output_shape))
pad_shape = input_shape
dtype = x.dtype
device = backend.get_device(x)
xp = device.xp
with device:
x = xp.conj(util.flip(x, axes=range(-ndim, 0)))
x = x.reshape((batch_size, 1, input_channel) + input_shape)
y = y.reshape((batch_size, output_channel, 1) + output_shape)
x_pad = util.resize(x, (batch_size, 1, input_channel) + pad_shape,
oshift=[0] * x.ndim)
y_pad = util.resize(y, (batch_size, output_channel, 1) + pad_shape,
oshift=[0] * y.ndim)
if np.issubdtype(dtype, np.floating):
x_fft = xp.fft.rfftn(x_pad, axes=range(-ndim, 0), norm='ortho')
y_fft = xp.fft.rfftn(y_pad, axes=range(-ndim, 0), norm='ortho')
W_fft = xp.sum(x_fft * y_fft, axis=0)
W = xp.fft.irfftn(W_fft,
pad_shape,
axes=range(-ndim, 0),
norm='ortho').astype(dtype)
else:
x_fft = fourier.fft(x_pad, axes=range(-ndim, 0), center=False)
y_fft = fourier.fft(y_pad, axes=range(-ndim, 0), center=False)
W_fft = xp.sum(x_fft * y_fft, axis=0)
W = fourier.ifft(W_fft, axes=range(-ndim, 0), center=False)
if mode == 'full':
shift = [0, 0] + [m - 1 for m in input_shape]
elif mode == 'valid':
shift = [0, 0] + [p - 1 for p in output_shape]
W = util.resize(W, (output_channel, input_channel) + filter_shape,
ishift=shift)
W *= util.prod(pad_shape)**0.5
return W
def gridding(input, shape, width, kernel, coord):
"""Gridding of points specified by coordinates to array.
Args:
input (array): Input array.
shape (array of ints): Output shape.
width (float): Interpolation kernel width.
kernel (array): Interpolation kernel.
coord (array): Coordinate array of shape [..., ndim]
Returns:
output (array): Output array.
"""
ndim = coord.shape[-1]
batch_shape = shape[:-ndim]
batch_size = util.prod(batch_shape)
pts_shape = coord.shape[:-1]
npts = util.prod(pts_shape)
device = backend.get_device(input)
xp = device.xp
isreal = np.issubdtype(input.dtype, np.floating)
with device:
input = input.reshape([batch_size, npts])
coord = coord.reshape([npts, ndim])
output = xp.zeros([batch_size] + list(shape[-ndim:]),
dtype=input.dtype)
_gridding = _select_gridding(ndim, npts, device, isreal)
if device == backend.cpu_device:
_gridding(output, input, width, kernel, coord)
else: # pragma: no cover
_gridding(input, width, kernel, coord, output, size=npts)
return output.reshape(shape)
def interpolate(input, width, kernel, coord):
"""Interpolation from array to points specified by coordinates.
Args:
input (array): Input array of shape [..., ny, nx]
width (float): Interpolation kernel width.
kernel (array): Interpolation kernel.
coord (array): Coordinate array of shape [..., ndim]
Returns:
output (array): Output array of coord.shape[:-1]
"""
ndim = coord.shape[-1]
batch_shape = input.shape[:-ndim]
batch_size = util.prod(batch_shape)
pts_shape = coord.shape[:-1]
npts = util.prod(pts_shape)
device = backend.get_device(input)
xp = device.xp
isreal = np.issubdtype(input.dtype, np.floating)
coord = backend.to_device(coord, device)
kernel = backend.to_device(kernel, device)
with device:
input = input.reshape([batch_size] + list(input.shape[-ndim:]))
coord = coord.reshape([npts, ndim])
output = xp.zeros([batch_size, npts], dtype=input.dtype)
_interpolate = _select_interpolate(ndim, npts, device, isreal)
if device == backend.cpu_device:
_interpolate(output, input, width, kernel, coord)
else: # pragma: no cover
_interpolate(input, width, kernel, coord, output, size=npts)
return output.reshape(batch_shape + pts_shape)
def _fft_convolve_adjoint_input(W, y, mode='full'):
ndim = y.ndim - 2
batch_size = len(y)
output_channel, input_channel = W.shape[:2]
output_shape = y.shape[-ndim:]
filter_shape = W.shape[-ndim:]
if mode == 'full':
input_shape = tuple(p - n + 1
for p, n in zip(output_shape, filter_shape))
pad_shape = output_shape
elif mode == 'valid':
input_shape = tuple(p + n - 1
for p, n in zip(output_shape, filter_shape))
pad_shape = input_shape
dtype = y.dtype
device = backend.get_device(y)
xp = device.xp
with device:
y = y.reshape((batch_size, output_channel, 1) + output_shape)
W = xp.conj(util.flip(W, axes=range(-ndim, 0)))
y_pad = util.resize(y, (batch_size, output_channel, 1) + pad_shape,
oshift=[0] * y.ndim)
W_pad = util.resize(W, (output_channel, input_channel) + pad_shape,
oshift=[0] * W.ndim)
if np.issubdtype(dtype, np.floating):
y_fft = xp.fft.rfftn(y_pad, axes=range(-ndim, 0), norm='ortho')
W_fft = xp.fft.rfftn(W_pad, axes=range(-ndim, 0), norm='ortho')
x_fft = xp.sum(y_fft * W_fft, axis=-ndim - 2)
x = xp.fft.irfftn(x_fft,
pad_shape,
axes=range(-ndim, 0),
norm='ortho').astype(dtype)
else:
y_fft = fourier.fft(y_pad, axes=range(-ndim, 0), center=False)
W_fft = fourier.fft(W_pad, axes=range(-ndim, 0), center=False)
x_fft = xp.sum(y_fft * W_fft, axis=-ndim - 2)
x = fourier.ifft(x_fft, axes=range(-ndim, 0), center=False)
if mode == 'full':
shift = [0, 0] + [n - 1 for n in filter_shape]
elif mode == 'valid':
shift = [0] * x.ndim
x = util.resize(x, (batch_size, input_channel) + input_shape,
ishift=shift)
x *= util.prod(pad_shape)**0.5
return x
def _fft_convolve(x, W, mode='full'):
ndim = x.ndim - 2
batch_size = len(x)
output_channel, input_channel = W.shape[:2]
input_shape = x.shape[-ndim:]
filter_shape = W.shape[-ndim:]
if mode == 'full':
output_shape = tuple(m + n - 1
for m, n in zip(input_shape, filter_shape))
pad_shape = output_shape
elif mode == 'valid':
output_shape = tuple(m - n + 1
for m, n in zip(input_shape, filter_shape))
pad_shape = input_shape
dtype = x.dtype
device = backend.get_device(x)
xp = device.xp
with device:
x = x.reshape((batch_size, 1, input_channel) + input_shape)
x_pad = util.resize(x, (batch_size, 1, input_channel) + pad_shape,
oshift=[0] * x.ndim)
W_pad = util.resize(W, (output_channel, input_channel) + pad_shape,
oshift=[0] * W.ndim)
if np.issubdtype(dtype, np.floating):
x_fft = xp.fft.rfftn(x_pad, axes=range(-ndim, 0), norm='ortho')
W_fft = xp.fft.rfftn(W_pad, axes=range(-ndim, 0), norm='ortho')
y_fft = xp.sum(x_fft * W_fft, axis=-ndim - 1)
y = xp.fft.irfftn(y_fft,
pad_shape,
axes=range(-ndim, 0),
norm='ortho').astype(dtype)
else:
x_fft = fourier.fft(x_pad, axes=range(-ndim, 0), center=False)
W_fft = fourier.fft(W_pad, axes=range(-ndim, 0), center=False)
y_fft = xp.sum(x_fft * W_fft, axis=-ndim - 1)
y = fourier.ifft(y_fft, axes=range(-ndim, 0), center=False)
if mode == 'full':
shift = [0] * y.ndim
elif mode == 'valid':
shift = [0, 0] + [n - 1 for n in filter_shape]
y = util.resize(y, (batch_size, output_channel) + output_shape,
ishift=shift)
y *= util.prod(pad_shape)**0.5
return y
def _get_convolve_adjoint_filter_params(x, y, ndim, input_multi_channel,
output_multi_channel):
output_shape = y.shape[-ndim:]
input_shape = x.shape[-ndim:]
batch_shape = y.shape[:-ndim - output_multi_channel]
batch_size = util.prod(batch_shape)
if input_multi_channel:
input_channel = x.shape[-ndim - 1]
else:
input_channel = 1
if output_multi_channel:
output_channel = y.shape[-ndim - 1]
else:
output_channel = 1
return input_shape, output_shape, batch_shape, \
batch_size, input_channel, output_channel
def _get_convolve_params(data_shape, filt_shape, mode, strides, multi_channel):
D = len(filt_shape) - 2 * multi_channel
m = tuple(data_shape[-D:])
n = tuple(filt_shape[-D:])
b = tuple(data_shape[:-D - multi_channel])
B = util.prod(b)
if multi_channel:
if filt_shape[-D - 1] != data_shape[-D - 1]:
raise ValueError('Data channel mismatch, '
'got {} from data and {} from filt.'.format(
data_shape[-D - 1], filt_shape[-D - 1]))
c_i = filt_shape[-D - 1]
c_o = filt_shape[-D - 2]
else:
c_i = 1
c_o = 1
if strides is None:
s = (1, ) * D
else:
if len(strides) != D:
raise ValueError('Strides must have length {}.'.format(D))
s = tuple(strides)
if mode == 'full':
p = tuple(
(m_d + n_d - 1 + s_d - 1) // s_d for m_d, n_d, s_d in zip(m, n, s))
elif mode == 'valid':
if (any(m_d >= n_d for m_d, n_d in zip(m, n))
and any(m_d < n_d for m_d, n_d in zip(m, n))):
raise ValueError('In valid mode, either data or filter must be '
'at least as large as the other in every axis.')
p = tuple(
(m_d - n_d + 1 + s_d - 1) // s_d for m_d, n_d, s_d in zip(m, n, s))
else:
raise ValueError('Invalid mode, got {}'.format(mode))
return D, b, B, m, n, s, c_i, c_o, p
def _get_convolve_adjoint_input_params(W, y, input_multi_channel,
output_multi_channel):
ndim = W.ndim - input_multi_channel - output_multi_channel
output_shape = y.shape[-ndim:]
filter_shape = W.shape[-ndim:]
batch_shape = y.shape[:-ndim - output_multi_channel]
batch_size = util.prod(batch_shape)
if input_multi_channel:
input_channel = W.shape[-ndim - 1]
else:
input_channel = 1
if output_multi_channel:
output_channel = y.shape[-ndim - 1]
else:
output_channel = 1
return ndim, output_shape, filter_shape, batch_shape, \
batch_size, input_channel, output_channel
def elitist_thresh(lamda, input, axes=None):
"""Elitist threshold.
Solves for
:math::
\text{argmin}_x \| x - y \|_2^2 + \lambda \| x \|_1^2
Args:
lamda (float, or array): Threshold parameter.
input (array): Input array.
axes (None or tuple of ints): Axes to perform threshold.
Returns:
array: Result.
References:
Kowalski, M. 2009. Sparse regression using mixed norms.
"""
shape = input.shape
axes = util._normalize_axes(axes, input.ndim)
remain_axes = tuple(set(range(input.ndim)) - set(axes))
length = util.prod([shape[a] for a in axes])
batch = input.size // length
input = input.transpose(remain_axes + axes)
input = input.reshape([batch, length])
thresh = find_elitist_thresh(lamda, input)
output = soft_thresh(thresh, input)
output = output.reshape([shape[a] for a in remain_axes + axes])
output = output.transpose(np.argsort(remain_axes + axes))
return output
def _vstack_params(shapes, axis):
if axis is None:
return _vstack_params([[util.prod(shape)] for shape in shapes], 0)
oshape = list(shapes[0])
ndim = len(oshape)
idx = shapes[0][axis]
indices = []
for shape in shapes[1:]:
if len(shape) != ndim:
raise Exception(
'Shapes must have the same lengths to concatenate.')
for i in range(ndim):
if i == axis:
oshape[i] += shape[i]
indices.append(idx)
idx += shape[i]
elif shape[i] != oshape[i]:
raise Exception(
'Shapes not along axis must be the same to concatenate.')
return oshape, indices
def blocks_to_array(input, oshape, blk_shape, blk_strides):
"""Accumulate blocks into an array in a sliding window manner.
Args:
input (array): input array of shape [...] + num_blks + blk_shape
oshape (tuple): output shape.
blk_shape (tuple): block shape of length D.
blk_strides (tuple): block strides of length D.
Returns:
array: array of shape oshape.
"""
if len(blk_shape) != len(blk_strides):
raise ValueError('blk_shape must have the same length as blk_strides.')
D = len(blk_shape)
num_blks = input.shape[-(2 * D):-D]
batch_shape = list(oshape[:-D])
batch_size = util.prod(batch_shape)
device = backend.get_device(input)
xp = device.xp
with device:
output = xp.zeros([batch_size] + list(oshape[-D:]), dtype=input.dtype)
input = input.reshape([batch_size] + list(input.shape[-2 * D:]))
if D == 1:
if device == backend.cpu_device:
_blocks_to_array1(output, input, batch_size, blk_shape[-1],
blk_strides[-1], num_blks[-1])
else: # pragma: no cover
if np.issubdtype(input.dtype, np.floating):
_blocks_to_array1_cuda(input,
batch_size,
blk_shape[-1],
blk_strides[-1],
num_blks[-1],
output,
size=batch_size * num_blks[-1] *
blk_shape[-1])
else:
_blocks_to_array1_cuda_complex(
input,
batch_size,
blk_shape[-1],
blk_strides[-1],
num_blks[-1],
output,
size=batch_size * num_blks[-1] * blk_shape[-1])
elif D == 2:
if device == backend.cpu_device:
_blocks_to_array2(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_strides[-1],
blk_strides[-2], num_blks[-1], num_blks[-2])
else: # pragma: no cover
if np.issubdtype(input.dtype, np.floating):
_blocks_to_array2_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_strides[-1],
blk_strides[-2],
num_blks[-1],
num_blks[-2],
output,
size=batch_size * num_blks[-1] *
num_blks[-2] * blk_shape[-1] *
blk_shape[-2])
else: # pragma: no cover
_blocks_to_array2_cuda_complex(
input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_strides[-1],
blk_strides[-2],
num_blks[-1],
num_blks[-2],
output,
size=batch_size * num_blks[-1] * num_blks[-2] *
blk_shape[-1] * blk_shape[-2])
elif D == 3:
if device == backend.cpu_device:
_blocks_to_array3(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3],
blk_strides[-1], blk_strides[-2],
blk_strides[-3], num_blks[-1], num_blks[-2],
num_blks[-3])
else: # pragma: no cover
if np.issubdtype(input.dtype, np.floating):
_blocks_to_array3_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
num_blks[-1],
num_blks[-2],
num_blks[-3],
output,
size=batch_size * num_blks[-1] *
num_blks[-2] * num_blks[-3] *
blk_shape[-1] * blk_shape[-2] *
blk_shape[-3])
else:
_blocks_to_array3_cuda_complex(
input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
num_blks[-1],
num_blks[-2],
num_blks[-3],
output,
size=batch_size * num_blks[-1] * num_blks[-2] *
num_blks[-3] * blk_shape[-1] * blk_shape[-2] *
blk_shape[-3])
elif D == 4:
if device == backend.cpu_device:
_blocks_to_array4(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3], blk_shape[-4],
blk_strides[-1], blk_strides[-2],
blk_strides[-3], blk_strides[-4],
num_blks[-1], num_blks[-2], num_blks[-3],
num_blks[-4])
else: # pragma: no cover
if np.issubdtype(input.dtype, np.floating):
_blocks_to_array4_cuda(
input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_shape[-4],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
blk_strides[-4],
num_blks[-1],
num_blks[-2],
num_blks[-3],
num_blks[-4],
output,
size=batch_size * num_blks[-1] * num_blks[-2] *
num_blks[-3] * num_blks[-4] * blk_shape[-1] *
blk_shape[-2] * blk_shape[-3] * blk_shape[-4])
else:
_blocks_to_array4_cuda_complex(
input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_shape[-4],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
blk_strides[-4],
num_blks[-1],
num_blks[-2],
num_blks[-3],
num_blks[-4],
output,
size=batch_size * num_blks[-1] * num_blks[-2] *
num_blks[-3] * num_blks[-4] * blk_shape[-1] *
blk_shape[-2] * blk_shape[-3] * blk_shape[-4])
else:
raise ValueError('Only support D <= 4, got {}'.format(D))
return output.reshape(oshape)
def array_to_blocks(input, blk_shape, blk_strides):
"""Extract blocks from an array in a sliding window manner.
Args:
input (array): input array of shape [..., N_1, ..., N_D]
blk_shape (tuple): block shape of length D, with D <= 4.
blk_strides (tuple): block strides of length D.
Returns:
array: array of shape [...] + num_blks + blk_shape, where
num_blks = (N - blk_shape + blk_strides) // blk_strides.
Example:
>>> input = np.array([0, 1, 2, 3, 4, 5])
>>> print(array_to_blocks(input, [2], [2]))
[[0, 1],
[2, 3],
[4, 5]]
"""
if len(blk_shape) != len(blk_strides):
raise ValueError('blk_shape must have the same length as blk_strides.')
D = len(blk_shape)
num_blks = [(i - b + s) // s
for i, b, s in zip(input.shape[-D:], blk_shape, blk_strides)]
batch_shape = list(input.shape[:-D])
batch_size = util.prod(batch_shape)
device = backend.get_device(input)
xp = device.xp
with device:
output = xp.zeros([batch_size] + num_blks + blk_shape,
dtype=input.dtype)
input = input.reshape([batch_size] + list(input.shape[-D:]))
if D == 1:
if device == backend.cpu_device:
_array_to_blocks1(output, input, batch_size, blk_shape[-1],
blk_strides[-1], num_blks[-1])
else: # pragma: no cover
_array_to_blocks1_cuda(input,
batch_size,
blk_shape[-1],
blk_strides[-1],
num_blks[-1],
output,
size=batch_size * num_blks[-1] *
blk_shape[-1])
elif D == 2:
if device == backend.cpu_device:
_array_to_blocks2(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_strides[-1],
blk_strides[-2], num_blks[-1], num_blks[-2])
else: # pragma: no cover
_array_to_blocks2_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_strides[-1],
blk_strides[-2],
num_blks[-1],
num_blks[-2],
output,
size=batch_size * num_blks[-1] *
num_blks[-2] * blk_shape[-1] *
blk_shape[-2])
elif D == 3:
if device == backend.cpu_device:
_array_to_blocks3(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3],
blk_strides[-1], blk_strides[-2],
blk_strides[-3], num_blks[-1], num_blks[-2],
num_blks[-3])
else: # pragma: no cover
_array_to_blocks3_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
num_blks[-1],
num_blks[-2],
num_blks[-3],
output,
size=batch_size * num_blks[-1] *
num_blks[-2] * num_blks[-3] *
blk_shape[-1] * blk_shape[-2] *
blk_shape[-3])
elif D == 4:
if device == backend.cpu_device:
_array_to_blocks4(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3], blk_shape[-4],
blk_strides[-1], blk_strides[-2],
blk_strides[-3], blk_strides[-4],
num_blks[-1], num_blks[-2], num_blks[-3],
num_blks[-4])
else: # pragma: no cover
_array_to_blocks4_cuda(
input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_shape[-4],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
blk_strides[-4],
num_blks[-1],
num_blks[-2],
num_blks[-3],
num_blks[-4],
output,
size=batch_size * num_blks[-1] * num_blks[-2] *
num_blks[-3] * num_blks[-4] * blk_shape[-1] *
blk_shape[-2] * blk_shape[-3] * blk_shape[-4])
else:
raise ValueError('Only support D <= 4, got {}'.format(D))
return output.reshape(batch_shape + num_blks + blk_shape)
def gridding(input, coord, shape, kernel="spline", width=2, param=1):
r"""Gridding of points specified by coordinates to array.
Let :math:`y` be the input, :math:`x` be the output,
:math:`c` be the coordinates, :math:`W` be the kernel width,
and :math:`K` be the interpolation kernel, then the function computes,
.. math ::
x[i] = \sum_{j : \| i - c[j] \|_\infty \leq W / 2}
K\left(\frac{i - c[j]}{W / 2}\right) y[j]
There are two types of kernels: 'spline' and 'kaiser_bessel'.
'spline' uses the cardinal B-spline functions as kernels.
The order of the spline can be specified using param.
For example, param=1 performs linear interpolation.
Concretely, for param=0, :math:`K(x) = 1`,
for param=1, :math:`K(x) = 1 - |x|`, and
for param=2, :math:`K(x) = \frac{9}{8} (1 - |x|)^2`
for :math:`|x| > \frac{1}{3}`
and :math:`K(x) = \frac{3}{4} (1 - 3 x^2)` for :math:`|x| < \frac{1}{3}`.
These function expressions are derived from the reference wikipedia
page by shifting and scaling the range to -1 to 1.
When the coordinates specifies a uniformly spaced grid,
it is recommended to use the original scaling with width=param + 1
so that the interpolation weights add up to one.
'kaiser_bessel' uses the Kaiser-Bessel function as kernel.
Concretely, :math:`K(x) = I_0(\beta \sqrt{1 - x^2})`,
where :math:`I_0` is the modified Bessel function of the first kind.
The beta parameter can be specified with param.
The modified Bessel function of the first kind is approximated
using the power series, following the reference.
Args:
input (array): Input array.
coord (array): Coordinate array of shape [..., ndim]
width (float or tuple of floats): Interpolation kernel full-width.
kernel (str): Interpolation kernel, {"spline", "kaiser_bessel"}.
param (float or tuple of floats): Kernel parameter.
Returns:
output (array): Output array.
References:
https://en.wikipedia.org/wiki/Spline_wavelet#Cardinal_B-splines_of_small_orders
http://people.math.sfu.ca/~cbm/aands/page_378.htm
"""
ndim = coord.shape[-1]
batch_shape = shape[:-ndim]
batch_size = util.prod(batch_shape)
pts_shape = coord.shape[:-1]
npts = util.prod(pts_shape)
xp = backend.get_array_module(input)
isreal = np.issubdtype(input.dtype, np.floating)
input = input.reshape([batch_size, npts])
coord = coord.reshape([npts, ndim])
output = xp.zeros([batch_size] + list(shape[-ndim:]), dtype=input.dtype)
if np.isscalar(param):
param = xp.array([param] * ndim, coord.dtype)
else:
param = xp.array(param, coord.dtype)
if np.isscalar(width):
width = xp.array([width] * ndim, coord.dtype)
else:
width = xp.array(width, coord.dtype)
if xp == np:
_gridding[kernel][ndim - 1](output, input, coord, width, param)
else: # pragma: no cover
if isreal:
_gridding_cuda[kernel][ndim - 1](
input, coord, width, param, output, size=npts)
else:
_gridding_cuda_complex[kernel][ndim - 1](
input, coord, width, param, output, size=npts)
return output.reshape(shape)
def blocks_to_array(input, oshape, blk_shape, blk_strides):
"""Accumulate blocks into an array in a sliding window manner.
Args:
input (array): input array of shape [...] + num_blks + blk_shape
oshape (tuple): output shape.
blk_shape (tuple): block shape of length ndim.
blk_strides (tuple): block strides of length ndim.
Returns:
array: array of shape oshape.
"""
if len(blk_shape) != len(blk_strides):
raise ValueError('blk_shape must have the same length as blk_strides.')
ndim = len(blk_shape)
num_blks = input.shape[-(2 * ndim):-ndim]
batch_shape = list(oshape[:-ndim])
batch_size = util.prod(batch_shape)
xp = backend.get_array_module(input)
output = xp.zeros([batch_size] + list(oshape[-ndim:]), dtype=input.dtype)
input = input.reshape([batch_size] + list(input.shape[-2 * ndim:]))
if ndim == 1:
if xp == np:
_blocks_to_array1(output, input, batch_size, blk_shape[-1],
blk_strides[-1], num_blks[-1])
else: # pragma: no cover
_blocks_to_array1_cuda(input,
batch_size,
blk_shape[-1],
blk_strides[-1],
num_blks[-1],
output,
size=output.size)
elif ndim == 2:
if xp == np:
_blocks_to_array2(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_strides[-1], blk_strides[-2],
num_blks[-1], num_blks[-2])
else: # pragma: no cover
_blocks_to_array2_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_strides[-1],
blk_strides[-2],
num_blks[-1],
num_blks[-2],
output,
size=output.size)
elif ndim == 3:
if xp == np:
_blocks_to_array3(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3], blk_strides[-1],
blk_strides[-2], blk_strides[-3], num_blks[-1],
num_blks[-2], num_blks[-3])
else: # pragma: no cover
_blocks_to_array3_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
num_blks[-1],
num_blks[-2],
num_blks[-3],
output,
size=output.size)
else:
raise ValueError('Only support ndim=1, 2, or 3, got {}'.format(ndim))
return output.reshape(oshape)
def array_to_blocks(input, blk_shape, blk_strides):
"""Extract blocks from an array in a sliding window manner.
Args:
input (array): input array of shape [..., N_1, ..., N_ndim]
blk_shape (tuple): block shape of length ndim, with ndim={1, 2, 3}.
blk_strides (tuple): block strides of length ndim.
Returns:
array: array of shape [...] + num_blks + blk_shape, where
num_blks = (N - blk_shape + blk_strides) // blk_strides.
Example:
>>> input = np.array([0, 1, 2, 3, 4, 5])
>>> print(array_to_blocks(input, [2], [2]))
[[0, 1],
[2, 3],
[4, 5]]
"""
if len(blk_shape) != len(blk_strides):
raise ValueError('blk_shape must have the same length as blk_strides.')
ndim = len(blk_shape)
num_blks = [(i - b + s) // s
for i, b, s in zip(input.shape[-ndim:], blk_shape, blk_strides)
]
batch_shape = list(input.shape[:-ndim])
batch_size = util.prod(batch_shape)
xp = backend.get_array_module(input)
output = xp.zeros([batch_size] + num_blks + list(blk_shape),
dtype=input.dtype)
input = input.reshape([batch_size] + list(input.shape[-ndim:]))
if ndim == 1:
if xp == np:
_array_to_blocks1(output, input, batch_size, blk_shape[-1],
blk_strides[-1], num_blks[-1])
else: # pragma: no cover
_array_to_blocks1_cuda(input,
batch_size,
blk_shape[-1],
blk_strides[-1],
num_blks[-1],
output,
size=output.size)
elif ndim == 2:
if xp == np:
_array_to_blocks2(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_strides[-1], blk_strides[-2],
num_blks[-1], num_blks[-2])
else: # pragma: no cover
_array_to_blocks2_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_strides[-1],
blk_strides[-2],
num_blks[-1],
num_blks[-2],
output,
size=output.size)
elif ndim == 3:
if xp == np:
_array_to_blocks3(output, input, batch_size, blk_shape[-1],
blk_shape[-2], blk_shape[-3], blk_strides[-1],
blk_strides[-2], blk_strides[-3], num_blks[-1],
num_blks[-2], num_blks[-3])
else: # pragma: no cover
_array_to_blocks3_cuda(input,
batch_size,
blk_shape[-1],
blk_shape[-2],
blk_shape[-3],
blk_strides[-1],
blk_strides[-2],
blk_strides[-3],
num_blks[-1],
num_blks[-2],
num_blks[-3],
output,
size=output.size)
else:
raise ValueError('Only support ndim=1, 2, or 3, got {}'.format(ndim))
return output.reshape(batch_shape + num_blks + list(blk_shape))