diff --git a/examples/ufunc_ops/ufunc_basics.py b/examples/oplib/ufunc_basics.py similarity index 77% rename from examples/ufunc_ops/ufunc_basics.py rename to examples/oplib/ufunc_basics.py index cd2fb583579..2d7f91428a0 100644 --- a/examples/ufunc_ops/ufunc_basics.py +++ b/examples/oplib/ufunc_basics.py @@ -6,8 +6,8 @@ # Trigonometric functions can be computed, along with their gradients. -cos = odl.ufunc_ops.cos() -sin = odl.ufunc_ops.sin() +cos = odl.oplib.cos() +sin = odl.oplib.sin() # Compute cosine and its gradient @@ -18,16 +18,15 @@ # Other functions include the square, exponential, etc # Higher order derivatives are obtained via the gradient of the gradient, etc. -square = odl.ufunc_ops.square() +square = odl.oplib.square() print('[x^2](3) = {}, [d/dx x^2](3) = {}, ' '[d^2/dx^2 x^2](3) = {}, [d^3/dx^3 x^2](3) = {}' - ''.format(square(3), square.gradient(3), - square.gradient.gradient(3), + ''.format(square(3), square.gradient(3), square.gradient.gradient(3), square.gradient.gradient.gradient(3))) # Can also define ufuncs on vector-spaces, then they act pointwise. r3 = odl.rn(3) -exp_r3 = odl.ufunc_ops.exp(r3) +exp_r3 = odl.oplib.exp(r3) print('e^[1, 2, 3] = {}'.format(exp_r3([1, 2, 3]))) diff --git a/examples/ufunc_ops/ufunc_composition.py b/examples/oplib/ufunc_composition.py similarity index 97% rename from examples/ufunc_ops/ufunc_composition.py rename to examples/oplib/ufunc_composition.py index 7c2306244a2..71b9a01e8f3 100644 --- a/examples/ufunc_ops/ufunc_composition.py +++ b/examples/oplib/ufunc_composition.py @@ -13,7 +13,7 @@ import odl # Create square functional. It's domain is by default the real numbers. -square = odl.ufunc_ops.square() +square = odl.oplib.square() # Create L2 norm functionals space = odl.rn(3) diff --git a/examples/ufunc_ops/ufunc_solvers.py b/examples/oplib/ufunc_solvers.py similarity index 96% rename from examples/ufunc_ops/ufunc_solvers.py rename to examples/oplib/ufunc_solvers.py index 56a5b2e43f3..748346b4d03 100644 --- a/examples/ufunc_ops/ufunc_solvers.py +++ b/examples/oplib/ufunc_solvers.py @@ -10,7 +10,7 @@ # Create space and functionals r2 = odl.rn(2) rosenbrock = odl.solvers.RosenbrockFunctional(r2, scale=2.0) -log = odl.ufunc_ops.log() +log = odl.oplib.log() # Create goal functional by composing log with rosenbrock and add 0.1 to # avoid singularity at 0 diff --git a/examples/space/auto_weighting.py b/examples/space/auto_weighting.py new file mode 100644 index 00000000000..ec25fb57c4b --- /dev/null +++ b/examples/space/auto_weighting.py @@ -0,0 +1,57 @@ +"""Example demonstrating the usage of the ``auto_weighting`` decorator.""" + +import odl +from odl.space.space_utils import auto_weighting + + +class ScalingOp(odl.Operator): + + """Operator that scales input by a constant.""" + + def __init__(self, dom, ran, c): + super(ScalingOp, self).__init__(dom, ran, linear=True) + self.c = c + + def _call(self, x): + return self.c * x + + @property + @auto_weighting + def adjoint(self): + return ScalingOp(self.range, self.domain, self.c) + + +rn = odl.rn(2) # Constant weight 1 +discr = odl.uniform_discr(0, 4, 2) # Constant weight 2 +print('*** Spaces ***') +print('Rn =', rn, '- weight =', rn.weighting.const) +print('discr =', discr, '- weight =', discr.weighting.const) +print('') + +op1 = ScalingOp(rn, rn, 2) # Same weightings, no scaling in adjoint +op2 = ScalingOp(discr, discr, 2) # Same weightings, no scaling in adjoint +op3 = ScalingOp(rn, discr, 2) # Different weightings, adjoint scales + +# Look at output of ajoint +print('*** Ajoint evaluation ***') +print('Rn -> Rn adjoint at one :', op1.adjoint(op1.range.one())) +print('discr -> discr adjoint at one:', op2.adjoint(op2.range.one())) +print('Rn -> discr adjoint at one :', op3.adjoint(op3.range.one())) +print('') + +# Check adjointness +print('*** Check adjointness ***') +inner1_dom = op1.domain.one().inner(op1.adjoint(op1.range.one())) +inner1_ran = op1(op1.domain.one()).inner(op1.range.one()) +print('Rn -> Rn: = {}, = {}' + ''.format(inner1_ran, inner1_dom)) + +inner2_dom = op2.domain.one().inner(op2.adjoint(op2.range.one())) +inner2_ran = op2(op2.domain.one()).inner(op2.range.one()) +print('discr -> discr: = {}, = {}' + ''.format(inner2_ran, inner2_dom)) + +inner3_dom = op3.domain.one().inner(op3.adjoint(op3.range.one())) +inner3_ran = op3(op3.domain.one()).inner(op3.range.one()) +print('Rn -> discr: = {}, = {}' + ''.format(inner3_ran, inner3_dom)) diff --git a/odl/__init__.py b/odl/__init__.py index 835f4027e86..9a0899f83ee 100644 --- a/odl/__init__.py +++ b/odl/__init__.py @@ -65,7 +65,7 @@ from . import solvers from . import tomo from . import trafos -from . import ufunc_ops +from . import oplib from . import util # Add `test` function to global namespace so users can run `odl.test()` diff --git a/odl/operator/tensor_ops.py b/odl/operator/tensor_ops.py index a015670b1d3..166b2342e7f 100644 --- a/odl/operator/tensor_ops.py +++ b/odl/operator/tensor_ops.py @@ -18,8 +18,8 @@ from odl.space import ProductSpace, tensor_space from odl.space.base_tensors import TensorSpace from odl.space.weighting import ArrayWeighting -from odl.util import ( - signature_string, indent, dtype_repr, moveaxis, writable_array) +from odl.util import signature_string, indent, dtype_repr, writable_array +from odl.util.npy_compat import moveaxis __all__ = ('PointwiseNorm', 'PointwiseInner', 'PointwiseSum', 'MatrixOperator', diff --git a/odl/oplib/__init__.py b/odl/oplib/__init__.py new file mode 100644 index 00000000000..5a3ed191c7d --- /dev/null +++ b/odl/oplib/__init__.py @@ -0,0 +1,24 @@ +# Copyright 2014-2017 The ODL contributors +# +# This file is part of ODL. +# +# This Source Code Form is subject to the terms of the Mozilla Public License, +# v. 2.0. If a copy of the MPL was not distributed with this file, You can +# obtain one at https://mozilla.org/MPL/2.0/. + +"""Library of operators. + +This submodule is a place for operators that either do not fit in any other +places or require "advanced" features of ODL that would make them hard +to put into any other submodule due to circular dependencies. +""" + +from __future__ import absolute_import + +__all__ = ('convolution', 'ufunc_ops') + +from .convolution import * +__all__ += convolution.__all__ + +from .ufunc_ops import * +__all__ += ufunc_ops.__all__ diff --git a/odl/oplib/convolution.py b/odl/oplib/convolution.py new file mode 100644 index 00000000000..5bdc94b6b1f --- /dev/null +++ b/odl/oplib/convolution.py @@ -0,0 +1,868 @@ +# Copyright 2014-2017 The ODL contributors +# +# This file is part of ODL. +# +# This Source Code Form is subject to the terms of the Mozilla Public License, +# v. 2.0. If a copy of the MPL was not distributed with this file, You can +# obtain one at https://mozilla.org/MPL/2.0/. + +"""Discretized continuous convolution and fully discrete convolution.""" + +from __future__ import division +import numpy as np + +from odl.discr import DiscreteLpElement +from odl.operator import Operator +from odl.space import tensor_space +from odl.space.base_tensors import TensorSpace, Tensor +from odl.space.space_utils import auto_weighting +from odl.trafos.backends import PYFFTW_AVAILABLE +from odl.util import ( + is_real_dtype, is_floating_dtype, dtype_str, writable_array) +from odl.util.npy_compat import roll + + +__all__ = ('DiscreteConvolution', 'convolve', 'correlate') + + +class DiscreteConvolution(Operator): + + """Fully discrete convolution with a given kernel.""" + + def __init__(self, domain, kernel, range=None, axis=None, impl='fft', + **kwargs): + """Initialize a new instance. + + Parameters + ---------- + domain : `TensorSpace` + Space on which the convolution is defined. If ``domain`` is + a `DiscreteLp`, it must be uniformly discretized. + kernel : array-like + The kernel with which input elements are convolved. It must + have the same number of dimensions as ``domain``, and its + shape can be at most equal to ``domain.shape``. In axes + with size 1, broadcasting is applied. + + **Important:** + + - The kernel must **always** have the same number of dimensions + as ``domain``, even for convolutions along axes. + - In axes where no convolution is performed, the shape of the + kernel must either be 1 (broadcasting along these axes), or + equal to the domain shape (stacked kernel, only supported + for `impl='fft'`). + - The ``'fft'`` implementation needs the kernel to have + floating-point ``dtype``, hence the smallest possible + float data type is used in that case to store the kernel. + - If the convolution kernel is complex, the ``domain`` **must** + be a complex space. Real-to-complex convolutions are not + allowed and need to be defined by composition with + `ComplexEmbedding` instead. + + See Examples for further clarification. + + range : `TensorSpace`, optional + Space of output elements of the convolution. Must be of the + same shape as ``domain``. If not given, the range is equal to + ``domain.astype(result_dtype)``, where ``result_dtype`` is + the data type of the convolution result. If ``impl='real'``, + integer dtypes are preserved, while for ``impl='fft'``, + the smallest possible floating-point type is chosen. + axis : int or sequence of ints, optional + Coordinate axis or axes in which to take the convolution. + ``None`` means all input axes. + impl : {'fft', 'real'} + Implementation of the convolution as FFT-based or using + direct summation. The fastest available FFT backend is + chosen automatically. Real space convolution is based on + `scipy.signal.convolve`. + See Notes for further information on the backends. + + Other Parameters + ---------------- + padding : int or sequence of ints, optional + Zero-padding used before Fourier transform in the FFT backend. + Does not apply for ``impl='real'``. A sequence is applied per + axis, with padding values corresponding to ``axis`` entries + as provided. + Default: ``min(kernel.shape - 1, 64)`` + padded_shape : sequence of ints, optional + Apply zero-padding with this target shape. Cannot be used + together with ``padding``. + variant : {'forward', 'adjoint'} + Convolution variant, mostly for internal use in the adjoint + operator. + Default: ``'forward'`` + cache_kernel_ft : bool, optional + If ``True``, store the Fourier transform of the kernel for + later reuse. + Default: ``False`` + + Notes + ----- + - In general, if nonzero padding is used, the out-of-place call + (no ``out`` parameter) is expected to be faster for this operator + for a variety of reasons: + + * The ``impl='real'`` backend does not support an `out` argument + and will thus create a new array anyway. Providing `out` will + require an additional copy from the new array to `out`. + * For ``impl='fft'``, the intermediate padding and unpadding steps + make it impossible to use an `out` object efficiently since + size and/or dtype do not match the requirements. + In addition, the Numpy FFT backend does not support an + `out` parameter. + + - Providing ``out`` will be faster and more memory-efficient + for ``impl='fft'`` and ``pyfftw`` backend with ``padding=0``, + since this is the only case when the final inverse FFT can be + written to ``out`` directly. + + This fact can be exploited by padding the input beforehand and + then using no padding in this operator. + + - ``scipy.convolve`` does not support an ``axis`` parameter. + However, a convolution along axes with a lower-dimensional + kernel can be achieved by adding empty dimensions. For example, + to convolve along axis 0 we can do the following :: + + ker_1d = np.array([-1.0, 1.0]) + ker_axis0 = ker_1d[:, None] + conv = DiscreteConvolution(space_2d, ker_axis0, impl='real') + + Not possible with this approach is a convolution with a + *different* kernel in each column ("stacked" kernel). + + - The NumPy FFT backend always uses ``'float64'/'complex128'`` + internally, so different data types will simply result in + additional casting, not speedup or different precision. + + Examples + -------- + Convolve in all axes: + + >>> space = odl.rn((3, 3)) + >>> kernel = [[0, 0, 0], # A discrete Dirac delta + ... [0, 1, 0], + ... [0, 0, 0]] + >>> conv = DiscreteConvolution(space, kernel) + >>> x = space.element([[1, 2, 3], + ... [2, 4, 6], + ... [-3, -6, -9]]) + >>> conv(x) + rn((3, 3)).element( + [[ 1., 2., 3.], + [ 2., 4., 6.], + [-3., -6., -9.]] + ) + + For even-sized kernels, the convolution is performed in a + "backwards" manner, i.e., the lower indices are affected by + implicit zero-padding: + + >>> kernel = [[1, 1], # 2x2 blurring kernel + ... [1, 1]] + >>> conv = DiscreteConvolution(space, kernel) + >>> x = space.element([[1, 2, 3], + ... [2, 4, 6], + ... [-3, -6, -9]]) + >>> conv(x) + rn((3, 3)).element( + [[ 1., 3., 5.], + [ 3., 9., 15.], + [ -1., -3., -5.]] + ) + + Convolution in selected axes can be done either with broadcasting + or with "stacked kernels": + + >>> kernel_1d = [1, -1] # backward difference kernel + >>> kernel = np.array(kernel_1d)[None, :] # broadcasting in axis 0 + >>> conv = DiscreteConvolution(space, kernel, axis=1) + >>> x = space.element([[1, 2, 3], + ... [2, 4, 6], + ... [-3, -6, -9]]) + >>> conv(x) + rn((3, 3)).element( + [[ 1., 1., 1.], + [ 2., 2., 2.], + [-3., -3., -3.]] + ) + >>> kernel_stack = [[1, -1], # separate kernel per row + ... [2, -2], + ... [3, -3]] + >>> conv = DiscreteConvolution(space, kernel_stack, axis=1) + >>> conv(x) + rn((3, 3)).element( + [[ 1., 1., 1.], + [ 4., 4., 4.], + [-9., -9., -9.]] + ) + """ + # Avoid name clash with range iterator + import builtins + range, ran = builtins.range, range + + if not isinstance(domain, TensorSpace): + raise TypeError('`domain` must be a TensorSpace, got {}' + ''.format(type(domain))) + + if not isinstance(kernel, Tensor): + try: + kernel = np.asarray(kernel, dtype=domain.dtype) + except TypeError: + raise ValueError( + 'cannot use a complex `kernel` with a real `domain`; ' + 'use a complex `domain` instead') + + ker_space = tensor_space(kernel.shape, kernel.dtype) + kernel = ker_space.element(kernel) + + if ran is None: + if str(impl).lower() == 'fft': + # Need a floating point dtype for FFT + result_dtype = np.result_type(domain.dtype, np.float16) + else: + result_dtype = domain.dtype + ran = domain.astype(result_dtype) + + # Disallow real-to-complex convolutions + if domain.is_real and ran.is_complex: + raise ValueError('cannot combine `domain` with real dtype {} ' + 'and `range` with complex dtype {}' + ''.format(dtype_str(domain.dtype), + dtype_str(ran.dtype))) + + super(DiscreteConvolution, self).__init__(domain, ran, linear=True) + self.__kernel = kernel + + ndim = self.domain.ndim + if kernel.ndim != ndim: + raise ValueError('`kernel` must have {} (=ndim) dimensions, but ' + 'got a {}-dimensional kernel' + ''.format(ndim, kernel.ndim)) + + if axis is None: + self.__axes = tuple(range(ndim)) + else: + try: + iter(axis) + except TypeError: + self.__axes = (int(axis),) + else: + self.__axes = tuple(int(ax) for ax in axis) + + if not all(-ndim <= ax < ndim for ax in self.axes): + raise ValueError('`axis` must (all) satisfy -{n} <= axis < {n}, ' + 'got {}'.format(axis, n=ndim)) + + for i in range(ndim): + if i in self.axes and kernel.shape[i] > self.domain.shape[i]: + raise ValueError( + 'kernel size in convolution axis {} can at most be equal ' + 'to domain size {}, but got size {}' + ''.format(i, self.domain.shape[i], kernel.shape[i])) + elif (i not in self.axes and + kernel.shape[i] not in (1, self.domain.shape[i])): + raise ValueError( + 'kernel size in non-convolution axis {} must be either 1 ' + '(broadcasting) or equal to domain size {}, but got size ' + '{}'.format(i, self.domain.shape[i], kernel.shape[i])) + + self.__impl = str(impl).lower() + if self.impl == 'real': + self.__real_impl = 'scipy' + self.__fft_impl = None + elif self.impl == 'fft': + self.__real_impl = None + self.__fft_impl = 'pyfftw' if PYFFTW_AVAILABLE else 'numpy' + else: + raise ValueError('unknown `impl` {!r}'.format(impl)) + + if self.impl == 'real': + for i in range(ndim): + if i not in self.axes and kernel.shape[i] != 1: + raise ValueError( + "for `impl='real', all non-convolution axes must " + 'have size 1, but got size {} in axis {}' + ''.format(kernel.shape[i], i)) + + # Handle padding and padded_shape + padding = kwargs.pop('padding', None) + padded_shape = kwargs.pop('padded_shape', None) + if padding is not None and padded_shape is not None: + raise TypeError('cannot give both `padding` and `padded_shape`') + + if padded_shape is not None and len(padded_shape) != ndim: + raise ValueError('`padded_shape` contains an invalid number of ' + 'entries: need {} (=ndim), got {}' + ''.format(ndim, len(padded_shape))) + + if padding is None: + full_padding = np.minimum(np.array(kernel.shape) - 1, 64) + padding = [full_padding[i] if i in self.axes else 0 + for i in range(ndim)] + else: + try: + iter(padding) + except TypeError: + padding = [int(padding) if i in self.axes else 0 + for i in range(ndim)] + else: + padding = [int(p) for p in padding] + if len(padding) == len(self.axes): + padding_lst = [0] * self.domain.ndim + for ax, pad in zip(self.axes, padding): + padding_lst[ax] = pad + padding = padding_lst + + if len(padding) != ndim: + raise ValueError('`padding` contains an invalid number of ' + 'entries: need {} (=ndim), got {}' + ''.format(ndim, len(padding))) + + if padded_shape is None: + padded_shape = tuple(np.array(self.domain.shape) + padding) + + self.__padded_shape = tuple(padded_shape) + + for i in range(ndim): + if self.padded_shape[i] < self.domain.shape[i]: + raise ValueError( + '`padded_shape` in axis {} must be larger than or equal ' + 'to domain size {}, but got {}' + ''.format(i, self.domain.shape[i], self.padded_shape[i])) + + self.__cache_kernel_ft = bool(kwargs.pop('cache_kernel_ft', False)) + self._kernel_ft = None + + variant = kwargs.pop('variant', 'forward') + self.__variant, variant_in = str(variant).lower(), variant + if self.variant not in ('forward', 'adjoint'): + raise ValueError('`variant` {!r} not understood' + ''.format(variant_in)) + + # Caching + self._adjoint = None + + if kwargs: + raise TypeError('got unexpected kwargs {}'.format(kwargs)) + + @property + def kernel(self): + """The `Tensor` used as kernel in the convolution.""" + return self.__kernel + + @property + def axes(self): + """The dimensions along which the convolution is taken.""" + return self.__axes + + @property + def impl(self): + """Implementation variant, ``'fft' or 'real'``.""" + return self.__impl + + @property + def real_impl(self): + """Backend for real-space conv., or ``None`` if not applicable.""" + return self.__real_impl + + @property + def fft_impl(self): + """Backend used for FFTs, or ``None`` if not applicable.""" + return self.__fft_impl + + @property + def padded_shape(self): + """Domain shape after padding for FFT-based convolution.""" + return self.__padded_shape + + @property + def variant(self): + """Convolution variant, ``'forward'`` or ``'adjoint'``.""" + return self.__variant + + @property + def cache_kernel_ft(self): + """If ``True``, the kernel FT is cached for later reuse.""" + return self.__cache_kernel_ft + + def _call(self, x, out=None): + """Perform convolution of ``f`` with `kernel`.""" + if self.impl == 'real' and self.real_impl == 'scipy': + res = self._call_scipy_convolve(x) + elif self.impl == 'fft' and self.fft_impl == 'numpy': + res = self._call_numpy_fft(x) + elif self.impl == 'fft' and self.fft_impl == 'pyfftw': + res = self._call_pyfftw(x, out=out) + else: + raise RuntimeError('bad `impl` {!r} or `fft_impl` {!r}' + ''.format(self.impl, self.fft_impl)) + if out is None: + out = res + else: + out[:] = res + return out + + def _call_scipy_convolve(self, x): + """Perform real-space convolution using ``scipy.signal.convolve``.""" + import scipy.signal + + if self.variant == 'forward': + conv = scipy.signal.convolve(x, self.kernel, mode='same', + method='direct') + elif self.variant == 'adjoint': + # We need the 'upper same' part of the correlation, not the + # lower one as in the 'same' mode. Unfortunately that mode + # doesn't exist so we implement it ourselves. + conv_full = scipy.signal.correlate(x, self.kernel, mode='full', + method='direct') + slc = [] + for i in range(self.domain.ndim): + if i in self.axes: + n_extra = self.kernel.shape[i] - 1 + right = n_extra // 2 + left = n_extra - right + slc.append(slice(left, conv_full.shape[i] - right)) + else: + slc.append(slice(None)) + conv = conv_full[slc] + + else: + raise RuntimeError('bad `variant` {!r}'.format(self.variant)) + + return conv + + def _call_numpy_fft(self, x): + """Perform FFT-based convolution using NumPy's backend.""" + # Use real-to-complex FFT if possible, it's faster + if (is_real_dtype(self.kernel.dtype) and + is_real_dtype(self.domain.dtype)): + fft = np.fft.rfftn + ifft = np.fft.irfftn + else: + fft = np.fft.fftn + ifft = np.fft.ifftn + + # Pad the input with zeros + paddings = [] + for i in range(self.domain.ndim): + diff = self.padded_shape[i] - x.shape[i] + left = diff // 2 + right = diff - left + paddings.append((left, right)) + + if any(p != (0, 0) for p in paddings): + x_prep = np.pad(x, paddings, 'constant') + else: + x_prep = np.asarray(x) + + # Perform FFTs of x and kernel (or retrieve from cache) + x_ft = fft(x_prep, axes=self.axes) + + if self._kernel_ft is not None: + kernel_ft = self._kernel_ft + else: + if self.variant == 'forward': + kernel = self.kernel + elif self.variant == 'adjoint': + # Flip kernel in conv axes for adjoint + slc = [slice(None, None, -1) if i in self.axes else slice(None) + for i in range(self.domain.ndim)] + kernel = self.kernel[slc] + else: + raise RuntimeError('bad `variant` {!r}'.format(self.variant)) + + # Prepare kernel, preserving length-1 axes for broadcasting + ker_padded_shp = [1 if self.kernel.shape[i] == 1 + else self.padded_shape[i] + for i in range(self.domain.ndim)] + + kernel_prep = _prepare_for_fft(kernel, ker_padded_shp, self.axes, + self.variant) + + kernel_ft = fft(kernel_prep, axes=self.axes) + if self.cache_kernel_ft: + self._kernel_ft = kernel_ft + + # Multiply `x_ft` with `kernel_ft` and transform back. Note that + # both have dtype 'float64' since that's what `numpy.fft` always uses. + x_ft *= kernel_ft + # `irfft` needs an explicit shape, otherwise the result shape may not + # be the same as the original one + s = [x_prep.shape[i] + for i in range(self.domain.ndim) if i in self.axes] + ifft_x = ifft(x_ft, axes=self.axes, s=s) + + # Unpad to get the "relevant" part + slc = [slice(l, n - r) for (l, r), n in zip(paddings, x_prep.shape)] + return ifft_x[slc] + + def _call_pyfftw(self, x, out=None): + """Perform FFT-based convolution using the pyfftw backend.""" + import multiprocessing + import pyfftw + + # Pad the input with zeros + paddings = [] + for i in range(self.domain.ndim): + diff = self.padded_shape[i] - x.shape[i] + left = diff // 2 + right = diff - left + paddings.append((left, right)) + + if any(p != (0, 0) for p in paddings): + x_prep = np.pad(x, paddings, 'constant') + else: + x_prep = np.asarray(x) + x_prep_shape = x_prep.shape + + # Real-to-halfcomplex only if both domain and kernel are eligible + use_halfcx = (is_real_dtype(self.domain.dtype) and + is_real_dtype(self.kernel.dtype)) + + def fft_out_array(arr, use_halfcx): + """Make an output array for FFTW with suitable dtype and shape.""" + ft_dtype = np.result_type(arr.dtype, 1j) + ft_shape = list(arr.shape) + if use_halfcx: + ft_shape[self.axes[-1]] = ft_shape[self.axes[-1]] // 2 + 1 + return np.empty(ft_shape, ft_dtype) + + # Perform FFT of `x`. Use 'FFTW_ESTIMATE', since other options destroy + # the input and would require a copy. + x_ft = fft_out_array(x_prep, use_halfcx) + if not use_halfcx and x_ft.dtype != x_prep.dtype: + # Need to perform C2C transform, hence a cast + x_prep = x_prep.astype(x_ft.dtype) + elif x_prep.dtype == 'float16': + # No native support for half floats + x_prep = x_prep.astype('float32') + + plan_x = pyfftw.FFTW(x_prep, x_ft, axes=self.axes, + direction='FFTW_FORWARD', + flags=['FFTW_ESTIMATE'], + threads=multiprocessing.cpu_count()) + plan_x(x_prep, x_ft) + plan_x = None # can be gc'ed + x_prep = None + + # Perform FFT of kernel if necessary + if self._kernel_ft is not None: + kernel_ft = self._kernel_ft + else: + if self.variant == 'forward': + kernel = self.kernel + elif self.variant == 'adjoint': + # Flip kernel in conv axes for adjoint + slc = [slice(None, None, -1) if i in self.axes else slice(None) + for i in range(self.domain.ndim)] + kernel = self.kernel[slc] + else: + raise RuntimeError('bad `variant` {!r}'.format(self.variant)) + + # Prepare kernel, preserving length-1 axes for broadcasting + if not is_floating_dtype(self.kernel.dtype): + flt_dtype = np.result_type(self.kernel.dtype, np.float16) + kernel = np.asarray(kernel, dtype=flt_dtype) + + ker_padded_shp = [1 if self.kernel.shape[i] == 1 + else self.padded_shape[i] + for i in range(self.domain.ndim)] + kernel_prep = _prepare_for_fft(kernel, ker_padded_shp, self.axes, + self.variant) + kernel = None # can be gc'ed + + kernel_ft = fft_out_array(kernel_prep, use_halfcx) + if not use_halfcx and kernel_ft.dtype != kernel_prep.dtype: + # Need to perform C2C transform, hence a cast + kernel_prep = kernel_prep.astype(kernel_ft.dtype) + elif kernel_prep.dtype == 'float16': + # No native support + kernel_prep = kernel_prep.astype('float32') + + plan_kernel = pyfftw.FFTW(kernel_prep, kernel_ft, axes=self.axes, + direction='FFTW_FORWARD', + flags=['FFTW_ESTIMATE'], + threads=multiprocessing.cpu_count()) + plan_kernel(kernel_prep, kernel_ft) + plan_kernel = None # can be gc'ed + kernel_prep = None + + if self.cache_kernel_ft: + self._kernel_ft = kernel_ft + + # Multiply x_ft with kernel_ft and transform back. Some care + # is required with respect to dtypes, in particular when + # x_ft.dtype < kernel_ft.dtype. + if x_ft.dtype < kernel_ft.dtype: + x_ft = x_ft * kernel_ft + else: + x_ft *= kernel_ft + + # Perform inverse FFT + if all(p == (0, 0) for p in paddings) and out is not None: + # No padding used, can write directly to `out` + with writable_array(out) as out_arr: + plan_ift = pyfftw.FFTW(x_ft, out_arr, axes=self.axes, + direction='FFTW_BACKWARD', + flags=['FFTW_ESTIMATE'], + threads=multiprocessing.cpu_count()) + plan_ift(x_ft, out_arr) + return out + else: + if use_halfcx: + x_ift_dtype = np.empty(0, dtype=x_ft.dtype).real.dtype + else: + x_ift_dtype = x_ft.dtype + x_ift = np.empty(x_prep_shape, x_ift_dtype) + plan_ift = pyfftw.FFTW(x_ft, x_ift, axes=self.axes, + direction='FFTW_BACKWARD', + flags=['FFTW_ESTIMATE'], + threads=multiprocessing.cpu_count()) + + plan_ift(x_ft, x_ift) + x_ft = None # can be gc'ed + + # Unpad to get the "relevant" part + slc = [slice(l, n - r) + for (l, r), n in zip(paddings, x_prep_shape)] + + return x_ift[slc] + + @property + @auto_weighting + def adjoint(self): + """Adjoint of the convolution operator. + + The adjoint convolution is a convolution with the adjoint + kernel, which is the (complex conjugate of the) original kernel, + (roughly) flipped in the convolution axes. See Notes. + + Examples + -------- + >>> conv = DiscreteConvolution(odl.rn(3), [1, -1]) + >>> conv.adjoint.kernel + rn(2).element([ 1., -1.]) + """ + if self._adjoint is None: + adj_variant = 'adjoint' if self.variant == 'forward' else 'forward' + self._adjoint = DiscreteConvolution( + self.range, self.kernel, range=self.domain, axis=self.axes, + impl=self.impl, padded_shape=self.padded_shape, + variant=adj_variant) + + return self._adjoint + + +def _prepare_for_fft(kernel, padded_shape, axes=None, variant='forward'): + """Return a kernel with desired shape with middle entry at index 0. + + This function applies the appropriate steps to prepare a kernel for + FFT-based convolution. It first pads the kernel with zeros *to the + right* up to ``padded_shape``, and then rolls the entries such that + the old middle element, i.e., the one at ``(kernel.shape - 1) // 2``, + lies at index 0. + + Parameters + ---------- + kernel : array-like + The kernel to be prepared for FFT convolution. + padded_shape : sequence of ints + The target shape to be reached by zero-padding. + axes : sequence of ints, optional + Dimensions in which to perform shifting. ``None`` means all axes. + variant : {'forward', 'adjoint'} + Convolution variant for which the kernel should be prepared. + + Returns + ------- + prepared : `numpy.ndarray` + The zero-padded and rolled kernel ready for FFT. + + Examples + -------- + >>> kernel = np.array([[1, 2, 3], + ... [4, 5, 6]]) # middle element is 2 + >>> _prepare_for_fft(kernel, padded_shape=(4, 4)) + array([[2, 3, 0, 1], + [5, 6, 0, 4], + [0, 0, 0, 0], + [0, 0, 0, 0]]) + >>> _prepare_for_fft(kernel, padded_shape=(5, 5)) + array([[2, 3, 0, 0, 1], + [5, 6, 0, 0, 4], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0], + [0, 0, 0, 0, 0]]) + """ + kernel = np.asarray(kernel) + if kernel.flags.f_contiguous and not kernel.flags.c_contiguous: + order = 'F' + else: + order = 'C' + + padded = np.zeros(padded_shape, kernel.dtype, order) + + if axes is None: + axes = list(range(kernel.ndim)) + + if any(padded_shape[i] != kernel.shape[i] for i in range(kernel.ndim) + if i not in axes): + raise ValueError( + '`padded_shape` can only differ from `kernel.shape` in `axes`; ' + 'got `padded_shape={}`, `kernel.shape={}`, `axes={}`' + ''.format(padded_shape, kernel.shape, axes)) + + orig_slc = [slice(n) for n in kernel.shape] + padded[orig_slc] = kernel + # This shift makes sure that the middle element is shifted to index 0, + # where "middle" means + # - actually middle (n-1)/2 for an odd number n + # - one below middle (n/2 - 1) for even n and 'forward' + # - one above middle (n/2) for even n and 'adjoint' + shifts = [] + for ax in axes: + if kernel.shape[ax] % 2 == 0 and variant == 'forward': + shifts.append(-(kernel.shape[ax] // 2 - 1)) + elif kernel.shape[ax] % 2 == 0 and variant == 'adjoint': + shifts.append(-(kernel.shape[ax] // 2)) + else: + shifts.append(-((kernel.shape[ax] - 1) // 2)) + return roll(padded, shifts, axis=axes) + + +def convolve(x, y, out=None, **kwargs): + """Return the convolution of ``x`` and ``y``. + + This is a convenience function for quickly computing a convolution + without having to explicitly create a `DiscreteConvolution` instance. + + Parameters + ---------- + x : array-like + Array or discrete function that is supposed to be convolved with + ``y``. Its type determines the return type. + y : array-like + The kernel with which ``x`` is convolved. It must have the same + number of dimensions as ``x``, and its shape can be at most equal + to ``x.shape``. In axes with size 1, broadcasting is applied. + + **Important:** + + - The kernel must **always** have the same number of dimensions + as ``x``, even for convolutions along axes. + - In axes where no convolution is performed, the shape of the + kernel must either be 1 (broadcasting along these axes), or + equal to the ``x.shape`` (stacked kernel, only supported + for `impl='fft'`). + - The ``'fft'`` implementation needs the kernel to have + floating-point ``dtype``, hence the smallest possible + float data type is used in that case to store the kernel. + - If the convolution kernel is complex, ``x`` will be cast to + complex dtype, and the returned object will be complex as well. + + See Examples for further clarification. + + out : `numpy.ndarray` or `Tensor`, optional + Object to which the result of the convolution should be written. + Its shape and data type must be compatible with the result of the + convolution, which can be determined by :: + + res_dtype = np.result_type(x.dtype, y.dtype, np.float16) + + axis : int or sequence of ints, optional + Coordinate axis or axes in which to take the convolution. + ``None`` means all input axes. + impl : {'fft', 'real'} + Implementation of the convolution as FFT-based or using + direct summation. The fastest available FFT backend is + chosen automatically. Real space convolution is based on + `scipy.signal.convolve`. + See Notes for further information on the backends. + padding : int or sequence of ints, optional + Zero-padding used before Fourier transform in the FFT backend. + Does not apply for ``impl='real'``. A sequence is applied per + axis, with padding values corresponding to ``axis`` entries + as provided. + Default: ``min(kernel.shape - 1, 64)`` + padded_shape : sequence of ints, optional + Apply zero-padding with this target shape. Cannot be used + together with ``padding``. + + Returns + ------- + convolved : `Tensor` + The convolution of ``x`` and ``y``. If ``x`` is of type + `DiscreteLpElement`, so is the result. Otherwise the returned + type is `NumpyTensor`. + + See Also + -------- + correlate + DiscreteConvolution + """ + y = np.asarray(y) + y_is_complex = issubclass(y.dtype.type, np.complexfloating) + + if not isinstance(x, (DiscreteLpElement, Tensor)): + x = np.asarray(x) + x = tensor_space(x.shape, dtype=x.dtype).element(x) + + dom_dtype = np.promote_types(x.dtype, y.dtype) if y_is_complex else x.dtype + if y_is_complex: + dom_dtype = np.promote_types(x.dtype, y.dtype) + else: + dom_dtype = x.dtype + + domain = x.space.astype(dom_dtype) + + conv = DiscreteConvolution(domain, y, **kwargs) + if out is None: + out = conv(x) + elif isinstance(out, np.ndarray): + res = conv.range.element(out) + if res.data is not out: + raise TypeError('`out` {!r} is not compatible with the range {!r}' + 'of the convolution'.format(out, conv.range)) + conv(x, out=res) + else: + res = conv.range.element(out) + if out is not res and out is not getattr(res, 'tensor', None): + raise TypeError('`out` {!r} is not compatible with the range {!r}' + 'of the convolution'.format(out, conv.range)) + conv(x, out=res) + + return out + + +def correlate(x, y, out=None, **kwargs): + """Return the cross-correlation of ``x`` and ``y``. + + This function computes the cross-correlation defined in continuum as + + .. math:: + [x \star y](t) = \int x(t + s)\, y(s)\, \,\mathrm{d}s. + + The order of ``x`` and ``y`` matters, i.e., this operation is + not commutative. + + For details on the function arguments, see `convolve`. + + See Also + -------- + convolve + DiscreteConvolution + """ + if 'variant' in kwargs: + raise TypeError('cannot use `variant` argument in `correlate`') + + kwargs['variant'] = 'adjoint' + return convolve(x, y, out, **kwargs) + + +if __name__ == '__main__': + from odl.util import run_doctests + run_doctests() diff --git a/odl/ufunc_ops/ufunc_ops.py b/odl/oplib/ufunc_ops.py similarity index 99% rename from odl/ufunc_ops/ufunc_ops.py rename to odl/oplib/ufunc_ops.py index af88b1f408a..7d08621efd3 100644 --- a/odl/ufunc_ops/ufunc_ops.py +++ b/odl/oplib/ufunc_ops.py @@ -96,7 +96,7 @@ def _is_integer_only_ufunc(name): -------- >>> import odl >>> space = odl.{space!r} ->>> op = odl.ufunc_ops.{name}(space) +>>> op = odl.oplib.{name}(space) >>> print(op({arg})) {result!s} """ @@ -382,14 +382,14 @@ def __repr__(self): RAW_UFUNC_FACTORY_FUNCTIONAL_DOCSTRING = """ Create functional with domain/range as real numbers: ->>> func = odl.ufunc_ops.{name}() +>>> func = odl.oplib.{name}() """ RAW_UFUNC_FACTORY_OPERATOR_DOCSTRING = """ Create operator that acts pointwise on a `TensorSpace` >>> space = odl.rn(3) ->>> op = odl.ufunc_ops.{name}(space) +>>> op = odl.oplib.{name}(space) """ diff --git a/odl/solvers/functional/default_functionals.py b/odl/solvers/functional/default_functionals.py index c382268e20f..ff12e079e6f 100644 --- a/odl/solvers/functional/default_functionals.py +++ b/odl/solvers/functional/default_functionals.py @@ -27,7 +27,8 @@ proximal_const_func, proximal_box_constraint, proximal_convex_conj_kl, proximal_convex_conj_kl_cross_entropy, combine_proximals, proximal_convex_conj) -from odl.util import conj_exponent, moveaxis +from odl.util import conj_exponent +from odl.util.npy_compat import moveaxis __all__ = ('ZeroFunctional', 'ConstantFunctional', 'ScalingFunctional', diff --git a/odl/space/space_utils.py b/odl/space/space_utils.py index 7beb9dff672..e574c2c5f9b 100644 --- a/odl/space/space_utils.py +++ b/odl/space/space_utils.py @@ -13,7 +13,8 @@ from odl.set import RealNumbers, ComplexNumbers from odl.space.entry_points import tensor_space_impl - +from odl.space.weighting import ArrayWeighting, ConstWeighting +from odl.util import OptionalArgDecorator __all__ = ('vector', 'tensor_space', 'cn', 'rn') @@ -281,6 +282,216 @@ def rn(shape, dtype=None, impl='numpy', **kwargs): return rn +class auto_weighting(OptionalArgDecorator): + + """Make an unweighted adjoint automatically account for weightings. + + Depending on the weightings, the correction is achieved by composing + the unweighted operator with either `ScalingOperator` or + `ConstantOperator`. The following rules are applied for the domain + weighting ``w``, the range weighting ``v`` and the provided unweighted + adjoint ``B^*``: + + - If both ``w`` and ``v`` are arrays, return :: + + (1 / w) * (B^*) * v + + - If ``w`` is an array and ``v`` a constant, return :: + + (v / w) * (B^*) + + - If ``w`` is a constant and ``v`` an array, return :: + + (B^*) * (w / v) + + - If both ``w`` and ``v`` are constants, return :: + + (B^*) * (v / w) + + if ``B.range.size < B.domain.size``, otherwise :: + + (v / w) * (B^*) + + - Ignore constants 1.0. + + To avoid the inconvenience of dealing with `OperatorComp` objects, + the given operator is monkey-patched instead of composed. + + Parameters + ---------- + unweighted_adjoint : `Operator` + Unweighted variant of the adjoint. It will be patched with a + new ``_call()`` method. + The weightings of ``domain`` and ``range`` of the operator + must be `ArrayWeighting` or `ConstWeighting`. + optimize : bool, optional + If ``True``, merge and move around constant weightings for + highest expected efficiency. + + Notes + ----- + Consider a linear operator :math:`A: X_w \\to Y_v` between spaces with + weights :math:`w` and :math:`v`, respectively, along with the same + operator :math:`B: X \\to Y` defined between the unweighted variants of + the spaces. (This means that :math:`B f = A f` for all + :math:`f \\in X \cong X_w`). + + Then, the adjoint of :math:`A` is related to the adjoint of :math:`B` + as follows: + + .. math:: + \\langle Af, g \\rangle_{Y_v} = + \\langle Bf, v \cdot g \\rangle_Y = + \\langle f, B^*(v \cdot g) \\rangle_X = + \\langle f, w^{-1}\, B^*(v \cdot g) \\rangle_{X_w}. + + Thus, from the existing unweighted adjoint :math:`B^*` one can compute + the weighted one as :math:`A^* = w^{-1}\, B^*(v\, \cdot)`. + Depending on the types of weighting, this expression can be simplified + further, e.g., a constant weight can be absorbed into the other weight. + """ + + @staticmethod + def _wrapper(unweighted_adjoint, optimize=True): + """Return the weighted variant of the unweighted adjoint.""" + # Support decorating the `adjoint` property directly + import inspect + from functools import wraps + from odl.operator.operator import Operator + + if (inspect.isfunction(unweighted_adjoint) and + unweighted_adjoint.__name__ == 'adjoint'): + # We need this level of indirection since `self` needs to + # be filled in with the instance, but we decorate at class + # level. + @wraps(unweighted_adjoint) + def weighted_adjoint(self): + adj = unweighted_adjoint(self) + if not isinstance(adj, Operator): + raise TypeError('`adjoint` did not return an `Operator`') + if adj is self: + raise TypeError( + 'returning `self` in an `adjoint` property using ' + '`auto_weighting` is not allowed') + + # This is for cached adjoints: don't double-wrap + if hasattr(adj, '_call_unweighted'): + return adj + else: + return auto_weighting._instance_wrapper(adj, optimize) + + return weighted_adjoint + + else: + raise TypeError( + "`auto_weighting` can only be applied to 'adjoint' methods " + '(@auto_weighting decorator)') + + @staticmethod + def _instance_wrapper(unweighted_adjoint, optimize=True): + """Wrapper for `Operator` instances.""" + # Use notions of the original operator, not the adjoint + dom_weighting = unweighted_adjoint.range.weighting + ran_weighting = unweighted_adjoint.domain.weighting + + if isinstance(dom_weighting, ArrayWeighting): + dom_w_type = 'array' + dom_w = dom_weighting.array + elif isinstance(dom_weighting, ConstWeighting): + dom_w_type = 'const' + dom_w = dom_weighting.const + else: + raise TypeError( + 'weighting of `unweighted_adjoint.range` must be of ' + 'type `ArrayWeighting` or `ConstWeighting`, got {}' + ''.format(type(dom_weighting))) + + if isinstance(ran_weighting, ArrayWeighting): + ran_w_type = 'array' + ran_w = ran_weighting.array + elif isinstance(ran_weighting, ConstWeighting): + ran_w_type = 'const' + ran_w = ran_weighting.const + else: + raise TypeError( + 'weighting of `unweighted_adjoint.domain` must be of ' + 'type `ArrayWeighting` or `ConstWeighting`, got {}' + ''.format(type(ran_weighting))) + + # Compute the effective weights and mark constants 1.0 as to be + # skipped + if not optimize: + new_dom_w, new_ran_w = dom_w, ran_w + skip_dom = dom_w_type == 'const' and dom_w == 1.0 + skip_ran = ran_w_type == 'const' and ran_w == 1.0 + elif dom_w_type == 'array' and ran_w_type == 'array': + new_dom_w, new_ran_w = dom_w, ran_w + skip_dom = skip_ran = False + elif dom_w_type == 'array' and ran_w_type == 'const': + new_dom_w = dom_w / ran_w + new_ran_w = 1.0 + skip_dom = False + skip_ran = True + elif dom_w_type == 'const' and ran_w_type == 'array': + new_dom_w = 1.0 + new_ran_w = ran_w / dom_w + skip_dom = True + skip_ran = False + elif dom_w_type == 'const' and ran_w_type == 'const': + if unweighted_adjoint.domain.size < unweighted_adjoint.range.size: + new_dom_w = 1.0 + new_ran_w = ran_w / dom_w + skip_dom = True + skip_ran = False + else: + new_dom_w = dom_w / ran_w + new_ran_w = 1.0 + skip_dom = False + skip_ran = True + + # Define the new `_call` depending on original signature + self = unweighted_adjoint + + # Monkey-patching starts here + if self._call_has_out and self._call_out_optional: + def _call(x, out=None): + if not skip_ran: + x = new_ran_w * x + out = self._call_unweighted(x, out=out) + if not skip_dom: + out /= new_dom_w + return out + + self._call_unweighted = self._call_in_place + self._call_in_place = self._call_out_of_place = _call + + elif self._call_has_out and not self._call_out_optional: + def _call(x, out): + if not skip_ran: + x = new_ran_w * x + self._call_unweighted(x, out=out) + if not skip_dom: + out /= new_dom_w + return out + + self._call_unweighted = self._call_in_place + self._call_in_place = _call + + else: + def _call(x): + if not skip_ran: + x = new_ran_w * x + out = self._call_unweighted(x) + if not skip_dom: + out /= new_dom_w + return out + + self._call_unweighted = self._call_out_of_place + self._call_out_of_place = _call + + return self + + if __name__ == '__main__': from odl.util.testutils import run_doctests run_doctests() diff --git a/odl/test/oplib/convolution_test.py b/odl/test/oplib/convolution_test.py new file mode 100644 index 00000000000..22f758dd35d --- /dev/null +++ b/odl/test/oplib/convolution_test.py @@ -0,0 +1,457 @@ +# Copyright 2014-2017 The ODL contributors +# +# This file is part of ODL. +# +# This Source Code Form is subject to the terms of the Mozilla Public License, +# v. 2.0. If a copy of the MPL was not distributed with this file, You can +# obtain one at https://mozilla.org/MPL/2.0/. + +from __future__ import division +import numpy as np +import pytest +import scipy.signal + +import odl +from odl.oplib import DiscreteConvolution +from odl.util.testutils import ( + simple_fixture, noise_elements, all_equal, all_almost_equal, noise_element) + + +# --- pytest fixtures --- # + +size_1d = simple_fixture('size', [3, 16]) +shape_2d = simple_fixture('shape', [(3, 4), (3, 8), (16, 4)]) +ker_kind = simple_fixture('kind', ['smooth', 'diff1', 'diff2']) +conv_type = simple_fixture('type', ['full', + 'bcast0', 'bcast1', + 'stack0', 'stack1']) +conv_impl = simple_fixture('impl', ['fft', 'real']) + + +@pytest.fixture(scope='module') +def kernel_1d(ker_kind): + if ker_kind == 'smooth': + kernel = [1, 1] + elif ker_kind == 'diff1': + kernel = [1, -1] + elif ker_kind == 'diff2': + kernel = [1, -2, 1] + else: + assert False + + return ker_kind, np.array(kernel) + + +# --- DiscreteConvolution --- # + + +def test_dconv_init_and_properties(): + """Check init, error handling and props of ``DiscreteConvolution``.""" + # 1D case + r3 = odl.rn(3, dtype='float32') + conv = DiscreteConvolution(r3, [1, 1]) + assert conv.domain == conv.range == r3 + assert conv.is_linear + assert conv.padded_shape == (4,) # adds kernel_size - 1 by default + assert conv.kernel.space == odl.rn(2, dtype='float32') + assert conv.axes == (0,) + assert conv.impl == 'fft' + assert conv.real_impl is None + + ran = odl.rn(3, dtype=float) + conv = DiscreteConvolution(r3, [1, 1], range=ran) + assert conv.range == ran + + conv = DiscreteConvolution(r3, [1, 1], axis=0) + assert conv.axes == (0,) + + conv = DiscreteConvolution(r3, [1, 1], impl='real') + assert conv.impl == 'real' + assert conv.real_impl == 'scipy' + + conv = DiscreteConvolution(r3, [1, 1], impl='fft') + assert conv.impl == 'fft' + + conv = DiscreteConvolution(r3, [1, 1], padding=4) + assert conv.padded_shape == (7,) + + conv = DiscreteConvolution(r3, [1, 1], padding=(4,)) + assert conv.padded_shape == (7,) + + conv = DiscreteConvolution(r3, [1, 1], padded_shape=(8,)) + assert conv.padded_shape == (8,) + + with pytest.raises(ValueError): + DiscreteConvolution(r3, [[1, 1], + [1, 1]]) # kernel has too many dims + with pytest.raises(ValueError): + DiscreteConvolution(r3, [1, 1, 1, 1]) # kernel too large + with pytest.raises(ValueError): + DiscreteConvolution(r3, [1j, 1j]) # complex kernel with real domain + with pytest.raises(ValueError): + # Real to complex not supported + DiscreteConvolution(r3, [1, 1], range=odl.cn(3, dtype=complex)) + with pytest.raises(ValueError): + DiscreteConvolution(r3, [1, 1], axis=1) # axis out of bounds + with pytest.raises(ValueError): + DiscreteConvolution(r3, [1, 1], impl='pyfftw') # bad impl + with pytest.raises(ValueError): + DiscreteConvolution(r3, [1, 1], padding=(2, 2)) # padding too long + with pytest.raises(ValueError): + # padded_shape too long + DiscreteConvolution(r3, [1, 1], padded_shape=(8, 8)) + with pytest.raises(ValueError): + # padded_shape cannot be smaller than original shape + DiscreteConvolution(r3, [1, 1], padded_shape=(2,)) + with pytest.raises(TypeError): + # cannot give both padding and padded_shape + DiscreteConvolution(r3, [1, 1], padding=3, padded_shape=(9,)) + with pytest.raises(TypeError): + DiscreteConvolution(r3, [1, 1], arg=0) # bad kwarg + + # 2D case, full kernel + rn = odl.rn((3, 4)) + kernel_full = [[1, 1], + [-1, -1]] + + conv = DiscreteConvolution(rn, kernel_full) + assert conv.domain == conv.range == rn + assert conv.padded_shape == (4, 5) # + kernel.shape - 1 + assert conv.kernel.space == odl.rn((2, 2)) + assert conv.axes == (0, 1) + + ran = odl.rn((3, 4), dtype=float) + conv = DiscreteConvolution(rn, kernel_full, range=ran) + assert conv.range == ran + + conv = DiscreteConvolution(rn, kernel_full, axis=(1, 0)) + assert conv.axes == (1, 0) + + conv = DiscreteConvolution(rn, kernel_full, padding=4) + assert conv.padded_shape == (7, 8) + + conv = DiscreteConvolution(rn, kernel_full, padding=(4, 8)) + assert conv.padded_shape == (7, 12) + + conv = DiscreteConvolution(rn, kernel_full, padded_shape=(8, 8)) + assert conv.padded_shape == (8, 8) + + with pytest.raises(ValueError): + DiscreteConvolution(rn, [1, 1]) # kernel has too few dims + with pytest.raises(ValueError): + kernel_too_large_0 = [[1, 1]] * 5 + DiscreteConvolution(rn, kernel_too_large_0) + with pytest.raises(ValueError): + kernel_too_large_1 = [[1, 1, 1, 1, 1]] * 2 + DiscreteConvolution(rn, kernel_too_large_1) + with pytest.raises(ValueError): + DiscreteConvolution(rn, kernel_full, axis=0) # too short in axis 1 + with pytest.raises(ValueError): + DiscreteConvolution(rn, kernel_full, axis=1) # too short in axis 0 + with pytest.raises(ValueError): + DiscreteConvolution(rn, kernel_full, axis=(0, 1, 2)) # axis oob + with pytest.raises(ValueError): + # padding too long + DiscreteConvolution(rn, kernel_full, padding=(2, 2, 2)) + with pytest.raises(ValueError): + # padded_shape too long + DiscreteConvolution(rn, kernel_full, padded_shape=(8, 8, 8)) + + # 2D case, conv along axis 0, broadcast along axis 1 + rn = odl.rn((3, 4)) + kernel_bcast = [[1], + [1]] + + conv = DiscreteConvolution(rn, kernel_bcast) + assert conv.kernel.space == odl.rn((2, 1)) + assert conv.padded_shape == (4, 4) # + kernel_shape - 1 + assert conv.axes == (0, 1) + + conv = DiscreteConvolution(rn, kernel_bcast, axis=0) + assert conv.padded_shape == (4, 4) # + kernel_shape - 1 in conv axis + assert conv.axes == (0,) + + conv = DiscreteConvolution(rn, kernel_bcast, axis=0, impl='real') + assert conv.impl == 'real' + + conv = DiscreteConvolution(rn, kernel_bcast, axis=0, padding=4) + assert conv.padded_shape == (7, 4) # pad only in conv axis + + conv = DiscreteConvolution(rn, kernel_bcast, axis=0, padded_shape=(8, 4)) + assert conv.padded_shape == (8, 4) + + # 2D case, conv along axis 0, stack along axis 1 + rn = odl.rn((3, 4)) + kernel_stack = [[1, 1, 1, 1], + [1, 1, 1, 1]] + conv = DiscreteConvolution(rn, kernel_stack, axis=0) + assert conv.kernel.space == odl.rn((2, 4)) + assert conv.padded_shape == (4, 4) # pad only in conv axis + + conv = DiscreteConvolution(rn, kernel_stack, axis=0, padding=4) + assert conv.padded_shape == (7, 4) # pad only in conv axis + + with pytest.raises(ValueError): + # stacked kernels not supported by 'real' impl + DiscreteConvolution(rn, kernel_stack, axis=0, impl='real') + + # 2D case, conv along axis 1, broadcast along axis 0 + rn = odl.rn((3, 4)) + kernel_bcast = [[1, 1]] + + conv = DiscreteConvolution(rn, kernel_bcast) + assert conv.kernel.space == odl.rn((1, 2)) + assert conv.padded_shape == (3, 5) # + kernel_shape - 1 + assert conv.axes == (0, 1) + + conv = DiscreteConvolution(rn, kernel_bcast, axis=1) + assert conv.padded_shape == (3, 5) # + kernel_shape - 1 in conv axis + assert conv.axes == (1,) + + conv = DiscreteConvolution(rn, kernel_bcast, axis=1, impl='real') + assert conv.impl == 'real' + + conv = DiscreteConvolution(rn, kernel_bcast, axis=1, padding=4) + assert conv.padded_shape == (3, 8) # pad only in conv axis + + conv = DiscreteConvolution(rn, kernel_bcast, axis=1, padded_shape=(3, 8)) + assert conv.padded_shape == (3, 8) + + # 2D case, conv along axis 1, stack along axis 0 + rn = odl.rn((3, 4)) + kernel_stack = [[1, 1], + [1, 1], + [1, 1]] + conv = DiscreteConvolution(rn, kernel_stack, axis=1) + assert conv.kernel.space == odl.rn((3, 2)) + assert conv.padded_shape == (3, 5) # pad only in conv axis + + conv = DiscreteConvolution(rn, kernel_stack, axis=1, padding=4) + assert conv.padded_shape == (3, 8) # pad only in conv axis + + with pytest.raises(ValueError): + # stacked kernels not supported by 'real' impl + DiscreteConvolution(rn, kernel_stack, axis=1, impl='real') + + +def test_dconv_1d(size_1d, kernel_1d, floating_dtype, conv_impl): + """Check discrete convolution in 1d.""" + ker_kind, kernel = kernel_1d + rn = odl.tensor_space(size_1d, dtype=floating_dtype) + conv = DiscreteConvolution(rn, kernel, impl=conv_impl) + + inp_arr, inp = noise_elements(rn) + + if ker_kind == 'smooth': + expected = scipy.signal.convolve(inp_arr, [1, 1], mode='same') + elif ker_kind == 'diff1': + # Backward diff, need to add first element in the beginning + expected = np.concatenate([[inp_arr[0]], np.diff(inp_arr, n=1)]) + elif ker_kind == 'diff2': + # Second diff, need to add the truncated convolution at the beginning + # and the end + expected = np.concatenate([[-2 * inp_arr[0] + inp_arr[1]], + np.diff(inp_arr, n=2), + [inp_arr[-2] - 2 * inp_arr[-1]]]) + else: + assert False + + # Make sure we don't compare with too high precision + expected = expected.astype(conv.range.dtype, copy=False) + + conv_res = conv(inp) + assert all_almost_equal(conv_res, expected.astype(conv_res.dtype)) + + out = conv.range.element() + conv(inp, out=out) + assert all_almost_equal(out, expected.astype(conv_res.dtype)) + + +def test_dconv_2d(shape_2d, kernel_1d, conv_type, conv_impl, floating_dtype): + """Check discrete convolution in 2d.""" + ker_kind, ker_1d = kernel_1d + + if conv_type.startswith('stack') and conv_impl == 'real': + pytest.skip('stacked kernels not supported in real-space convolution') + + if floating_dtype == 'float16' and conv_impl == 'real': + pytest.xfail('bug in scipy.signal.convolve for half float') + + if conv_type == 'full': + axis = None + else: + axis = int(conv_type[-1]) + + rn = odl.tensor_space(shape_2d, dtype=floating_dtype) + if conv_type == 'full': + kernel = np.outer(ker_1d, ker_1d) + elif conv_type == 'bcast0': + kernel = ker_1d[:, None] + elif conv_type == 'bcast1': + kernel = ker_1d[None, :] + elif conv_type == 'stack0': + kernel = ker_1d[:, None] * np.arange(shape_2d[1])[None, :] + elif conv_type == 'stack1': + kernel = ker_1d[None, :] * np.arange(shape_2d[0])[:, None] + else: + assert False + + conv = DiscreteConvolution(rn, kernel, impl=conv_impl, axis=axis) + + inp_arr, inp = noise_elements(rn) + + # Reference impl of stacked convolutions + def conv_stack_0(arr_2d, ker_1d): + cols = [ + i * scipy.signal.convolve(arr_2d[:, i], ker_1d, mode='same') + for i in range(arr_2d.shape[1])] + return np.hstack([col[:, None] for col in cols]) + + def conv_stack_1(arr_2d, ker_1d): + rows = [ + i * scipy.signal.convolve(arr_2d[i], ker_1d, mode='same') + for i in range(arr_2d.shape[0])] + return np.vstack([row[None, :] for row in rows]) + + if ker_kind == 'smooth': + if conv_type == 'full': + expected = scipy.signal.convolve(inp_arr, [[1, 1], + [1, 1]], mode='same') + elif conv_type == 'bcast0': + expected = scipy.signal.convolve(inp_arr, [[1], + [1]], mode='same') + elif conv_type == 'bcast1': + expected = scipy.signal.convolve(inp_arr, [[1, 1]], mode='same') + elif conv_type == 'stack0': + expected = conv_stack_0(inp_arr, [1, 1]) + elif conv_type == 'stack1': + expected = conv_stack_1(inp_arr, [1, 1]) + + elif ker_kind == 'diff1': + if conv_type == 'full': + padded = np.pad(inp_arr, (1, 0), mode='constant') + expected = np.diff(np.diff(padded, n=1, axis=0), n=1, axis=1) + elif conv_type == 'bcast0': + padded = np.pad(inp_arr, [(1, 0), (0, 0)], mode='constant') + expected = np.diff(padded, n=1, axis=0) + elif conv_type == 'bcast1': + padded = np.pad(inp_arr, [(0, 0), (1, 0)], mode='constant') + expected = np.diff(padded, n=1, axis=1) + elif conv_type == 'stack0': + padded = np.pad(inp_arr, [(1, 0), (0, 0)], mode='constant') + expected = (np.diff(padded, n=1, axis=0) * + np.arange(inp_arr.shape[1])[None, :]) + elif conv_type == 'stack1': + padded = np.pad(inp_arr, [(0, 0), (1, 0)], mode='constant') + expected = (np.diff(padded, n=1, axis=1) * + np.arange(inp_arr.shape[0])[:, None]) + + elif ker_kind == 'diff2': + if conv_type == 'full': + padded = np.pad(inp_arr, (1, 1), mode='constant') + expected = np.diff(np.diff(padded, n=2, axis=0), n=2, axis=1) + elif conv_type == 'bcast0': + padded = np.pad(inp_arr, [(1, 1), (0, 0)], mode='constant') + expected = np.diff(padded, n=2, axis=0) + elif conv_type == 'bcast1': + padded = np.pad(inp_arr, [(0, 0), (1, 1)], mode='constant') + expected = np.diff(padded, n=2, axis=1) + elif conv_type == 'stack0': + padded = np.pad(inp_arr, [(1, 1), (0, 0)], mode='constant') + expected = (np.diff(padded, n=2, axis=0) * + np.arange(inp_arr.shape[1])[None, :]) + elif conv_type == 'stack1': + padded = np.pad(inp_arr, [(0, 0), (1, 1)], mode='constant') + expected = (np.diff(padded, n=2, axis=1) * + np.arange(inp_arr.shape[0])[:, None]) + + else: + assert False + + # Make sure we don't compare with too high precision + expected = expected.astype(conv.range.dtype, copy=False) + + conv_res = conv(inp) + assert all_almost_equal(conv_res, expected.astype(conv_res.dtype)) + + out = conv.range.element() + conv(inp, out=out) + assert all_almost_equal(out, expected.astype(conv_res.dtype)) + + +def test_dconv_no_padding_out(): + """Test special case of no padding with direct writing to out.""" + rn = odl.rn((3, 4)) + conv = DiscreteConvolution(rn, [[1, 1], + [1, 1]], impl='fft') + out = conv.range.element() + x = noise_element(conv.domain) + conv_x = conv(x) + conv(x, out=out) + assert all_equal(conv_x, out) + + cn = odl.cn((3, 4)) + conv = DiscreteConvolution(cn, [[1, 1], + [1, 1]], impl='fft') + out = conv.range.element() + x = noise_element(conv.domain) + conv_x = conv(x) + conv(x, out=out) + assert all_equal(conv_x, out) + + +def test_dconv_kernel_caching(): + """Check if the kernel caching in DiscreteConvolution works.""" + rn = odl.rn((3, 4)) + conv = DiscreteConvolution(rn, [[1, 1], + [1, 1]], cache_kernel_ft=True, impl='fft') + + conv(rn.one()) + assert conv._kernel_ft is not None + + +def test_dconv_adjoint(shape_2d, kernel_1d, conv_type, floating_dtype, + conv_impl): + """Check if the adjoint of DiscreteConvolution is correct.""" + ker_kind, ker_1d = kernel_1d + + if conv_type.startswith('stack') and conv_impl == 'real': + pytest.skip('stacked kernels not supported in real-space convolution') + + if floating_dtype == 'float16' and conv_impl == 'real': + pytest.xfail('bug in scipy.signal.convolve for half float') + + if conv_type == 'full': + axis = None + else: + axis = int(conv_type[-1]) + + spc = odl.tensor_space(shape_2d, dtype=floating_dtype) + spc_w = odl.tensor_space(shape_2d, dtype=floating_dtype, weighting=2.0) + if conv_type == 'full': + kernel = np.outer(ker_1d, ker_1d) + elif conv_type == 'bcast0': + kernel = ker_1d[:, None] + elif conv_type == 'bcast1': + kernel = ker_1d[None, :] + elif conv_type == 'stack0': + kernel = ker_1d[:, None] * np.arange(shape_2d[1])[None, :] + elif conv_type == 'stack1': + kernel = ker_1d[None, :] * np.arange(shape_2d[0])[:, None] + else: + assert False + + # Use enough padding so the adjoint convolution isn't too far off + kwargs = {'padding': 4} if conv_impl == 'fft' else {} + conv = DiscreteConvolution(spc, kernel, axis=axis, range=spc_w, + impl=conv_impl, **kwargs) + dom_el = noise_element(conv.domain) + ran_el = noise_element(conv.range) + # Don't be too harsh when comparing + rtol = 10 * np.prod(shape_2d) * np.finfo(floating_dtype).resolution + assert (conv(dom_el).inner(ran_el) == + pytest.approx(dom_el.inner(conv.adjoint(ran_el)), rel=rtol)) + + +if __name__ == '__main__': + odl.util.test_file(__file__) diff --git a/odl/test/space/space_utils_test.py b/odl/test/space/space_utils_test.py index 661325021ae..ed52d5378c5 100644 --- a/odl/test/space/space_utils_test.py +++ b/odl/test/space/space_utils_test.py @@ -1,4 +1,4 @@ -# Copyright 2014-2017 The ODL contributors +# Copyright 2014-2019 The ODL contributors # # This file is part of ODL. # @@ -7,12 +7,19 @@ # obtain one at https://mozilla.org/MPL/2.0/. from __future__ import division + import numpy as np +import pytest import odl from odl import vector from odl.space.npy_tensors import NumpyTensor -from odl.util.testutils import all_equal +from odl.space.space_utils import auto_weighting +from odl.util.testutils import all_equal, simple_fixture, noise_element + +auto_weighting_optimize = simple_fixture('optimize', [True, False]) +call_variant = simple_fixture('call_variant', ['oop', 'ip', 'dual']) +weighting = simple_fixture('weighting', [1.0, 2.0, [1.0, 2.0]]) def test_vector_numpy(): @@ -76,5 +83,169 @@ def test_vector_numpy(): assert x.shape == (0,) +def test_auto_weighting(call_variant, weighting, auto_weighting_optimize): + """Test the auto_weighting decorator for different adjoint variants.""" + rn = odl.rn(2) + rn_w = odl.rn(2, weighting=weighting) + + class ScalingOpBase(odl.Operator): + + def __init__(self, dom, ran, c): + super(ScalingOpBase, self).__init__(dom, ran, linear=True) + self.c = c + + if call_variant == 'oop': + + class ScalingOp(ScalingOpBase): + + def _call(self, x): + return self.c * x + + @property + @auto_weighting(optimize=auto_weighting_optimize) + def adjoint(self): + return ScalingOp(self.range, self.domain, self.c) + + elif call_variant == 'ip': + + class ScalingOp(ScalingOpBase): + + def _call(self, x, out): + out[:] = self.c * x + return out + + @property + @auto_weighting(optimize=auto_weighting_optimize) + def adjoint(self): + return ScalingOp(self.range, self.domain, self.c) + + elif call_variant == 'dual': + + class ScalingOp(ScalingOpBase): + + def _call(self, x, out=None): + if out is None: + out = self.c * x + else: + out[:] = self.c * x + return out + + @property + @auto_weighting(optimize=auto_weighting_optimize) + def adjoint(self): + return ScalingOp(self.range, self.domain, self.c) + + else: + assert False + + op1 = ScalingOp(rn, rn_w, 1.5) + op2 = ScalingOp(rn_w, rn, 1.5) + + for op in [op1, op2]: + dom_el = noise_element(op.domain) + ran_el = noise_element(op.range) + assert pytest.approx(op(dom_el).inner(ran_el), + dom_el.inner(op.adjoint(ran_el))) + + +def test_auto_weighting_noarg(): + """Test the auto_weighting decorator without the optimize argument.""" + rn = odl.rn(2) + rn_w = odl.rn(2, weighting=2) + + class ScalingOp(odl.Operator): + + def __init__(self, dom, ran, c): + super(ScalingOp, self).__init__(dom, ran, linear=True) + self.c = c + + def _call(self, x): + return self.c * x + + @property + @auto_weighting + def adjoint(self): + return ScalingOp(self.range, self.domain, self.c) + + op1 = ScalingOp(rn, rn, 1.5) + op2 = ScalingOp(rn_w, rn_w, 1.5) + op3 = ScalingOp(rn, rn_w, 1.5) + op4 = ScalingOp(rn_w, rn, 1.5) + + for op in [op1, op2, op3, op4]: + dom_el = noise_element(op.domain) + ran_el = noise_element(op.range) + assert pytest.approx(op(dom_el).inner(ran_el), + dom_el.inner(op.adjoint(ran_el))) + + +def test_auto_weighting_cached_adjoint(): + """Check if auto_weighting plays well with adjoint caching.""" + rn = odl.rn(2) + rn_w = odl.rn(2, weighting=2) + + class ScalingOp(odl.Operator): + + def __init__(self, dom, ran, c): + super(ScalingOp, self).__init__(dom, ran, linear=True) + self.c = c + self._adjoint = None + + def _call(self, x): + return self.c * x + + @property + @auto_weighting + def adjoint(self): + if self._adjoint is None: + self._adjoint = ScalingOp(self.range, self.domain, self.c) + return self._adjoint + + op = ScalingOp(rn, rn_w, 1.5) + dom_el = noise_element(op.domain) + op_eval_before = op(dom_el) + + adj = op.adjoint + adj_again = op.adjoint + assert adj_again is adj + + # Check that original op is intact + assert not hasattr(op, '_call_unweighted') # op shouldn't be mutated + op_eval_after = op(dom_el) + assert all_equal(op_eval_before, op_eval_after) + + dom_el = noise_element(op.domain) + ran_el = noise_element(op.range) + op(dom_el) + op.adjoint(ran_el) + assert pytest.approx(op(dom_el).inner(ran_el), + dom_el.inner(op.adjoint(ran_el))) + + +def test_auto_weighting_raise_on_return_self(): + """Check that auto_weighting raises when adjoint returns self.""" + rn = odl.rn(2) + + class InvalidScalingOp(odl.Operator): + + def __init__(self, dom, ran, c): + super(InvalidScalingOp, self).__init__(dom, ran, linear=True) + self.c = c + self._adjoint = None + + def _call(self, x): + return self.c * x + + @property + @auto_weighting + def adjoint(self): + return self + + # This would be a vaild situation for adjont just returning self + op = InvalidScalingOp(rn, rn, 1.5) + with pytest.raises(TypeError): + op.adjoint + + if __name__ == '__main__': odl.util.test_file(__file__) diff --git a/odl/ufunc_ops/__init__.py b/odl/ufunc_ops/__init__.py deleted file mode 100644 index 0dd0ece6807..00000000000 --- a/odl/ufunc_ops/__init__.py +++ /dev/null @@ -1,16 +0,0 @@ -# Copyright 2014-2017 The ODL contributors -# -# This file is part of ODL. -# -# This Source Code Form is subject to the terms of the Mozilla Public License, -# v. 2.0. If a copy of the MPL was not distributed with this file, You can -# obtain one at https://mozilla.org/MPL/2.0/. - -"""Universal functions as `Operator` and `Functional`.""" - -from __future__ import absolute_import - -__all__ = () - -from .ufunc_ops import * -__all__ = ufunc_ops.__all__ diff --git a/odl/util/npy_compat.py b/odl/util/npy_compat.py index 54306cc10dc..3dda88740b4 100644 --- a/odl/util/npy_compat.py +++ b/odl/util/npy_compat.py @@ -12,10 +12,12 @@ import numpy as np -__all__ = ('moveaxis', 'flip') +__all__ = ('moveaxis', 'flip', 'roll') + + +# --- Numpy 1.11 --- # -# TODO: Remove when Numpy 1.11 is an ODL dependency def moveaxis(a, source, destination): """Move axes of an array to new positions. @@ -28,31 +30,33 @@ def moveaxis(a, source, destination): -------- numpy.moveaxis """ - import numpy - if hasattr(numpy, 'moveaxis'): - return numpy.moveaxis(a, source, destination) + major, minor, _ = [int(s) for s in np.version.short_version.split('.')] + if (major, minor) >= (1, 11): + return np.moveaxis(a, source, destination) + else: + try: + source = list(source) + except TypeError: + source = [source] + try: + destination = list(destination) + except TypeError: + destination = [destination] - try: - source = list(source) - except TypeError: - source = [source] - try: - destination = list(destination) - except TypeError: - destination = [destination] + source = [ax + a.ndim if ax < 0 else ax for ax in source] + destination = [ax + a.ndim if ax < 0 else ax for ax in destination] - source = [ax + a.ndim if ax < 0 else ax for ax in source] - destination = [ax + a.ndim if ax < 0 else ax for ax in destination] + order = [n for n in range(a.ndim) if n not in source] - order = [n for n in range(a.ndim) if n not in source] + for dest, src in sorted(zip(destination, source)): + order.insert(dest, src) - for dest, src in sorted(zip(destination, source)): - order.insert(dest, src) + return a.transpose(order) - return a.transpose(order) + +# --- Numpy 1.12 --- # -# TODO: Remove when Numpy 1.12 is an ODL dependency def flip(a, axis): """Reverse the order of elements in an array along the given axis. @@ -62,15 +66,60 @@ def flip(a, axis): -------- numpy.flip """ - if not hasattr(a, 'ndim'): - a = np.asarray(a) - indexer = [slice(None)] * a.ndim - try: - indexer[axis] = slice(None, None, -1) - except IndexError: - raise ValueError('axis={} is invalid for the {}-dimensional input ' - 'array'.format(axis, a.ndim)) - return a[tuple(indexer)] + major, minor, _ = [int(s) for s in np.version.short_version.split('.')] + if (major, minor) >= (1, 12): + return np.flip(a, axis) + else: + if not hasattr(a, 'ndim'): + a = np.asarray(a) + indexer = [slice(None)] * a.ndim + try: + indexer[axis] = slice(None, None, -1) + except IndexError: + raise ValueError('axis={} is invalid for the {}-dimensional input ' + 'array'.format(axis, a.ndim)) + return a[tuple(indexer)] + + +# --- Numpy 1.13 --- # + + +def roll(a, shift, axis=None): + """Roll array elements along a given axis. + + Elements that roll beyond the last position are re-introduced at + the first. + + This function is a backport of `numpy.roll` introduced in NumPy 1.13. + + See Also + -------- + numpy.roll + """ + major, minor, _ = [int(s) for s in np.version.short_version.split('.')] + if (major, minor) >= (1, 13): + return np.roll(a, shift, axis) + else: + if axis is None: + return roll(a.ravel(), shift, 0).reshape(a.shape) + elif np.isscalar(axis): + return np.roll(a, shift, axis) + else: + axis = tuple(axis) + if axis == (): + return a.copy() + + if np.isscalar(shift): + shift = [shift] * len(axis) + elif len(shift) != len(axis): + raise ValueError('`shift` must be integer or of the same ' + 'length as `axis`') + + rolled = np.roll(a, shift[0], axis[0]) + for sh, ax in zip(shift[1:], axis[1:]): + rolled = np.roll(rolled, sh, ax) + + return rolled if __name__ == '__main__':