def divergence_ravel_offsets(
g,
direction="forward",
deltas=None,
grad_axis="last",
offsets=None,
use_corners=False,
):
xp, on_gpu = get_array_module(g)
g = xp.asanyarray(g)
# Note: direction here is opposite of the corresponding gradient_periodic
if direction.lower() == "forward":
n_roll = 1
elif direction.lower() == "backward":
n_roll = -1
else:
raise ValueError("direction must be 'forward' or 'backward'")
otype = g.dtype.char
if otype not in ["f", "d", "F", "D", "m", "M"]:
otype = "d"
if grad_axis in [0, "first"]:
grad_axis = 0
fshape = g.shape[1:]
elif grad_axis in [-1, "last"]:
grad_axis = -1
fshape = g.shape[:-1]
else:
raise ValueError("Unsupported grad_axis: {}".format(grad_axis) +
"... must be first or last axis")
if offsets is None:
offsets = compute_offsets(fshape, use_corners)
n = len(offsets)
if grad_axis in [0, "first"]:
g = g.reshape((n, prod(fshape)), order="F")
elif grad_axis in [-1, "last"]:
g = g.reshape((prod(fshape), n), order="F")
f = xp.empty(g.shape, dtype=otype)
if deltas is not None:
deltas = np.asanyarray(deltas)
if len(deltas) != n:
raise ValueError("deltas array length must match f.ndim")
for n, off in enumerate(offsets):
if grad_axis == 0:
f[n, ...] = xp.roll(g[n, ...], n_roll * off, axis=0) - g[n, ...]
if deltas is not None:
f[n, ...] /= deltas[n]
elif grad_axis == -1:
f[..., n] = xp.roll(g[..., n], n_roll * off, axis=0) - g[..., n]
if deltas is not None:
f[..., n] /= deltas[n]
div = -f.sum(axis=grad_axis)
return div.reshape(fshape, order="F")
def block_inverse_transform(
x,
inverse_transform,
block_size=(8, 8),
reshape_input=False,
transform_kwargs={},
):
""" block-wise dctn
x.shape must be an integer multiple of by block_size
any axis where block_size = 1 will not be transformed
if reshape_output is false, it assumes the input is from block_dctn with
reshape_output=False.
"""
from skimage.util.shape import view_as_blocks
axes = np.where(np.asarray(block_size) > 1)[0]
if reshape_input:
ndim = x.ndim
if ndim != len(block_size):
raise ValueError("block_size must match size of x")
xview = view_as_blocks(x, block_size)
out_shape = x.shape
else:
ndim = x.ndim // 2
if ndim != len(block_size):
raise ValueError("block_size must match size of x")
xview = x
out_shape = np.asarray(x.shape[:ndim]) * np.asarray(x.shape[ndim:])
dct_axes = axes + ndim
y = inverse_transform(xview, axes=dct_axes, **transform_kwargs)
if not np.iscomplexobj(y):
out = np.zeros(out_shape)
else:
out = np.zeros(out_shape, np.result_type(y.dtype, np.complex64))
y_slices = [slice(None)] * (2 * ndim)
out_slices = [slice(None)] * ndim
y_axes = np.arange(ndim)
if prod(block_size) > prod(y.shape[:ndim]):
# size of each block exceeds the number of blocks
for ells in product(*([range(s) for s in y.shape[:ndim]])):
for n, ax, b in zip(ells, y_axes, block_size):
out_slices[ax] = slice(n * b, (n + 1) * b)
y_slices[ax] = slice(n, n + 1)
out[tuple(out_slices)] = y[tuple(y_slices)]
else:
# number of blocks exceeds the size of each block
for ells in product(*([range(s) for s in block_size])):
for n, ax, b in zip(ells, y_axes, block_size):
out_slices[ax] = slice(n, None, b)
y_slices[ndim + ax] = slice(n, n + 1)
out[tuple(out_slices)] = np.squeeze(y[tuple(y_slices)])
return out
def forward(self, x):
""" image gradient """
xp = self.xp_in
nreps = int(x.size / prod(self.arr_shape))
if nreps == 1:
x = x.reshape(self.arr_shape, order=self.order)
g = self.grad_func(
self._prior_subtract(x),
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
else:
if self.order == "C":
g = xp.zeros((nreps, ) + self.grad_shape, dtype=x.dtype)
for r in range(nreps):
xr = x[r, ...].reshape(self.arr_shape, order=self.order)
g[r, ...] = self.grad_func(
self._prior_subtract(xr),
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
else:
g = xp.zeros(self.grad_shape + (nreps, ), dtype=x.dtype)
for r in range(nreps):
xr = x[..., r].reshape(self.arr_shape, order=self.order)
g[..., r] = self.grad_func(
self._prior_subtract(xr),
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
return g
def adjoint(self, g):
""" image divergence """
xp = self.xp_in
nreps = int(g.size / prod(self.grad_shape))
# TODO: fix gradient_adjoint for N-dimensional case
# TODO: prior case-> add prior back?
# return gradient_adjoint(g, grad_axis='last')
if nreps == 1:
g = g.reshape(self.grad_shape, order=self.order)
d = self.div_func(
g,
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
else:
if self.order == "C":
d = xp.zeros((nreps, ) + self.arr_shape, dtype=g.dtype)
for r in range(nreps):
d[r, ...] = self.div_func(
g[r, ...],
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
else:
d = xp.zeros(self.arr_shape + (nreps, ), dtype=g.dtype)
for r in range(nreps):
d[..., r] = self.div_func(
g[..., r],
deltas=self.grid_size,
direction="forward",
grad_axis=self.grad_axis,
)
return -d # adjoint of grad is -div
def forward(self, x):
xp = self.xp
if self.force_real_image:
x = x.real
if self.im_mask is None:
size_1rep = self.nargin
else:
size_1rep = prod(self.shape_in)
if x.size < size_1rep:
raise ValueError("data, x, too small to transform.")
elif x.size == size_1rep:
nreps = 1
# 1D or single repetition nD
x = x.reshape(self.shape_in, order=self.order)
y = self._forward_single_rep(x)
else:
if self.order == "C":
nreps = x.shape[0]
shape_tmp = (nreps, ) + self.shape_in
else:
nreps = x.shape[-1]
shape_tmp = self.shape_in + (nreps, )
x = x.reshape(shape_tmp, order=self.order)
if self.sample_mask is not None:
# number of samples by number of repetitions
if self.order == "C":
shape_y = (nreps, self.nargout)
else:
shape_y = (self.nargout, nreps)
y = xp.zeros(
shape_y,
dtype=xp.result_type(x, np.complex64),
order=self.order,
)
else:
if self.order == "C":
shape_y = (nreps, ) + self.shape_in
else:
shape_y = self.shape_in + (nreps, )
y = xp.zeros(
shape_y,
dtype=xp.result_type(x, np.complex64),
order=self.order,
)
if self.order == "C":
for rep in range(nreps):
y[rep, ...] = self._forward_single_rep(
x[rep, ...]).reshape(y.shape[1:], order=self.order)
else:
for rep in range(nreps):
y[..., rep] = self._forward_single_rep(
x[..., rep]).reshape(y.shape[:-1], order=self.order)
if self.ortho:
y /= self.scale_ortho
if y.dtype != self.result_dtype:
y = y.astype(self.result_dtype)
return y
def compute_offsets(shape, use_corners=False):
ndim = len(shape)
if ndim == 1:
offsets = [1]
elif ndim == 2:
nx = shape[0]
strides = [1, nx]
if use_corners:
offsets = [1, nx, nx - 1, nx + 1]
else:
offsets = strides
elif ndim == 3:
nx, ny = shape[:2]
strides = [1, nx, nx * ny]
if not use_corners:
offsets = strides
else:
offsets = []
for xoff in [-1, 0, 1]:
for yoff in [-1, 0, 1]:
for zoff in [-1, 0, 1]:
if xoff == 0 and yoff == 0 and zoff == 0:
continue
offset = (xoff * strides[0] + yoff * strides[1] +
zoff * strides[2])
if offset < 0:
# can skip all negative offsets by symmetry
continue
offsets.append(offset)
offsets = np.sort(offsets).tolist()
elif ndim == 4:
nx, ny, nz = shape[:3]
strides = [1, nx, nx * ny, nx * ny * nz]
if not use_corners:
offsets = strides
else:
offsets = []
for xoff in [-1, 0, 1]:
for yoff in [-1, 0, 1]:
for zoff in [-1, 0, 1]:
for toff in [-1, 0, 1]:
if (xoff == 0 and yoff == 0 and zoff == 0
and toff == 0):
continue
offset = (xoff * strides[0] + yoff * strides[1] +
zoff * strides[2] + toff * strides[3])
if offset < 0:
# can skip all negative offsets by symmetry
continue
offsets.append(offset)
offsets = np.sort(offsets).tolist()
else:
if not use_corners:
offsets = [1]
for d in range(1, len(shape)):
offsets.append(prod(shape[:d]))
else:
raise ValueError(">4D currently unsupported for use_corners case")
return offsets
def block_transform(x,
transform,
block_size=(8, 8),
reshape_output=False,
transform_kwargs={}):
""" block-wise dctn
x.shape must be an integer multiple of by block_size
any axis where block_size = 1 will not be transformed
if reshape_output is false, the output will be shape
# x.shape//block_size + block_size
otherwise, the output is reshaped to match the shape of the input
"""
from skimage.util.shape import view_as_blocks
if x.ndim != len(block_size):
raise ValueError("block_size must match size of x")
ndim = x.ndim
block_axes = np.where(np.asarray(block_size) > 1)[0]
dct_axes = block_axes + ndim
xview = view_as_blocks(x, block_size)
y = transform(xview, axes=dct_axes, **transform_kwargs)
if not reshape_output:
return y
else:
# TODO: implement faster inverse of view_as_blocks
# tmp2 = idctn(tmp, axes=(0, 1))
out = np.zeros(x.shape)
y_axes = np.arange(ndim)
y_slices = [slice(None)] * (y.ndim)
out_slices = [slice(None)] * ndim
if prod(block_size) > prod(y.shape[:ndim]):
# size of each block exceeds the number of blocks
for ells in product(*([range(s) for s in y.shape[:ndim]])):
for n, ax, b in zip(ells, y_axes, block_size):
out_slices[ax] = slice(n * b, (n + 1) * b)
y_slices[ax] = slice(n, n + 1)
out[tuple(out_slices)] = y[tuple(y_slices)]
else:
# number of blocks exceeds the size of each block
for ells in product(*([range(s) for s in block_size])):
for n, ax, b in zip(ells, y_axes, block_size):
out_slices[ax] = slice(n, None, b)
y_slices[ndim + ax] = slice(n, n + 1)
out[tuple(out_slices)] = np.squeeze(y[tuple(y_slices)])
return out
def compute_Q_v2(G, copy_X=True):
"""Alternative version of compute_Q.
experimental: not recommended over compute_Q()
"""
from mrrt.nufft._nufft import nufft_adj
ones = np.ones(G.kspace.shape[0], G.Gnufft._cplx_dtype)
sf = np.sqrt(prod(G.Gnufft.Kd))
return sf * fftn(nufft_adj(G.Gnufft, ones, copy_X=True, return_psf=True))
def _norm(self, x):
# forward transform, immediately followed by inverse transform
# slightly faster than calling self.adjoint(self.forward(x))
xp = self.xp
if self.force_real_image:
x = x.real
if self.im_mask is None:
size_1rep = self.nargin
else:
size_1rep = prod(self.shape_in)
if x.size < size_1rep:
raise ValueError("data, x, too small to transform.")
elif x.size == size_1rep:
nreps = 1
# 1D or single repetition nD
x = x.reshape(self.shape_in, order=self.order)
y = self._norm_single_rep(x)
else:
nreps = x.size // size_1rep
if self.order == "C":
x = x.reshape((nreps, ) + self.shape_in, order=self.order)
y = xp.zeros_like(x)
for rep in range(nreps):
y[rep, ...] = self._norm_single_rep(x[rep, ...]).reshape(
y.shape[1:], order=self.order)
else:
x = x.reshape(self.shape_in + (nreps, ), order=self.order)
y = xp.zeros_like(x)
for rep in range(nreps):
y[..., rep] = self._norm_single_rep(x[..., rep]).reshape(
y.shape[:-1], order=self.order)
if y.dtype != self.result_dtype:
y = y.astype(self.result_dtype)
if not self.nd_input:
if self.squeeze_reps_in and (nreps == 1):
y = xp.ravel(y, order=self.order)
else:
if self.order == "C":
y = y.reshape((nreps, -1), order=self.order)
else:
y = y.reshape((-1, nreps), order=self.order)
return y
def __init__(
self,
arr_shape,
order="F",
arr_dtype=np.float32,
use_FFT_shifts=True,
ortho=False,
debug=False,
dct_axes=None,
fftshift_axes=None,
**kwargs,
):
"""
Parameters
----------
arr_shape : int
shape of the array
order : {'C','F'}, optional
array ordering that will be assumed if inputs/outputs need to be
reshaped
arr_dtype : numpy.dtype, optional
dtype for the array
use_FFT_shifts : bool, optional
If False, do not apply any FFT shifts
dct_axes : tuple or None, optional
Specify a subset of the axes to transform. The default is to
transform all axes.
fftshift_axes : tuple or None, optional
Specify a subset of the axes to fftshift. The default is to
shift all axes.
debug : bool, optional
Additional Parameters
---------------------
nd_input : bool, optional
nd_output : bool, optional
"""
if isinstance(arr_shape, (np.ndarray, list)):
# retrieve shape from array
arr_shape = tuple(arr_shape)
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.use_FFT_shifts = use_FFT_shifts
self.debug = debug
nargin = prod(arr_shape)
nargout = nargin
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
raise NotImplementedError("GPU version not implemented")
# can specify a subset of the axes to perform the FFT/FFTshifts over
self.dct_axes = dct_axes
if self.dct_axes is None:
self.dct_axes = tuple(np.arange(self.ndim))
if fftshift_axes is None:
self.fftshift_axes = self.dct_axes
else:
self.fftshift_axes = fftshift_axes
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.complex64)
self.ortho = ortho
if self.ortho:
# sqrt of product of shape along axes where FFT is performed
if self.dct_axes is None:
self.scale_ortho = sqrt(nargin)
else:
self.scale_ortho = sqrt(
prod(
np.asarray(self.arr_shape)[np.asarray(self.fft_axes)]))
else:
self.scale_ortho = None
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if nd_output:
shape_out = self.arr_shape
else:
shape_out = (nargout, 1)
shape_in = self.arr_shape
# mask_out = self.sample_mask
mask_in = kwargs.pop("mask_in", None)
mask_out = kwargs.pop("mask_out", None)
self.dctn = dctn
self.idctn = idctn
matvec = self.forward
matvec_adj = self.adjoint
super(DCT_Operator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=matvec,
matvec_transp=matvec_adj,
matvec_adj=matvec_adj,
nd_input=nd_input,
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=mask_in,
mask_out=mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
def __init__(
self,
arr_shape,
order="F",
arr_dtype=np.float32,
use_fft_shifts=True,
sample_mask=None,
ortho=False,
force_real_image=False,
debug=False,
preplan_pyfftw=True,
pyfftw_threads=None,
fft_axes=None,
fftshift_axes=None,
planner_effort="FFTW_ESTIMATE",
disable_warnings=False,
im_mask=None,
rel_fov=None,
**kwargs,
):
"""Cartesian MRI Operator (with partial FFT and coil maps).
Parameters
----------
arr_shape : int
shape of the array
order : {'C','F'}, optional
array ordering that will be assumed if inputs/outputs need to be
reshaped
arr_dtype : numpy.dtype, optional
dtype for the array
sample_mask : array_like, optional
boolean mask of which FFT coefficients to keep
ortho : bool, optional
if True, change the normalizeation to the orthogonal case
preplan_pyfftw : bool, optional
if True, precompute the pyFFTW plan upon object creation
pyfftw_threads : int, optional
number of threads to be used by pyFFTW. defaults to
multiprocessing.cpu_count() // 2.
use_fft_shifts : bool, optional
If False, do not apply any FFT shifts
fft_axes : tuple or None, optional
Specify a subset of the axes to transform. The default is to
transform all axes.
fftshift_axes : tuple or None, optional
Specify a subset of the axes to fftshift. The default is to
shift all axes.
im_mask : ndarray or None, optional
Image domain mask
force_real_image : bool, optional
debug : bool, optional
Additional Parameters
---------------------
nd_input : bool, optional
nd_output : bool, optional
"""
if isinstance(arr_shape, (np.ndarray, list)):
# retrieve shape from array
arr_shape = tuple(arr_shape)
if not isinstance(arr_shape, tuple):
raise ValueError("expected array_shape to be a tuple or list")
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.use_fft_shifts = use_fft_shifts
self.disable_warnings = disable_warnings
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
xp = cupy
on_gpu = True
else:
xp = np
on_gpu = False
self._on_gpu = on_gpu
if sample_mask is not None:
# masking faster if continguity of mask matches
if self.order == "F":
sample_mask = xp.asfortranarray(sample_mask)
elif self.order == "C":
sample_mask = xp.ascontiguousarray(sample_mask)
else:
raise ValueError("order must be C or F")
self.sample_mask = sample_mask
self.force_real_image = force_real_image
self.debug = debug
if self.sample_mask is not None:
if sample_mask.shape != arr_shape:
raise ValueError("sample mask shape must match arr_shape")
# make sure it is boolean
self.sample_mask = self.sample_mask > 0
# prestore raveled mask indices to save time later during masking
# self.sample_mask_idx = xp.where(
# self.sample_mask.ravel(order=self.order)
# )
# can specify a subset of the axes to perform the FFT/FFTshifts over
self.fft_axes = fft_axes
if self.fft_axes is None:
self.fft_axes = tuple(range(self.ndim))
if fftshift_axes is None:
self.fftshift_axes = self.fft_axes
else:
self.fftshift_axes = fftshift_axes
# configure scaling (e.g. unitary operator or not)
self.ortho = ortho
if self.fft_axes is None:
Ntrans = prod(self.arr_shape)
else:
Ntrans = prod(
np.asarray(self.arr_shape)[np.asarray(self.fft_axes)])
if self.ortho:
# sqrt of product of shape along axes where FFT is performed
self.scale_ortho = sqrt(Ntrans)
self.gpu_scale_inverse = 1 # self.scale_ortho / Ntrans
# self.gpu_scale_forward = self.scale_ortho
else:
self.scale_ortho = None
self.gpu_scale_inverse = 1 / Ntrans
if "mask_out" in kwargs:
raise ValueError("This operator specifies `mask_out` via the "
"parameter `sample_mask")
if ("mask_in" in kwargs) or ("mask_out" in kwargs):
raise ValueError("This operator specifies `mask_in` via the "
"parameter `im_mask")
if im_mask is not None:
if im_mask.shape != arr_shape:
raise ValueError("im_mask shape mismatch")
if order != "F":
raise ValueError("only order='F' supported for im_mask case")
nargin = im_mask.sum()
self.im_mask = im_mask
else:
nargin = prod(arr_shape)
self.im_mask = None
if sample_mask is not None:
nargout = sample_mask.sum()
else:
nargout = nargin
nargout = int(nargout)
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.complex64)
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if (self.sample_mask is not None) and nd_output:
raise ValueError("cannot have both nd_output and sample_mask")
if nd_output:
shape_out = self.arr_shape
else:
shape_out = (nargout, 1)
shape_in = self.arr_shape
self.have_pyfftw = config.have_pyfftw
if self.on_gpu:
self.preplan_pyfftw = False
else:
self.preplan_pyfftw = preplan_pyfftw if self.have_pyfftw else False
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
# raise ValueError("Implementation Incomplete")
if self.on_gpu:
self.fftn = partial(fftn, xp=cupy)
self.ifftn = partial(ifftn, xp=cupy)
else:
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
else:
self.fftn = partial(fftn, xp=np)
self.ifftn = partial(ifftn, xp=np)
self.fftshift = xp.fft.fftshift
self.ifftshift = xp.fft.ifftshift
self.rel_fov = rel_fov
self.mask = None # TODO: implement or remove (expected by CUDA code)
matvec = self.forward
matvec_adj = self.adjoint
self.norm_available = True
self.norm = self._norm
super(FFT_Operator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=matvec,
matvec_transp=matvec_adj,
matvec_adj=matvec_adj,
nd_input=nd_input or (im_mask is not None),
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=im_mask,
mask_out=None, # mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
def __init__(
self,
arr_shape,
order="F",
arr_dtype=np.float32,
filterbank=None,
mode="symmetric",
level=None,
prior=None,
force_real=False,
random_shift=False,
autopad=False,
autopad_mode="symmetric",
axes=None,
**kwargs,
):
""" TODO: fix docstring
Parameters
----------
arr_shape : int
shape of the array to filter
degree : int
degree of the directional derivatives to apply
order : {'C','F'}
array ordering that will be assumed if inputs/outputs need to be
reshaped
arr_dtype : numpy.dtype
dtype for the filter coefficients
prior : array_like
subtract this prior before computing the DWT
force_real : bool
subtract this prior before computing the DWT
random_shift : bool
Random shifts can be introduced to achieve translation invariance
as described in [1]_
autopad : bool
If True, the array will be padded from ``arr_shape`` up to the
nearest integer multiple of 2**``level`` prior to the transform.
The padding will be removed upon the adjoint transform.
autopad_mode : str
The mode for `numpy.pad` to use when ``autopad`` is True.
References
----------
..[1] MAT Figueiredo and RD Nowak.
An EM Algorithm for Wavelet-Based Image Restoration.
IEEE Trans. Image Process. 2003; 12(8):906-916
"""
if isinstance(arr_shape, np.ndarray):
# retrieve shape from array
arr_shape = arr_shape.shape
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.prior = prior
self.force_real = force_real
self.random_shift = random_shift
# determine axes and the shape of the axes to be transformed
self.axes, self.axes_shape, self.ndim_transform = _prep_axes_nd(
self.arr_shape, axes)
if self.random_shift:
self.shifts = np.zeros(self.ndim, dtype=np.int)
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
xp = cupy
on_gpu = True
else:
xp = np
on_gpu = False
self._on_gpu = on_gpu
self.autopad = autopad
self.autopad_mode = autopad_mode
if not is_nonstring_sequence(filterbank):
filterbank = [filterbank] * len(self.axes)
self.filterbank = _prep_filterbank(filterbank, self.axes, xp=xp)
max_levels = []
for ax, fb in zip(self.axes, self.filterbank):
max_levels.append(sep.dwt_max_level(self.arr_shape, fb, axes=ax))
if level is None or xp.isscalar(level):
level = [level] * len(self.axes)
if len(level) != len(self.axes):
raise ValueError("level must match the length of the axes list")
for n in range(len(level)):
if level[n] is None:
# default to 4 levels unless otherwise specified
level[n] = min(4, max_levels[n])
elif not self.autopad:
if level[n] > max_levels[n]:
level[n] = max_levels[n]
warnings.warn(
"level exceeds max level for the size of the "
"input along axis {}. Reducing level to {} for this "
"axis".format(self.axes[n], level[n]))
self.level = level
mode = _modes_per_axis(mode, self.axes)
if self.autopad:
min_factors = np.asarray([
fb.sizes[0][0] * fb.bw_decimation[0]**level
for fb in self.filterbank
])
arr_shape = np.asarray(arr_shape)
pad_shape = arr_shape.copy() # includes non-transformed axes
pad_shape[np.asarray(self.axes)] = min_factors * np.ceil(
np.asarray(self.axes_shape) / min_factors)
pad_shape = pad_shape.astype(int)
self.pad_width = [(0, s) for s in (pad_shape - arr_shape)]
self.pad_shape = tuple(pad_shape)
else:
self.pad_shape = self.arr_shape
self.mode = mode
self.truncate_coefficients = False
# TODO: define per-axis shift_range instead
self.shift_range = 1 << np.min(self.level)
if True:
# determine the shape of the coefficient arrays
# TODO: determine without running the transform
coeffs_tmp = sep.fswavedecn(
xp.ones(self.pad_shape, dtype=np.float32),
filterbank=self.filterbank,
mode=self.mode,
level=self.level,
axes=self.axes,
)
self.coeff_arr_shape = coeffs_tmp._coeffs.shape
nargin = prod(self.arr_shape)
nargout = prod(self.coeff_arr_shape)
# output of DWT may be complex, depending on input type
self.result_dtype = np.result_type(arr_dtype, np.float32)
self.coeffs = None
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if nd_output:
raise ValueError(
"nd_output = True is not supported. All framelet coefficients "
"will be stacked into a 1d array")
self.mask_in = kwargs.pop("mask_in", None)
if self.mask_in is not None:
nargin = self.mask_in.sum()
self.mask_out = kwargs.pop("mask_out", None)
if self.mask_out is not None:
nargout = self.mask_out.sum()
super(Framelet_Operator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=self.forward,
matvec_transp=self.adjoint,
matvec_adj=self.adjoint,
nd_input=nd_input,
nd_output=nd_output,
shape_in=self.arr_shape,
shape_out=self.coeff_arr_shape,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=self.mask_in,
mask_out=self.mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
def __init__(self,
kspace,
pixel_basis="dirac",
dx=None,
fov=None,
Nd=None):
""" Initialize a pixel-basis and its transform.
Parameters
----------
kspace : array
2D array of k-space coordinates :math:`[n_{samples}, n_{dim}]`.
pixel_basis : {'dirac', 'rect'}
basis type
dx : array_like, optional
image pixel spacings
fov : array_like, optional
image field of view (i.e. `Nd` * `dx`)
Nd : array_like, optional
image matrix size along each dimension
Attributes
----------
type : str
basis type
dx : array or None
pixel spacing in image domain
fov : array or None
image field of view (i.e. `Nd` * `dx`)
Nd : array or None
image matrix size along each dimension
transform : array
Fourier transform of the basis evaluated at the k-space sample
locations.
Notes
-----
A continuous function, :math:`f(\vec{x})`, is represented by a
discrete set of pixel values :math:`x_{i}` in basis,
:math:`b(\vec{x})`
.. math::
f\left(\vec{x}\right)=\sum_{i=1}^{N}\,x_{i}b\left(\vec{x}-x_{i}\right)
units for `dx`, `fov` must match (e.g. mm).
`kspace` units (e.g. mm^-1) should be the inverse of those for `dx`,
`fov`.
References
----------
.. [1] Pruessmann KP, Weiger M, Scheidegger MB, Boesiger P. SENSE:
Sensitivity Encoding for Fast MRI. Magn. Reson. Med. 1999;
42:952-962.
.. [2] Sutton, BP. Physics Based Iterative Reconstruction for MRI:
Compensating and Estimating Field Inhomogeneity and T2\* Relaxation.
Doctoral Dissertation. University of Michigan. 2003.
.. [3] Barrett HH, Myers KJ. "7.1: Objects and Images" In *Foundations
of Image Science*. 2003. Wiley-Interscience.
"""
self.type = pixel_basis
self.dx = dx
self.fov = fov
self.Nd = Nd
if self.dx is not None:
self.dx = np.atleast_1d(self.dx)
if self.fov is not None:
self.fov = np.atleast_1d(self.fov)
if self.Nd is not None:
self.Nd = np.atleast_1d(self.Nd)
else:
self.Nd = np.round(self.fov / self.dx).astype(int)
self.kspace = kspace
if kspace.ndim != 2:
raise ValueError("kspace must be 2D: [nsamples x ndim]")
nk, ndim = kspace.shape
if self.type == "dirac":
self.transform = None
elif self.type == "rect":
if self.dx is None:
if (self.fov is None) or (self.Nd is None):
raise ValueError("must specify dx or both fov and Nd")
self.dx = self.fov / self.Nd
if len(self.dx) != ndim:
if len(self.dx) == 1:
self.dx = np.asarray(self.dx.tolist() * ndim)
else:
raise ValueError("len(dx) != ndim")
self.transform = np.ones((nk, ))
for idx in range(ndim):
# allow for zero-sized pixels (Dirac impulses)
if self.dx[idx] > 0:
# dx* sinc(dx*kx) term in Eq. 6.15 of Brad Sutton's thesis
self.transform = self.transform * (
self.dx[idx] * np.sinc(self.dx[idx] * kspace[:, idx]))
# make the average scaling close to 1
self.transform /= self.transform.mean()
elif self.type in ["sinc", "dirac*dx"]:
# simply provide an appropriate "scale factor" dx to relate Fourier
# integral and summation
self.transform = np.ones(
(nk, )) * prod(f / n for f, n in zip(self.fov, self.Nd))
# make the average scaling close to 1
self.transform /= self.transform.mean()
else:
raise ValueError("unknown PixelBasis type: {}".format(pixel_basis))
def __init__(
self,
arr_shape,
order="F",
arr_dtype=np.float32,
use_fft_shifts=True,
sample_mask=None,
ortho=False,
coil_sensitivities=None,
force_real_image=False,
debug=False,
preplan_pyfftw=True,
pyfftw_threads=None,
fft_axes=None,
fftshift_axes=None,
planner_effort="FFTW_ESTIMATE",
loop_over_coils=False,
preserve_memory=False,
disable_warnings=False,
im_mask=None,
pixel_basis="dirac",
rel_fov=None,
**kwargs,
):
"""Cartesian MRI Operator (with partial FFT and coil maps).
Parameters
----------
arr_shape : int
shape of the array
order : {'C','F'}, optional
array ordering that will be assumed if inputs/outputs need to be
reshaped
arr_dtype : numpy.dtype, optional
dtype for the array
sample_mask : array_like, optional
boolean mask of which FFT coefficients to keep
coil_sensitivities : array, optional
Array of coil sensitivities.
ortho : bool, optional
if True, change the normalizeation to the orthogonal case
preplan_pyfftw : bool, optional
if True, precompute the pyFFTW plan upon object creation
pyfftw_threads : int, optional
number of threads to be used by pyFFTW. defaults to
multiprocessing.cpu_count() // 2.
use_fft_shifts : bool, optional
If False, do not apply any FFT shifts
fft_axes : tuple or None, optional
Specify a subset of the axes to transform. The default is to
transform all axes.
fftshift_axes : tuple or None, optional
Specify a subset of the axes to fftshift. The default is to
shift all axes.
im_mask : ndarray or None, optional
Image domain mask
force_real_image : bool, optional
loop_over_coils : bool, optional
If True, memory required is lower, but speed will be slower.
preserve_memory : bool, optional
If False, conjugate copy of coils won't be precomputed
debug : bool, optional
Additional Parameters
---------------------
nd_input : bool, optional
nd_output : bool, optional
"""
if isinstance(arr_shape, (np.ndarray, list)):
# retrieve shape from array
arr_shape = tuple(arr_shape)
if not isinstance(arr_shape, tuple):
raise ValueError("expected array_shape to be a tuple or list")
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.use_fft_shifts = use_fft_shifts
self.disable_warnings = disable_warnings
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
xp = cupy
on_gpu = True
else:
xp = np
on_gpu = False
self._on_gpu = on_gpu
if sample_mask is not None:
# masking faster if continguity of mask matches
if self.order == "F":
sample_mask = xp.asfortranarray(sample_mask)
elif self.order == "C":
sample_mask = xp.ascontiguousarray(sample_mask)
else:
raise ValueError("order must be C or F")
self.sample_mask = sample_mask
self.force_real_image = force_real_image
self.debug = debug
if self.sample_mask is not None:
if sample_mask.shape != arr_shape:
raise ValueError("sample mask shape must match arr_shape")
# make sure it is boolean
self.sample_mask = self.sample_mask > 0
# # prestore raveled mask indices to save time later during masking
# self.sample_mask_idx = xp.where(
# self.sample_mask.ravel(order=self.order)
# )
self.preserve_memory = preserve_memory
self.loop_over_coils = loop_over_coils
# can specify a subset of the axes to perform the FFT/FFTshifts over
self.fft_axes = fft_axes
if self.fft_axes is None:
if self.order == "C":
# last ndim axes
self.fft_axes = tuple([-ax for ax in range(self.ndim, 0, -1)])
else:
# first ndim axes
self.fft_axes = tuple(range(self.ndim))
if fftshift_axes is None:
self.fftshift_axes = self.fft_axes
else:
self.fftshift_axes = fftshift_axes
# configure scaling (e.g. unitary operator or not)
self.ortho = ortho
if self.fft_axes is None:
Ntrans = prod(self.arr_shape)
else:
Ntrans = prod(
np.asarray(self.arr_shape)[np.asarray(self.fft_axes)])
if self.ortho:
# sqrt of product of shape along axes where FFT is performed
self.scale_ortho = sqrt(Ntrans)
self.gpu_scale_inverse = 1 # self.scale_ortho / Ntrans
# self.gpu_scale_forward = self.scale_ortho
else:
self.scale_ortho = None
self.gpu_scale_inverse = 1 / Ntrans
if "mask_out" in kwargs:
raise ValueError("This operator specifies `mask_out` via the "
"parameter `sample_mask")
if ("mask_in" in kwargs) or ("mask_out" in kwargs):
raise ValueError("This operator specifies `mask_in` via the "
"parameter `im_mask")
if coil_sensitivities is not None:
if coil_sensitivities.ndim == self.ndim + 1:
# self.arr_shape + (Ncoils, )
if self.order == "C":
Ncoils = coil_sensitivities.shape[0]
else:
Ncoils = coil_sensitivities.shape[-1]
Nmaps = 1
elif coil_sensitivities.ndim == self.ndim + 2:
# case with multiple maps (e.g. ESPIRIT soft-SENSE)
# self.arr_shape + (Ncoils, Nmaps)
if self.order == "C":
Ncoils = coil_sensitivities.shape[1]
Nmaps = coil_sensitivities.shape[0]
else:
Ncoils = coil_sensitivities.shape[-2]
Nmaps = coil_sensitivities.shape[-1]
else:
# determine based on size
Ncoils = coil_sensitivities.size / prod(self.arr_shape)
if (Ncoils % 1) != 0:
raise ValueError("sensitivity map size mismatch")
self.Ncoils = int(Ncoils)
self.Nmaps = Nmaps
if self.order == "C":
cmap_shape = (self.Nmaps, self.Ncoils) + self.arr_shape
if not coil_sensitivities.flags.c_contiguous:
msg = (
"Converting coil_sensitivities to be C contiguous"
" (requires a copy). To avoid the copy, convert to "
"Fortran contiguous order prior to calling "
"MRI_Cartesian (see np.ascontiguousarray)")
if not self.disable_warnings:
warnings.warn(msg)
coil_sensitivities = xp.ascontiguousarray(
coil_sensitivities)
else:
cmap_shape = self.arr_shape + (self.Ncoils, self.Nmaps)
if not coil_sensitivities.flags.f_contiguous:
msg = (
"Converting coil_sensitivities to be Fortan contiguous"
" (requires a copy). To avoid the copy, convert to "
"Fortran contiguous order prior to calling "
"MRI_Cartesian (see np.asfortranarray)")
if not self.disable_warnings:
warnings.warn(msg)
coil_sensitivities = xp.asfortranarray(coil_sensitivities)
if tuple(coil_sensitivities.shape) != tuple(cmap_shape):
coil_sensitivities = coil_sensitivities.reshape(
cmap_shape, order=self.order)
if not self.preserve_memory:
self.coil_sensitivities_conj = xp.conj(coil_sensitivities)
else:
self.Ncoils = 1
self.Nmaps = 1
if self.Ncoils == 1:
# TODO: currently has a shape bug if Ncoils == 1
# and loop_over_coils = False.
self.loop_over_coils = True
self.coil_sensitivities = coil_sensitivities
if im_mask is not None:
if im_mask.shape != arr_shape:
raise ValueError("im_mask shape mismatch")
if order != "F":
raise ValueError("only order='F' supported for im_mask case")
nargin = xp.count_nonzero(im_mask)
self.im_mask = im_mask
else:
nargin = prod(arr_shape)
self.im_mask = None
nargin *= self.Nmaps
# nargout = # of k-space samples
if sample_mask is not None:
nargout = xp.count_nonzero(sample_mask) * self.Ncoils
else:
nargout = nargin // self.Nmaps * self.Ncoils
nargout = int(nargout)
self.idx_orig = None
self.idx_conj = None
self.sample_mask_conj = None
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.complex64)
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if (self.sample_mask is not None) and nd_output:
raise ValueError("cannot have both nd_output and sample_mask")
if nd_output:
if self.Ncoils == 1:
shape_out = self.arr_shape
else:
if Nmaps == 1:
if self.order == "C":
shape_out = (self.Ncoils, ) + self.arr_shape
else:
shape_out = self.arr_shape + (self.Ncoils, )
else:
if self.order == "C":
shape_out = (self.Nmaps, self.Ncoils) + self.arr_shape
else:
shape_out = self.arr_shape + (self.Ncoils, self.Nmaps)
else:
shape_out = (nargout, 1)
if self.Nmaps == 1:
shape_in = self.arr_shape
else:
if self.order == "C":
shape_in = (self.Nmaps, ) + self.arr_shape
else:
shape_in = self.arr_shape + (self.Nmaps, )
if self.order == "C":
self.shape_inM = (self.Nmaps, ) + self.arr_shape
else:
self.shape_inM = self.arr_shape + (self.Nmaps, )
self.shape_in1 = self.arr_shape
self.have_pyfftw = config.have_pyfftw
if self.on_gpu:
self.preplan_pyfftw = False
else:
self.preplan_pyfftw = preplan_pyfftw if self.have_pyfftw else False
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
# raise ValueError("Implementation Incomplete")
if self.on_gpu:
self.fftn = partial(fftn, xp=cupy)
self.ifftn = partial(ifftn, xp=cupy)
else:
if self.preplan_pyfftw:
self._preplan_fft(pyfftw_threads, planner_effort)
else:
self.fftn = partial(fftn, xp=np)
self.ifftn = partial(ifftn, xp=np)
self.fftshift = xp.fft.fftshift
self.ifftshift = xp.fft.ifftshift
self.rel_fov = rel_fov
self.mask = None # TODO: implement or remove (expected by CUDA code)
matvec = self.forward
matvec_adj = self.adjoint
self.norm_available = False
self.norm = self._norm
super(MRI_Cartesian, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=matvec,
matvec_transp=matvec_adj,
matvec_adj=matvec_adj,
nd_input=nd_input or (im_mask is not None),
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=im_mask,
mask_out=None, # mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
self._init_pixel_basis(pixel_basis=pixel_basis)
def forward(self, x):
xp = self.xp
if self.force_real_image:
if x.dtype in [np.complex64, np.complex128]:
x.imag[:] = 0
if self.im_mask is None:
size_1rep = self.nargin
else:
size_1rep = prod(self.shape_inM)
if x.size < size_1rep:
raise ValueError("data, x, too small to transform.")
elif x.size == size_1rep:
nreps = 1
# 1D or single repetition nD
x = x.reshape(self.shape_inM, order=self.order)
if self.order == "C":
y = self._forward_single_rep(x[0, ...], i_map=0)
for i_map in range(1, self.Nmaps):
y += self._forward_single_rep(x[i_map, ...], i_map=i_map)
else:
y = self._forward_single_rep(x[..., 0], i_map=0)
for i_map in range(1, self.Nmaps):
y += self._forward_single_rep(x[..., i_map], i_map=i_map)
else:
if self.order == "C":
nreps = x.shape[0]
x_shape = (nreps, ) + self.shape_inM
if self.sample_mask is not None:
y_shape = (nreps, self.Ncoils, self.nargout)
else:
y_shape = (nreps, self.Ncoils) + self.shape_in1
else:
nreps = x.shape[-1]
x_shape = self.shape_inM + (nreps, )
if self.sample_mask is not None:
y_shape = (self.nargout, self.Ncoils, nreps)
else:
y_shape = self.shape_in1 + (self.Ncoils, nreps)
x = x.reshape(x_shape, order=self.order)
if self.sample_mask is not None:
# number of samples by number of repetitions
y = xp.zeros(
y_shape,
dtype=xp.result_type(x, xp.complex64),
order=self.order,
)
else:
y = xp.zeros(
y_shape,
dtype=xp.result_type(x, xp.complex64),
order=self.order,
)
if self.order == "C":
for rep in range(nreps):
y[rep, ...] = self._forward_single_rep(
x[rep, 0, ...], i_map=0).reshape(y.shape[1:],
order=self.order)
for i_map in range(1, self.Nmaps):
y[rep, ...] += self._forward_single_rep(
x[rep, i_map,
...], i_map=i_map).reshape(y.shape[1:],
order=self.order)
else:
for rep in range(nreps):
y[..., rep] = self._forward_single_rep(
x[..., 0, rep], i_map=0).reshape(y.shape[:-1],
order=self.order)
for i_map in range(1, self.Nmaps):
y[..., rep] += self._forward_single_rep(
x[..., i_map,
rep], i_map=i_map).reshape(y.shape[:-1],
order=self.order)
if self.ortho:
y /= self.scale_ortho
if y.dtype != self.result_dtype:
y = y.astype(self.result_dtype)
return y
def __init__(
self,
Nd,
omega,
mask=None,
squeeze_reps=True,
loc_in="cpu",
loc_out="cpu",
order="F",
**kwargs,
):
if loc_in == "cpu":
xp = np
else:
import cupy
xp = cupy
# masked boolean if mask is True everywhere or no mask is provided
if mask is not None:
mask = xp.asarray(mask, dtype=bool)
self.masked = (
(mask is not None)
and (not isinstance(mask, tuple))
and (mask.size != xp.count_nonzero(mask))
)
if isinstance(mask, tuple):
self.__mask = None
else:
self.__mask = mask
self.__Nd = Nd
nargout = omega.shape[0]
if self.masked:
nargin = xp.count_nonzero(mask)
else:
nargin = prod(self.__Nd)
if order != "F":
raise ValueError("NUFFT_Operator only supports order='F'.")
# initialize NUFFT
NufftBase.__init__(self, Nd=Nd, omega=omega, order=order, **kwargs)
# set appropriate operations as initialized by NufftBase
self.__matvec = self.fft
self.__matvec_transp = self.adj
self.__matvec_adj = self.adj
# construct the linear operator
LinearOperatorMulti.__init__(
self,
nargin,
nargout,
symmetric=False,
hermetian=False,
matvec=self.__matvec,
matvec_adj=self.__matvec_adj,
matvec_transp=self.__matvec_adj,
nd_input=False, # True,
nd_output=False,
loc_in=loc_in,
loc_out=loc_out,
shape_in=self.Nd,
shape_out=(nargout, 1),
order=order,
matvec_allows_repetitions=True,
# MRI_Operator code assumes squeeze_reps_* are False
squeeze_reps=squeeze_reps,
mask_in=self.mask,
mask_out=None,
dtype=self._cplx_dtype,
)
def __init__(
self,
arr_shape,
weight=1,
norm=1,
prior=None,
order="F",
arr_dtype=np.float32,
tv_type="iso",
grid_size=None,
nd_input=True,
nd_output=True,
squeeze_reps=True,
axes=None,
**kwargs,
):
"""
Parameters
----------
arr_shape : int
shape of the array to filter
degree : int
degree of the directional derivatives to apply
order : {'C','F'}
array ordering that will be assumed if inputs/outputs need to be
reshaped
prior : array_like
If provided, subtract the prior before computing the TV.
arr_dtype : numpy.dtype
dtype for the filter coefficients
tv_type : {'iso', 'aniso'}
Select whether isotropic or anisotropic (l1) TV is used
grid_size : array, optional
size of the grid along each dimension of the image. defaults to
ones. (e.g. needed if the voxels are anisotropic.)
"""
if isinstance(arr_shape, np.ndarray):
# retrieve shape from array
arr_shape = arr_shape.shape
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
self.grid_size = grid_size
self.tv_type = tv_type
if tv_type == "aniso":
# TODO
raise ValueError("Not implemented")
self.offsets = compute_offsets(tuple(arr_shape), use_corners=False)
if axes is not None:
if np.isscalar(axes):
axes = (axes, )
axes = np.asarray(axes)
if len(axes) > self.ndim:
raise ValueError("invalid number of axes")
if np.any(np.abs(axes) > self.ndim):
raise ValueError("invalid axis")
axes = axes % self.ndim
self.offsets = np.asarray(self.offsets)[axes].tolist()
self.num_offsets = len(self.offsets)
if self.order == "C":
self.grad_shape = (self.num_offsets, ) + tuple(arr_shape)
else:
self.grad_shape = tuple(arr_shape) + (self.num_offsets, )
shape_out = self.grad_shape
shape_in = self.arr_shape
nargin = prod(self.arr_shape)
nargout = prod(self.grad_shape)
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.float32)
self.weight = weight
self.prior = prior
self.norm_p = norm
self._limit = 1e-6
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
import cupy
xp = cupy
else:
xp = np
if self.order == "C":
grad_axis = 0
else:
grad_axis = -1
self.grad_axis = grad_axis
if axes is not None:
# customized offsets
if True:
# non-periodic, slightly faster
self.grad_func = functools.partial(gradient_ravel_offsets,
offsets=self.offsets)
self.div_func = functools.partial(divergence_ravel_offsets,
offsets=self.offsets)
else:
# periodic
self.grad_func = functools.partial(gradient_periodic,
axes=axes)
self.div_func = functools.partial(divergence_periodic,
axes=axes)
else:
self.grad_func = functools.partial(gradient_periodic)
self.div_func = functools.partial(divergence_periodic)
self.mask_in = kwargs.pop("mask_in", None)
if self.mask_in is not None:
nargin = self.mask_in.sum()
if isinstance(nargin, cupy_ndarray_type):
nargin = nargin.get()
nd_input = True # TODO: fix LinOp to remove need for this. why does DWT case not need it?
self.mask_out = kwargs.pop("mask_out", None)
if self.mask_out is not None:
# probably wrong if order = 'C'
if order == "C":
stack_axis = 0
else:
stack_axis = -1
self.mask_out = xp.stack([self.mask_out] * self.num_offsets,
stack_axis)
nargout = self.mask_out.sum()
if isinstance(nargout, cupy_ndarray_type):
nargout = nargout.get()
nd_output = True
super(TV_Operator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=self.forward,
matvec_transp=self.adjoint,
matvec_adj=self.adjoint,
nd_input=nd_input,
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=True,
squeeze_reps=squeeze_reps,
mask_in=self.mask_in,
mask_out=self.mask_out,
symmetric=False,
hermetian=False,
dtype=self.result_dtype,
**kwargs,
)
def __init__(
self,
arr_shape,
axis,
m,
order="F",
arr_dtype=np.float32,
debug=False,
**kwargs,
):
"""
Parameters
----------
arr_shape : int
shape of the array
axis : int
The axis to transform
m : np.ndarray
Orthogonal matrix representing the transform.
order : {'F', 'C'}
The memory layout ordering of the data.
arr_dtype: np.dtype
The dtype of the array to be transformed.
debug : bool, optional
Additional Parameters
---------------------
nd_input : bool, optional
nd_output : bool, optional
"""
if isinstance(arr_shape, (np.ndarray, list)):
# retrieve shape from array
arr_shape = tuple(arr_shape)
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
if axis > self.ndim or axis < (-self.ndim + 1):
raise ValueError("invalid axis")
else:
axis = axis % self.ndim
self.axis = axis
self.debug = debug
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
if not config.have_cupy:
raise ImportError("CuPy not found")
xp = cupy
else:
xp = np
if m == "dct":
# default is a DCT transform
m = dct_matrix(self.arr_shape[self.axis], xp=xp)
else:
m = xp.asarray(m)
self.m = m
if self.m.dtype.kind == "c":
m_inv = xp.conj(self.m).H
else:
m_inv = m.T
self.m_inv = m_inv
nargin = prod(arr_shape)
nargout = nargin
self.result_dtype = xp.result_type(arr_dtype, m.dtype)
# self.ortho = True
matvec_allows_repetitions = kwargs.pop("matvec_allows_repetitions",
True)
squeeze_reps = kwargs.pop("squeeze_reps", True)
nd_input = kwargs.pop("nd_input", False)
nd_output = kwargs.pop("nd_output", False)
if nd_output:
shape_out = self.arr_shape
else:
shape_out = (nargout, 1)
shape_in = self.arr_shape
# mask_out = self.sample_mask
mask_in = kwargs.pop("mask_in", None)
mask_out = kwargs.pop("mask_out", None)
matvec = self.forward
matvec_adj = self.adjoint
super(OrthoMatrixOperator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=matvec,
matvec_transp=matvec_adj,
matvec_adj=matvec_adj,
nd_input=nd_input,
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=matvec_allows_repetitions,
squeeze_reps=squeeze_reps,
mask_in=mask_in,
mask_out=mask_out,
symmetric=False, # TODO: set properly
hermetian=False, # TODO: set properly
dtype=self.result_dtype,
**kwargs,
)
def __init__(
self,
arr_shape,
norm=1,
prior=None,
order="F",
arr_dtype=np.float32,
grid_size=None,
nd_input=True,
nd_output=True,
squeeze_reps=True,
use_corners=False,
custom_offsets=None,
axes=None,
**kwargs,
):
"""
Parameters
----------
arr_shape : int
shape of the array to filter
degree : int
degree of the directional derivatives to apply
order : {'C','F'}
array ordering that will be assumed if inputs/outputs need to be
reshaped
prior : array_like
subtract this prior before computing the finite difference
arr_dtype : numpy.dtype
dtype for the filter coefficients
grid_size : array, optional
size of the grid along each dimension of the image. defaults to
ones. (e.g. needed if the voxels are anisotropic.)
use_corners : bool, optional
If True, use additional "diagonal" directions as opposed to only
differences along the primary axes. When True, the number of
directions will be (3**ndim - 1)/2. When False, ndim directions
are used.
custom_offsets : list of int or None
can be used to override the default offsets
axes : list of int or None
Can be used to specify a subset of axes over which to take the
finite difference. If None, all axes are used. This option is
only compatible with `use_corners = False`. If `custom_offsets`
is provided, this argument is ignored.
"""
if isinstance(arr_shape, np.ndarray):
# retrieve shape from array
arr_shape = arr_shape.shape
self.arr_shape = arr_shape
self.ndim = len(arr_shape)
self.order = order
if axes is not None:
if np.isscalar(axes):
axes = (axes, )
self.axes = axes
else:
self.axes = None
if custom_offsets is None:
self.offsets = compute_offsets(tuple(arr_shape),
use_corners=use_corners)
if use_corners and axes is not None:
raise ValueError(
"use_corners not supported with customized axes")
elif axes is not None:
self.offsets = [self.offsets[ax] for ax in axes]
else:
self.offsets = custom_offsets
self.num_offsets = len(self.offsets)
self.grid_size = grid_size
if self.grid_size is not None:
if len(self.grid_size) != self.num_offsets:
raise ValueError(
"grid_size array must match the number of offsets")
# else:
# if len(self.grid_size) != len(self.axes):
# raise ValueError(
# "grid_size array must match the number of "
# "axes specified")
if self.order == "C":
self.grad_shape = (self.num_offsets, ) + tuple(arr_shape)
else:
self.grad_shape = tuple(arr_shape) + (self.num_offsets, )
shape_out = self.grad_shape
shape_in = self.arr_shape
nargin = prod(self.arr_shape)
nargout = prod(self.grad_shape)
# output of FFTs will be complex, regardless of input type
self.result_dtype = np.result_type(arr_dtype, np.float32)
self.prior = prior
self.norm = norm
self._limit = 1e-6
if "loc_in" in kwargs and kwargs["loc_in"] == "gpu":
import cupy
xp = cupy
else:
xp = np
self.mask_in = kwargs.pop("mask_in", None)
if self.mask_in is not None:
nargin = self.mask_in.sum()
nd_input = True # TODO: fix LinOp to remove need for this. why does DWT case not need it?
self.mask_out = kwargs.pop("mask_out", None)
if self.mask_out is not None:
# TODO: probably wrong if order = 'C'
if self.order == "C":
stack_axis = 0
else:
stack_axis = -1
self.mask_out = xp.stack([self.mask_out] * self.num_offsets,
stack_axis)
nargout = self.mask_out.sum()
nd_output = True
self.grad_func = gradient_periodic
self.div_func = divergence_periodic
if self.order == "C":
grad_axis = 0
else:
grad_axis = -1
self.grad_axis = grad_axis
if custom_offsets is None and (not use_corners):
# standard grad/div along all axes
# This verion uses periodic boundary conditions
self.grad_func = functools.partial(gradient_periodic,
axes=self.axes)
self.div_func = functools.partial(divergence_periodic,
axes=self.axes)
else:
# customized offsets
self.grad_func = functools.partial(gradient_ravel_offsets,
offsets=self.offsets)
self.div_func = functools.partial(divergence_ravel_offsets,
offsets=self.offsets)
super(FiniteDifferenceOperator, self).__init__(
nargin=nargin,
nargout=nargout,
matvec=self.forward,
matvec_transp=self.adjoint,
matvec_adj=self.adjoint,
nd_input=nd_input,
nd_output=nd_output,
shape_in=shape_in,
shape_out=shape_out,
order=self.order,
matvec_allows_repetitions=True,
squeeze_reps=squeeze_reps,
mask_in=self.mask_in,
mask_out=self.mask_out,
symmetric=False,
hermetian=False,
dtype=self.result_dtype,
**kwargs,
)
def test_prod(seq):
assert utils.prod(seq) == np.prod(seq)