def _ifftc(input, oshape=None, axes=None, norm='ortho'): ndim = input.ndim axes = util._normalize_axes(axes, ndim) xp = backend.get_array_module(input) if oshape is None: oshape = input.shape tmp = util.resize(input, oshape) tmp = xp.fft.ifftshift(tmp, axes=axes) tmp = xp.fft.ifftn(tmp, axes=axes, norm=norm) output = xp.fft.fftshift(tmp, axes=axes) return output
def _ifftc(input, oshape=None, axes=None, norm='ortho'): ndim = input.ndim axes = util._normalize_axes(axes, ndim) device = backend.get_device(input) xp = device.xp if oshape is None: oshape = input.shape with device: tmp = util.resize(input, oshape) tmp = xp.fft.ifftshift(tmp, axes=axes) tmp = xp.fft.ifftn(tmp, axes=axes, norm=norm) output = xp.fft.fftshift(tmp, axes=axes) return output
def Gradient(ishape, axes=None): """Linear operator that computes numerical gradient. Args: ishape (tuple of ints): Input shape. """ I = Identity(ishape) axes = util._normalize_axes(axes, len(ishape)) ndim = len(ishape) G = Vstack([ I - Circshift(ishape, [0] * i + [1] + [0] * (ndim - i - 1)) for i in range(ndim) ]) G.repr_str = 'Gradient' return G
def FiniteDifference(ishape, axes=None): """Linear operator that computes finite difference gradient. Args: ishape (tuple of ints): Input shape. """ I = Identity(ishape) axes = util._normalize_axes(axes, len(ishape)) ndim = len(ishape) linops = [] for i in range(ndim): D = I - Circshift(ishape, [0] * i + [1] + [0] * (ndim - i - 1)) R = Reshape([1] + list(ishape), ishape) linops.append(R * D) G = Vstack(linops, axis=0) return G
def l2_proj(eps, input, axes=None): """Projection onto L2 ball. Args: eps (float, or array): L2 ball scaling. input (array) Returns: array: Result. """ axes = util._normalize_axes(axes, input.ndim) xp = backend.get_array_module(input) norm = xp.sum(xp.abs(input)**2, axis=axes, keepdims=True)**0.5 mask = norm < eps output = mask * input + (1 - mask) * (eps * input / (norm + mask)) return output
def FiniteDifference(ishape, axes=None): """Linear operator that computes finite difference gradient. Args: ishape (tuple of ints): Input shape. axes (tuple or list): Axes to circularly shift. All axes are used if None. """ I = Identity(ishape) ndim = len(ishape) axes = util._normalize_axes(axes, ndim) linops = [] for i in axes: D = I - Circshift(ishape, [1], axes=[i]) R = Reshape([1] + list(ishape), ishape) linops.append(R * D) G = Vstack(linops, axis=0) return G
def elitist_thresh(lamda, input, axes=None): """Elitist threshold. Solves for :math:: \text{argmin}_x \| x - y \|_2^2 + \lambda \| x \|_1^2 Args: lamda (float, or array): Threshold parameter. input (array): Input array. axes (None or tuple of ints): Axes to perform threshold. Returns: array: Result. References: Kowalski, M. 2009. Sparse regression using mixed norms. """ shape = input.shape axes = util._normalize_axes(axes, input.ndim) remain_axes = tuple(set(range(input.ndim)) - set(axes)) length = util.prod([shape[a] for a in axes]) batch = input.size // length input = input.transpose(remain_axes + axes) input = input.reshape([batch, length]) thresh = find_elitist_thresh(lamda, input) output = soft_thresh(thresh, input) output = output.reshape([shape[a] for a in remain_axes + axes]) output = output.transpose(np.argsort(remain_axes + axes)) return output