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)