From 426a7e546b5a14f62743b3bcff37c9835dcdc8df Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 18 Mar 2024 12:19:37 +0100 Subject: [PATCH 01/52] Copy the numpy-tensors module, for migration to pytorch. No changes made yet. --- odl/space/pytorch_tensors.py | 2363 ++++++++++++++++++++++++++++++++++ 1 file changed, 2363 insertions(+) create mode 100644 odl/space/pytorch_tensors.py diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py new file mode 100644 index 00000000000..d80bac82087 --- /dev/null +++ b/odl/space/pytorch_tensors.py @@ -0,0 +1,2363 @@ +# Copyright 2014-2020 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/. + +"""NumPy implementation of tensor spaces.""" + +from __future__ import absolute_import, division, print_function +from future.utils import native + +import ctypes +from builtins import object +from functools import partial + +import numpy as np + +from odl.set.sets import ComplexNumbers, RealNumbers +from odl.set.space import LinearSpaceTypeError +from odl.space.base_tensors import Tensor, TensorSpace +from odl.space.weighting import ( + ArrayWeighting, ConstWeighting, CustomDist, CustomInner, CustomNorm, + Weighting) +from odl.util import ( + dtype_str, is_floating_dtype, is_numeric_dtype, is_real_dtype, nullcontext, + signature_string, writable_array) + +__all__ = ('NumpyTensorSpace',) + + +_BLAS_DTYPES = (np.dtype('float32'), np.dtype('float64'), + np.dtype('complex64'), np.dtype('complex128')) + +# Define size thresholds to switch implementations +THRESHOLD_SMALL = 100 +THRESHOLD_MEDIUM = 50000 + + +class NumpyTensorSpace(TensorSpace): + + """Set of tensors of arbitrary data type, implemented with NumPy. + + A tensor is, in the most general sense, a multi-dimensional array + that allows operations per entry (keep the rank constant), + reductions / contractions (reduce the rank) and broadcasting + (raises the rank). + For non-numeric data type like ``object``, the range of valid + operations is rather limited since such a set of tensors does not + define a vector space. + Any numeric data type, on the other hand, is considered valid for + a tensor space, although certain operations - like division with + integer dtype - are not guaranteed to yield reasonable results. + + Under these restrictions, all basic vector space operations are + supported by this class, along with reductions based on arithmetic + or comparison, and element-wise mathematical functions ("ufuncs"). + + This class is implemented using `numpy.ndarray`'s as back-end. + + See the `Wikipedia article on tensors`_ for further details. + See also [Hac2012] "Part I Algebraic Tensors" for a rigorous + treatment of tensors with a definition close to this one. + + Note also that this notion of tensors is the same as in popular + Deep Learning frameworks. + + References + ---------- + [Hac2012] Hackbusch, W. *Tensor Spaces and Numerical Tensor Calculus*. + Springer, 2012. + + .. _Wikipedia article on tensors: https://en.wikipedia.org/wiki/Tensor + """ + + def __init__(self, shape, dtype=None, **kwargs): + r"""Initialize a new instance. + + Parameters + ---------- + shape : positive int or sequence of positive ints + Number of entries per axis for elements in this space. A + single integer results in a space with rank 1, i.e., 1 axis. + dtype : + Data type of each element. Can be provided in any + way the `numpy.dtype` function understands, e.g. + as built-in type or as a string. For ``None``, + the `default_dtype` of this space (``float64``) is used. + exponent : positive float, optional + Exponent of the norm. For values other than 2.0, no + inner product is defined. + + This option has no impact if either ``dist``, ``norm`` or + ``inner`` is given, or if ``dtype`` is non-numeric. + + Default: 2.0 + + Other Parameters + ---------------- + weighting : optional + Use weighted inner product, norm, and dist. The following + types are supported as ``weighting``: + + ``None``: no weighting, i.e. weighting with ``1.0`` (default). + + `Weighting`: Use this weighting as-is. Compatibility + with this space's elements is not checked during init. + + ``float``: Weighting by a constant. + + array-like: Pointwise weighting by an array. + + This option cannot be combined with ``dist``, + ``norm`` or ``inner``. It also cannot be used in case of + non-numeric ``dtype``. + + dist : callable, optional + Distance function defining a metric on the space. + It must accept two `NumpyTensor` arguments and return + a non-negative real number. See ``Notes`` for + mathematical requirements. + + By default, ``dist(x, y)`` is calculated as ``norm(x - y)``. + + This option cannot be combined with ``weight``, + ``norm`` or ``inner``. It also cannot be used in case of + non-numeric ``dtype``. + + norm : callable, optional + The norm implementation. It must accept a + `NumpyTensor` argument, return a non-negative real number. + See ``Notes`` for mathematical requirements. + + By default, ``norm(x)`` is calculated as ``inner(x, x)``. + + This option cannot be combined with ``weight``, + ``dist`` or ``inner``. It also cannot be used in case of + non-numeric ``dtype``. + + inner : callable, optional + The inner product implementation. It must accept two + `NumpyTensor` arguments and return an element of the field + of the space (usually real or complex number). + See ``Notes`` for mathematical requirements. + + This option cannot be combined with ``weight``, + ``dist`` or ``norm``. It also cannot be used in case of + non-numeric ``dtype``. + + kwargs : + Further keyword arguments are passed to the weighting + classes. + + See Also + -------- + odl.space.space_utils.rn : constructor for real tensor spaces + odl.space.space_utils.cn : constructor for complex tensor spaces + odl.space.space_utils.tensor_space : + constructor for tensor spaces of arbitrary scalar data type + + Notes + ----- + - A distance function or metric on a space :math:`\mathcal{X}` + is a mapping + :math:`d:\mathcal{X} \times \mathcal{X} \to \mathbb{R}` + satisfying the following conditions for all space elements + :math:`x, y, z`: + + * :math:`d(x, y) \geq 0`, + * :math:`d(x, y) = 0 \Leftrightarrow x = y`, + * :math:`d(x, y) = d(y, x)`, + * :math:`d(x, y) \leq d(x, z) + d(z, y)`. + + - A norm on a space :math:`\mathcal{X}` is a mapping + :math:`\| \cdot \|:\mathcal{X} \to \mathbb{R}` + satisfying the following conditions for all + space elements :math:`x, y`: and scalars :math:`s`: + + * :math:`\| x\| \geq 0`, + * :math:`\| x\| = 0 \Leftrightarrow x = 0`, + * :math:`\| sx\| = |s| \cdot \| x \|`, + * :math:`\| x+y\| \leq \| x\| + + \| y\|`. + + - An inner product on a space :math:`\mathcal{X}` over a field + :math:`\mathbb{F} = \mathbb{R}` or :math:`\mathbb{C}` is a + mapping + :math:`\langle\cdot, \cdot\rangle: \mathcal{X} \times + \mathcal{X} \to \mathbb{F}` + satisfying the following conditions for all + space elements :math:`x, y, z`: and scalars :math:`s`: + + * :math:`\langle x, y\rangle = + \overline{\langle y, x\rangle}`, + * :math:`\langle sx + y, z\rangle = s \langle x, z\rangle + + \langle y, z\rangle`, + * :math:`\langle x, x\rangle = 0 \Leftrightarrow x = 0`. + + Examples + -------- + Explicit initialization with the class constructor: + + >>> space = NumpyTensorSpace(3, float) + >>> space + rn(3) + >>> space.shape + (3,) + >>> space.dtype + dtype('float64') + + A more convenient way is to use factory functions: + + >>> space = odl.rn(3, weighting=[1, 2, 3]) + >>> space + rn(3, weighting=[1, 2, 3]) + >>> space = odl.tensor_space((2, 3), dtype=int) + >>> space + tensor_space((2, 3), dtype=int) + """ + super(NumpyTensorSpace, self).__init__(shape, dtype) + if self.dtype.char not in self.available_dtypes(): + raise ValueError('`dtype` {!r} not supported' + ''.format(dtype_str(dtype))) + + dist = kwargs.pop('dist', None) + norm = kwargs.pop('norm', None) + inner = kwargs.pop('inner', None) + weighting = kwargs.pop('weighting', None) + exponent = kwargs.pop('exponent', getattr(weighting, 'exponent', 2.0)) + + if (not is_numeric_dtype(self.dtype) and + any(x is not None for x in (dist, norm, inner, weighting))): + raise ValueError('cannot use any of `weighting`, `dist`, `norm` ' + 'or `inner` for non-numeric `dtype` {}' + ''.format(dtype)) + if exponent != 2.0 and any(x is not None for x in (dist, norm, inner)): + raise ValueError('cannot use any of `dist`, `norm` or `inner` ' + 'for exponent != 2') + # Check validity of option combination (0 or 1 may be provided) + num_extra_args = sum(a is not None + for a in (dist, norm, inner, weighting)) + if num_extra_args > 1: + raise ValueError('invalid combination of options `weighting`, ' + '`dist`, `norm` and `inner`') + + # Set the weighting + if weighting is not None: + if isinstance(weighting, Weighting): + if weighting.impl != 'numpy': + raise ValueError("`weighting.impl` must be 'numpy', " + '`got {!r}'.format(weighting.impl)) + if weighting.exponent != exponent: + raise ValueError('`weighting.exponent` conflicts with ' + '`exponent`: {} != {}' + ''.format(weighting.exponent, exponent)) + self.__weighting = weighting + else: + self.__weighting = _weighting(weighting, exponent) + + # Check (afterwards) that the weighting input was sane + if isinstance(self.weighting, NumpyTensorSpaceArrayWeighting): + if self.weighting.array.dtype == object: + raise ValueError('invalid `weighting` argument: {}' + ''.format(weighting)) + elif not np.can_cast(self.weighting.array.dtype, self.dtype): + raise ValueError( + 'cannot cast from `weighting` data type {} to ' + 'the space `dtype` {}' + ''.format(dtype_str(self.weighting.array.dtype), + dtype_str(self.dtype))) + if self.weighting.array.shape != self.shape: + raise ValueError('array-like weights must have same ' + 'shape {} as this space, got {}' + ''.format(self.shape, + self.weighting.array.shape)) + + elif dist is not None: + self.__weighting = NumpyTensorSpaceCustomDist(dist) + elif norm is not None: + self.__weighting = NumpyTensorSpaceCustomNorm(norm) + elif inner is not None: + self.__weighting = NumpyTensorSpaceCustomInner(inner) + else: + # No weighting, i.e., weighting with constant 1.0 + self.__weighting = NumpyTensorSpaceConstWeighting(1.0, exponent) + + # Make sure there are no leftover kwargs + if kwargs: + raise TypeError('got unknown keyword arguments {}'.format(kwargs)) + + @property + def impl(self): + """Name of the implementation back-end: ``'numpy'``.""" + return 'numpy' + + @property + def default_order(self): + """Default storage order for new elements in this space: ``'C'``.""" + return 'C' + + @property + def weighting(self): + """This space's weighting scheme.""" + return self.__weighting + + @property + def is_weighted(self): + """Return ``True`` if the space is not weighted by constant 1.0.""" + return not ( + isinstance(self.weighting, NumpyTensorSpaceConstWeighting) and + self.weighting.const == 1.0) + + @property + def exponent(self): + """Exponent of the norm and the distance.""" + return self.weighting.exponent + + def element(self, inp=None, data_ptr=None, order=None): + """Create a new element. + + Parameters + ---------- + inp : `array-like`, optional + Input used to initialize the new element. + + If ``inp`` is `None`, an empty element is created with no + guarantee of its state (memory allocation only). + The new element will use ``order`` as storage order if + provided, otherwise `default_order`. + + Otherwise, a copy is avoided whenever possible. This requires + correct `shape` and `dtype`, and if ``order`` is provided, + also contiguousness in that ordering. If any of these + conditions is not met, a copy is made. + + data_ptr : int, optional + Pointer to the start memory address of a contiguous Numpy array + or an equivalent raw container with the same total number of + bytes. For this option, ``order`` must be either ``'C'`` or + ``'F'``. + The option is also mutually exclusive with ``inp``. + order : {None, 'C', 'F'}, optional + Storage order of the returned element. For ``'C'`` and ``'F'``, + contiguous memory in the respective ordering is enforced. + The default ``None`` enforces no contiguousness. + + Returns + ------- + element : `NumpyTensor` + The new element, created from ``inp`` or from scratch. + + Examples + -------- + Without arguments, an uninitialized element is created. With an + array-like input, the element can be initialized: + + >>> space = odl.rn(3) + >>> empty = space.element() + >>> empty.shape + (3,) + >>> empty.space + rn(3) + >>> x = space.element([1, 2, 3]) + >>> x + rn(3).element([ 1., 2., 3.]) + + If the input already is a `numpy.ndarray` of correct `dtype`, it + will merely be wrapped, i.e., both array and space element access + the same memory, such that mutations will affect both: + + >>> arr = np.array([1, 2, 3], dtype=float) + >>> elem = odl.rn(3).element(arr) + >>> elem[0] = 0 + >>> elem + rn(3).element([ 0., 2., 3.]) + >>> arr + array([ 0., 2., 3.]) + + Elements can also be constructed from a data pointer, resulting + again in shared memory: + + >>> int_space = odl.tensor_space((2, 3), dtype=int) + >>> arr = np.array([[1, 2, 3], + ... [4, 5, 6]], dtype=int, order='F') + >>> ptr = arr.ctypes.data + >>> y = int_space.element(data_ptr=ptr, order='F') + >>> y + tensor_space((2, 3), dtype=int).element( + [[1, 2, 3], + [4, 5, 6]] + ) + >>> y[0, 1] = -1 + >>> arr + array([[ 1, -1, 3], + [ 4, 5, 6]]) + """ + if order is not None and str(order).upper() not in ('C', 'F'): + raise ValueError("`order` {!r} not understood".format(order)) + + if inp is None and data_ptr is None: + if order is None: + arr = np.empty(self.shape, dtype=self.dtype, + order=self.default_order) + else: + arr = np.empty(self.shape, dtype=self.dtype, order=order) + + return self.element_type(self, arr) + + elif inp is None and data_ptr is not None: + if order is None: + raise ValueError('`order` cannot be None for element ' + 'creation from pointer') + + ctype_array_def = ctypes.c_byte * self.nbytes + as_ctype_array = ctype_array_def.from_address(data_ptr) + as_numpy_array = np.ctypeslib.as_array(as_ctype_array) + arr = as_numpy_array.view(dtype=self.dtype) + arr = arr.reshape(self.shape, order=order) + return self.element_type(self, arr) + + elif inp is not None and data_ptr is None: + if inp in self and order is None: + # Short-circuit for space elements and no enforced ordering + return inp + + # Try to not copy but require dtype and order if given + # (`order=None` is ok as np.array argument) + arr = np.array(inp, copy=False, dtype=self.dtype, ndmin=self.ndim, + order=order) + # Make sure the result is writeable, if not make copy. + # This happens for e.g. results of `np.broadcast_to()`. + if not arr.flags.writeable: + arr = arr.copy() + if arr.shape != self.shape: + raise ValueError('shape of `inp` not equal to space shape: ' + '{} != {}'.format(arr.shape, self.shape)) + return self.element_type(self, arr) + + else: + raise TypeError('cannot provide both `inp` and `data_ptr`') + + def zero(self): + """Return a tensor of all zeros. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.zero() + >>> x + rn(3).element([ 0., 0., 0.]) + """ + return self.element(np.zeros(self.shape, dtype=self.dtype, + order=self.default_order)) + + def one(self): + """Return a tensor of all ones. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.one() + >>> x + rn(3).element([ 1., 1., 1.]) + """ + return self.element(np.ones(self.shape, dtype=self.dtype, + order=self.default_order)) + + @staticmethod + def available_dtypes(): + """Return the set of data types available in this implementation. + + Notes + ----- + This is all dtypes available in Numpy. See ``numpy.sctypes`` + for more information. + + The available dtypes may depend on the specific system used. + """ + all_dtypes = [] + for lst in np.sctypes.values(): + for dtype in lst: + if dtype not in (object, np.void): + all_dtypes.append(np.dtype(dtype)) + # Need to add these manually since np.sctypes['others'] will only + # contain one of them (depending on Python version) + all_dtypes.extend([np.dtype('S'), np.dtype('U')]) + return tuple(sorted(set(all_dtypes))) + + @staticmethod + def default_dtype(field=None): + """Return the default data type of this class for a given field. + + Parameters + ---------- + field : `Field`, optional + Set of numbers to be represented by a data type. + Currently supported : `RealNumbers`, `ComplexNumbers` + The default ``None`` means `RealNumbers` + + Returns + ------- + dtype : `numpy.dtype` + Numpy data type specifier. The returned defaults are: + + ``RealNumbers()`` : ``np.dtype('float64')`` + + ``ComplexNumbers()`` : ``np.dtype('complex128')`` + """ + if field is None or field == RealNumbers(): + return np.dtype('float64') + elif field == ComplexNumbers(): + return np.dtype('complex128') + else: + raise ValueError('no default data type defined for field {}' + ''.format(field)) + + def _lincomb(self, a, x1, b, x2, out): + """Implement the linear combination of ``x1`` and ``x2``. + + Compute ``out = a*x1 + b*x2`` using optimized + BLAS routines if possible. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + a, b : `TensorSpace.field` element + Scalars to multiply ``x1`` and ``x2`` with. + x1, x2 : `NumpyTensor` + Summands in the linear combination. + out : `NumpyTensor` + Tensor to which the result is written. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([0, 1, 1]) + >>> y = space.element([0, 0, 1]) + >>> out = space.element() + >>> result = space.lincomb(1, x, 2, y, out) + >>> result + rn(3).element([ 0., 1., 3.]) + >>> result is out + True + """ + _lincomb_impl(a, x1, b, x2, out) + + def _dist(self, x1, x2): + """Return the distance between ``x1`` and ``x2``. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Elements whose mutual distance is calculated. + + Returns + ------- + dist : `float` + Distance between the elements. + + Examples + -------- + Different exponents result in difference metrics: + + >>> space_2 = odl.rn(3, exponent=2) + >>> x = space_2.element([-1, -1, 2]) + >>> y = space_2.one() + >>> space_2.dist(x, y) + 3.0 + + >>> space_1 = odl.rn(3, exponent=1) + >>> x = space_1.element([-1, -1, 2]) + >>> y = space_1.one() + >>> space_1.dist(x, y) + 5.0 + + Weighting is supported, too: + + >>> space_1_w = odl.rn(3, exponent=1, weighting=[2, 1, 1]) + >>> x = space_1_w.element([-1, -1, 2]) + >>> y = space_1_w.one() + >>> space_1_w.dist(x, y) + 7.0 + """ + return self.weighting.dist(x1, x2) + + def _norm(self, x): + """Return the norm of ``x``. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + x : `NumpyTensor` + Element whose norm is calculated. + + Returns + ------- + norm : `float` + Norm of the element. + + Examples + -------- + Different exponents result in difference norms: + + >>> space_2 = odl.rn(3, exponent=2) + >>> x = space_2.element([3, 0, 4]) + >>> space_2.norm(x) + 5.0 + >>> space_1 = odl.rn(3, exponent=1) + >>> x = space_1.element([3, 0, 4]) + >>> space_1.norm(x) + 7.0 + + Weighting is supported, too: + + >>> space_1_w = odl.rn(3, exponent=1, weighting=[2, 1, 1]) + >>> x = space_1_w.element([3, 0, 4]) + >>> space_1_w.norm(x) + 10.0 + """ + return self.weighting.norm(x) + + def _inner(self, x1, x2): + """Return the inner product of ``x1`` and ``x2``. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Elements whose inner product is calculated. + + Returns + ------- + inner : `field` `element` + Inner product of the elements. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([1, 0, 3]) + >>> y = space.one() + >>> space.inner(x, y) + 4.0 + + Weighting is supported, too: + + >>> space_w = odl.rn(3, weighting=[2, 1, 1]) + >>> x = space_w.element([1, 0, 3]) + >>> y = space_w.one() + >>> space_w.inner(x, y) + 5.0 + """ + return self.weighting.inner(x1, x2) + + def _multiply(self, x1, x2, out): + """Compute the entry-wise product ``out = x1 * x2``. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Factors in the product. + out : `NumpyTensor` + Element to which the result is written. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([1, 0, 3]) + >>> y = space.element([-1, 1, -1]) + >>> space.multiply(x, y) + rn(3).element([-1., 0., -3.]) + >>> out = space.element() + >>> result = space.multiply(x, y, out=out) + >>> result + rn(3).element([-1., 0., -3.]) + >>> result is out + True + """ + np.multiply(x1.data, x2.data, out=out.data) + + def _divide(self, x1, x2, out): + """Compute the entry-wise quotient ``x1 / x2``. + + This function is part of the subclassing API. Do not + call it directly. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Dividend and divisor in the quotient. + out : `NumpyTensor` + Element to which the result is written. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([2, 0, 4]) + >>> y = space.element([1, 1, 2]) + >>> space.divide(x, y) + rn(3).element([ 2., 0., 2.]) + >>> out = space.element() + >>> result = space.divide(x, y, out=out) + >>> result + rn(3).element([ 2., 0., 2.]) + >>> result is out + True + """ + np.divide(x1.data, x2.data, out=out.data) + + def __eq__(self, other): + """Return ``self == other``. + + Returns + ------- + equals : bool + True if ``other`` is an instance of ``type(self)`` + with the same `NumpyTensorSpace.shape`, `NumpyTensorSpace.dtype` + and `NumpyTensorSpace.weighting`, otherwise False. + + Examples + -------- + >>> space = odl.rn(3) + >>> same_space = odl.rn(3, exponent=2) + >>> same_space == space + True + + Different `shape`, `exponent` or `dtype` all result in different + spaces: + + >>> diff_space = odl.rn((3, 4)) + >>> diff_space == space + False + >>> diff_space = odl.rn(3, exponent=1) + >>> diff_space == space + False + >>> diff_space = odl.rn(3, dtype='float32') + >>> diff_space == space + False + >>> space == object + False + """ + if other is self: + return True + + return (super(NumpyTensorSpace, self).__eq__(other) and + self.weighting == other.weighting) + + def __hash__(self): + """Return ``hash(self)``.""" + return hash((super(NumpyTensorSpace, self).__hash__(), + self.weighting)) + + @property + def byaxis(self): + """Return the subspace defined along one or several dimensions. + + Examples + -------- + Indexing with integers or slices: + + >>> space = odl.rn((2, 3, 4)) + >>> space.byaxis[0] + rn(2) + >>> space.byaxis[1:] + rn((3, 4)) + + Lists can be used to stack spaces arbitrarily: + + >>> space.byaxis[[2, 1, 2]] + rn((4, 3, 4)) + """ + space = self + + class NpyTensorSpacebyaxis(object): + + """Helper class for indexing by axis.""" + + def __getitem__(self, indices): + """Return ``self[indices]``.""" + try: + iter(indices) + except TypeError: + newshape = space.shape[indices] + else: + newshape = tuple(space.shape[i] for i in indices) + + if isinstance(space.weighting, ArrayWeighting): + new_array = np.asarray(space.weighting.array[indices]) + weighting = NumpyTensorSpaceArrayWeighting( + new_array, space.weighting.exponent) + else: + weighting = space.weighting + + return type(space)(newshape, space.dtype, weighting=weighting) + + def __repr__(self): + """Return ``repr(self)``.""" + return repr(space) + '.byaxis' + + return NpyTensorSpacebyaxis() + + def __repr__(self): + """Return ``repr(self)``.""" + if self.ndim == 1: + posargs = [self.size] + else: + posargs = [self.shape] + + if self.is_real: + ctor_name = 'rn' + elif self.is_complex: + ctor_name = 'cn' + else: + ctor_name = 'tensor_space' + + if (ctor_name == 'tensor_space' or + not is_numeric_dtype(self.dtype) or + self.dtype != self.default_dtype(self.field)): + optargs = [('dtype', dtype_str(self.dtype), '')] + if self.dtype in (float, complex, int, bool): + optmod = '!s' + else: + optmod = '' + else: + optargs = [] + optmod = '' + + inner_str = signature_string(posargs, optargs, mod=['', optmod]) + weight_str = self.weighting.repr_part + if weight_str: + inner_str += ', ' + weight_str + + return '{}({})'.format(ctor_name, inner_str) + + @property + def element_type(self): + """Type of elements in this space: `NumpyTensor`.""" + return NumpyTensor + + +class NumpyTensor(Tensor): + + """Representation of a `NumpyTensorSpace` element.""" + + def __init__(self, space, data): + """Initialize a new instance.""" + Tensor.__init__(self, space) + self.__data = data + + @property + def data(self): + """The `numpy.ndarray` representing the data of ``self``.""" + return self.__data + + def asarray(self, out=None): + """Extract the data of this array as a ``numpy.ndarray``. + + This method is invoked when calling `numpy.asarray` on this + tensor. + + Parameters + ---------- + out : `numpy.ndarray`, optional + Array in which the result should be written in-place. + Has to be contiguous and of the correct dtype. + + Returns + ------- + asarray : `numpy.ndarray` + Numpy array with the same data type as ``self``. If + ``out`` was given, the returned object is a reference + to it. + + Examples + -------- + >>> space = odl.rn(3, dtype='float32') + >>> x = space.element([1, 2, 3]) + >>> x.asarray() + array([ 1., 2., 3.], dtype=float32) + >>> np.asarray(x) is x.asarray() + True + >>> out = np.empty(3, dtype='float32') + >>> result = x.asarray(out=out) + >>> out + array([ 1., 2., 3.], dtype=float32) + >>> result is out + True + >>> space = odl.rn((2, 3)) + >>> space.one().asarray() + array([[ 1., 1., 1.], + [ 1., 1., 1.]]) + """ + if out is None: + return self.data + else: + out[:] = self.data + return out + + def astype(self, dtype): + """Return a copy of this element with new ``dtype``. + + Parameters + ---------- + dtype : + Scalar data type of the returned space. Can be provided + in any way the `numpy.dtype` constructor understands, e.g. + as built-in type or as a string. Data types with non-trivial + shapes are not allowed. + + Returns + ------- + newelem : `NumpyTensor` + Version of this element with given data type. + """ + return self.space.astype(dtype).element(self.data.astype(dtype)) + + @property + def data_ptr(self): + """A raw pointer to the data container of ``self``. + + Examples + -------- + >>> import ctypes + >>> space = odl.tensor_space(3, dtype='uint16') + >>> x = space.element([1, 2, 3]) + >>> arr_type = ctypes.c_uint16 * 3 # C type "array of 3 uint16" + >>> buffer = arr_type.from_address(x.data_ptr) + >>> arr = np.frombuffer(buffer, dtype='uint16') + >>> arr + array([1, 2, 3], dtype=uint16) + + In-place modification via pointer: + + >>> arr[0] = 42 + >>> x + tensor_space(3, dtype='uint16').element([42, 2, 3]) + """ + return self.data.ctypes.data + + def __eq__(self, other): + """Return ``self == other``. + + Returns + ------- + equals : bool + True if all entries of ``other`` are equal to this + the entries of ``self``, False otherwise. + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([1, 2, 3]) + >>> y = space.element([1, 2, 3]) + >>> x == y + True + + >>> y = space.element([-1, 2, 3]) + >>> x == y + False + >>> x == object + False + + Space membership matters: + + >>> space2 = odl.tensor_space(3, dtype='int64') + >>> y = space2.element([1, 2, 3]) + >>> x == y or y == x + False + """ + if other is self: + return True + elif other not in self.space: + return False + else: + return np.array_equal(self.data, other.data) + + def copy(self): + """Return an identical (deep) copy of this tensor. + + Parameters + ---------- + None + + Returns + ------- + copy : `NumpyTensor` + The deep copy + + Examples + -------- + >>> space = odl.rn(3) + >>> x = space.element([1, 2, 3]) + >>> y = x.copy() + >>> y == x + True + >>> y is x + False + """ + return self.space.element(self.data.copy()) + + def __copy__(self): + """Return ``copy(self)``. + + This implements the (shallow) copy interface of the ``copy`` + module of the Python standard library. + + See Also + -------- + copy + + Examples + -------- + >>> from copy import copy + >>> space = odl.rn(3) + >>> x = space.element([1, 2, 3]) + >>> y = copy(x) + >>> y == x + True + >>> y is x + False + """ + return self.copy() + + def __getitem__(self, indices): + """Return ``self[indices]``. + + Parameters + ---------- + indices : index expression + Integer, slice or sequence of these, defining the positions + of the data array which should be accessed. + + Returns + ------- + values : `NumpyTensorSpace.dtype` or `NumpyTensor` + The value(s) at the given indices. Note that the returned + object is a writable view into the original tensor, except + for the case when ``indices`` is a list. + + Examples + -------- + For one-dimensional spaces, indexing is as in linear arrays: + + >>> space = odl.rn(3) + >>> x = space.element([1, 2, 3]) + >>> x[0] + 1.0 + >>> x[1:] + rn(2).element([ 2., 3.]) + + In higher dimensions, the i-th index expression accesses the + i-th axis: + + >>> space = odl.rn((2, 3)) + >>> x = space.element([[1, 2, 3], + ... [4, 5, 6]]) + >>> x[0, 1] + 2.0 + >>> x[:, 1:] + rn((2, 2)).element( + [[ 2., 3.], + [ 5., 6.]] + ) + + Slices can be assigned to, except if lists are used for indexing: + + >>> y = x[:, ::2] # view into x + >>> y[:] = -9 + >>> x + rn((2, 3)).element( + [[-9., 2., -9.], + [-9., 5., -9.]] + ) + >>> y = x[[0, 1], [1, 2]] # not a view, won't modify x + >>> y + rn(2).element([ 2., -9.]) + >>> y[:] = 0 + >>> x + rn((2, 3)).element( + [[-9., 2., -9.], + [-9., 5., -9.]] + ) + """ + # Lazy implementation: index the array and deal with it + if isinstance(indices, NumpyTensor): + indices = indices.data + arr = self.data[indices] + + if np.isscalar(arr): + if self.space.field is not None: + return self.space.field.element(arr) + else: + return arr + else: + if is_numeric_dtype(self.dtype): + weighting = self.space.weighting + else: + weighting = None + space = type(self.space)( + arr.shape, dtype=self.dtype, exponent=self.space.exponent, + weighting=weighting) + return space.element(arr) + + def __setitem__(self, indices, values): + """Implement ``self[indices] = values``. + + Parameters + ---------- + indices : index expression + Integer, slice or sequence of these, defining the positions + of the data array which should be written to. + values : scalar, array-like or `NumpyTensor` + The value(s) that are to be assigned. + + If ``index`` is an integer, ``value`` must be a scalar. + + If ``index`` is a slice or a sequence of slices, ``value`` + must be broadcastable to the shape of the slice. + + Examples + -------- + For 1d spaces, entries can be set with scalars or sequences of + correct shape: + + >>> space = odl.rn(3) + >>> x = space.element([1, 2, 3]) + >>> x[0] = -1 + >>> x[1:] = (0, 1) + >>> x + rn(3).element([-1., 0., 1.]) + + It is also possible to use tensors of other spaces for + casting and assignment: + + >>> space = odl.rn((2, 3)) + >>> x = space.element([[1, 2, 3], + ... [4, 5, 6]]) + >>> x[0, 1] = -1 + >>> x + rn((2, 3)).element( + [[ 1., -1., 3.], + [ 4., 5., 6.]] + ) + >>> short_space = odl.tensor_space((2, 2), dtype='short') + >>> y = short_space.element([[-1, 2], + ... [0, 0]]) + >>> x[:, :2] = y + >>> x + rn((2, 3)).element( + [[-1., 2., 3.], + [ 0., 0., 6.]] + ) + + The Numpy assignment and broadcasting rules apply: + + >>> x[:] = np.array([[0, 0, 0], + ... [1, 1, 1]]) + >>> x + rn((2, 3)).element( + [[ 0., 0., 0.], + [ 1., 1., 1.]] + ) + >>> x[:, 1:] = [7, 8] + >>> x + rn((2, 3)).element( + [[ 0., 7., 8.], + [ 1., 7., 8.]] + ) + >>> x[:, ::2] = -2. + >>> x + rn((2, 3)).element( + [[-2., 7., -2.], + [-2., 7., -2.]] + ) + """ + if isinstance(indices, type(self)): + indices = indices.data + if isinstance(values, type(self)): + values = values.data + + self.data[indices] = values + + @property + def real(self): + """Real part of ``self``. + + Returns + ------- + real : `NumpyTensor` + Real part of this element as a member of a + `NumpyTensorSpace` with corresponding real data type. + + Examples + -------- + Get the real part: + + >>> space = odl.cn(3) + >>> x = space.element([1 + 1j, 2, 3 - 3j]) + >>> x.real + rn(3).element([ 1., 2., 3.]) + + Set the real part: + + >>> space = odl.cn(3) + >>> x = space.element([1 + 1j, 2, 3 - 3j]) + >>> zero = odl.rn(3).zero() + >>> x.real = zero + >>> x + cn(3).element([ 0.+1.j, 0.+0.j, 0.-3.j]) + + Other array-like types and broadcasting: + + >>> x.real = 1.0 + >>> x + cn(3).element([ 1.+1.j, 1.+0.j, 1.-3.j]) + >>> x.real = [2, 3, 4] + >>> x + cn(3).element([ 2.+1.j, 3.+0.j, 4.-3.j]) + """ + if self.space.is_real: + return self + elif self.space.is_complex: + real_space = self.space.astype(self.space.real_dtype) + return real_space.element(self.data.real) + else: + raise NotImplementedError('`real` not defined for non-numeric ' + 'dtype {}'.format(self.dtype)) + + @real.setter + def real(self, newreal): + """Setter for the real part. + + This method is invoked by ``x.real = other``. + + Parameters + ---------- + newreal : array-like or scalar + Values to be assigned to the real part of this element. + """ + self.real.data[:] = newreal + + @property + def imag(self): + """Imaginary part of ``self``. + + Returns + ------- + imag : `NumpyTensor` + Imaginary part this element as an element of a + `NumpyTensorSpace` with real data type. + + Examples + -------- + Get the imaginary part: + + >>> space = odl.cn(3) + >>> x = space.element([1 + 1j, 2, 3 - 3j]) + >>> x.imag + rn(3).element([ 1., 0., -3.]) + + Set the imaginary part: + + >>> space = odl.cn(3) + >>> x = space.element([1 + 1j, 2, 3 - 3j]) + >>> zero = odl.rn(3).zero() + >>> x.imag = zero + >>> x + cn(3).element([ 1.+0.j, 2.+0.j, 3.+0.j]) + + Other array-like types and broadcasting: + + >>> x.imag = 1.0 + >>> x + cn(3).element([ 1.+1.j, 2.+1.j, 3.+1.j]) + >>> x.imag = [2, 3, 4] + >>> x + cn(3).element([ 1.+2.j, 2.+3.j, 3.+4.j]) + """ + if self.space.is_real: + return self.space.zero() + elif self.space.is_complex: + real_space = self.space.astype(self.space.real_dtype) + return real_space.element(self.data.imag) + else: + raise NotImplementedError('`imag` not defined for non-numeric ' + 'dtype {}'.format(self.dtype)) + + @imag.setter + def imag(self, newimag): + """Setter for the imaginary part. + + This method is invoked by ``x.imag = other``. + + Parameters + ---------- + newimag : array-like or scalar + Values to be assigned to the imaginary part of this element. + + Raises + ------ + ValueError + If the space is real, i.e., no imagninary part can be set. + """ + if self.space.is_real: + raise ValueError('cannot set imaginary part in real spaces') + self.imag.data[:] = newimag + + def conj(self, out=None): + """Return the complex conjugate of ``self``. + + Parameters + ---------- + out : `NumpyTensor`, optional + Element to which the complex conjugate is written. + Must be an element of ``self.space``. + + Returns + ------- + out : `NumpyTensor` + The complex conjugate element. If ``out`` was provided, + the returned object is a reference to it. + + Examples + -------- + >>> space = odl.cn(3) + >>> x = space.element([1 + 1j, 2, 3 - 3j]) + >>> x.conj() + cn(3).element([ 1.-1.j, 2.-0.j, 3.+3.j]) + >>> out = space.element() + >>> result = x.conj(out=out) + >>> result + cn(3).element([ 1.-1.j, 2.-0.j, 3.+3.j]) + >>> result is out + True + + In-place conjugation: + + >>> result = x.conj(out=x) + >>> x + cn(3).element([ 1.-1.j, 2.-0.j, 3.+3.j]) + >>> result is x + True + """ + if self.space.is_real: + if out is None: + return self + else: + out[:] = self + return out + + if not is_numeric_dtype(self.space.dtype): + raise NotImplementedError('`conj` not defined for non-numeric ' + 'dtype {}'.format(self.dtype)) + + if out is None: + return self.space.element(self.data.conj()) + else: + if out not in self.space: + raise LinearSpaceTypeError('`out` {!r} not in space {!r}' + ''.format(out, self.space)) + self.data.conj(out.data) + return out + + def __ipow__(self, other): + """Return ``self **= other``.""" + try: + if other == int(other): + return super(NumpyTensor, self).__ipow__(other) + except TypeError: + pass + + np.power(self.data, other, out=self.data) + return self + + def __int__(self): + """Return ``int(self)``.""" + return int(self.data) + + def __long__(self): + """Return ``long(self)``. + + This method is only useful in Python 2. + """ + return long(self.data) + + def __float__(self): + """Return ``float(self)``.""" + return float(self.data) + + def __complex__(self): + """Return ``complex(self)``.""" + if self.size != 1: + raise TypeError('only size-1 tensors can be converted to ' + 'Python scalars') + return complex(self.data.ravel()[0]) + + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + """Interface to Numpy's ufunc machinery. + + This method is called by Numpy version 1.13 and higher as a single + point for the ufunc dispatch logic. An object implementing + ``__array_ufunc__`` takes over control when a `numpy.ufunc` is + called on it, allowing it to use custom implementations and + output types. + + This includes handling of in-place arithmetic like + ``npy_array += custom_obj``. In this case, the custom object's + ``__array_ufunc__`` takes precedence over the baseline + `numpy.ndarray` implementation. It will be called with + ``npy_array`` as ``out`` argument, which ensures that the + returned object is a Numpy array. For this to work properly, + ``__array_ufunc__`` has to accept Numpy arrays as ``out`` arguments. + + See the `corresponding NEP`_ and the `interface documentation`_ + for further details. See also the `general documentation on + Numpy ufuncs`_. + + .. note:: + This basic implementation casts inputs and + outputs to Numpy arrays and evaluates ``ufunc`` on those. + For `numpy.ndarray` based data storage, this incurs no + significant overhead compared to direct usage of Numpy arrays. + + For other (in particular non-local) implementations, e.g., + GPU arrays or distributed memory, overhead is significant due + to copies to CPU main memory. In those classes, the + ``__array_ufunc__`` mechanism should be overridden to use + native implementations if possible. + + .. note:: + When using operations that alter the shape (like ``reduce``), + or the data type (can be any of the methods), + the resulting array is wrapped in a space of the same + type as ``self.space``, propagating space properties like + `exponent` or `weighting` as closely as possible. + + Parameters + ---------- + ufunc : `numpy.ufunc` + Ufunc that should be called on ``self``. + method : str + Method on ``ufunc`` that should be called on ``self``. + Possible values: + + ``'__call__'``, ``'accumulate'``, ``'at'``, ``'outer'``, + ``'reduce'``, ``'reduceat'`` + + input1, ..., inputN : + Positional arguments to ``ufunc.method``. + kwargs : + Keyword arguments to ``ufunc.method``. + + Returns + ------- + ufunc_result : `Tensor`, `numpy.ndarray` or tuple + Result of the ufunc evaluation. If no ``out`` keyword argument + was given, the result is a `Tensor` or a tuple + of such, depending on the number of outputs of ``ufunc``. + If ``out`` was provided, the returned object or tuple entries + refer(s) to ``out``. + + Examples + -------- + We apply `numpy.add` to ODL tensors: + + >>> r3 = odl.rn(3) + >>> x = r3.element([1, 2, 3]) + >>> y = r3.element([-1, -2, -3]) + >>> x.__array_ufunc__(np.add, '__call__', x, y) + rn(3).element([ 0., 0., 0.]) + >>> np.add(x, y) # same mechanism for Numpy >= 1.13 + rn(3).element([ 0., 0., 0.]) + + As ``out``, a Numpy array or an ODL tensor can be given (wrapped + in a sequence): + + >>> out = r3.element() + >>> res = x.__array_ufunc__(np.add, '__call__', x, y, out=(out,)) + >>> out + rn(3).element([ 0., 0., 0.]) + >>> res is out + True + >>> out_arr = np.empty(3) + >>> res = x.__array_ufunc__(np.add, '__call__', x, y, out=(out_arr,)) + >>> out_arr + array([ 0., 0., 0.]) + >>> res is out_arr + True + + With multiple dimensions: + + >>> r23 = odl.rn((2, 3)) + >>> x = y = r23.one() + >>> x.__array_ufunc__(np.add, '__call__', x, y) + rn((2, 3)).element( + [[ 2., 2., 2.], + [ 2., 2., 2.]] + ) + + The ``ufunc.accumulate`` method retains the original `shape` and + `dtype`. The latter can be changed with the ``dtype`` parameter: + + >>> x = r3.element([1, 2, 3]) + >>> x.__array_ufunc__(np.add, 'accumulate', x) + rn(3).element([ 1., 3., 6.]) + >>> np.add.accumulate(x) # same mechanism for Numpy >= 1.13 + rn(3).element([ 1., 3., 6.]) + >>> x.__array_ufunc__(np.add, 'accumulate', x, dtype=complex) + cn(3).element([ 1.+0.j, 3.+0.j, 6.+0.j]) + + For multi-dimensional tensors, an optional ``axis`` parameter + can be provided: + + >>> z = r23.one() + >>> z.__array_ufunc__(np.add, 'accumulate', z, axis=1) + rn((2, 3)).element( + [[ 1., 2., 3.], + [ 1., 2., 3.]] + ) + + The ``ufunc.at`` method operates in-place. Here we add the second + operand ``[5, 10]`` to ``x`` at indices ``[0, 2]``: + + >>> x = r3.element([1, 2, 3]) + >>> x.__array_ufunc__(np.add, 'at', x, [0, 2], [5, 10]) + >>> x + rn(3).element([ 6., 2., 13.]) + + For outer-product-type operations, i.e., operations where the result + shape is the sum of the individual shapes, the ``ufunc.outer`` + method can be used: + + >>> x = odl.rn(2).element([0, 3]) + >>> y = odl.rn(3).element([1, 2, 3]) + >>> x.__array_ufunc__(np.add, 'outer', x, y) + rn((2, 3)).element( + [[ 1., 2., 3.], + [ 4., 5., 6.]] + ) + >>> y.__array_ufunc__(np.add, 'outer', y, x) + rn((3, 2)).element( + [[ 1., 4.], + [ 2., 5.], + [ 3., 6.]] + ) + + Using ``ufunc.reduce`` produces a scalar, which can be avoided with + ``keepdims=True``: + + >>> x = r3.element([1, 2, 3]) + >>> x.__array_ufunc__(np.add, 'reduce', x) + 6.0 + >>> x.__array_ufunc__(np.add, 'reduce', x, keepdims=True) + rn(1).element([ 6.]) + + In multiple dimensions, ``axis`` can be provided for reduction over + selected axes: + + >>> z = r23.element([[1, 2, 3], + ... [4, 5, 6]]) + >>> z.__array_ufunc__(np.add, 'reduce', z, axis=1) + rn(2).element([ 6., 15.]) + + Finally, ``add.reduceat`` is a combination of ``reduce`` and + ``at`` with rather flexible and complex semantics (see the + `reduceat documentation`_ for details): + + >>> x = r3.element([1, 2, 3]) + >>> x.__array_ufunc__(np.add, 'reduceat', x, [0, 1]) + rn(2).element([ 1., 5.]) + + References + ---------- + .. _corresponding NEP: + https://docs.scipy.org/doc/numpy/neps/ufunc-overrides.html + + .. _interface documentation: + https://docs.scipy.org/doc/numpy/reference/arrays.classes.html\ +#numpy.class.__array_ufunc__ + + .. _general documentation on Numpy ufuncs: + https://docs.scipy.org/doc/numpy/reference/ufuncs.html + + .. _reduceat documentation: + https://docs.scipy.org/doc/numpy/reference/generated/\ +numpy.ufunc.reduceat.html + """ + # Remark: this method differs from the parent implementation only + # in the propagation of additional space properties. + + # --- Process `out` --- # + + # Unwrap out if provided. The output parameters are all wrapped + # in one tuple, even if there is only one. + out_tuple = kwargs.pop('out', ()) + + # Check number of `out` args, depending on `method` + if method == '__call__' and len(out_tuple) not in (0, ufunc.nout): + raise ValueError( + "ufunc {}: need 0 or {} `out` arguments for " + "`method='__call__'`, got {}" + ''.format(ufunc.__name__, ufunc.nout, len(out_tuple))) + elif method != '__call__' and len(out_tuple) not in (0, 1): + raise ValueError( + 'ufunc {}: need 0 or 1 `out` arguments for `method={!r}`, ' + 'got {}'.format(ufunc.__name__, method, len(out_tuple))) + + # We allow our own tensors, the data container type and + # `numpy.ndarray` objects as `out` (see docs for reason for the + # latter) + valid_types = (type(self), type(self.data), np.ndarray) + if not all(isinstance(o, valid_types) or o is None + for o in out_tuple): + return NotImplemented + + # Assign to `out` or `out1` and `out2`, respectively + out = out1 = out2 = None + if len(out_tuple) == 1: + out = out_tuple[0] + elif len(out_tuple) == 2: + out1 = out_tuple[0] + out2 = out_tuple[1] + + # --- Process `inputs` --- # + + # Convert inputs that are ODL tensors to Numpy arrays so that the + # native Numpy ufunc is called later + inputs = tuple( + inp.asarray() if isinstance(inp, type(self)) else inp + for inp in inputs) + + # --- Get some parameters for later --- # + + # Arguments for `writable_array` and/or space constructors + out_dtype = kwargs.get('dtype', None) + if out_dtype is None: + array_kwargs = {} + else: + array_kwargs = {'dtype': out_dtype} + + exponent = self.space.exponent + weighting = self.space.weighting + + # --- Evaluate ufunc --- # + + if method == '__call__': + if ufunc.nout == 1: + # Make context for output (trivial one returns `None`) + if out is None: + out_ctx = nullcontext() + else: + out_ctx = writable_array(out, **array_kwargs) + + # Evaluate ufunc + with out_ctx as out_arr: + kwargs['out'] = out_arr + res = ufunc(*inputs, **kwargs) + + # Wrap result if necessary (lazily) + if out is None: + if is_floating_dtype(res.dtype): + # Weighting contains exponent + spc_kwargs = {'weighting': weighting} + else: + # No `exponent` or `weighting` applicable + spc_kwargs = {} + out_space = type(self.space)(self.shape, res.dtype, + **spc_kwargs) + out = out_space.element(res) + + return out + + elif ufunc.nout == 2: + # Make contexts for outputs (trivial ones return `None`) + if out1 is not None: + out1_ctx = writable_array(out1, **array_kwargs) + else: + out1_ctx = nullcontext() + if out2 is not None: + out2_ctx = writable_array(out2, **array_kwargs) + else: + out2_ctx = nullcontext() + + # Evaluate ufunc + with out1_ctx as out1_arr, out2_ctx as out2_arr: + kwargs['out'] = (out1_arr, out2_arr) + res1, res2 = ufunc(*inputs, **kwargs) + + # Wrap results if necessary (lazily) + # We don't use exponents or weightings since we don't know + # how to map them to the spaces + if out1 is None: + out1_space = type(self.space)(self.shape, res1.dtype) + out1 = out1_space.element(res1) + if out2 is None: + out2_space = type(self.space)(self.shape, res2.dtype) + out2 = out2_space.element(res2) + + return out1, out2 + + else: + raise NotImplementedError('nout = {} not supported' + ''.format(ufunc.nout)) + + else: # method != '__call__' + # Make context for output (trivial one returns `None`) + if out is None: + out_ctx = nullcontext() + else: + out_ctx = writable_array(out, **array_kwargs) + + # Evaluate ufunc method + with out_ctx as out_arr: + if method != 'at': + # No kwargs allowed for 'at' + kwargs['out'] = out_arr + res = getattr(ufunc, method)(*inputs, **kwargs) + + # Shortcut for scalar or no return value + if np.isscalar(res) or res is None: + # The first occurs for `reduce` with all axes, + # the second for in-place stuff (`at` currently) + return res + + # Wrap result if necessary (lazily) + if out is None: + if is_floating_dtype(res.dtype): + if res.shape != self.shape: + # Don't propagate weighting if shape changes + weighting = NumpyTensorSpaceConstWeighting(1.0, + exponent) + spc_kwargs = {'weighting': weighting} + else: + spc_kwargs = {} + + out_space = type(self.space)(res.shape, res.dtype, + **spc_kwargs) + out = out_space.element(res) + + return out + + +def _blas_is_applicable(*args): + """Whether BLAS routines can be applied or not. + + BLAS routines are available for single and double precision + float or complex data only. If the arrays are non-contiguous, + BLAS methods are usually slower, and array-writing routines do + not work at all. Hence, only contiguous arrays are allowed. + + Parameters + ---------- + x1,...,xN : `NumpyTensor` + The tensors to be tested for BLAS conformity. + + Returns + ------- + blas_is_applicable : bool + ``True`` if all mentioned requirements are met, ``False`` otherwise. + """ + if any(x.dtype != args[0].dtype for x in args[1:]): + return False + elif any(x.dtype not in _BLAS_DTYPES for x in args): + return False + elif not (all(x.flags.f_contiguous for x in args) or + all(x.flags.c_contiguous for x in args)): + return False + elif any(x.size > np.iinfo('int32').max for x in args): + # Temporary fix for 32 bit int overflow in BLAS + # TODO: use chunking instead + return False + else: + return True + + +def _lincomb_impl(a, x1, b, x2, out): + """Optimized implementation of ``out[:] = a * x1 + b * x2``.""" + # Lazy import to improve `import odl` time + import scipy.linalg + + size = native(x1.size) + + if size < THRESHOLD_SMALL: + # Faster for small arrays + out.data[:] = a * x1.data + b * x2.data + return + + elif (size < THRESHOLD_MEDIUM or + not _blas_is_applicable(x1.data, x2.data, out.data)): + + def fallback_axpy(x1, x2, n, a): + """Fallback axpy implementation avoiding copy.""" + if a != 0: + x2 /= a + x2 += x1 + x2 *= a + return x2 + + def fallback_scal(a, x, n): + """Fallback scal implementation.""" + x *= a + return x + + def fallback_copy(x1, x2, n): + """Fallback copy implementation.""" + x2[...] = x1[...] + return x2 + + axpy, scal, copy = (fallback_axpy, fallback_scal, fallback_copy) + x1_arr = x1.data + x2_arr = x2.data + out_arr = out.data + + else: + # Need flat data for BLAS, otherwise in-place does not work. + # Raveling must happen in fixed order for non-contiguous out, + # otherwise 'A' is applied to arrays, which makes the outcome + # dependent on their respective contiguousness. + if out.data.flags.f_contiguous: + ravel_order = 'F' + else: + ravel_order = 'C' + + x1_arr = x1.data.ravel(order=ravel_order) + x2_arr = x2.data.ravel(order=ravel_order) + out_arr = out.data.ravel(order=ravel_order) + axpy, scal, copy = scipy.linalg.blas.get_blas_funcs( + ['axpy', 'scal', 'copy'], arrays=(x1_arr, x2_arr, out_arr)) + + if x1 is x2 and b != 0: + # x1 is aligned with x2 -> out = (a+b)*x1 + _lincomb_impl(a + b, x1, 0, x1, out) + elif out is x1 and out is x2: + # All the vectors are aligned -> out = (a+b)*out + if (a + b) != 0: + scal(a + b, out_arr, size) + else: + out_arr[:] = 0 + elif out is x1: + # out is aligned with x1 -> out = a*out + b*x2 + if a != 1: + scal(a, out_arr, size) + if b != 0: + axpy(x2_arr, out_arr, size, b) + elif out is x2: + # out is aligned with x2 -> out = a*x1 + b*out + if b != 1: + scal(b, out_arr, size) + if a != 0: + axpy(x1_arr, out_arr, size, a) + else: + # We have exhausted all alignment options, so x1 is not x2 is not out + # We now optimize for various values of a and b + if b == 0: + if a == 0: # Zero assignment -> out = 0 + out_arr[:] = 0 + else: # Scaled copy -> out = a*x1 + copy(x1_arr, out_arr, size) + if a != 1: + scal(a, out_arr, size) + + else: # b != 0 + if a == 0: # Scaled copy -> out = b*x2 + copy(x2_arr, out_arr, size) + if b != 1: + scal(b, out_arr, size) + + elif a == 1: # No scaling in x1 -> out = x1 + b*x2 + copy(x1_arr, out_arr, size) + axpy(x2_arr, out_arr, size, b) + else: # Generic case -> out = a*x1 + b*x2 + copy(x2_arr, out_arr, size) + if b != 1: + scal(b, out_arr, size) + axpy(x1_arr, out_arr, size, a) + + +def _weighting(weights, exponent): + """Return a weighting whose type is inferred from the arguments.""" + if np.isscalar(weights): + weighting = NumpyTensorSpaceConstWeighting(weights, exponent) + elif weights is None: + weighting = NumpyTensorSpaceConstWeighting(1.0, exponent) + else: # last possibility: make an array + arr = np.asarray(weights) + weighting = NumpyTensorSpaceArrayWeighting(arr, exponent) + return weighting + + +def npy_weighted_inner(weights): + """Weighted inner product on `TensorSpace`'s as free function. + + Parameters + ---------- + weights : scalar or `array-like` + Weights of the inner product. A scalar is interpreted as a + constant weight, a 1-dim. array as a weighting vector. + + Returns + ------- + inner : `callable` + Inner product function with given weight. Constant weightings + are applicable to spaces of any size, for arrays the sizes + of the weighting and the space must match. + + See Also + -------- + NumpyTensorSpaceConstWeighting + NumpyTensorSpaceArrayWeighting + """ + return _weighting(weights, exponent=2.0).inner + + +def npy_weighted_norm(weights, exponent=2.0): + """Weighted norm on `TensorSpace`'s as free function. + + Parameters + ---------- + weights : scalar or `array-like` + Weights of the norm. A scalar is interpreted as a + constant weight, a 1-dim. array as a weighting vector. + exponent : positive `float` + Exponent of the norm. + + Returns + ------- + norm : `callable` + Norm function with given weight. Constant weightings + are applicable to spaces of any size, for arrays the sizes + of the weighting and the space must match. + + See Also + -------- + NumpyTensorSpaceConstWeighting + NumpyTensorSpaceArrayWeighting + """ + return _weighting(weights, exponent=exponent).norm + + +def npy_weighted_dist(weights, exponent=2.0): + """Weighted distance on `TensorSpace`'s as free function. + + Parameters + ---------- + weights : scalar or `array-like` + Weights of the distance. A scalar is interpreted as a + constant weight, a 1-dim. array as a weighting vector. + exponent : positive `float` + Exponent of the norm. + + Returns + ------- + dist : `callable` + Distance function with given weight. Constant weightings + are applicable to spaces of any size, for arrays the sizes + of the weighting and the space must match. + + See Also + -------- + NumpyTensorSpaceConstWeighting + NumpyTensorSpaceArrayWeighting + """ + return _weighting(weights, exponent=exponent).dist + + +def _norm_default(x): + """Default Euclidean norm implementation.""" + # Lazy import to improve `import odl` time + import scipy.linalg + + if _blas_is_applicable(x.data): + nrm2 = scipy.linalg.blas.get_blas_funcs('nrm2', dtype=x.dtype) + norm = partial(nrm2, n=native(x.size)) + else: + norm = np.linalg.norm + return norm(x.data.ravel()) + + +def _pnorm_default(x, p): + """Default p-norm implementation.""" + return np.linalg.norm(x.data.ravel(), ord=p) + + +def _pnorm_diagweight(x, p, w): + """Diagonally weighted p-norm implementation.""" + # Ravel both in the same order (w is a numpy array) + order = 'F' if all(a.flags.f_contiguous for a in (x.data, w)) else 'C' + + # This is faster than first applying the weights and then summing with + # BLAS dot or nrm2 + xp = np.abs(x.data.ravel(order)) + if p == float('inf'): + xp *= w.ravel(order) + return np.max(xp) + else: + xp = np.power(xp, p, out=xp) + xp *= w.ravel(order) + return np.sum(xp) ** (1 / p) + + +def _inner_default(x1, x2): + """Default Euclidean inner product implementation.""" + # Ravel both in the same order + order = 'F' if all(a.data.flags.f_contiguous for a in (x1, x2)) else 'C' + + if is_real_dtype(x1.dtype): + if x1.size > THRESHOLD_MEDIUM: + # This is as fast as BLAS dotc + return np.tensordot(x1, x2, [range(x1.ndim)] * 2) + else: + # Several times faster for small arrays + return np.dot(x1.data.ravel(order), + x2.data.ravel(order)) + else: + # x2 as first argument because we want linearity in x1 + return np.vdot(x2.data.ravel(order), + x1.data.ravel(order)) + + +# TODO: implement intermediate weighting schemes with arrays that are +# broadcast, i.e. between scalar and full-blown in dimensionality? + + +class NumpyTensorSpaceArrayWeighting(ArrayWeighting): + + """Weighting of a `NumpyTensorSpace` by an array. + + This class defines a weighting by an array that has the same shape + as the tensor space. Since the space is not known to this class, + no checks of shape or data type are performed. + See ``Notes`` for mathematical details. + """ + + def __init__(self, array, exponent=2.0): + r"""Initialize a new instance. + + Parameters + ---------- + array : `array-like`, one-dim. + Weighting array of the inner product, norm and distance. + All its entries must be positive, however this is not + verified during initialization. + exponent : positive `float` + Exponent of the norm. For values other than 2.0, no inner + product is defined. + + Notes + ----- + - For exponent 2.0, a new weighted inner product with array + :math:`W` is defined as + + .. math:: + \langle A, B\rangle_W := + \langle W \odot A, B\rangle = + \langle w \odot a, b\rangle = + b^{\mathrm{H}} (w \odot a), + + where :math:`a, b, w` are the "flattened" counterparts of + tensors :math:`A, B, W`, respectively, :math:`b^{\mathrm{H}}` + stands for transposed complex conjugate and :math:`w \odot a` + for element-wise multiplication. + + - For other exponents, only norm and dist are defined. In the + case of exponent :math:`\infty`, the weighted norm is + + .. math:: + \| A\|_{W, \infty} := + \| W \odot A\|_{\infty} = + \| w \odot a\|_{\infty}, + + otherwise it is (using point-wise exponentiation) + + .. math:: + \| A\|_{W, p} := + \| W^{1/p} \odot A\|_{p} = + \| w^{1/p} \odot a\|_{\infty}. + + - Note that this definition does **not** fulfill the limit + property in :math:`p`, i.e. + + .. math:: + \| A\|_{W, p} \not\to + \| A\|_{W, \infty} \quad (p \to \infty) + + unless all weights are equal to 1. + + - The array :math:`W` may only have positive entries, otherwise + it does not define an inner product or norm, respectively. This + is not checked during initialization. + """ + if isinstance(array, NumpyTensor): + array = array.data + elif not isinstance(array, np.ndarray): + array = np.asarray(array) + super(NumpyTensorSpaceArrayWeighting, self).__init__( + array, impl='numpy', exponent=exponent) + + def __hash__(self): + """Return ``hash(self)``.""" + return hash((type(self), self.array.tobytes(), self.exponent)) + + def inner(self, x1, x2): + """Return the weighted inner product of ``x1`` and ``x2``. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Tensors whose inner product is calculated. + + Returns + ------- + inner : float or complex + The inner product of the two provided vectors. + """ + if self.exponent != 2.0: + raise NotImplementedError('no inner product defined for ' + 'exponent != 2 (got {})' + ''.format(self.exponent)) + else: + inner = _inner_default(x1 * self.array, x2) + if is_real_dtype(x1.dtype): + return float(inner) + else: + return complex(inner) + + def norm(self, x): + """Return the weighted norm of ``x``. + + Parameters + ---------- + x : `NumpyTensor` + Tensor whose norm is calculated. + + Returns + ------- + norm : float + The norm of the provided tensor. + """ + if self.exponent == 2.0: + norm_squared = self.inner(x, x).real # TODO: optimize?! + if norm_squared < 0: + norm_squared = 0.0 # Compensate for numerical error + return float(np.sqrt(norm_squared)) + else: + return float(_pnorm_diagweight(x, self.exponent, self.array)) + + +class NumpyTensorSpaceConstWeighting(ConstWeighting): + + """Weighting of a `NumpyTensorSpace` by a constant. + + See ``Notes`` for mathematical details. + """ + + def __init__(self, const, exponent=2.0): + r"""Initialize a new instance. + + Parameters + ---------- + const : positive float + Weighting constant of the inner product, norm and distance. + exponent : positive float + Exponent of the norm. For values other than 2.0, the inner + product is not defined. + + Notes + ----- + - For exponent 2.0, a new weighted inner product with constant + :math:`c` is defined as + + .. math:: + \langle a, b\rangle_c := + c \, \langle a, b\rangle_c = + c \, b^{\mathrm{H}} a, + + where :math:`b^{\mathrm{H}}` standing for transposed complex + conjugate. + + - For other exponents, only norm and dist are defined. In the + case of exponent :math:`\infty`, the weighted norm is defined + as + + .. math:: + \| a \|_{c, \infty} := + c\, \| a \|_{\infty}, + + otherwise it is + + .. math:: + \| a \|_{c, p} := + c^{1/p}\, \| a \|_{p}. + + - Note that this definition does **not** fulfill the limit + property in :math:`p`, i.e. + + .. math:: + \| a\|_{c, p} \not\to + \| a \|_{c, \infty} \quad (p \to \infty) + + unless :math:`c = 1`. + + - The constant must be positive, otherwise it does not define an + inner product or norm, respectively. + """ + super(NumpyTensorSpaceConstWeighting, self).__init__( + const, impl='numpy', exponent=exponent) + + def inner(self, x1, x2): + """Return the weighted inner product of ``x1`` and ``x2``. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Tensors whose inner product is calculated. + + Returns + ------- + inner : float or complex + The inner product of the two provided tensors. + """ + if self.exponent != 2.0: + raise NotImplementedError('no inner product defined for ' + 'exponent != 2 (got {})' + ''.format(self.exponent)) + else: + inner = self.const * _inner_default(x1, x2) + if x1.space.field is None: + return inner + else: + return x1.space.field.element(inner) + + def norm(self, x): + """Return the weighted norm of ``x``. + + Parameters + ---------- + x1 : `NumpyTensor` + Tensor whose norm is calculated. + + Returns + ------- + norm : float + The norm of the tensor. + """ + if self.exponent == 2.0: + return float(np.sqrt(self.const) * _norm_default(x)) + elif self.exponent == float('inf'): + return float(self.const * _pnorm_default(x, self.exponent)) + else: + return float((self.const ** (1 / self.exponent) * + _pnorm_default(x, self.exponent))) + + def dist(self, x1, x2): + """Return the weighted distance between ``x1`` and ``x2``. + + Parameters + ---------- + x1, x2 : `NumpyTensor` + Tensors whose mutual distance is calculated. + + Returns + ------- + dist : float + The distance between the tensors. + """ + if self.exponent == 2.0: + return float(np.sqrt(self.const) * _norm_default(x1 - x2)) + elif self.exponent == float('inf'): + return float(self.const * _pnorm_default(x1 - x2, self.exponent)) + else: + return float((self.const ** (1 / self.exponent) * + _pnorm_default(x1 - x2, self.exponent))) + + +class NumpyTensorSpaceCustomInner(CustomInner): + + """Class for handling a user-specified inner product.""" + + def __init__(self, inner): + """Initialize a new instance. + + Parameters + ---------- + inner : `callable` + The inner product implementation. It must accept two + `Tensor` arguments, return an element from their space's + field (real or complex number) and satisfy the following + conditions for all vectors ``x, y, z`` and scalars ``s``: + + - `` = conj()`` + - `` = s * + `` + - `` = 0`` if and only if ``x = 0`` + """ + super(NumpyTensorSpaceCustomInner, self).__init__(inner, impl='numpy') + + +class NumpyTensorSpaceCustomNorm(CustomNorm): + + """Class for handling a user-specified norm. + + Note that this removes ``inner``. + """ + + def __init__(self, norm): + """Initialize a new instance. + + Parameters + ---------- + norm : `callable` + The norm implementation. It must accept a `Tensor` + argument, return a `float` and satisfy the following + conditions for all any two elements ``x, y`` and scalars + ``s``: + + - ``||x|| >= 0`` + - ``||x|| = 0`` if and only if ``x = 0`` + - ``||s * x|| = |s| * ||x||`` + - ``||x + y|| <= ||x|| + ||y||`` + """ + super(NumpyTensorSpaceCustomNorm, self).__init__(norm, impl='numpy') + + +class NumpyTensorSpaceCustomDist(CustomDist): + + """Class for handling a user-specified distance in `TensorSpace`. + + Note that this removes ``inner`` and ``norm``. + """ + + def __init__(self, dist): + """Initialize a new instance. + + Parameters + ---------- + dist : `callable` + The distance function defining a metric on `TensorSpace`. It + must accept two `Tensor` arguments, return a `float` and + fulfill the following mathematical conditions for any three + elements ``x, y, z``: + + - ``dist(x, y) >= 0`` + - ``dist(x, y) = 0`` if and only if ``x = y`` + - ``dist(x, y) = dist(y, x)`` + - ``dist(x, y) <= dist(x, z) + dist(z, y)`` + """ + super(NumpyTensorSpaceCustomDist, self).__init__(dist, impl='numpy') + + +if __name__ == '__main__': + from odl.util.testutils import run_doctests + run_doctests() From 7bc54d77805bb7690782ace5cc5bc244016bc79f Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 19 Mar 2024 15:07:06 +0100 Subject: [PATCH 02/52] A basic PyTorch pendant to NumpyTensorSpace. Only superficially tested so far. --- odl/space/__init__.py | 2 + odl/space/entry_points.py | 5 +- odl/space/pytorch_tensors.py | 853 +++++++---------------------------- 3 files changed, 168 insertions(+), 692 deletions(-) diff --git a/odl/space/__init__.py b/odl/space/__init__.py index 59368edebf7..e66e59d6fa0 100644 --- a/odl/space/__init__.py +++ b/odl/space/__init__.py @@ -12,10 +12,12 @@ from . import base_tensors, entry_points, weighting from .npy_tensors import * +from .pytorch_tensors import * from .pspace import * from .space_utils import * __all__ = () __all__ += npy_tensors.__all__ +__all__ += pytorch_tensors.__all__ __all__ += pspace.__all__ __all__ += space_utils.__all__ diff --git a/odl/space/entry_points.py b/odl/space/entry_points.py index fe1fc7644f8..e571869938b 100644 --- a/odl/space/entry_points.py +++ b/odl/space/entry_points.py @@ -23,12 +23,15 @@ from __future__ import print_function, division, absolute_import from odl.space.npy_tensors import NumpyTensorSpace +from odl.space.pytorch_tensors import PytorchTensorSpace # We don't expose anything to odl.space __all__ = () IS_INITIALIZED = False -TENSOR_SPACE_IMPLS = {'numpy': NumpyTensorSpace} +TENSOR_SPACE_IMPLS = {'numpy': NumpyTensorSpace, + 'pytorch': PytorchTensorSpace + } def _initialize_if_needed(): diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index d80bac82087..ca843d7b8ad 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -1,4 +1,4 @@ -# Copyright 2014-2020 The ODL contributors +# Copyright 2024 The ODL contributors # # This file is part of ODL. # @@ -6,7 +6,7 @@ # 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/. -"""NumPy implementation of tensor spaces.""" +"""PyTorch implementation of tensor spaces.""" from __future__ import absolute_import, division, print_function from future.utils import native @@ -16,6 +16,7 @@ from functools import partial import numpy as np +import torch from odl.set.sets import ComplexNumbers, RealNumbers from odl.set.space import LinearSpaceTypeError @@ -27,20 +28,22 @@ dtype_str, is_floating_dtype, is_numeric_dtype, is_real_dtype, nullcontext, signature_string, writable_array) -__all__ = ('NumpyTensorSpace',) +__all__ = ('PytorchTensorSpace',) -_BLAS_DTYPES = (np.dtype('float32'), np.dtype('float64'), - np.dtype('complex64'), np.dtype('complex128')) +_PYTORCH_DTYPES = {np.dtype('float32'): torch.float32, + np.dtype('float64'): torch.float64, + np.dtype('complex64'): torch.complex64, + np.dtype('complex128'): torch.complex128} # Define size thresholds to switch implementations THRESHOLD_SMALL = 100 THRESHOLD_MEDIUM = 50000 -class NumpyTensorSpace(TensorSpace): +class PytorchTensorSpace(TensorSpace): - """Set of tensors of arbitrary data type, implemented with NumPy. + """Set of tensors of arbitrary data type, implemented with Pytorch. A tensor is, in the most general sense, a multi-dimensional array that allows operations per entry (keep the rank constant), @@ -57,7 +60,7 @@ class NumpyTensorSpace(TensorSpace): supported by this class, along with reductions based on arithmetic or comparison, and element-wise mathematical functions ("ufuncs"). - This class is implemented using `numpy.ndarray`'s as back-end. + This class is implemented using `torch.Tensor`'s as back-end. See the `Wikipedia article on tensors`_ for further details. See also [Hac2012] "Part I Algebraic Tensors" for a rigorous @@ -84,7 +87,7 @@ def __init__(self, shape, dtype=None, **kwargs): single integer results in a space with rank 1, i.e., 1 axis. dtype : Data type of each element. Can be provided in any - way the `numpy.dtype` function understands, e.g. + way the `torch.dtype` function understands, e.g. as built-in type or as a string. For ``None``, the `default_dtype` of this space (``float64``) is used. exponent : positive float, optional @@ -117,7 +120,7 @@ def __init__(self, shape, dtype=None, **kwargs): dist : callable, optional Distance function defining a metric on the space. - It must accept two `NumpyTensor` arguments and return + It must accept two `PytorchTensor` arguments and return a non-negative real number. See ``Notes`` for mathematical requirements. @@ -129,7 +132,7 @@ def __init__(self, shape, dtype=None, **kwargs): norm : callable, optional The norm implementation. It must accept a - `NumpyTensor` argument, return a non-negative real number. + `PytorchTensor` argument, return a non-negative real number. See ``Notes`` for mathematical requirements. By default, ``norm(x)`` is calculated as ``inner(x, x)``. @@ -140,7 +143,7 @@ def __init__(self, shape, dtype=None, **kwargs): inner : callable, optional The inner product implementation. It must accept two - `NumpyTensor` arguments and return an element of the field + `PytorchTensor` arguments and return an element of the field of the space (usually real or complex number). See ``Notes`` for mathematical requirements. @@ -201,25 +204,16 @@ def __init__(self, shape, dtype=None, **kwargs): -------- Explicit initialization with the class constructor: - >>> space = NumpyTensorSpace(3, float) + >>> space = PytorchTensorSpace(3, float) >>> space rn(3) >>> space.shape (3,) >>> space.dtype dtype('float64') - - A more convenient way is to use factory functions: - - >>> space = odl.rn(3, weighting=[1, 2, 3]) - >>> space - rn(3, weighting=[1, 2, 3]) - >>> space = odl.tensor_space((2, 3), dtype=int) - >>> space - tensor_space((2, 3), dtype=int) """ - super(NumpyTensorSpace, self).__init__(shape, dtype) - if self.dtype.char not in self.available_dtypes(): + super(PytorchTensorSpace, self).__init__(shape, dtype) + if self.dtype not in self.available_dtypes(): raise ValueError('`dtype` {!r} not supported' ''.format(dtype_str(dtype))) @@ -234,6 +228,8 @@ def __init__(self, shape, dtype=None, **kwargs): raise ValueError('cannot use any of `weighting`, `dist`, `norm` ' 'or `inner` for non-numeric `dtype` {}' ''.format(dtype)) + else: + self._torch_dtype = _PYTORCH_DTYPES[self.dtype] if exponent != 2.0 and any(x is not None for x in (dist, norm, inner)): raise ValueError('cannot use any of `dist`, `norm` or `inner` ' 'for exponent != 2') @@ -247,8 +243,8 @@ def __init__(self, shape, dtype=None, **kwargs): # Set the weighting if weighting is not None: if isinstance(weighting, Weighting): - if weighting.impl != 'numpy': - raise ValueError("`weighting.impl` must be 'numpy', " + if weighting.impl != 'torch': + raise ValueError("`weighting.impl` must be 'torch', " '`got {!r}'.format(weighting.impl)) if weighting.exponent != exponent: raise ValueError('`weighting.exponent` conflicts with ' @@ -259,7 +255,7 @@ def __init__(self, shape, dtype=None, **kwargs): self.__weighting = _weighting(weighting, exponent) # Check (afterwards) that the weighting input was sane - if isinstance(self.weighting, NumpyTensorSpaceArrayWeighting): + if isinstance(self.weighting, PytorchTensorSpaceArrayWeighting): if self.weighting.array.dtype == object: raise ValueError('invalid `weighting` argument: {}' ''.format(weighting)) @@ -276,14 +272,14 @@ def __init__(self, shape, dtype=None, **kwargs): self.weighting.array.shape)) elif dist is not None: - self.__weighting = NumpyTensorSpaceCustomDist(dist) + self.__weighting = PytorchTensorSpaceCustomDist(dist) elif norm is not None: - self.__weighting = NumpyTensorSpaceCustomNorm(norm) + self.__weighting = PytorchTensorSpaceCustomNorm(norm) elif inner is not None: - self.__weighting = NumpyTensorSpaceCustomInner(inner) + self.__weighting = PytorchTensorSpaceCustomInner(inner) else: # No weighting, i.e., weighting with constant 1.0 - self.__weighting = NumpyTensorSpaceConstWeighting(1.0, exponent) + self.__weighting = PytorchTensorSpaceConstWeighting(1.0, exponent) # Make sure there are no leftover kwargs if kwargs: @@ -296,7 +292,7 @@ def impl(self): @property def default_order(self): - """Default storage order for new elements in this space: ``'C'``.""" + """Default (and only) storage order for new elements in this space: ``'C'``.""" return 'C' @property @@ -308,7 +304,7 @@ def weighting(self): def is_weighted(self): """Return ``True`` if the space is not weighted by constant 1.0.""" return not ( - isinstance(self.weighting, NumpyTensorSpaceConstWeighting) and + isinstance(self.weighting, PytorchTensorSpaceConstWeighting) and self.weighting.const == 1.0) @property @@ -326,8 +322,8 @@ def element(self, inp=None, data_ptr=None, order=None): If ``inp`` is `None`, an empty element is created with no guarantee of its state (memory allocation only). - The new element will use ``order`` as storage order if - provided, otherwise `default_order`. + All tensors use row-major storage (corrsponding to + `order='C'` in NumPy). Otherwise, a copy is avoided whenever possible. This requires correct `shape` and `dtype`, and if ``order`` is provided, @@ -335,19 +331,14 @@ def element(self, inp=None, data_ptr=None, order=None): conditions is not met, a copy is made. data_ptr : int, optional - Pointer to the start memory address of a contiguous Numpy array + Pointer to the start memory address of a contiguous PyTorch array or an equivalent raw container with the same total number of - bytes. For this option, ``order`` must be either ``'C'`` or - ``'F'``. + bytes. The option is also mutually exclusive with ``inp``. - order : {None, 'C', 'F'}, optional - Storage order of the returned element. For ``'C'`` and ``'F'``, - contiguous memory in the respective ordering is enforced. - The default ``None`` enforces no contiguousness. Returns ------- - element : `NumpyTensor` + element : `PytorchTensor` The new element, created from ``inp`` or from scratch. Examples @@ -355,7 +346,7 @@ def element(self, inp=None, data_ptr=None, order=None): Without arguments, an uninitialized element is created. With an array-like input, the element can be initialized: - >>> space = odl.rn(3) + >>> space = odl.rn(3) # TODO adapt / test >>> empty = space.element() >>> empty.shape (3,) @@ -365,11 +356,11 @@ def element(self, inp=None, data_ptr=None, order=None): >>> x rn(3).element([ 1., 2., 3.]) - If the input already is a `numpy.ndarray` of correct `dtype`, it + If the input already is a `torch.Tensor` of correct `dtype`, it will merely be wrapped, i.e., both array and space element access the same memory, such that mutations will affect both: - >>> arr = np.array([1, 2, 3], dtype=float) + >>> arr = torch.Tensor([1, 2, 3], dtype=float) # TODO test >>> elem = odl.rn(3).element(arr) >>> elem[0] = 0 >>> elem @@ -381,8 +372,8 @@ def element(self, inp=None, data_ptr=None, order=None): again in shared memory: >>> int_space = odl.tensor_space((2, 3), dtype=int) - >>> arr = np.array([[1, 2, 3], - ... [4, 5, 6]], dtype=int, order='F') + >>> arr = torch.Tensor([[1, 2, 3], + ... [4, 5, 6]], dtype=int, order='F') >>> ptr = arr.ctypes.data >>> y = int_space.element(data_ptr=ptr, order='F') >>> y @@ -395,15 +386,11 @@ def element(self, inp=None, data_ptr=None, order=None): array([[ 1, -1, 3], [ 4, 5, 6]]) """ - if order is not None and str(order).upper() not in ('C', 'F'): - raise ValueError("`order` {!r} not understood".format(order)) + if order is not None and str(order).upper() not in ('C'): + raise ValueError(f"Only row-major order supported ('C'), not '{order}'.") if inp is None and data_ptr is None: - if order is None: - arr = np.empty(self.shape, dtype=self.dtype, - order=self.default_order) - else: - arr = np.empty(self.shape, dtype=self.dtype, order=order) + arr = torch.empty(self.shape, dtype=self._torch_dtype) return self.element_type(self, arr) @@ -415,23 +402,18 @@ def element(self, inp=None, data_ptr=None, order=None): ctype_array_def = ctypes.c_byte * self.nbytes as_ctype_array = ctype_array_def.from_address(data_ptr) as_numpy_array = np.ctypeslib.as_array(as_ctype_array) - arr = as_numpy_array.view(dtype=self.dtype) + arr = as_numpy_array.view(dtype=self._torch_dtype) arr = arr.reshape(self.shape, order=order) - return self.element_type(self, arr) + return self.element_type(self, torch.Tensor(arr)) elif inp is not None and data_ptr is None: if inp in self and order is None: # Short-circuit for space elements and no enforced ordering return inp - # Try to not copy but require dtype and order if given - # (`order=None` is ok as np.array argument) - arr = np.array(inp, copy=False, dtype=self.dtype, ndmin=self.ndim, - order=order) - # Make sure the result is writeable, if not make copy. - # This happens for e.g. results of `np.broadcast_to()`. - if not arr.flags.writeable: - arr = arr.copy() + # TODO avoid copy when it's not necessary + arr = torch.tensor(inp, dtype=self._torch_dtype) + if arr.shape != self.shape: raise ValueError('shape of `inp` not equal to space shape: ' '{} != {}'.format(arr.shape, self.shape)) @@ -445,26 +427,24 @@ def zero(self): Examples -------- - >>> space = odl.rn(3) + >>> space = odl.rn(3) # TODO adapt >>> x = space.zero() >>> x rn(3).element([ 0., 0., 0.]) """ - return self.element(np.zeros(self.shape, dtype=self.dtype, - order=self.default_order)) + return self.element(torch.zeros(self.shape, dtype=self._torch_dtype)) def one(self): """Return a tensor of all ones. Examples -------- - >>> space = odl.rn(3) + >>> space = odl.rn(3) # TODO adapt >>> x = space.one() >>> x rn(3).element([ 1., 1., 1.]) """ - return self.element(np.ones(self.shape, dtype=self.dtype, - order=self.default_order)) + return self.element(torch.ones(self.shape, dtype=self._torch_dtype)) @staticmethod def available_dtypes(): @@ -472,20 +452,10 @@ def available_dtypes(): Notes ----- - This is all dtypes available in Numpy. See ``numpy.sctypes`` - for more information. - - The available dtypes may depend on the specific system used. + Currently only a conservative selection of the types supported + by Pytorch. """ - all_dtypes = [] - for lst in np.sctypes.values(): - for dtype in lst: - if dtype not in (object, np.void): - all_dtypes.append(np.dtype(dtype)) - # Need to add these manually since np.sctypes['others'] will only - # contain one of them (depending on Python version) - all_dtypes.extend([np.dtype('S'), np.dtype('U')]) - return tuple(sorted(set(all_dtypes))) + return [np.float32, np.float64, np.complex64, np.complex128] @staticmethod def default_dtype(field=None): @@ -500,17 +470,20 @@ def default_dtype(field=None): Returns ------- - dtype : `numpy.dtype` - Numpy data type specifier. The returned defaults are: + dtype : `torch.dtype` + Pytorch data type specifier. The returned defaults are: ``RealNumbers()`` : ``np.dtype('float64')`` ``ComplexNumbers()`` : ``np.dtype('complex128')`` """ + # Note that we're using the NumPy versions of the types, rather + # than the equivalent Pytorch ones. This is for compatibility + # with the rest of ODL, which is not aware of Pytorch. if field is None or field == RealNumbers(): - return np.dtype('float64') + return np.float64 elif field == ComplexNumbers(): - return np.dtype('complex128') + return np.complex128 else: raise ValueError('no default data type defined for field {}' ''.format(field)) @@ -528,14 +501,14 @@ def _lincomb(self, a, x1, b, x2, out): ---------- a, b : `TensorSpace.field` element Scalars to multiply ``x1`` and ``x2`` with. - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Summands in the linear combination. - out : `NumpyTensor` + out : `PytorchTensor` Tensor to which the result is written. Examples -------- - >>> space = odl.rn(3) + >>> space = odl.rn(3) # TODO adapt >>> x = space.element([0, 1, 1]) >>> y = space.element([0, 0, 1]) >>> out = space.element() @@ -545,7 +518,7 @@ def _lincomb(self, a, x1, b, x2, out): >>> result is out True """ - _lincomb_impl(a, x1, b, x2, out) + torch.add(input=a*x1.data, other=x2.data, alpha=b, out=out.data) def _dist(self, x1, x2): """Return the distance between ``x1`` and ``x2``. @@ -555,7 +528,7 @@ def _dist(self, x1, x2): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Elements whose mutual distance is calculated. Returns @@ -597,7 +570,7 @@ def _norm(self, x): Parameters ---------- - x : `NumpyTensor` + x : `PytorchTensor` Element whose norm is calculated. Returns @@ -635,7 +608,7 @@ def _inner(self, x1, x2): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Elements whose inner product is calculated. Returns @@ -669,9 +642,9 @@ def _multiply(self, x1, x2, out): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Factors in the product. - out : `NumpyTensor` + out : `PytorchTensor` Element to which the result is written. Examples @@ -688,7 +661,7 @@ def _multiply(self, x1, x2, out): >>> result is out True """ - np.multiply(x1.data, x2.data, out=out.data) + torch.mul(x1.data, x2.data, out=out.data) def _divide(self, x1, x2, out): """Compute the entry-wise quotient ``x1 / x2``. @@ -698,9 +671,9 @@ def _divide(self, x1, x2, out): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Dividend and divisor in the quotient. - out : `NumpyTensor` + out : `PytorchTensor` Element to which the result is written. Examples @@ -717,7 +690,7 @@ def _divide(self, x1, x2, out): >>> result is out True """ - np.divide(x1.data, x2.data, out=out.data) + torch.div(x1.data, x2.data, out=out.data) def __eq__(self, other): """Return ``self == other``. @@ -726,8 +699,8 @@ def __eq__(self, other): ------- equals : bool True if ``other`` is an instance of ``type(self)`` - with the same `NumpyTensorSpace.shape`, `NumpyTensorSpace.dtype` - and `NumpyTensorSpace.weighting`, otherwise False. + with the same `PytorchTensorSpace.shape`, `PytorchTensorSpace.dtype` + and `PytorchTensorSpace.weighting`, otherwise False. Examples -------- @@ -754,12 +727,12 @@ def __eq__(self, other): if other is self: return True - return (super(NumpyTensorSpace, self).__eq__(other) and + return (super(PytorchTensorSpace, self).__eq__(other) and self.weighting == other.weighting) def __hash__(self): """Return ``hash(self)``.""" - return hash((super(NumpyTensorSpace, self).__hash__(), + return hash((super(PytorchTensorSpace, self).__hash__(), self.weighting)) @property @@ -770,7 +743,7 @@ def byaxis(self): -------- Indexing with integers or slices: - >>> space = odl.rn((2, 3, 4)) + >>> space = odl.rn((2, 3, 4)) # TODO adapt >>> space.byaxis[0] rn(2) >>> space.byaxis[1:] @@ -783,7 +756,7 @@ def byaxis(self): """ space = self - class NpyTensorSpacebyaxis(object): + class PytorchTensorSpacebyaxis(object): """Helper class for indexing by axis.""" @@ -798,7 +771,7 @@ def __getitem__(self, indices): if isinstance(space.weighting, ArrayWeighting): new_array = np.asarray(space.weighting.array[indices]) - weighting = NumpyTensorSpaceArrayWeighting( + weighting = PytorchTensorSpaceArrayWeighting( new_array, space.weighting.exponent) else: weighting = space.weighting @@ -809,7 +782,7 @@ def __repr__(self): """Return ``repr(self)``.""" return repr(space) + '.byaxis' - return NpyTensorSpacebyaxis() + return PytorchTensorSpacebyaxis() def __repr__(self): """Return ``repr(self)``.""" @@ -819,7 +792,7 @@ def __repr__(self): posargs = [self.shape] if self.is_real: - ctor_name = 'rn' + ctor_name = 'rn' # TODO adapt elif self.is_complex: ctor_name = 'cn' else: @@ -846,13 +819,13 @@ def __repr__(self): @property def element_type(self): - """Type of elements in this space: `NumpyTensor`.""" - return NumpyTensor + """Type of elements in this space: `PytorchTensor`.""" + return PytorchTensor -class NumpyTensor(Tensor): +class PytorchTensor(Tensor): - """Representation of a `NumpyTensorSpace` element.""" + """Representation of a `PytorchTensorSpace` element.""" def __init__(self, space, data): """Initialize a new instance.""" @@ -861,31 +834,31 @@ def __init__(self, space, data): @property def data(self): - """The `numpy.ndarray` representing the data of ``self``.""" + """The `torch.Tensor` representing the data of ``self``.""" return self.__data def asarray(self, out=None): - """Extract the data of this array as a ``numpy.ndarray``. + """Extract the data of this array as a ``torch.Tensor``. - This method is invoked when calling `numpy.asarray` on this + This method is invoked when calling `torch.tensor` on this tensor. Parameters ---------- - out : `numpy.ndarray`, optional + out : `np.ndarray`, optional Array in which the result should be written in-place. Has to be contiguous and of the correct dtype. Returns ------- - asarray : `numpy.ndarray` - Numpy array with the same data type as ``self``. If + asarray : `torch.Tensor` + Pytorch array with the same data type as ``self``. If ``out`` was given, the returned object is a reference to it. Examples -------- - >>> space = odl.rn(3, dtype='float32') + >>> space = odl.rn(3, dtype='float32') # TODO adapt >>> x = space.element([1, 2, 3]) >>> x.asarray() array([ 1., 2., 3.], dtype=float32) @@ -903,9 +876,9 @@ def asarray(self, out=None): [ 1., 1., 1.]]) """ if out is None: - return self.data + return self.data.cpu().numpy() else: - out[:] = self.data + out[:] = self.data.cpu().numpy() return out def astype(self, dtype): @@ -921,7 +894,7 @@ def astype(self, dtype): Returns ------- - newelem : `NumpyTensor` + newelem : `PytorchTensor` Version of this element with given data type. """ return self.space.astype(dtype).element(self.data.astype(dtype)) @@ -933,7 +906,7 @@ def data_ptr(self): Examples -------- >>> import ctypes - >>> space = odl.tensor_space(3, dtype='uint16') + >>> space = odl.tensor_space(3, dtype='uint16') # TODO check example >>> x = space.element([1, 2, 3]) >>> arr_type = ctypes.c_uint16 * 3 # C type "array of 3 uint16" >>> buffer = arr_type.from_address(x.data_ptr) @@ -943,11 +916,11 @@ def data_ptr(self): In-place modification via pointer: - >>> arr[0] = 42 + >>> arr[0] = 42 # TODO doubtful if this actually works >>> x tensor_space(3, dtype='uint16').element([42, 2, 3]) """ - return self.data.ctypes.data + return self.data.data_ptr() def __eq__(self, other): """Return ``self == other``. @@ -984,7 +957,7 @@ def __eq__(self, other): elif other not in self.space: return False else: - return np.array_equal(self.data, other.data) + return torch.equal(self.data, other.data) def copy(self): """Return an identical (deep) copy of this tensor. @@ -995,12 +968,12 @@ def copy(self): Returns ------- - copy : `NumpyTensor` + copy : `PytorchTensor` The deep copy Examples -------- - >>> space = odl.rn(3) + >>> space = odl.rn(3) # TODO adapt >>> x = space.element([1, 2, 3]) >>> y = x.copy() >>> y == x @@ -1008,7 +981,7 @@ def copy(self): >>> y is x False """ - return self.space.element(self.data.copy()) + return self.space.element(self.data.clone()) def __copy__(self): """Return ``copy(self)``. @@ -1044,7 +1017,7 @@ def __getitem__(self, indices): Returns ------- - values : `NumpyTensorSpace.dtype` or `NumpyTensor` + values : `PytorchTensorSpace.dtype` or `PytorchTensor` The value(s) at the given indices. Note that the returned object is a writable view into the original tensor, except for the case when ``indices`` is a list. @@ -1094,11 +1067,11 @@ def __getitem__(self, indices): ) """ # Lazy implementation: index the array and deal with it - if isinstance(indices, NumpyTensor): + if isinstance(indices, PytorchTensor): indices = indices.data arr = self.data[indices] - if np.isscalar(arr): + if arr.shape == (): # scalar if self.space.field is not None: return self.space.field.element(arr) else: @@ -1121,7 +1094,7 @@ def __setitem__(self, indices, values): indices : index expression Integer, slice or sequence of these, defining the positions of the data array which should be written to. - values : scalar, array-like or `NumpyTensor` + values : scalar, array-like or `PytorchTensor` The value(s) that are to be assigned. If ``index`` is an integer, ``value`` must be a scalar. @@ -1198,9 +1171,9 @@ def real(self): Returns ------- - real : `NumpyTensor` + real : `PytorchTensor` Real part of this element as a member of a - `NumpyTensorSpace` with corresponding real data type. + `PytorchTensorSpace` with corresponding real data type. Examples -------- @@ -1257,9 +1230,9 @@ def imag(self): Returns ------- - imag : `NumpyTensor` + imag : `PytorchTensor` Imaginary part this element as an element of a - `NumpyTensorSpace` with real data type. + `PytorchTensorSpace` with real data type. Examples -------- @@ -1322,13 +1295,13 @@ def conj(self, out=None): Parameters ---------- - out : `NumpyTensor`, optional + out : `PytorchTensor`, optional Element to which the complex conjugate is written. Must be an element of ``self.space``. Returns ------- - out : `NumpyTensor` + out : `PytorchTensor` The complex conjugate element. If ``out`` was provided, the returned object is a reference to it. @@ -1377,11 +1350,11 @@ def __ipow__(self, other): """Return ``self **= other``.""" try: if other == int(other): - return super(NumpyTensor, self).__ipow__(other) + return super(PytorchTensor, self).__ipow__(other) except TypeError: pass - np.power(self.data, other, out=self.data) + torch.pow(self.data, other, out=self.data) return self def __int__(self): @@ -1406,500 +1379,22 @@ def __complex__(self): 'Python scalars') return complex(self.data.ravel()[0]) - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - """Interface to Numpy's ufunc machinery. - - This method is called by Numpy version 1.13 and higher as a single - point for the ufunc dispatch logic. An object implementing - ``__array_ufunc__`` takes over control when a `numpy.ufunc` is - called on it, allowing it to use custom implementations and - output types. - - This includes handling of in-place arithmetic like - ``npy_array += custom_obj``. In this case, the custom object's - ``__array_ufunc__`` takes precedence over the baseline - `numpy.ndarray` implementation. It will be called with - ``npy_array`` as ``out`` argument, which ensures that the - returned object is a Numpy array. For this to work properly, - ``__array_ufunc__`` has to accept Numpy arrays as ``out`` arguments. - - See the `corresponding NEP`_ and the `interface documentation`_ - for further details. See also the `general documentation on - Numpy ufuncs`_. - - .. note:: - This basic implementation casts inputs and - outputs to Numpy arrays and evaluates ``ufunc`` on those. - For `numpy.ndarray` based data storage, this incurs no - significant overhead compared to direct usage of Numpy arrays. - - For other (in particular non-local) implementations, e.g., - GPU arrays or distributed memory, overhead is significant due - to copies to CPU main memory. In those classes, the - ``__array_ufunc__`` mechanism should be overridden to use - native implementations if possible. - - .. note:: - When using operations that alter the shape (like ``reduce``), - or the data type (can be any of the methods), - the resulting array is wrapped in a space of the same - type as ``self.space``, propagating space properties like - `exponent` or `weighting` as closely as possible. - - Parameters - ---------- - ufunc : `numpy.ufunc` - Ufunc that should be called on ``self``. - method : str - Method on ``ufunc`` that should be called on ``self``. - Possible values: - - ``'__call__'``, ``'accumulate'``, ``'at'``, ``'outer'``, - ``'reduce'``, ``'reduceat'`` - - input1, ..., inputN : - Positional arguments to ``ufunc.method``. - kwargs : - Keyword arguments to ``ufunc.method``. - - Returns - ------- - ufunc_result : `Tensor`, `numpy.ndarray` or tuple - Result of the ufunc evaluation. If no ``out`` keyword argument - was given, the result is a `Tensor` or a tuple - of such, depending on the number of outputs of ``ufunc``. - If ``out`` was provided, the returned object or tuple entries - refer(s) to ``out``. - - Examples - -------- - We apply `numpy.add` to ODL tensors: - - >>> r3 = odl.rn(3) - >>> x = r3.element([1, 2, 3]) - >>> y = r3.element([-1, -2, -3]) - >>> x.__array_ufunc__(np.add, '__call__', x, y) - rn(3).element([ 0., 0., 0.]) - >>> np.add(x, y) # same mechanism for Numpy >= 1.13 - rn(3).element([ 0., 0., 0.]) - - As ``out``, a Numpy array or an ODL tensor can be given (wrapped - in a sequence): - - >>> out = r3.element() - >>> res = x.__array_ufunc__(np.add, '__call__', x, y, out=(out,)) - >>> out - rn(3).element([ 0., 0., 0.]) - >>> res is out - True - >>> out_arr = np.empty(3) - >>> res = x.__array_ufunc__(np.add, '__call__', x, y, out=(out_arr,)) - >>> out_arr - array([ 0., 0., 0.]) - >>> res is out_arr - True - - With multiple dimensions: - - >>> r23 = odl.rn((2, 3)) - >>> x = y = r23.one() - >>> x.__array_ufunc__(np.add, '__call__', x, y) - rn((2, 3)).element( - [[ 2., 2., 2.], - [ 2., 2., 2.]] - ) - - The ``ufunc.accumulate`` method retains the original `shape` and - `dtype`. The latter can be changed with the ``dtype`` parameter: - - >>> x = r3.element([1, 2, 3]) - >>> x.__array_ufunc__(np.add, 'accumulate', x) - rn(3).element([ 1., 3., 6.]) - >>> np.add.accumulate(x) # same mechanism for Numpy >= 1.13 - rn(3).element([ 1., 3., 6.]) - >>> x.__array_ufunc__(np.add, 'accumulate', x, dtype=complex) - cn(3).element([ 1.+0.j, 3.+0.j, 6.+0.j]) - - For multi-dimensional tensors, an optional ``axis`` parameter - can be provided: - - >>> z = r23.one() - >>> z.__array_ufunc__(np.add, 'accumulate', z, axis=1) - rn((2, 3)).element( - [[ 1., 2., 3.], - [ 1., 2., 3.]] - ) - - The ``ufunc.at`` method operates in-place. Here we add the second - operand ``[5, 10]`` to ``x`` at indices ``[0, 2]``: - - >>> x = r3.element([1, 2, 3]) - >>> x.__array_ufunc__(np.add, 'at', x, [0, 2], [5, 10]) - >>> x - rn(3).element([ 6., 2., 13.]) - - For outer-product-type operations, i.e., operations where the result - shape is the sum of the individual shapes, the ``ufunc.outer`` - method can be used: - - >>> x = odl.rn(2).element([0, 3]) - >>> y = odl.rn(3).element([1, 2, 3]) - >>> x.__array_ufunc__(np.add, 'outer', x, y) - rn((2, 3)).element( - [[ 1., 2., 3.], - [ 4., 5., 6.]] - ) - >>> y.__array_ufunc__(np.add, 'outer', y, x) - rn((3, 2)).element( - [[ 1., 4.], - [ 2., 5.], - [ 3., 6.]] - ) - - Using ``ufunc.reduce`` produces a scalar, which can be avoided with - ``keepdims=True``: - - >>> x = r3.element([1, 2, 3]) - >>> x.__array_ufunc__(np.add, 'reduce', x) - 6.0 - >>> x.__array_ufunc__(np.add, 'reduce', x, keepdims=True) - rn(1).element([ 6.]) - - In multiple dimensions, ``axis`` can be provided for reduction over - selected axes: - - >>> z = r23.element([[1, 2, 3], - ... [4, 5, 6]]) - >>> z.__array_ufunc__(np.add, 'reduce', z, axis=1) - rn(2).element([ 6., 15.]) - Finally, ``add.reduceat`` is a combination of ``reduce`` and - ``at`` with rather flexible and complex semantics (see the - `reduceat documentation`_ for details): - - >>> x = r3.element([1, 2, 3]) - >>> x.__array_ufunc__(np.add, 'reduceat', x, [0, 1]) - rn(2).element([ 1., 5.]) - - References - ---------- - .. _corresponding NEP: - https://docs.scipy.org/doc/numpy/neps/ufunc-overrides.html - - .. _interface documentation: - https://docs.scipy.org/doc/numpy/reference/arrays.classes.html\ -#numpy.class.__array_ufunc__ - - .. _general documentation on Numpy ufuncs: - https://docs.scipy.org/doc/numpy/reference/ufuncs.html - - .. _reduceat documentation: - https://docs.scipy.org/doc/numpy/reference/generated/\ -numpy.ufunc.reduceat.html - """ - # Remark: this method differs from the parent implementation only - # in the propagation of additional space properties. - - # --- Process `out` --- # - - # Unwrap out if provided. The output parameters are all wrapped - # in one tuple, even if there is only one. - out_tuple = kwargs.pop('out', ()) - - # Check number of `out` args, depending on `method` - if method == '__call__' and len(out_tuple) not in (0, ufunc.nout): - raise ValueError( - "ufunc {}: need 0 or {} `out` arguments for " - "`method='__call__'`, got {}" - ''.format(ufunc.__name__, ufunc.nout, len(out_tuple))) - elif method != '__call__' and len(out_tuple) not in (0, 1): - raise ValueError( - 'ufunc {}: need 0 or 1 `out` arguments for `method={!r}`, ' - 'got {}'.format(ufunc.__name__, method, len(out_tuple))) - - # We allow our own tensors, the data container type and - # `numpy.ndarray` objects as `out` (see docs for reason for the - # latter) - valid_types = (type(self), type(self.data), np.ndarray) - if not all(isinstance(o, valid_types) or o is None - for o in out_tuple): - return NotImplemented - - # Assign to `out` or `out1` and `out2`, respectively - out = out1 = out2 = None - if len(out_tuple) == 1: - out = out_tuple[0] - elif len(out_tuple) == 2: - out1 = out_tuple[0] - out2 = out_tuple[1] - - # --- Process `inputs` --- # - - # Convert inputs that are ODL tensors to Numpy arrays so that the - # native Numpy ufunc is called later - inputs = tuple( - inp.asarray() if isinstance(inp, type(self)) else inp - for inp in inputs) - - # --- Get some parameters for later --- # - - # Arguments for `writable_array` and/or space constructors - out_dtype = kwargs.get('dtype', None) - if out_dtype is None: - array_kwargs = {} - else: - array_kwargs = {'dtype': out_dtype} - - exponent = self.space.exponent - weighting = self.space.weighting - - # --- Evaluate ufunc --- # - - if method == '__call__': - if ufunc.nout == 1: - # Make context for output (trivial one returns `None`) - if out is None: - out_ctx = nullcontext() - else: - out_ctx = writable_array(out, **array_kwargs) - - # Evaluate ufunc - with out_ctx as out_arr: - kwargs['out'] = out_arr - res = ufunc(*inputs, **kwargs) - - # Wrap result if necessary (lazily) - if out is None: - if is_floating_dtype(res.dtype): - # Weighting contains exponent - spc_kwargs = {'weighting': weighting} - else: - # No `exponent` or `weighting` applicable - spc_kwargs = {} - out_space = type(self.space)(self.shape, res.dtype, - **spc_kwargs) - out = out_space.element(res) - - return out - - elif ufunc.nout == 2: - # Make contexts for outputs (trivial ones return `None`) - if out1 is not None: - out1_ctx = writable_array(out1, **array_kwargs) - else: - out1_ctx = nullcontext() - if out2 is not None: - out2_ctx = writable_array(out2, **array_kwargs) - else: - out2_ctx = nullcontext() - - # Evaluate ufunc - with out1_ctx as out1_arr, out2_ctx as out2_arr: - kwargs['out'] = (out1_arr, out2_arr) - res1, res2 = ufunc(*inputs, **kwargs) - - # Wrap results if necessary (lazily) - # We don't use exponents or weightings since we don't know - # how to map them to the spaces - if out1 is None: - out1_space = type(self.space)(self.shape, res1.dtype) - out1 = out1_space.element(res1) - if out2 is None: - out2_space = type(self.space)(self.shape, res2.dtype) - out2 = out2_space.element(res2) - - return out1, out2 - - else: - raise NotImplementedError('nout = {} not supported' - ''.format(ufunc.nout)) - - else: # method != '__call__' - # Make context for output (trivial one returns `None`) - if out is None: - out_ctx = nullcontext() - else: - out_ctx = writable_array(out, **array_kwargs) - - # Evaluate ufunc method - with out_ctx as out_arr: - if method != 'at': - # No kwargs allowed for 'at' - kwargs['out'] = out_arr - res = getattr(ufunc, method)(*inputs, **kwargs) - - # Shortcut for scalar or no return value - if np.isscalar(res) or res is None: - # The first occurs for `reduce` with all axes, - # the second for in-place stuff (`at` currently) - return res - - # Wrap result if necessary (lazily) - if out is None: - if is_floating_dtype(res.dtype): - if res.shape != self.shape: - # Don't propagate weighting if shape changes - weighting = NumpyTensorSpaceConstWeighting(1.0, - exponent) - spc_kwargs = {'weighting': weighting} - else: - spc_kwargs = {} - - out_space = type(self.space)(res.shape, res.dtype, - **spc_kwargs) - out = out_space.element(res) - - return out - - -def _blas_is_applicable(*args): - """Whether BLAS routines can be applied or not. - - BLAS routines are available for single and double precision - float or complex data only. If the arrays are non-contiguous, - BLAS methods are usually slower, and array-writing routines do - not work at all. Hence, only contiguous arrays are allowed. - - Parameters - ---------- - x1,...,xN : `NumpyTensor` - The tensors to be tested for BLAS conformity. - - Returns - ------- - blas_is_applicable : bool - ``True`` if all mentioned requirements are met, ``False`` otherwise. - """ - if any(x.dtype != args[0].dtype for x in args[1:]): - return False - elif any(x.dtype not in _BLAS_DTYPES for x in args): - return False - elif not (all(x.flags.f_contiguous for x in args) or - all(x.flags.c_contiguous for x in args)): - return False - elif any(x.size > np.iinfo('int32').max for x in args): - # Temporary fix for 32 bit int overflow in BLAS - # TODO: use chunking instead - return False - else: - return True - - -def _lincomb_impl(a, x1, b, x2, out): - """Optimized implementation of ``out[:] = a * x1 + b * x2``.""" - # Lazy import to improve `import odl` time - import scipy.linalg - - size = native(x1.size) - - if size < THRESHOLD_SMALL: - # Faster for small arrays - out.data[:] = a * x1.data + b * x2.data - return - - elif (size < THRESHOLD_MEDIUM or - not _blas_is_applicable(x1.data, x2.data, out.data)): - - def fallback_axpy(x1, x2, n, a): - """Fallback axpy implementation avoiding copy.""" - if a != 0: - x2 /= a - x2 += x1 - x2 *= a - return x2 - - def fallback_scal(a, x, n): - """Fallback scal implementation.""" - x *= a - return x - - def fallback_copy(x1, x2, n): - """Fallback copy implementation.""" - x2[...] = x1[...] - return x2 - - axpy, scal, copy = (fallback_axpy, fallback_scal, fallback_copy) - x1_arr = x1.data - x2_arr = x2.data - out_arr = out.data - - else: - # Need flat data for BLAS, otherwise in-place does not work. - # Raveling must happen in fixed order for non-contiguous out, - # otherwise 'A' is applied to arrays, which makes the outcome - # dependent on their respective contiguousness. - if out.data.flags.f_contiguous: - ravel_order = 'F' - else: - ravel_order = 'C' - - x1_arr = x1.data.ravel(order=ravel_order) - x2_arr = x2.data.ravel(order=ravel_order) - out_arr = out.data.ravel(order=ravel_order) - axpy, scal, copy = scipy.linalg.blas.get_blas_funcs( - ['axpy', 'scal', 'copy'], arrays=(x1_arr, x2_arr, out_arr)) - - if x1 is x2 and b != 0: - # x1 is aligned with x2 -> out = (a+b)*x1 - _lincomb_impl(a + b, x1, 0, x1, out) - elif out is x1 and out is x2: - # All the vectors are aligned -> out = (a+b)*out - if (a + b) != 0: - scal(a + b, out_arr, size) - else: - out_arr[:] = 0 - elif out is x1: - # out is aligned with x1 -> out = a*out + b*x2 - if a != 1: - scal(a, out_arr, size) - if b != 0: - axpy(x2_arr, out_arr, size, b) - elif out is x2: - # out is aligned with x2 -> out = a*x1 + b*out - if b != 1: - scal(b, out_arr, size) - if a != 0: - axpy(x1_arr, out_arr, size, a) - else: - # We have exhausted all alignment options, so x1 is not x2 is not out - # We now optimize for various values of a and b - if b == 0: - if a == 0: # Zero assignment -> out = 0 - out_arr[:] = 0 - else: # Scaled copy -> out = a*x1 - copy(x1_arr, out_arr, size) - if a != 1: - scal(a, out_arr, size) - - else: # b != 0 - if a == 0: # Scaled copy -> out = b*x2 - copy(x2_arr, out_arr, size) - if b != 1: - scal(b, out_arr, size) - - elif a == 1: # No scaling in x1 -> out = x1 + b*x2 - copy(x1_arr, out_arr, size) - axpy(x2_arr, out_arr, size, b) - else: # Generic case -> out = a*x1 + b*x2 - copy(x2_arr, out_arr, size) - if b != 1: - scal(b, out_arr, size) - axpy(x1_arr, out_arr, size, a) def _weighting(weights, exponent): """Return a weighting whose type is inferred from the arguments.""" - if np.isscalar(weights): - weighting = NumpyTensorSpaceConstWeighting(weights, exponent) + if np.isscalar(weights) or weights.shape == (): + weighting = PytorchTensorSpaceConstWeighting(weights, exponent) elif weights is None: - weighting = NumpyTensorSpaceConstWeighting(1.0, exponent) + weighting = PytorchTensorSpaceConstWeighting(1.0, exponent) else: # last possibility: make an array - arr = np.asarray(weights) - weighting = NumpyTensorSpaceArrayWeighting(arr, exponent) + arr = torch.tensor(weights) + weighting = PytorchTensorSpaceArrayWeighting(arr, exponent) return weighting -def npy_weighted_inner(weights): +def pytorch_weighted_inner(weights): """Weighted inner product on `TensorSpace`'s as free function. Parameters @@ -1917,13 +1412,13 @@ def npy_weighted_inner(weights): See Also -------- - NumpyTensorSpaceConstWeighting - NumpyTensorSpaceArrayWeighting + PytorchTensorSpaceConstWeighting + PytorchTensorSpaceArrayWeighting """ return _weighting(weights, exponent=2.0).inner -def npy_weighted_norm(weights, exponent=2.0): +def pytorch_weighted_norm(weights, exponent=2.0): """Weighted norm on `TensorSpace`'s as free function. Parameters @@ -1943,13 +1438,13 @@ def npy_weighted_norm(weights, exponent=2.0): See Also -------- - NumpyTensorSpaceConstWeighting - NumpyTensorSpaceArrayWeighting + PytorchTensorSpaceConstWeighting + PytorchTensorSpaceArrayWeighting """ return _weighting(weights, exponent=exponent).norm -def npy_weighted_dist(weights, exponent=2.0): +def pytorch_weighted_dist(weights, exponent=2.0): """Weighted distance on `TensorSpace`'s as free function. Parameters @@ -1969,73 +1464,49 @@ def npy_weighted_dist(weights, exponent=2.0): See Also -------- - NumpyTensorSpaceConstWeighting - NumpyTensorSpaceArrayWeighting + PytorchTensorSpaceConstWeighting + PytorchTensorSpaceArrayWeighting """ return _weighting(weights, exponent=exponent).dist def _norm_default(x): """Default Euclidean norm implementation.""" - # Lazy import to improve `import odl` time - import scipy.linalg - if _blas_is_applicable(x.data): - nrm2 = scipy.linalg.blas.get_blas_funcs('nrm2', dtype=x.dtype) - norm = partial(nrm2, n=native(x.size)) - else: - norm = np.linalg.norm - return norm(x.data.ravel()) + return x.data.norm(p=2) def _pnorm_default(x, p): """Default p-norm implementation.""" - return np.linalg.norm(x.data.ravel(), ord=p) + return x.data.norm(p=p) def _pnorm_diagweight(x, p, w): """Diagonally weighted p-norm implementation.""" - # Ravel both in the same order (w is a numpy array) - order = 'F' if all(a.flags.f_contiguous for a in (x.data, w)) else 'C' - - # This is faster than first applying the weights and then summing with - # BLAS dot or nrm2 - xp = np.abs(x.data.ravel(order)) + xp = torch.abs(x.data) if p == float('inf'): - xp *= w.ravel(order) - return np.max(xp) + xp *= w + return torch.max(xp) else: - xp = np.power(xp, p, out=xp) - xp *= w.ravel(order) - return np.sum(xp) ** (1 / p) + torch.pow(xp, p, out=xp) + xp *= w + return torch.sum(xp) ** (1 / p) def _inner_default(x1, x2): """Default Euclidean inner product implementation.""" - # Ravel both in the same order - order = 'F' if all(a.data.flags.f_contiguous for a in (x1, x2)) else 'C' if is_real_dtype(x1.dtype): - if x1.size > THRESHOLD_MEDIUM: - # This is as fast as BLAS dotc - return np.tensordot(x1, x2, [range(x1.ndim)] * 2) - else: - # Several times faster for small arrays - return np.dot(x1.data.ravel(order), - x2.data.ravel(order)) + return torch.dot(x1.data, x2.data) else: # x2 as first argument because we want linearity in x1 - return np.vdot(x2.data.ravel(order), - x1.data.ravel(order)) - + return torch.vdot(x2.data, x1.data) -# TODO: implement intermediate weighting schemes with arrays that are -# broadcast, i.e. between scalar and full-blown in dimensionality? -class NumpyTensorSpaceArrayWeighting(ArrayWeighting): +class PytorchTensorSpaceArrayWeighting(ArrayWeighting): - """Weighting of a `NumpyTensorSpace` by an array. + """Weighting of a `PytorchTensorSpace` by an array. This class defines a weighting by an array that has the same shape as the tensor space. Since the space is not known to this class, @@ -2100,23 +1571,23 @@ def __init__(self, array, exponent=2.0): it does not define an inner product or norm, respectively. This is not checked during initialization. """ - if isinstance(array, NumpyTensor): + if isinstance(array, PytorchTensor): array = array.data - elif not isinstance(array, np.ndarray): - array = np.asarray(array) - super(NumpyTensorSpaceArrayWeighting, self).__init__( - array, impl='numpy', exponent=exponent) + elif not isinstance(array, torch.Tensor): + array = torch.tensor(array) + super(PytorchTensorSpaceArrayWeighting, self).__init__( + array, impl='pytorch', exponent=exponent) def __hash__(self): """Return ``hash(self)``.""" - return hash((type(self), self.array.tobytes(), self.exponent)) + return hash((type(self), hash(self.array), self.exponent)) def inner(self, x1, x2): """Return the weighted inner product of ``x1`` and ``x2``. Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Tensors whose inner product is calculated. Returns @@ -2140,7 +1611,7 @@ def norm(self, x): Parameters ---------- - x : `NumpyTensor` + x : `PytorchTensor` Tensor whose norm is calculated. Returns @@ -2157,9 +1628,9 @@ def norm(self, x): return float(_pnorm_diagweight(x, self.exponent, self.array)) -class NumpyTensorSpaceConstWeighting(ConstWeighting): +class PytorchTensorSpaceConstWeighting(ConstWeighting): - """Weighting of a `NumpyTensorSpace` by a constant. + """Weighting of a `PytorchTensorSpace` by a constant. See ``Notes`` for mathematical details. """ @@ -2214,7 +1685,7 @@ def __init__(self, const, exponent=2.0): - The constant must be positive, otherwise it does not define an inner product or norm, respectively. """ - super(NumpyTensorSpaceConstWeighting, self).__init__( + super(PytorchTensorSpaceConstWeighting, self).__init__( const, impl='numpy', exponent=exponent) def inner(self, x1, x2): @@ -2222,7 +1693,7 @@ def inner(self, x1, x2): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Tensors whose inner product is calculated. Returns @@ -2246,7 +1717,7 @@ def norm(self, x): Parameters ---------- - x1 : `NumpyTensor` + x1 : `PytorchTensor` Tensor whose norm is calculated. Returns @@ -2267,7 +1738,7 @@ def dist(self, x1, x2): Parameters ---------- - x1, x2 : `NumpyTensor` + x1, x2 : `PytorchTensor` Tensors whose mutual distance is calculated. Returns @@ -2284,7 +1755,7 @@ def dist(self, x1, x2): _pnorm_default(x1 - x2, self.exponent))) -class NumpyTensorSpaceCustomInner(CustomInner): +class PytorchTensorSpaceCustomInner(CustomInner): """Class for handling a user-specified inner product.""" @@ -2303,10 +1774,10 @@ def __init__(self, inner): - `` = s * + `` - `` = 0`` if and only if ``x = 0`` """ - super(NumpyTensorSpaceCustomInner, self).__init__(inner, impl='numpy') + super(PytorchTensorSpaceCustomInner, self).__init__(inner, impl='numpy') -class NumpyTensorSpaceCustomNorm(CustomNorm): +class PytorchTensorSpaceCustomNorm(CustomNorm): """Class for handling a user-specified norm. @@ -2329,10 +1800,10 @@ def __init__(self, norm): - ``||s * x|| = |s| * ||x||`` - ``||x + y|| <= ||x|| + ||y||`` """ - super(NumpyTensorSpaceCustomNorm, self).__init__(norm, impl='numpy') + super(PytorchTensorSpaceCustomNorm, self).__init__(norm, impl='numpy') -class NumpyTensorSpaceCustomDist(CustomDist): +class PytorchTensorSpaceCustomDist(CustomDist): """Class for handling a user-specified distance in `TensorSpace`. @@ -2355,7 +1826,7 @@ def __init__(self, dist): - ``dist(x, y) = dist(y, x)`` - ``dist(x, y) <= dist(x, z) + dist(z, y)`` """ - super(NumpyTensorSpaceCustomDist, self).__init__(dist, impl='numpy') + super(PytorchTensorSpaceCustomDist, self).__init__(dist, impl='numpy') if __name__ == '__main__': From 3ff5a9750507f84381a868d07c8225fa54fa25b4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Thu, 21 Mar 2024 17:17:32 +0100 Subject: [PATCH 03/52] Some numpy->pytorch changes that slipped past the previous commit. --- odl/space/pytorch_tensors.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index ca843d7b8ad..7bb3e1b5efc 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -287,8 +287,8 @@ def __init__(self, shape, dtype=None, **kwargs): @property def impl(self): - """Name of the implementation back-end: ``'numpy'``.""" - return 'numpy' + """Name of the implementation back-end: ``'pytorch'``.""" + return 'pytorch' @property def default_order(self): @@ -1686,7 +1686,7 @@ def __init__(self, const, exponent=2.0): inner product or norm, respectively. """ super(PytorchTensorSpaceConstWeighting, self).__init__( - const, impl='numpy', exponent=exponent) + const, impl='pytorch', exponent=exponent) def inner(self, x1, x2): """Return the weighted inner product of ``x1`` and ``x2``. @@ -1774,7 +1774,7 @@ def __init__(self, inner): - `` = s * + `` - `` = 0`` if and only if ``x = 0`` """ - super(PytorchTensorSpaceCustomInner, self).__init__(inner, impl='numpy') + super(PytorchTensorSpaceCustomInner, self).__init__(inner, impl='pytorch') class PytorchTensorSpaceCustomNorm(CustomNorm): @@ -1800,7 +1800,7 @@ def __init__(self, norm): - ``||s * x|| = |s| * ||x||`` - ``||x + y|| <= ||x|| + ||y||`` """ - super(PytorchTensorSpaceCustomNorm, self).__init__(norm, impl='numpy') + super(PytorchTensorSpaceCustomNorm, self).__init__(norm, impl='pytorch') class PytorchTensorSpaceCustomDist(CustomDist): @@ -1826,7 +1826,7 @@ def __init__(self, dist): - ``dist(x, y) = dist(y, x)`` - ``dist(x, y) <= dist(x, z) + dist(z, y)`` """ - super(PytorchTensorSpaceCustomDist, self).__init__(dist, impl='numpy') + super(PytorchTensorSpaceCustomDist, self).__init__(dist, impl='pytorch') if __name__ == '__main__': From 0e0e9b2564a5d20c6df49d91b904eafa628c5c68 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 2 Apr 2024 17:31:53 +0200 Subject: [PATCH 04/52] An example for an operator with Torch implementation. --- examples/operator/torch_convo_operator.py | 65 +++++++++++++++++++++++ 1 file changed, 65 insertions(+) create mode 100644 examples/operator/torch_convo_operator.py diff --git a/examples/operator/torch_convo_operator.py b/examples/operator/torch_convo_operator.py new file mode 100644 index 00000000000..7c7ce21d54a --- /dev/null +++ b/examples/operator/torch_convo_operator.py @@ -0,0 +1,65 @@ +"""Create a convolution operator by wrapping a library.""" + +import odl +import numpy as np +import torch + + +class Convolution(odl.Operator): + """Operator calculating the convolution of a kernel with a function. + + The operator inherits from ``odl.Operator`` to be able to be used with ODL. + """ + + def __init__(self, kernel, domain, range): + """Initialize a convolution operator with a known kernel.""" + + # Store the kernel + self.kernel = kernel + + # Initialize the Operator class by calling its __init__ method. + # This sets properties such as domain and range and allows the other + # operator convenience functions to work. + super(Convolution, self).__init__( + domain=domain, range=range, linear=True) + + def _call(self, x): + """Implement calling the operator by calling PyTorch.""" + return self.range.element(torch.conv2d( input=x.data.unsqueeze(0) + , weight=self.kernel.unsqueeze(0).unsqueeze(0) + , stride=(1,1) + , padding="same" + ).squeeze(0) + ) + + @property + def adjoint(self): + """Implement ``self.adjoint``. + + For a convolution operator, the adjoint is given by the convolution + with a kernel with flipped axes. In particular, if the kernel is + symmetric the operator is self-adjoint. + """ + return Convolution( torch.flip(self.kernel, dims=(0,1)) + , domain=self.range, range=self.domain ) + + +# Define the space on which the problem should be solved +# Here the square [-1, 1] x [-1, 1] discretized on a 100x100 grid +space = odl.uniform_discr([-1, -1], [1, 1], [100, 100], impl='pytorch', dtype=np.float32) + +# Convolution kernel, a small centered rectangle +kernel = torch.ones((5,5)) + +# Create convolution operator +A = Convolution(kernel, domain=space, range=space) + +# Create phantom (the "unknown" solution) +phantom = odl.phantom.shepp_logan(space, modified=True) + +# Apply convolution to phantom to create data +g = A.adjoint(phantom) + +# Display the results using the show method +phantom.show('phantom') +g.show('convolved phantom', force_show=True) From b74e8e90a34f832014b71d46baa523ae1e39c81b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 3 Apr 2024 18:13:11 +0200 Subject: [PATCH 05/52] Fix inconsistency about backend/impl naming. --- odl/space/pytorch_tensors.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 7bb3e1b5efc..6b60eb6343a 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -243,8 +243,8 @@ def __init__(self, shape, dtype=None, **kwargs): # Set the weighting if weighting is not None: if isinstance(weighting, Weighting): - if weighting.impl != 'torch': - raise ValueError("`weighting.impl` must be 'torch', " + if weighting.impl != 'pytorch': + raise ValueError("`weighting.impl` must be 'pytorch', " '`got {!r}'.format(weighting.impl)) if weighting.exponent != exponent: raise ValueError('`weighting.exponent` conflicts with ' From 2b92872823a2919a6305e72a5ffcbf7571e4ca50 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 8 Apr 2024 14:06:14 +0200 Subject: [PATCH 06/52] Adapt the deconvolution example to something more comparable with compact kernels. --- examples/solvers/deconvolution_1d.py | 41 +++++++++++++++++----------- 1 file changed, 25 insertions(+), 16 deletions(-) diff --git a/examples/solvers/deconvolution_1d.py b/examples/solvers/deconvolution_1d.py index b60d6d23d55..d6ee6bb53b4 100644 --- a/examples/solvers/deconvolution_1d.py +++ b/examples/solvers/deconvolution_1d.py @@ -27,11 +27,13 @@ def opnorm(self): # Discretization -discr_space = odl.uniform_discr(0, 10, 500, impl='numpy') +discr_space = odl.uniform_discr(-5, 5, 500, impl='numpy') # Complicated functions to check performance -kernel = discr_space.element(lambda x: np.exp(x / 2) * np.cos(x * 1.172)) -phantom = discr_space.element(lambda x: x ** 2 * np.sin(x) ** 2 * (x > 5)) +kernel = discr_space.element(lambda x: np.exp(-x**2 * 2) * np.cos(x * 1.172)) + +# phantom = discr_space.element(lambda x: (x+5) ** 2 * np.sin(x+5) ** 2 * (x > 0)) +phantom = discr_space.element(lambda x: np.cos(0*x) * (x > -1) * (x < 1)) # Create operator conv = Convolution(kernel) @@ -41,21 +43,28 @@ def opnorm(self): omega = 1 / conv.opnorm() ** 2 -# Display callback -def callback(x): - plt.plot(conv(x)) - +def test_with_plot(conv, phantom, solver, **extra_args): + fig, axs = plt.subplots(2) + fig.suptitle("CGN") + axs[0].set_title("x") + axs[1].set_title("k*x") + axs[0].plot(phantom) + axs[1].plot(conv(phantom)) + def plot_callback(x): + axs[0].plot(conv(x), '--') + axs[1].plot(conv(x), '--') + solver(conv, discr_space.zero(), phantom, iterations, callback=plot_callback, **extra_args) # Test CGN -plt.figure() -plt.plot(phantom) -odl.solvers.conjugate_gradient_normal(conv, discr_space.zero(), phantom, - iterations, callback) +test_with_plot(conv, phantom, odl.solvers.conjugate_gradient_normal) -# Landweber -plt.figure() -plt.plot(phantom) -odl.solvers.landweber(conv, discr_space.zero(), phantom, - iterations, omega, callback) +# test_with_plot(conv, phantom, odl.solvers.landweber, omega=omega) +# # Landweber +# lw_fig, lw_axs = plt.subplots(1) +# lw_fig.suptitle("Landweber") +# lw_axs.plot(phantom) +# odl.solvers.landweber(conv, discr_space.zero(), phantom, +# iterations, omega, lambda x: lw_axs.plot(conv(x))) +# plt.show() From ca336cc119f1df094c12f2bb8e36501287584954 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 9 Apr 2024 10:38:34 +0200 Subject: [PATCH 07/52] More proper result-showing in examples for solvers. --- examples/operator/convolution_operator.py | 2 +- examples/solvers/pdhg_denoising.py | 2 +- 2 files changed, 2 insertions(+), 2 deletions(-) diff --git a/examples/operator/convolution_operator.py b/examples/operator/convolution_operator.py index bf344a183e4..319c357dcc5 100644 --- a/examples/operator/convolution_operator.py +++ b/examples/operator/convolution_operator.py @@ -56,4 +56,4 @@ def adjoint(self): # Display the results using the show method kernel.show('kernel') phantom.show('phantom') -g.show('convolved phantom') +g.show('convolved phantom', force_show=True) diff --git a/examples/solvers/pdhg_denoising.py b/examples/solvers/pdhg_denoising.py index ed2662d3cf9..1bb6cffffeb 100644 --- a/examples/solvers/pdhg_denoising.py +++ b/examples/solvers/pdhg_denoising.py @@ -26,7 +26,7 @@ space = odl.uniform_discr([0, 0], shape, shape) # Original image -orig = space.element(image) +orig = space.element(image.copy()) # Add noise image += 0.1 * odl.phantom.white_noise(orig.space) From ee003d3370bf77c61bca63c311ca7f111e91ca72 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 9 Apr 2024 11:05:16 +0200 Subject: [PATCH 08/52] PyTorch version of the deconvolution example. Results match the NumPy version, though not exactly because a compact kernel is used with direct convolution whereas the NumPy version uses FFT convolution. --- examples/solvers/deconvolution_1d_pytorch.py | 89 ++++++++++++++++++++ 1 file changed, 89 insertions(+) create mode 100644 examples/solvers/deconvolution_1d_pytorch.py diff --git a/examples/solvers/deconvolution_1d_pytorch.py b/examples/solvers/deconvolution_1d_pytorch.py new file mode 100644 index 00000000000..7eff71e4614 --- /dev/null +++ b/examples/solvers/deconvolution_1d_pytorch.py @@ -0,0 +1,89 @@ +"""Example of a deconvolution problem with different solvers (CPU).""" + +import numpy as np +import torch +import matplotlib.pyplot as plt +import scipy.signal +import odl + + +class Convolution(odl.Operator): + def __init__(self, kernel, domain, range, adjkernel=None): + self.kernel = kernel + self.adjkernel = torch.flip(kernel, dims=(0,)) if adjkernel is None else adjkernel + self.norm = float(torch.sum(torch.abs(self.kernel))) + super(Convolution, self).__init__( + domain=domain, range=range, linear=True) + + def _call(self, x): + return self.range.element( + torch.conv1d( input=x.data.unsqueeze(0) + , weight=self.kernel.unsqueeze(0).unsqueeze(0) + , stride=1 + , padding="same" + ).squeeze(0) + ) + + @property + def adjoint(self): + return Convolution( self.adjkernel + , domain=self.range, range=self.domain + , adjkernel = self.kernel + ) + + def opnorm(self): + return self.norm + + +resolution = 50 + +# Discretization +discr_space = odl.uniform_discr(-5, 5, resolution*10, impl='pytorch', dtype=np.float32) + +# Complicated functions to check performance +def mk_kernel(): + q = 1.172 + # Select main lobe and one side lobe on each side + r = np.ceil(3*np.pi/(2*q)) + # Quantised to resolution + nr = int(np.ceil(r*resolution)) + r = nr / resolution + x = torch.linspace(-r, r, nr*2 + 1) + return torch.exp(-x**2 * 2) * np.cos(x * q) +kernel = mk_kernel() + +phantom = discr_space.element(lambda x: np.ones_like(x) ** 2 * (x > -1) * (x < 1)) +# phantom = discr_space.element(lambda x: x ** 2 * np.sin(x) ** 2 * (x > 5)) + +# Create operator +conv = Convolution(kernel, domain=discr_space, range=discr_space) + +# Dampening parameter for landweber +iterations = 100 +omega = 1 / conv.opnorm() ** 2 + + + +def test_with_plot(conv, phantom, solver, **extra_args): + fig, axs = plt.subplots(2) + fig.suptitle("CGN") + axs[0].set_title("x") + axs[1].set_title("k*x") + axs[0].plot(phantom) + axs[1].plot(conv(phantom)) + def plot_callback(x): + axs[0].plot(conv(x), '--') + axs[1].plot(conv(x), '--') + solver(conv, discr_space.zero(), phantom, iterations, callback=plot_callback, **extra_args) + +# Test CGN +test_with_plot(conv, phantom, odl.solvers.conjugate_gradient_normal) + +# # Landweber +# lw_fig, lw_axs = plt.subplots(1) +# lw_fig.suptitle("Landweber") +# lw_axs.plot(phantom) +# odl.solvers.landweber(conv, discr_space.zero(), phantom, +# iterations, omega, lambda x: lw_axs.plot(conv(x))) + +plt.show() From 6f6df0b60b4212a5dbe35e3761949ed3337fa864 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 9 Apr 2024 13:45:39 +0200 Subject: [PATCH 09/52] More consistent naming for the PyTorch versions of examples. --- .../{torch_convo_operator.py => convolution_operator_pytorch.py} | 0 1 file changed, 0 insertions(+), 0 deletions(-) rename examples/operator/{torch_convo_operator.py => convolution_operator_pytorch.py} (100%) diff --git a/examples/operator/torch_convo_operator.py b/examples/operator/convolution_operator_pytorch.py similarity index 100% rename from examples/operator/torch_convo_operator.py rename to examples/operator/convolution_operator_pytorch.py From 595e1665789ce34aa377a464c22ebd165fd55f97 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 9 Apr 2024 22:26:18 +0200 Subject: [PATCH 10/52] Generalise some utility to support PyTorch in addition to NumPy. --- odl/util/utility.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/odl/util/utility.py b/odl/util/utility.py index 1df8f24896d..651c68bf45b 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -17,6 +17,7 @@ from itertools import product import numpy as np +import torch __all__ = ( 'REPR_PRECISION', @@ -33,6 +34,7 @@ 'is_real_dtype', 'is_real_floating_dtype', 'is_complex_floating_dtype', + 'uses_pytorch', 'real_dtype', 'complex_dtype', 'is_string', @@ -497,6 +499,13 @@ def complex_dtype(dtype, default=None): else: return np.dtype((complex_base_dtype, dtype.shape)) +def uses_pytorch(obj): + if isinstance(obj, torch.Tensor): + return True + elif getattr(obj, "impl", None)=="pytorch": + return True + else: + return False def is_string(obj): """Return ``True`` if ``obj`` behaves like a string, ``False`` else.""" @@ -628,8 +637,12 @@ def writable_array(obj, **kwargs): [2, 4, 6] """ arr = None + torch_impl = uses_pytorch(obj) try: - arr = np.asarray(obj, **kwargs) + if torch_impl: + arr = torch.tensor(obj, **kwargs) + else: + arr = np.asarray(obj, **kwargs) yield arr finally: if arr is not None: From f650814a49c3ae629c6bb9acfe89c51e62861543 Mon Sep 17 00:00:00 2001 From: Justus Sagemuller Date: Fri, 12 Apr 2024 17:26:22 +0200 Subject: [PATCH 11/52] Implementation-consistent types for the `asarray` method. It is debateable whether this is what `asarray` is there for. Perhaps it would be better to simply expose `.data` for this purpose and keep `.asarray` NumPy- specific. However, since `__array__` is already there, it seems to fulfill that purpose as well. --- odl/discr/discr_space.py | 2 +- odl/space/base_tensors.py | 15 +++++++++------ odl/space/pytorch_tensors.py | 4 ++-- 3 files changed, 12 insertions(+), 9 deletions(-) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index 033fb5e0c95..6fab8f61c6b 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -1528,7 +1528,7 @@ def show(self, title=None, method='', coords=None, indices=None, # Squeeze grid and values according to the index expression part = self.space.partition[indices].squeeze() - values = self.asarray()[indices].squeeze() + values = np.array(self)[indices].squeeze() return show_discrete_data(values, part, title=title, method=method, force_show=force_show, fig=fig, diff --git a/odl/space/base_tensors.py b/odl/space/base_tensors.py index 3396b6d1142..c43afc3c86d 100644 --- a/odl/space/base_tensors.py +++ b/odl/space/base_tensors.py @@ -510,19 +510,20 @@ class Tensor(LinearSpaceElement): """Abstract class for representation of `TensorSpace` elements.""" def asarray(self, out=None): - """Extract the data of this tensor as a Numpy array. + """Extract the data of this tensor as an array. This could be a NumPy array + or a PyTorch tensor, depending on what implementation backend is used. This method should be overridden by subclasses. Parameters ---------- - out : `numpy.ndarray`, optional + out : `array_like`, optional Array to write the result to. Returns ------- - asarray : `numpy.ndarray` - Numpy array of the same data type and shape as the space. + asarray : `array_like` + Array of the same type, data type and shape as the space. If ``out`` was given, the returned object is a reference to it. """ @@ -652,6 +653,8 @@ def __bool__(self): def __array__(self, dtype=None): """Return a Numpy array from this tensor. + (Contrast with the `asarray` method, which may give other types of array, + not just NumPy.) Parameters ---------- @@ -663,9 +666,9 @@ def __array__(self, dtype=None): array : `numpy.ndarray` """ if dtype is None: - return self.asarray() + return np.array(self.asarray()) else: - return self.asarray().astype(dtype, copy=AVOID_UNNECESSARY_COPY) + return np.array(self.asarray()).astype(dtype, copy=AVOID_UNNECESSARY_COPY) def __array_wrap__(self, array): """Return a new tensor wrapping the ``array``. diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 6b60eb6343a..985ef26830c 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -876,9 +876,9 @@ def asarray(self, out=None): [ 1., 1., 1.]]) """ if out is None: - return self.data.cpu().numpy() + return self.data else: - out[:] = self.data.cpu().numpy() + out[:] = self.data return out def astype(self, dtype): From 45386abc474bc693788ed635a1cfe50fc41fcdfc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 12 Apr 2024 18:52:07 +0200 Subject: [PATCH 12/52] A generic way of obtaining a compatible scalar dtype for various things. --- odl/util/utility.py | 22 ++++++++++++++++++++++ 1 file changed, 22 insertions(+) diff --git a/odl/util/utility.py b/odl/util/utility.py index 651c68bf45b..0dfd97cf9ec 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -27,6 +27,7 @@ 'array_str', 'dtype_repr', 'dtype_str', + 'dtype_type', 'cache_arguments', 'is_numeric_dtype', 'is_int_dtype', @@ -321,6 +322,27 @@ def dtype_str(dtype): else: return '{}'.format(dtype) +def dtype_type(dtype): + """Obtain a Python type corresponding to the given NumPy or PyTorch + dtype. This can be used for constructing values of a suitable type + for storing in an array of either backend.""" + if isinstance(dtype, str) or isinstance(dtype, type): + dtype = np.dtype(dtype) + + if hasattr(dtype, 'dtype'): + return dtype_type(dtype.dtype) + elif dtype == np.dtype(int): + return int + elif dtype == np.dtype(float): + return float + elif dtype == np.dtype(complex): + return complex + elif dtype == torch.float64: + return float + else: + raise ValueError(f"No suitable Python type available for {dtype}.") + + def cache_arguments(function): """Decorate function to cache the result with given arguments. From c3fe2c50a8a5409df77e36aebaf6ceb572da15bc Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 12 Apr 2024 18:53:43 +0200 Subject: [PATCH 13/52] Attempt at making gradient operators compatible with PyTorch. Does not seem to work yet. --- odl/discr/diff_ops.py | 31 ++++++++++++++++++++++++------- 1 file changed, 24 insertions(+), 7 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index e7ba9d7f168..c6e715d548b 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -11,11 +11,12 @@ from __future__ import absolute_import, division, print_function import numpy as np +import torch from odl.discr.discr_space import DiscretizedSpace from odl.operator.tensor_ops import PointwiseTensorFieldOperator from odl.space import ProductSpace -from odl.util import indent, signature_string, writable_array +from odl.util import indent, signature_string, writable_array, uses_pytorch, dtype_type __all__ = ('PartialDerivative', 'Gradient', 'Divergence', 'Laplacian') @@ -556,11 +557,14 @@ def _call(self, x, out=None): """Calculate the divergence of ``x``.""" if out is None: out = self.range.element() + # print(f"{type(out.data)=}") ndim = self.range.ndim dx = self.range.cell_sides + torch_impl = uses_pytorch(x[0]) - tmp = np.empty(out.shape, out.dtype, order=out.space.default_order) + tmp = self.range.element().asarray() + # print(f"{type(tmp)=}") with writable_array(out) as out_arr: for axis in range(ndim): finite_diff(x[axis], axis=axis, dx=dx[axis], @@ -884,7 +888,14 @@ def finite_diff(f, axis, dx=1.0, method='forward', out=None, >>> out is finite_diff(f, axis=0, out=out) True """ - f_arr = np.asarray(f) + torch_impl = uses_pytorch(f) + if torch_impl and out is not None: + assert(isinstance(out, torch.Tensor)), f"{type(out)=}" + + if torch_impl: + f_arr = torch.tensor(f) + else: + f_arr = np.asarray(f) ndim = f_arr.ndim if f_arr.shape[axis] < 2: @@ -909,7 +920,7 @@ def finite_diff(f, axis, dx=1.0, method='forward', out=None, raise ValueError('`pad_mode` {} not understood' ''.format(pad_mode)) - pad_const = f.dtype.type(pad_const) + pad_const = dtype_type(f)(pad_const) if out is None: out = np.empty_like(f_arr) @@ -932,19 +943,25 @@ def finite_diff(f, axis, dx=1.0, method='forward', out=None, out, out_in = np.swapaxes(out, 0, axis), out f_arr = np.swapaxes(f_arr, 0, axis) + def fd_subtraction(a, b): + if torch_impl: + out[1:-1] = a - b + else: + np.subtract(a, b, out=out[1:-1]) + # Interior of the domain of f if method == 'central': # 1D equivalent: out[1:-1] = (f[2:] - f[:-2])/2.0 - np.subtract(f_arr[2:], f_arr[:-2], out=out[1:-1]) + fd_subtraction(f_arr[2:], f_arr[:-2]) out[1:-1] /= 2.0 elif method == 'forward': # 1D equivalent: out[1:-1] = (f[2:] - f[1:-1]) - np.subtract(f_arr[2:], f_arr[1:-1], out=out[1:-1]) + fd_subtraction(f_arr[2:], f_arr[1:-1]) elif method == 'backward': # 1D equivalent: out[1:-1] = (f[1:-1] - f[:-2]) - np.subtract(f_arr[1:-1], f_arr[:-2], out=out[1:-1]) + fd_subtraction(f_arr[1:-1], f_arr[:-2]) # Boundaries if pad_mode == 'constant': From e244f0e208af6aa372a0931ca042095695ca5076 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Thu, 18 Apr 2024 00:32:20 +0200 Subject: [PATCH 14/52] More flexible plotting in 1D example. --- examples/solvers/deconvolution_1d_pytorch.py | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/examples/solvers/deconvolution_1d_pytorch.py b/examples/solvers/deconvolution_1d_pytorch.py index 7eff71e4614..c517baa2eac 100644 --- a/examples/solvers/deconvolution_1d_pytorch.py +++ b/examples/solvers/deconvolution_1d_pytorch.py @@ -67,13 +67,15 @@ def mk_kernel(): def test_with_plot(conv, phantom, solver, **extra_args): fig, axs = plt.subplots(2) fig.suptitle("CGN") + def plot_fn(ax_id, fn, *plot_args, **plot_kwargs): + axs[ax_id].plot(fn, *plot_args, **plot_kwargs) axs[0].set_title("x") axs[1].set_title("k*x") - axs[0].plot(phantom) - axs[1].plot(conv(phantom)) + plot_fn(0, phantom) + plot_fn(1, conv(phantom)) def plot_callback(x): - axs[0].plot(conv(x), '--') - axs[1].plot(conv(x), '--') + plot_fn(0, conv(x), '--') + plot_fn(1, conv(x), '--') solver(conv, discr_space.zero(), phantom, iterations, callback=plot_callback, **extra_args) # Test CGN From 5227dac5c2a50b214e3517133eb9f9aa18474153 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 19 Apr 2024 18:21:41 +0200 Subject: [PATCH 15/52] Methods for converting numbers to scalars. That is, values that are compatible for being multiplied with element-vectors of the space in question. --- odl/discr/discr_space.py | 12 ++++++++++++ odl/set/space.py | 12 ++++++++++++ odl/space/npy_tensors.py | 9 +++++++++ odl/space/pspace.py | 11 +++++++++++ odl/space/pytorch_tensors.py | 29 +++++++++++++++++++++++++++++ 5 files changed, 73 insertions(+) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index 6fab8f61c6b..741c2802440 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -251,6 +251,18 @@ def available_dtypes(self): """ return self.tspace.available_dtypes() + def is_suitable_scalar(self, s): + """Determine whether `s` has a type that can be scalar-multiplied with + elements of this space. + """ + return self.tspace.is_suitable_scalar(s) + + def as_suitable_scalar(self, s): + """Try to convert `s` to a type that can be scalar-multiplied with + elements of this space. + """ + return self.tspace.as_suitable_scalar(s) + # --- Derived properties @property diff --git a/odl/set/space.py b/odl/set/space.py index b1b4b380b8c..2d10651ab40 100644 --- a/odl/set/space.py +++ b/odl/set/space.py @@ -506,6 +506,18 @@ def __mul__(self, other): return ProductSpace(self, other) + def is_suitable_scalar(self, s): + """Determine whether `s` has a type that can be scalar-multiplied with + elements of this space. + """ + raise NotImplementedError(f'Abstract method not implemented for {type(self)}') + + def as_suitable_scalar(self, s): + """Try to convert `s` to a type that can be scalar-multiplied with + elements of this space. + """ + raise NotImplementedError(f'Abstract method not implemented for {type(self)}') + def __str__(self): """Return ``str(self)``.""" return repr(self) diff --git a/odl/space/npy_tensors.py b/odl/space/npy_tensors.py index c5a497ab231..f456f890cdc 100644 --- a/odl/space/npy_tensors.py +++ b/odl/space/npy_tensors.py @@ -879,6 +879,15 @@ def element_type(self): """Type of elements in this space: `NumpyTensor`.""" return NumpyTensor + def is_suitable_scalar(self, s): + return type(s) is self.dtype.type + + def as_suitable_scalar(self, s): + """Try to convert `s` to a type that can be scalar-multiplied with + numpy arrays. + """ + return self.dtype.type(s) + class NumpyTensor(Tensor): diff --git a/odl/space/pspace.py b/odl/space/pspace.py index 6273e19532a..c4628f72983 100644 --- a/odl/space/pspace.py +++ b/odl/space/pspace.py @@ -389,6 +389,17 @@ def dtype(self): else: raise AttributeError("`dtype`'s of subspaces not equal") + def is_suitable_scalar(self, s): + return all(space.is_suitable_scalar(s) for space in self.spaces) + + def as_suitable_scalar(self, s): + """Try to convert `s` to a type that can be scalar-multiplied with + elements of this space. + """ + s_sui = self.spaces[0].as_suitable_scalar(s) + assert(self.is_suitable_scalar(s_sui)) + return s_sui + @property def supported_num_operation_paradigms(self) -> NumOperationParadigmSupport: """Whether in-place operations an out-of-place operations are supported diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 985ef26830c..8ed49f408fe 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -295,6 +295,35 @@ def default_order(self): """Default (and only) storage order for new elements in this space: ``'C'``.""" return 'C' + def is_suitable_scalar(self, s): + if self._torch_dtype in [torch.complex64, torch.complex128]: + return type(s) is complex + else: + return type(s) is float + # Singleton-tensor version: + # if not isinstance(s, torch.Tensor): + # return False + # elif s.dtype != self._torch_dtype: + # return False + # elif s.shape != (): + # return False + # else: + # return True + + def as_suitable_scalar(self, s): + """Try to convert `s` to a type that can be scalar-multiplied with + torch tensors. + """ + if self._torch_dtype in [torch.complex64, torch.complex128]: + return complex(s) + # Arguably, this would be more appropriate: + # return torch.tensor(complex(s), dtype=self._torch_dtype) + # But this results in wrong PyTorch multiplication functions + # being called. + else: + return float(s) + # return torch.tensor(float(s), dtype=self._torch_dtype) + @property def weighting(self): """This space's weighting scheme.""" From eefca31fd0df7a1a4ccf382086a83cdb2d787b8b Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 12 Jun 2024 15:38:55 +0200 Subject: [PATCH 16/52] Multiplication operators with explicitly selected scalar types. --- odl/discr/discr_space.py | 4 ++++ odl/operator/default_ops.py | 7 ++++--- odl/solvers/nonsmooth/proximal_operators.py | 2 +- odl/space/pspace.py | 6 ++++++ odl/space/pytorch_tensors.py | 4 ++++ 5 files changed, 19 insertions(+), 4 deletions(-) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index 741c2802440..a6a5b4fb8af 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -790,6 +790,10 @@ def __ipow__(self, p): self.tensor.__ipow__(p) return self + def __rmul__(self, μ): + """Implement ``μ * self``.""" + return self.space.element(μ * self.tensor) + @property def real(self): """Real part of this element. diff --git a/odl/operator/default_ops.py b/odl/operator/default_ops.py index 448da71f2c2..0f1df24f9d9 100644 --- a/odl/operator/default_ops.py +++ b/odl/operator/default_ops.py @@ -320,13 +320,14 @@ def multiplicand(self): def _call(self, x, out=None): """Multiply ``x`` and write to ``out`` if given.""" + μ = x.space.as_suitable_scalar(self.multiplicand) if out is None: - return x * self.multiplicand + return x * μ elif not self.__range_is_field: if self.__domain_is_field: - out.lincomb(x, self.multiplicand) + out.lincomb(x, μ) else: - out.assign(self.multiplicand * x) + out.assign(x * μ) else: raise ValueError('can only use `out` with `LinearSpace` range') diff --git a/odl/solvers/nonsmooth/proximal_operators.py b/odl/solvers/nonsmooth/proximal_operators.py index 0d83472ff27..03b0e40e332 100644 --- a/odl/solvers/nonsmooth/proximal_operators.py +++ b/odl/solvers/nonsmooth/proximal_operators.py @@ -391,7 +391,7 @@ def quadratic_perturbation_prox_factory(sigma): return (MultiplyOperator(const, domain=u.space, range=u.space) * prox * (MultiplyOperator(const, domain=u.space, range=u.space) - - sigma * const * u)) + u.space.as_suitable_scalar(sigma * const) * u)) else: space = prox.domain return (MultiplyOperator(const, domain=space, range=space) * diff --git a/odl/space/pspace.py b/odl/space/pspace.py index c4628f72983..2131c187e3d 100644 --- a/odl/space/pspace.py +++ b/odl/space/pspace.py @@ -1567,6 +1567,12 @@ def show(self, title=None, indices=None, **kwargs): return tuple(figs) + def __rmul__(self, other): + if self.space.is_suitable_scalar(other): + return self.space.element([other*part for part in self.parts]) + else: + raise TypeError("Only multiplication with suitable scalar supported for product spaces.") + # --- Add arithmetic operators that broadcast --- # diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 8ed49f408fe..313f8c3f9ae 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -1386,6 +1386,10 @@ def __ipow__(self, other): torch.pow(self.data, other, out=self.data) return self + def __rmul__(self, other): + result = self.space.element(other * self.data) + return result + def __int__(self): """Return ``int(self)``.""" return int(self.data) From b6975f8c8abf31f04baf323a6468947a8554db36 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 12 Jun 2024 18:46:01 +0200 Subject: [PATCH 17/52] Sketch of a `ufuncs` version for PyTorch. This approach diverges from the plan to base everything on `__array_ufunc__` as the deprecation notes suggest, but that is probably not tenable if we want to properly support dissimilar backends. --- odl/space/base_tensors.py | 10 ++++- odl/util/ufuncs.py | 84 ++++++++++++++++++++++++++++++++++++--- 2 files changed, 87 insertions(+), 7 deletions(-) diff --git a/odl/space/base_tensors.py b/odl/space/base_tensors.py index c43afc3c86d..6c4678e1d5f 100644 --- a/odl/space/base_tensors.py +++ b/odl/space/base_tensors.py @@ -22,7 +22,7 @@ array_str, dtype_str, indent, is_complex_floating_dtype, is_floating_dtype, is_numeric_dtype, is_real_dtype, is_real_floating_dtype, safe_int_conv, signature_string, writable_array) -from odl.util.ufuncs import TensorSpaceUfuncs +from odl.util.ufuncs import NumpyTensorSpaceUfuncs, PytorchTensorSpaceUfuncs from odl.util.utility import TYPE_MAP_C2R, TYPE_MAP_R2C, nullcontext __all__ = ('TensorSpace',) @@ -892,7 +892,13 @@ def ufuncs(self): the minimum required version. Use Numpy ufuncs directly, e.g., ``np.sqrt(x)`` instead of ``x.ufuncs.sqrt()``. """ - return TensorSpaceUfuncs(self) + if self.impl == "numpy": + return NumpyTensorSpaceUfuncs(self) + elif self.impl == "pytorch": + return PytorchTensorSpaceUfuncs(self) + else: + raise NotImplementedError() + def show(self, title=None, method='', indices=None, force_show=False, fig=None, **kwargs): diff --git a/odl/util/ufuncs.py b/odl/util/ufuncs.py index 6926e642501..caec596ae70 100644 --- a/odl/util/ufuncs.py +++ b/odl/util/ufuncs.py @@ -26,10 +26,11 @@ from __future__ import print_function, division, absolute_import from builtins import object import numpy as np +import torch import re -__all__ = ('TensorSpaceUfuncs', 'ProductSpaceUfuncs') +__all__ = ('NumpyTensorSpaceUfuncs', 'ProductSpaceUfuncs') # Some are ignored since they don't cooperate with dtypes, needs fix @@ -64,6 +65,36 @@ """.format(name) UFUNCS.append((name, n_in, n_out, doc)) +TORCH_RAW_UFUNCS = ['absolute', 'add', 'arccos', 'arccosh', 'arcsin', 'arcsinh', + 'arctan', 'arctan2', 'arctanh', 'bitwise_and', 'bitwise_or', + 'bitwise_xor', 'ceil', 'conj', 'copysign', 'cos', 'cosh', + 'deg2rad', 'divide', 'equal', 'exp', 'exp2', 'expm1', 'floor', + 'floor_divide', 'fmax', 'fmin', 'fmod', 'greater', + 'greater_equal', 'hypot', 'isfinite', 'isinf', 'isnan', + 'less', 'less_equal', 'log', 'log10', 'log1p', + 'log2', 'logaddexp', 'logaddexp2', 'logical_and', 'logical_not', + 'logical_or', 'logical_xor', 'maximum', 'minimum', + 'multiply', 'negative', 'not_equal', + 'rad2deg', 'reciprocal', 'remainder', + 'sign', 'signbit', 'sin', 'sinh', 'sqrt', 'square', 'subtract', + 'tan', 'tanh', 'true_divide', 'trunc'] +# Add some standardized information +TORCH_UFUNCS = [] +for name in TORCH_RAW_UFUNCS: + ufunc = getattr(np, name) + n_in, n_out = ufunc.nin, ufunc.nout + descr = ufunc.__doc__.splitlines()[2] + # Numpy occasionally uses single ticks for doc, we only use them for links + descr = re.sub('`+', '``', descr) + doc = descr + """ + +See Also +-------- +torch.{} +""".format(name) + TORCH_UFUNCS.append((name, n_in, n_out, doc)) + + # TODO: add the following reductions (to the CUDA implementation): # ['var', 'trace', 'tensordot', 'std', 'ptp', 'mean', 'diff', 'cumsum', # 'cumprod', 'average'] @@ -72,7 +103,7 @@ # --- Wrappers for `Tensor` --- # -def wrap_ufunc_base(name, n_in, n_out, doc): +def wrap_ufunc_numpy(name, n_in, n_out, doc): """Return ufunc wrapper for implementation-agnostic ufunc classes.""" ufunc = getattr(np, name) if n_in == 1: @@ -111,8 +142,42 @@ def wrapper(self, x2, out=None, **kwargs): wrapper.__doc__ = doc return wrapper +def wrap_ufunc_pytorch(name, n_in, n_out, doc): + """Return ufunc wrapper for implementation-agnostic ufunc classes.""" + ufunc = getattr(torch, name) + + if n_in == 1: + def wrapper(self, out=None, **kwargs): + if out is None: + return self.elem.space.element(ufunc(self.elem.data, **kwargs)) + elif isinstance(out, type(self.elem)): + ufunc(self.elem.data, out=out.data, **kwargs) + return + raise NotImplementedError() + + elif n_in == 2: + def wrapper(self, x2, out=None, **kwargs): + if out is None: + return self.elem.space.element(ufunc(self.elem.data, **kwargs)) + elif isinstance(out, type(self.elem)): + selfdata = self.elem.data + if isinstance(x2, type(self.elem)): + x2 = x2.data + elif isinstance(x2, (float, int)): + x2 = torch.tensor(x2).to(selfdata.device) + ufunc(selfdata, x2, out=out.data, **kwargs) + return + raise NotImplementedError() -class TensorSpaceUfuncs(object): + else: + raise NotImplementedError + + wrapper.__name__ = wrapper.__qualname__ = name + wrapper.__doc__ = doc + return wrapper + + +class NumpyTensorSpaceUfuncs(object): """Ufuncs for `Tensor` objects. @@ -176,9 +241,18 @@ def max(self, axis=None, dtype=None, out=None, keepdims=False): # Add ufunc methods to ufunc class for name, n_in, n_out, doc in UFUNCS: - method = wrap_ufunc_base(name, n_in, n_out, doc) - setattr(TensorSpaceUfuncs, name, method) + method = wrap_ufunc_numpy(name, n_in, n_out, doc) + setattr(NumpyTensorSpaceUfuncs, name, method) + + +class PytorchTensorSpaceUfuncs(object): + def __init__(self, elem): + """Create ufunc wrapper for elem.""" + self.elem = elem +for name, n_in, n_out, doc in TORCH_UFUNCS: + method = wrap_ufunc_pytorch(name, n_in, n_out, doc) + setattr(PytorchTensorSpaceUfuncs, name, method) # --- Wrappers for `ProductSpaceElement` --- # From 40b8d4699d9135b4219f9dc057ad6857efd36b22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 17 Jun 2024 11:28:30 +0200 Subject: [PATCH 18/52] Propose using PyTorch convolution for finite-differences. --- odl/discr/diff_ops.py | 26 +++++++++++++++++++++++--- 1 file changed, 23 insertions(+), 3 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index c6e715d548b..1be4722b27d 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -555,14 +555,34 @@ def __init__(self, domain=None, range=None, method='forward', def _call(self, x, out=None): """Calculate the divergence of ``x``.""" - if out is None: - out = self.range.element() - # print(f"{type(out.data)=}") ndim = self.range.ndim dx = self.range.cell_sides + torch_impl = uses_pytorch(x[0]) + if out is None: + if torch_impl and len(x)==2: + dtype = x[0].data.dtype + + assert(self.method=='backward'), f"{self.method=}" + assert(self.pad_mode=='constant') + assert(self.pad_const==0) + + # Add singleton channel- and batch dimensions + horizconv_data = x[0].data[None,None] + horizconv_kern = torch.tensor([[[[-1,1]]]], dtype=dtype) + verticonv_data = x[1].data[None,None] + verticonv_kern = torch.tensor([[[[-1],[1]]]], dtype=dtype) + return self.range.element( + torch.conv2d(horizconv_data, horizconv_kern, padding='same')[0,0] + / dx[0] + + torch.conv2d(verticonv_data, verticonv_kern, padding='same')[0,0] + / dx[1] + ) + else: + out = self.range.element() + tmp = self.range.element().asarray() # print(f"{type(tmp)=}") with writable_array(out) as out_arr: From 7783af802d954937c4e18392bb0c7f67b3515f75 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 17 Jun 2024 16:00:08 +0200 Subject: [PATCH 19/52] Correct axis association of the convolution FDs. --- odl/discr/diff_ops.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index 1be4722b27d..e7b6e192e01 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -571,9 +571,9 @@ def _call(self, x, out=None): # Add singleton channel- and batch dimensions horizconv_data = x[0].data[None,None] - horizconv_kern = torch.tensor([[[[-1,1]]]], dtype=dtype) + horizconv_kern = torch.tensor([[[[-1],[1],[0]]]], dtype=dtype) verticonv_data = x[1].data[None,None] - verticonv_kern = torch.tensor([[[[-1],[1]]]], dtype=dtype) + verticonv_kern = torch.tensor([[[[-1,1,0]]]], dtype=dtype) return self.range.element( torch.conv2d(horizconv_data, horizconv_kern, padding='same')[0,0] / dx[0] From b8d0e04b8289242e445620aa094fe0ae5cf33df8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 17 Jun 2024 19:16:17 +0200 Subject: [PATCH 20/52] PyTorch version of finite-difference grad etc.. The numpy-style version was very slow. These operations can be expressed nicely in terms of convolutions, which PyTorch supports well. Still requires _not_ performing in-place update to get good performance. Some padding modes are not supported yet. --- odl/discr/diff_ops.py | 335 ++++++++++++++++++++++++++---------------- 1 file changed, 210 insertions(+), 125 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index e7b6e192e01..c499a048484 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -12,6 +12,7 @@ import numpy as np import torch +from math import prod from odl.discr.discr_space import DiscretizedSpace from odl.operator.tensor_ops import PointwiseTensorFieldOperator @@ -345,20 +346,25 @@ def __init__(self, domain=None, range=None, method='forward', def _call(self, x, out=None): """Calculate the spatial gradient of ``x``.""" - if out is None: - out = self.range.element() - x_arr = x.asarray() ndim = self.domain.ndim dx = self.domain.cell_sides - for axis in range(ndim): - with writable_array(out[axis]) as out_arr: - finite_diff(x_arr, axis=axis, dx=dx[axis], method=self.method, + if out is None: + return self.range.element([ + finite_diff(x_arr, axis=axis, dx=dx[axis], method=self.method, pad_mode=self.pad_mode, pad_const=self.pad_const, - out=out_arr) - return out + ) + for axis in range(ndim)]) + else: + for axis in range(ndim): + with writable_array(out[axis]) as out_arr: + finite_diff(x_arr, axis=axis, dx=dx[axis], method=self.method, + pad_mode=self.pad_mode, + pad_const=self.pad_const, + out=out_arr) + return out def derivative(self, point=None): """Return the derivative operator. @@ -809,113 +815,10 @@ def __str__(self): return '{}:\n{}'.format(self.__class__.__name__, indent(dom_ran_str)) -def finite_diff(f, axis, dx=1.0, method='forward', out=None, +def _finite_diff_numpy(f_arr, axis, dx=1.0, method='forward', out=None, pad_mode='constant', pad_const=0): - """Calculate the partial derivative of ``f`` along a given ``axis``. + """ NumPy-specific version of `finite_diff`. """ - In the interior of the domain of f, the partial derivative is computed - using first-order accurate forward or backward difference or - second-order accurate central differences. - - With padding the same method and thus accuracy is used on endpoints as - in the interior i.e. forward and backward differences use first-order - accuracy on edges while central differences use second-order accuracy at - edges. - - Without padding one-sided forward or backward differences are used at - the boundaries. The accuracy at the endpoints can then also be - triggered by the edge order. - - The returned array has the same shape as the input array ``f``. - - Per default forward difference with dx=1 and no padding is used. - - Parameters - ---------- - f : `array-like` - An N-dimensional array. - axis : int - The axis along which the partial derivative is evaluated. - dx : float, optional - Scalar specifying the distance between sampling points along ``axis``. - method : {'central', 'forward', 'backward'}, optional - Finite difference method which is used in the interior of the domain - of ``f``. - out : `numpy.ndarray`, optional - An N-dimensional array to which the output is written. Has to have - the same shape as the input array ``f``. - pad_mode : string, optional - The padding mode to use outside the domain. - - ``'constant'``: Fill with ``pad_const``. - - ``'symmetric'``: Reflect at the boundaries, not doubling the - outmost values. - - ``'periodic'``: Fill in values from the other side, keeping - the order. - - ``'order0'``: Extend constantly with the outmost values - (ensures continuity). - - ``'order1'``: Extend with constant slope (ensures continuity of - the first derivative). This requires at least 2 values along - each axis where padding is applied. - - ``'order2'``: Extend with second order accuracy (ensures continuity - of the second derivative). This requires at least 3 values along - each axis where padding is applied. - - pad_const : float, optional - For ``pad_mode == 'constant'``, ``f`` assumes ``pad_const`` for - indices outside the domain of ``f`` - - Returns - ------- - out : `numpy.ndarray` - N-dimensional array of the same shape as ``f``. If ``out`` was - provided, the returned object is a reference to it. - - Examples - -------- - >>> f = np.array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - - >>> finite_diff(f, axis=0) - array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.]) - - Without arguments the above defaults to: - - >>> finite_diff(f, axis=0, dx=1.0, method='forward', pad_mode='constant') - array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.]) - - Parameters can be changed one by one: - - >>> finite_diff(f, axis=0, dx=0.5) - array([ 2., 2., 2., 2., 2., 2., 2., 2., 2., -18.]) - >>> finite_diff(f, axis=0, pad_mode='order1') - array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]) - - Central differences and different edge orders: - - >>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order1') - array([ 0.5, 1. , 2. , 3. , 4. , 5. , 6. , 7. , 8. , 8.5]) - >>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order2') - array([-0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) - - In-place evaluation: - - >>> out = f.copy() - >>> out is finite_diff(f, axis=0, out=out) - True - """ - torch_impl = uses_pytorch(f) - if torch_impl and out is not None: - assert(isinstance(out, torch.Tensor)), f"{type(out)=}" - - if torch_impl: - f_arr = torch.tensor(f) - else: - f_arr = np.asarray(f) ndim = f_arr.ndim if f_arr.shape[axis] < 2: @@ -940,34 +843,30 @@ def finite_diff(f, axis, dx=1.0, method='forward', out=None, raise ValueError('`pad_mode` {} not understood' ''.format(pad_mode)) - pad_const = dtype_type(f)(pad_const) + pad_const = np.array([pad_const], dtype = f_arr.dtype) if out is None: out = np.empty_like(f_arr) else: - if out.shape != f.shape: + if out.shape != f_arr.shape: raise ValueError('expected output shape {}, got {}' ''.format(f.shape, out.shape)) + orig_shape = f_arr.shape - if f_arr.shape[axis] < 2 and pad_mode == 'order1': + if orig_shape[axis] < 2 and pad_mode == 'order1': raise ValueError("size of array to small to use 'order1', needs at " "least 2 elements along axis {}.".format(axis)) - if f_arr.shape[axis] < 3 and pad_mode == 'order2': + if orig_shape[axis] < 3 and pad_mode == 'order2': raise ValueError("size of array to small to use 'order2', needs at " "least 3 elements along axis {}.".format(axis)) - # create slice objects: initially all are [:, :, ..., :] - - # Swap axes so that the axis of interest is first. This is a O(1) - # operation and is done to simplify the code below. + # Swap axes so that the axis of interest is first. In NumPy (but not PyTorch), + # this is a O(1) operation and is done to simplify the code below. out, out_in = np.swapaxes(out, 0, axis), out f_arr = np.swapaxes(f_arr, 0, axis) def fd_subtraction(a, b): - if torch_impl: - out[1:-1] = a - b - else: - np.subtract(a, b, out=out[1:-1]) + np.subtract(a, b, out=out[1:-1]) # Interior of the domain of f if method == 'central': @@ -1166,6 +1065,192 @@ def fd_subtraction(a, b): return out_in +def _finite_diff_pytorch(f_arr, axis, dx=1.0, method='forward', + pad_mode='constant', pad_const=0): + """ PyTorch-specific version of `finite_diff`. Notice that this has no output argument. """ + + ndim = f_arr.ndim + + if f_arr.shape[axis] < 2: + raise ValueError('in axis {}: at least two elements required, got {}' + ''.format(axis, f_arr.shape[axis])) + + if axis < 0: + axis += ndim + if not (0 <= axis < ndim): + raise IndexError('`axis` {} outside the valid range 0 ... {}' + ''.format(axis, ndim - 1)) + + dx, dx_in = float(dx), dx + if dx <= 0 or not np.isfinite(dx): + raise ValueError("`dx` must be positive, got {}".format(dx_in)) + + method, method_in = str(method).lower(), method + if method not in _SUPPORTED_DIFF_METHODS: + raise ValueError('`method` {} was not understood'.format(method_in)) + + if pad_mode not in _SUPPORTED_PAD_MODES: + raise ValueError('`pad_mode` {} not understood' + ''.format(pad_mode)) + + orig_shape = f_arr.shape + + if orig_shape[axis] < 2 and pad_mode == 'order1': + raise ValueError("size of array to small to use 'order1', needs at " + "least 2 elements along axis {}.".format(axis)) + if orig_shape[axis] < 3 and pad_mode == 'order2': + raise ValueError("size of array to small to use 'order2', needs at " + "least 3 elements along axis {}.".format(axis)) + + # Reshape (in O(1)), so the axis of interest is the pænultimate, all previous + # axes are flattened into the batch dimension, and all subsequent axes flattened + # into the final dimension. This allows a batched 2D convolution of final size 1 + # to perform the differentiation in only the axis of interest. + f_arr = f_arr.reshape([ prod(orig_shape[:axis]) + , 1 + , orig_shape[axis] + , prod(orig_shape[axis+1:]) + ]) + + dtype = f_arr.dtype + + # Kernel for convolution that expresses the finite-difference operator on, at least, + # the interior of the domain of f + if method == 'central': + fd_kernel = torch.tensor([[[[-1],[0],[1]]]], dtype=dtype) / (2*dx) + elif method == 'forward': + fd_kernel = torch.tensor([[[[0],[-1],[1]]]], dtype=dtype) / dx + elif method == 'backward': + fd_kernel = torch.tensor([[[[-1],[1],[0]]]], dtype=dtype) / dx + + if pad_mode == 'constant': + if pad_const==0: + result = torch.conv2d(f_arr, fd_kernel, padding='same') + else: + padding_arr = torch.ones_like(f_arr[:,:,0:1,:]) * pad_const + result = torch.conv2d( torch.cat([padding_arr, f_arr, padding_arr], dim=-2) + , fd_kernel, padding='valid' ) + elif pad_mode == 'periodic': + result = torch.conv2d(f_arr, fd_kernel, padding='circular') + + else: + raise NotImplementedError(f'{pad_mode=} not implemented for PyTorch') + + return result.reshape(orig_shape) + + +def finite_diff(f, axis, dx=1.0, method='forward', out=None, + pad_mode='constant', pad_const=0): + """Calculate the partial derivative of ``f`` along a given ``axis``. + + In the interior of the domain of f, the partial derivative is computed + using first-order accurate forward or backward difference or + second-order accurate central differences. + + With padding the same method and thus accuracy is used on endpoints as + in the interior i.e. forward and backward differences use first-order + accuracy on edges while central differences use second-order accuracy at + edges. + + Without padding one-sided forward or backward differences are used at + the boundaries. The accuracy at the endpoints can then also be + triggered by the edge order. + + The returned array has the same shape as the input array ``f``. + + Per default forward difference with dx=1 and no padding is used. + + Parameters + ---------- + f : `array-like` + An N-dimensional array. + axis : int + The axis along which the partial derivative is evaluated. + dx : float, optional + Scalar specifying the distance between sampling points along ``axis``. + method : {'central', 'forward', 'backward'}, optional + Finite difference method which is used in the interior of the domain + of ``f``. + out : `numpy.ndarray`, optional + An N-dimensional array to which the output is written. Has to have + the same shape as the input array ``f``. + pad_mode : string, optional + The padding mode to use outside the domain. + + ``'constant'``: Fill with ``pad_const``. + + ``'symmetric'``: Reflect at the boundaries, not doubling the + outmost values. + + ``'periodic'``: Fill in values from the other side, keeping + the order. + + ``'order0'``: Extend constantly with the outmost values + (ensures continuity). + + ``'order1'``: Extend with constant slope (ensures continuity of + the first derivative). This requires at least 2 values along + each axis where padding is applied. + + ``'order2'``: Extend with second order accuracy (ensures continuity + of the second derivative). This requires at least 3 values along + each axis where padding is applied. + + pad_const : float, optional + For ``pad_mode == 'constant'``, ``f`` assumes ``pad_const`` for + indices outside the domain of ``f`` + + Returns + ------- + out : `numpy.ndarray` + N-dimensional array of the same shape as ``f``. If ``out`` was + provided, the returned object is a reference to it. + + Examples + -------- + >>> f = np.array([ 0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + + >>> finite_diff(f, axis=0) + array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.]) + + Without arguments the above defaults to: + + >>> finite_diff(f, axis=0, dx=1.0, method='forward', pad_mode='constant') + array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., -9.]) + + Parameters can be changed one by one: + + >>> finite_diff(f, axis=0, dx=0.5) + array([ 2., 2., 2., 2., 2., 2., 2., 2., 2., -18.]) + >>> finite_diff(f, axis=0, pad_mode='order1') + array([ 1., 1., 1., 1., 1., 1., 1., 1., 1., 1.]) + + Central differences and different edge orders: + + >>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order1') + array([ 0.5, 1. , 2. , 3. , 4. , 5. , 6. , 7. , 8. , 8.5]) + >>> finite_diff(0.5 * f ** 2, axis=0, method='central', pad_mode='order2') + array([-0., 1., 2., 3., 4., 5., 6., 7., 8., 9.]) + + In-place evaluation: + + >>> out = f.copy() + >>> out is finite_diff(f, axis=0, out=out) + True + """ + if uses_pytorch(f): + if out is None: + return _finite_diff_pytorch(torch.tensor(f), axis, dx=dx, method=method, + pad_mode=pad_mode, pad_const=pad_const) + else: + assert(isinstance(out, torch.Tensor)), f"{type(out)=}" + out[:] = _finite_diff_pytorch(torch.tensor(f), axis, dx=dx, method=method, + pad_mode=pad_mode, pad_const=pad_const) + else: + return _finite_diff_numpy(np.asarray(f), axis, dx=dx, method=method, out=out, + pad_mode=pad_mode, pad_const=pad_const) + + if __name__ == '__main__': from odl.util.testutils import run_doctests From 3392da0346bb8e1feef20d5017cbc98e4695e7c8 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 17 Jun 2024 20:10:51 +0200 Subject: [PATCH 21/52] Consistent use of PyTorch finite_diff also for divergence operator. --- odl/discr/diff_ops.py | 62 ++++++++++++++++++------------------------- 1 file changed, 26 insertions(+), 36 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index c499a048484..ebacb6f7cb9 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -567,42 +567,32 @@ def _call(self, x, out=None): torch_impl = uses_pytorch(x[0]) + def directional_derivative(axis, dd_out=None): + return finite_diff( x[axis], axis=axis, dx=dx[axis] + , method=self.method, pad_mode=self.pad_mode + , pad_const=self.pad_const + , out=dd_out ) + if out is None: - if torch_impl and len(x)==2: - dtype = x[0].data.dtype - - assert(self.method=='backward'), f"{self.method=}" - assert(self.pad_mode=='constant') - assert(self.pad_const==0) - - # Add singleton channel- and batch dimensions - horizconv_data = x[0].data[None,None] - horizconv_kern = torch.tensor([[[[-1],[1],[0]]]], dtype=dtype) - verticonv_data = x[1].data[None,None] - verticonv_kern = torch.tensor([[[[-1,1,0]]]], dtype=dtype) - return self.range.element( - torch.conv2d(horizconv_data, horizconv_kern, padding='same')[0,0] - / dx[0] - + torch.conv2d(verticonv_data, verticonv_kern, padding='same')[0,0] - / dx[1] - ) - else: - out = self.range.element() - - tmp = self.range.element().asarray() - # print(f"{type(tmp)=}") - with writable_array(out) as out_arr: - for axis in range(ndim): - finite_diff(x[axis], axis=axis, dx=dx[axis], - method=self.method, pad_mode=self.pad_mode, - pad_const=self.pad_const, - out=tmp) - if axis == 0: - out_arr[:] = tmp - else: - out_arr += tmp + result = directional_derivative(0) + for axis in range(1,len(x)): + result += directional_derivative(axis) - return out + return self.range.element(result) + + else: + assert(not torch_impl) + + tmp = self.range.element().asarray() + with writable_array(out) as out_arr: + for axis in range(ndim): + directional_derivative(axis, out=tmp) + if axis == 0: + out_arr[:] = tmp + else: + out_arr += tmp + + return out def derivative(self, point=None): """Return the derivative operator. @@ -1240,11 +1230,11 @@ def finite_diff(f, axis, dx=1.0, method='forward', out=None, """ if uses_pytorch(f): if out is None: - return _finite_diff_pytorch(torch.tensor(f), axis, dx=dx, method=method, + return _finite_diff_pytorch(f.data, axis, dx=dx, method=method, pad_mode=pad_mode, pad_const=pad_const) else: assert(isinstance(out, torch.Tensor)), f"{type(out)=}" - out[:] = _finite_diff_pytorch(torch.tensor(f), axis, dx=dx, method=method, + out[:] = _finite_diff_pytorch(f.data, axis, dx=dx, method=method, pad_mode=pad_mode, pad_const=pad_const) else: return _finite_diff_numpy(np.asarray(f), axis, dx=dx, method=method, out=out, From e47ff959bdd0c862151f9d781809579045ea1943 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 15:38:23 +0200 Subject: [PATCH 22/52] Abolish in-place updates for PyTorch in PDHG. This solver / backend combination now runs efficiently in simple tests. --- odl/operator/oputils.py | 24 ++++++++++---- .../nonsmooth/primal_dual_hybrid_gradient.py | 32 +++++++++++++++---- 2 files changed, 42 insertions(+), 14 deletions(-) diff --git a/odl/operator/oputils.py b/odl/operator/oputils.py index 74f478cc5eb..0cc1b2f8660 100644 --- a/odl/operator/oputils.py +++ b/odl/operator/oputils.py @@ -14,7 +14,7 @@ from future.utils import native from odl.space import ProductSpace from odl.space.base_tensors import TensorSpace -from odl.util import nd_iterator +from odl.util import nd_iterator, uses_pytorch from odl.util.testutils import noise_element __all__ = ( @@ -228,17 +228,27 @@ def calc_opnorm(x_norm): # initial guess of opnorm opnorm = calc_opnorm(x_norm) - # temporary to improve performance - tmp = op.range.element() + if uses_pytorch(x): + calc_in_place = False # In-place updates are not efficient in PyTorch + else: + calc_in_place = True + # temporary to improve performance in NumPy + tmp = op.range.element() # Use the power method to estimate opnorm for i in range(ncalls): if use_normal: - op(x, out=tmp) - op.adjoint(tmp, out=x) + if calc_in_place: + op(x, out=tmp) + op.adjoint(tmp, out=x) + else: + x = op.adjoint(op(x), out=x) else: - op(x, out=tmp) - x, tmp = tmp, x + if calc_in_place: + op(x, out=tmp) + x, tmp = tmp, x + else: + x = op(x) # Calculate x norm and verify it is valid x_norm = x.norm() diff --git a/odl/solvers/nonsmooth/primal_dual_hybrid_gradient.py b/odl/solvers/nonsmooth/primal_dual_hybrid_gradient.py index ae7aea3cdd9..9630c5b3b7b 100644 --- a/odl/solvers/nonsmooth/primal_dual_hybrid_gradient.py +++ b/odl/solvers/nonsmooth/primal_dual_hybrid_gradient.py @@ -14,7 +14,7 @@ from __future__ import print_function, division, absolute_import import numpy as np - +from odl.util import uses_pytorch from odl.operator import Operator @@ -263,14 +263,22 @@ def pdhg(x, f, g, L, niter, tau=None, sigma=None, **kwargs): dual_tmp = L.range.element() primal_tmp = L.domain.element() + if uses_pytorch(x): + calc_in_place = False # In-place updates are not efficient in PyTorch + else: + calc_in_place = True + for _ in range(niter): # Copy required for relaxation x_old.assign(x) # Gradient ascent in the dual variable y # Compute dual_tmp = y + sigma * L(x_relax) - L(x_relax, out=dual_tmp) - dual_tmp.lincomb(1, y, sigma, dual_tmp) + if calc_in_place: + L(x_relax, out=dual_tmp) + dual_tmp.lincomb(1, y, sigma, dual_tmp) + else: + dual_tmp = y + sigma*L(x_relax) # Apply the dual proximal if not proximal_constant: @@ -279,13 +287,20 @@ def pdhg(x, f, g, L, niter, tau=None, sigma=None, **kwargs): # Gradient descent in the primal variable x # Compute primal_tmp = x + (- tau) * L.derivative(x).adjoint(y) - L.derivative(x).adjoint(y, out=primal_tmp) - primal_tmp.lincomb(1, x, -tau, primal_tmp) + if calc_in_place: + L.derivative(x).adjoint(y, out=primal_tmp) + primal_tmp.lincomb(1, x, -tau, primal_tmp) + else: + primal_tmp = x - L.derivative(x).adjoint(y)*tau # Apply the primal proximal if not proximal_constant: proximal_primal_tau = proximal_primal(tau) - proximal_primal_tau(primal_tmp, out=x) + + if True or calc_in_place: + proximal_primal_tau(primal_tmp, out=x) + else: + x.assign(proximal_primal_tau(primal_tmp)) # Acceleration if gamma_primal is not None: @@ -299,7 +314,10 @@ def pdhg(x, f, g, L, niter, tau=None, sigma=None, **kwargs): sigma *= theta # Over-relaxation in the primal variable x - x_relax.lincomb(1 + theta, x, -theta, x_old) + if calc_in_place: + x_relax.lincomb(1 + theta, x, -theta, x_old) + else: + x_relax = x*(1+theta) - x_old*theta if callback is not None: callback(x) From 21d069432b4f57ea5c349d3fc31a5463825ec132 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 16:18:17 +0200 Subject: [PATCH 23/52] Update PDHG example and enable PyTorch in it. --- examples/solvers/pdhg_denoising.py | 13 +++++++------ 1 file changed, 7 insertions(+), 6 deletions(-) diff --git a/examples/solvers/pdhg_denoising.py b/examples/solvers/pdhg_denoising.py index 1bb6cffffeb..321891efe3e 100644 --- a/examples/solvers/pdhg_denoising.py +++ b/examples/solvers/pdhg_denoising.py @@ -11,28 +11,29 @@ """ import numpy as np +import torch import scipy.misc import odl +impl = 'numpy' +# impl = 'pytorch' + # Read test image: use only every second pixel, convert integer to float, # and rotate to get the image upright -image = np.rot90(scipy.misc.ascent()[::2, ::2], 3).astype('float') +image = np.rot90(scipy.datasets.ascent()[::2, ::2], 3).astype('float') shape = image.shape # Rescale max to 1 image /= image.max() # Discretized spaces -space = odl.uniform_discr([0, 0], shape, shape) +space = odl.uniform_discr([0, 0], shape, shape, impl=impl) # Original image orig = space.element(image.copy()) # Add noise -image += 0.1 * odl.phantom.white_noise(orig.space) - -# Data of noisy image -noisy = space.element(image) +noisy = space.element(image) + 0.1 * odl.phantom.white_noise(orig.space) # Gradient operator gradient = odl.Gradient(space) From 090b875c62e6d49578c42520963b468bf5158aa7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 18:10:24 +0200 Subject: [PATCH 24/52] Add the torch device as a parameter to tensor spaces. --- odl/discr/discr_space.py | 8 +++++++- odl/space/pytorch_tensors.py | 12 ++++++++++-- 2 files changed, 17 insertions(+), 3 deletions(-) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index a6a5b4fb8af..6ee8413d1a0 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -1609,8 +1609,14 @@ def uniform_discr_frompartition(partition, dtype=None, impl='numpy', **kwargs): else: weighting = partition.cell_volume + if impl=='pytorch': + tensor_impl_args = {} + for arg in ['torch_device']: + if arg in kwargs: + tensor_impl_args[arg] = kwargs.pop(arg) + tspace = tspace_type(partition.shape, dtype, exponent=exponent, - weighting=weighting) + weighting=weighting, **tensor_impl_args) return DiscretizedSpace(partition, tspace, **kwargs) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 313f8c3f9ae..b273f4fa302 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -101,6 +101,11 @@ def __init__(self, shape, dtype=None, **kwargs): Other Parameters ---------------- + torch_device : optional, PyTorch device identifier + Where to store and process data (i.e. arrays) representing elements + of this space. Should typically be a GPU (cuda) if available, else + CPU as also used by NumPy. + weighting : optional Use weighted inner product, norm, and dist. The following types are supported as ``weighting``: @@ -217,12 +222,15 @@ def __init__(self, shape, dtype=None, **kwargs): raise ValueError('`dtype` {!r} not supported' ''.format(dtype_str(dtype))) + torch_device = kwargs.pop('torch_device', "cpu") dist = kwargs.pop('dist', None) norm = kwargs.pop('norm', None) inner = kwargs.pop('inner', None) weighting = kwargs.pop('weighting', None) exponent = kwargs.pop('exponent', getattr(weighting, 'exponent', 2.0)) + self._torch_device = torch.device(torch_device) + if (not is_numeric_dtype(self.dtype) and any(x is not None for x in (dist, norm, inner, weighting))): raise ValueError('cannot use any of `weighting`, `dist`, `norm` ' @@ -419,7 +427,7 @@ def element(self, inp=None, data_ptr=None, order=None): raise ValueError(f"Only row-major order supported ('C'), not '{order}'.") if inp is None and data_ptr is None: - arr = torch.empty(self.shape, dtype=self._torch_dtype) + arr = torch.empty(self.shape, dtype=self._torch_dtype, device=self._torch_device) return self.element_type(self, arr) @@ -441,7 +449,7 @@ def element(self, inp=None, data_ptr=None, order=None): return inp # TODO avoid copy when it's not necessary - arr = torch.tensor(inp, dtype=self._torch_dtype) + arr = torch.tensor(inp, dtype=self._torch_dtype, device=self._torch_device) if arr.shape != self.shape: raise ValueError('shape of `inp` not equal to space shape: ' From 76caa8e1cf99a28ca5d70a620ffee0a6261e1da2 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 18:45:12 +0200 Subject: [PATCH 25/52] Refactor finite-difference kernels. Less code duplication. --- odl/discr/diff_ops.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index ebacb6f7cb9..152298b558f 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -1106,12 +1106,14 @@ def _finite_diff_pytorch(f_arr, axis, dx=1.0, method='forward', # Kernel for convolution that expresses the finite-difference operator on, at least, # the interior of the domain of f + def as_kernel(mat): + return torch.tensor(mat, dtype=dtype) if method == 'central': - fd_kernel = torch.tensor([[[[-1],[0],[1]]]], dtype=dtype) / (2*dx) + fd_kernel = as_kernel([[[[-1],[0],[1]]]]) / (2*dx) elif method == 'forward': - fd_kernel = torch.tensor([[[[0],[-1],[1]]]], dtype=dtype) / dx + fd_kernel = as_kernel([[[[0],[-1],[1]]]]) / dx elif method == 'backward': - fd_kernel = torch.tensor([[[[-1],[1],[0]]]], dtype=dtype) / dx + fd_kernel = as_kernel([[[[-1],[1],[0]]]]) / dx if pad_mode == 'constant': if pad_const==0: From f053134563b2af6851bfc39a528a68006b9a7e99 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 18:45:57 +0200 Subject: [PATCH 26/52] Use correct Torch device for FD convolutions. --- odl/discr/diff_ops.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/odl/discr/diff_ops.py b/odl/discr/diff_ops.py index 152298b558f..23e09678930 100644 --- a/odl/discr/diff_ops.py +++ b/odl/discr/diff_ops.py @@ -1107,7 +1107,7 @@ def _finite_diff_pytorch(f_arr, axis, dx=1.0, method='forward', # Kernel for convolution that expresses the finite-difference operator on, at least, # the interior of the domain of f def as_kernel(mat): - return torch.tensor(mat, dtype=dtype) + return torch.tensor(mat, dtype=dtype, device=f_arr.device) if method == 'central': fd_kernel = as_kernel([[[[-1],[0],[1]]]]) / (2*dx) elif method == 'forward': From 0845126a85ca020a83413e4f01605a716b055fd0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 18 Jun 2024 18:51:48 +0200 Subject: [PATCH 27/52] Make `tensor_impl_args` compatible (albeit empty) on NumPy. --- odl/discr/discr_space.py | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index 6ee8413d1a0..48fdb025608 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -1609,8 +1609,9 @@ def uniform_discr_frompartition(partition, dtype=None, impl='numpy', **kwargs): else: weighting = partition.cell_volume + tensor_impl_args = {} + if impl=='pytorch': - tensor_impl_args = {} for arg in ['torch_device']: if arg in kwargs: tensor_impl_args[arg] = kwargs.pop(arg) From e815813225ae75ecc7fa4084cfb7b7ddcb931e7e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 21 Jun 2024 12:17:25 +0200 Subject: [PATCH 28/52] GPU-compatible conversions to NumPy. --- odl/discr/discr_space.py | 16 ++++++++++++++++ odl/space/pytorch_tensors.py | 14 ++++++++++++++ 2 files changed, 30 insertions(+) diff --git a/odl/discr/discr_space.py b/odl/discr/discr_space.py index 48fdb025608..6fab39cacf8 100644 --- a/odl/discr/discr_space.py +++ b/odl/discr/discr_space.py @@ -972,6 +972,22 @@ def __setitem__(self, indices, values): values = values.tensor self.tensor.__setitem__(indices, values) + def __array__(self, dtype=None): + """Return a Numpy array from this tensor. + (Contrast with the `asarray` method, which may give other types of array, + not just NumPy.) + + Parameters + ---------- + dtype : + Specifier for the data type of the output array. + + Returns + ------- + array : `numpy.ndarray` + """ + return self.tensor.__array__() + def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): """Interface to Numpy's ufunc machinery. diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index b273f4fa302..0109c77841a 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -1202,6 +1202,20 @@ def __setitem__(self, indices, values): self.data[indices] = values + def __array__(self, dtype=None): + """Return a Numpy array from this tensor. + + Parameters + ---------- + dtype : + Specifier for the data type of the output array. + + Returns + ------- + array : `numpy.ndarray` + """ + return self.data.cpu().numpy() + @property def real(self): """Real part of ``self``. From e12c2db4613f914ec2c6ef3daa744a8d297429f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 21 Aug 2024 16:50:16 +0200 Subject: [PATCH 29/52] Refactor Fourier trafo classes. This is to make them more general with respect to backend, particularly towards PyTorch. --- odl/trafos/fourier.py | 72 ++++++++++++++++++++++++------------------- 1 file changed, 41 insertions(+), 31 deletions(-) diff --git a/odl/trafos/fourier.py b/odl/trafos/fourier.py index 15424f402f7..a4640655e0c 100644 --- a/odl/trafos/fourier.py +++ b/odl/trafos/fourier.py @@ -26,15 +26,32 @@ complex_dtype, conj_exponent, dtype_repr, is_complex_floating_dtype, is_real_dtype, normalized_axes_tuple, normalized_scalar_param_list) +from typing import Optional + __all__ = ('DiscreteFourierTransform', 'DiscreteFourierTransformInverse', 'FourierTransform', 'FourierTransformInverse') -_SUPPORTED_FOURIER_IMPLS = ('numpy',) -_DEFAULT_FOURIER_IMPL = 'numpy' +_SUPPORTED_FOURIER_IMPLS = {'numpy': ('numpy',)} +_DEFAULT_FOURIER_IMPL = {'numpy': 'numpy'} if PYFFTW_AVAILABLE: - _SUPPORTED_FOURIER_IMPLS += ('pyfftw',) - _DEFAULT_FOURIER_IMPL = 'pyfftw' + _SUPPORTED_FOURIER_IMPLS['numpy'] += ('pyfftw',) + _DEFAULT_FOURIER_IMPL['numpy'] = 'pyfftw' + + +def _select_fft_impl(impl_suggestion: Optional[str], domain_impl: str): + if impl_suggestion is None: + impl = _DEFAULT_FOURIER_IMPL.get(domain_impl) + if impl is None: + raise ValueError("There is no default FFT implementation for" + + " tensors with {domain_impl} implementation.") + else: + impl = impl_suggestion + impl, impl_in = str(impl).lower(), impl + if impl not in _SUPPORTED_FOURIER_IMPLS.get(domain_impl): + raise ValueError(f"`impl` '{impl_in}' not supported for" + + " tensors with {domain_impl} implementation.") + return impl class DiscreteFourierTransformBase(Operator): @@ -91,12 +108,7 @@ def __init__(self, inverse, domain, range=None, axes=None, sign='-', ''.format(range)) # Implementation - if impl is None: - impl = _DEFAULT_FOURIER_IMPL - impl, impl_in = str(impl).lower(), impl - if impl not in _SUPPORTED_FOURIER_IMPLS: - raise ValueError("`impl` '{}' not supported".format(impl_in)) - self.__impl = impl + self.__impl = _select_fft_impl(impl) # Axes if axes is None: @@ -125,11 +137,11 @@ def __init__(self, inverse, domain, range=None, axes=None, sign='-', domain.grid, shift=False, halfcomplex=halfcomplex, axes=axes).shape if range is None: - impl = domain.tspace.impl + domain_impl = domain.tspace.impl shape = np.atleast_1d(ran_shape) range = uniform_discr( - [0] * len(shape), shape - 1, shape, ran_dtype, impl, + [0] * len(shape), shape - 1, shape, ran_dtype, domain_impl, nodes_on_bdry=True, exponent=conj_exponent(domain.exponent)) else: @@ -171,10 +183,13 @@ def _call(self, x, out, **kwargs): Call pyfftw backend directly """ # TODO: Implement zero padding - if self.impl == 'numpy': - out[:] = self._call_numpy(x.asarray()) - else: - out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + match self.impl: + case 'numpy': + out[:] = self._call_numpy(x.asarray()) + case 'pyfftw': + out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + case _: + raise NotImplementedError(self.impl) @property def impl(self): @@ -804,22 +819,13 @@ def __init__(self, inverse, domain, range=None, impl=None, **kwargs): if not isinstance(domain, DiscretizedSpace): raise TypeError('domain {!r} is not a `DiscretizedSpace` instance' ''.format(domain)) - if domain.impl != 'numpy': - raise NotImplementedError( - 'Only Numpy-based data spaces are supported, got {}' - ''.format(domain.tspace)) # axes axes = kwargs.pop('axes', np.arange(domain.ndim)) self.__axes = normalized_axes_tuple(axes, domain.ndim) # Implementation - if impl is None: - impl = _DEFAULT_FOURIER_IMPL - impl, impl_in = str(impl).lower(), impl - if impl not in _SUPPORTED_FOURIER_IMPLS: - raise ValueError("`impl` '{}' not supported".format(impl_in)) - self.__impl = impl + self.__impl = _select_fft_impl(impl, domain.impl) # Handle half-complex yes/no and shifts halfcomplex = kwargs.pop('halfcomplex', True) @@ -905,11 +911,15 @@ def _call(self, x, out, **kwargs): Call pyfftw backend directly """ # TODO: Implement zero padding - if self.impl == 'numpy': - out[:] = self._call_numpy(x.asarray()) - else: - # 0-overhead assignment if asarray() does not copy - out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + match self.impl: + case 'numpy': + out[:] = self._call_numpy(x.asarray()) + case 'pyfftw': + # 0-overhead assignment if asarray() does not copy + out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + case _: + raise NotImplementedError(self.impl) + def _call_numpy(self, x): """Return ``self(x)`` for numpy back-end. From d8523e9142ae19a34ae7cb59e3c775cf963c9052 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 4 Oct 2024 12:07:10 +0200 Subject: [PATCH 30/52] Move the lookup dict for PyTorch dtypes to a more global level. --- odl/space/pytorch_tensors.py | 8 ++------ odl/util/utility.py | 5 +++++ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 0109c77841a..b159fabbad7 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -24,6 +24,7 @@ from odl.space.weighting import ( ArrayWeighting, ConstWeighting, CustomDist, CustomInner, CustomNorm, Weighting) +from odl.util.utility import _CORRESPONDING_PYTORCH_DTYPES from odl.util import ( dtype_str, is_floating_dtype, is_numeric_dtype, is_real_dtype, nullcontext, signature_string, writable_array) @@ -31,11 +32,6 @@ __all__ = ('PytorchTensorSpace',) -_PYTORCH_DTYPES = {np.dtype('float32'): torch.float32, - np.dtype('float64'): torch.float64, - np.dtype('complex64'): torch.complex64, - np.dtype('complex128'): torch.complex128} - # Define size thresholds to switch implementations THRESHOLD_SMALL = 100 THRESHOLD_MEDIUM = 50000 @@ -237,7 +233,7 @@ def __init__(self, shape, dtype=None, **kwargs): 'or `inner` for non-numeric `dtype` {}' ''.format(dtype)) else: - self._torch_dtype = _PYTORCH_DTYPES[self.dtype] + self._torch_dtype = _CORRESPONDING_PYTORCH_DTYPES[self.dtype] if exponent != 2.0 and any(x is not None for x in (dist, norm, inner)): raise ValueError('cannot use any of `dist`, `norm` or `inner` ' 'for exponent != 2') diff --git a/odl/util/utility.py b/odl/util/utility.py index 0dfd97cf9ec..1595ef1235c 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -521,6 +521,11 @@ def complex_dtype(dtype, default=None): else: return np.dtype((complex_base_dtype, dtype.shape)) +_CORRESPONDING_PYTORCH_DTYPES = {np.dtype('float32'): torch.float32, + np.dtype('float64'): torch.float64, + np.dtype('complex64'): torch.complex64, + np.dtype('complex128'): torch.complex128} + def uses_pytorch(obj): if isinstance(obj, torch.Tensor): return True From afeaf52e1dd74db232501c821677984c4741db2e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 4 Oct 2024 13:24:06 +0200 Subject: [PATCH 31/52] Propose a backend-agnostic ways of checking dtype compatibility. The implementation with a random check is an ugly hack, barely acceptable for this purpose. Arguably, hardcoding a full matrix of what is convertible would be a cleaner solution, but it would actuall be more problematic from a maintenance perspective because Torch might change what conversions are supported. The only proper solution would be to use a Torch function corresponding to np.can_cast, but torch.can_cast does not do that, it is more restrictive. --- odl/util/utility.py | 32 ++++++++++++++++++++++++++++++++ 1 file changed, 32 insertions(+) diff --git a/odl/util/utility.py b/odl/util/utility.py index 1595ef1235c..9a83ac72ae9 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -35,6 +35,7 @@ 'is_real_dtype', 'is_real_floating_dtype', 'is_complex_floating_dtype', + 'is_castable_to', 'uses_pytorch', 'real_dtype', 'complex_dtype', @@ -526,6 +527,37 @@ def complex_dtype(dtype, default=None): np.dtype('complex64'): torch.complex64, np.dtype('complex128'): torch.complex128} +@cache_arguments +def is_castable(from_dtype, to_dtype): + """Determine whether the type `from` is safely convertible to `to`. + Both should be either NumPy `dtype` or PyTorch `dtype`.""" + if isinstance(from_dtype, np.dtype) and isinstance(to_dtype, np.dtype): + return np.can_cast(from_dtype, to_dtype) + elif (isinstance(to_dtype, torch.dtype)): + from_dtype = _CORRESPONDING_PYTORCH_DTYPES.get(from_dtype, from_dtype) + # Torch does not provide a satisfying way to determine castability, + # so we find it out by experiment. + # This is somewhat expensive, so it is important that this function is + # memoised (cache_arguments). + try: + gen = torch.Generator() + gen.manual_seed(1232451) # Avoid nondeterministic behaviour + test_arr = torch.rand((1000,), generator=gen, dtype=from_dtype) + roundtripped = test_arr.to(to_dtype).to(from_dtype) + except TypeError: + return False + return torch.equal(roundtripped, test_arr) + +def is_castable_to(obj, dtype): + """Determine whether there is a safe way to cast `obj` to the type + specified by `dtype`, which can be either a NumPy dtype or a PyTorch + `dtype`.""" + if hasattr(obj, 'dtype'): + obj_dtype = obj.dtype + else: + obj_dtype = np.array([obj]).dtype + return is_castable(obj_dtype, dtype) + def uses_pytorch(obj): if isinstance(obj, torch.Tensor): return True From 21cfabd8f9fdd8a6be137e9003caa1dc79fd36c7 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Fri, 4 Oct 2024 13:59:05 +0200 Subject: [PATCH 32/52] Generalize array resizing to Torch. The size-filling required some nontrivial conversions. --- odl/util/numerics.py | 46 +++++++++++++++++++++++++++++--------------- 1 file changed, 31 insertions(+), 15 deletions(-) diff --git a/odl/util/numerics.py b/odl/util/numerics.py index d5f59fbb67b..54beeb11d79 100644 --- a/odl/util/numerics.py +++ b/odl/util/numerics.py @@ -11,6 +11,8 @@ from __future__ import absolute_import, division, print_function import numpy as np +import torch +from odl.util.utility import is_castable_to from odl.util.normalize import normalized_scalar_param_list, safe_int_conv __all__ = ( @@ -421,16 +423,31 @@ def resize_array(arr, newshp, offset=None, pad_mode='constant', pad_const=0, except TypeError: raise TypeError('`newshp` must be a sequence, got {!r}'.format(newshp)) + if isinstance(arr, np.ndarray): + impl = 'numpy' + elif isinstance(arr, torch.Tensor): + impl = 'pytorch' + else: + raise TypeError(f"Unknown how to resize array (?) of type {type(arr)}.") + if out is not None: - if not isinstance(out, np.ndarray): + if impl=='numpy' and not isinstance(out, np.ndarray): raise TypeError('`out` must be a `numpy.ndarray` instance, got ' '{!r}'.format(out)) + elif impl=='pytorch' and not isinstance(out, torch.Tensor): + raise TypeError('`out` must be a `torch.Tensor` instance, got ' + '{!r}'.format(out)) if out.shape != newshp: raise ValueError('`out` must have shape {}, got {}' ''.format(newshp, out.shape)) - order = 'C' if out.flags.c_contiguous else 'F' - arr = np.asarray(arr, dtype=out.dtype, order=order) + if impl=='pytorch': + if arr.dtype != out.dtype: + arr = torch.tensor(arr, dtype=out.dtype) + else: # NumPy + order = 'C' if out.flags.c_contiguous else 'F' + arr = np.asarray(arr, dtype=out.dtype, order=order) + if arr.ndim != out.ndim: raise ValueError('number of axes of `arr` and `out` do not match ' '({} != {})'.format(arr.ndim, out.ndim)) @@ -455,16 +472,14 @@ def resize_array(arr, newshp, offset=None, pad_mode='constant', pad_const=0, if pad_mode not in _SUPPORTED_RESIZE_PAD_MODES: raise ValueError("`pad_mode` '{}' not understood".format(pad_mode_in)) - if (pad_mode == 'constant' and - any(n_new > n_orig - for n_orig, n_new in zip(arr.shape, out.shape))): - try: - pad_const_scl = np.array([pad_const], out.dtype) - assert(pad_const_scl == np.array([pad_const])) - except Exception as e: - raise ValueError('`pad_const` {} cannot be safely cast to the data ' - 'type {} of the output array' - ''.format(pad_const, out.dtype)) + if pad_mode == 'constant': + incompatible_const_error = ValueError( + f'`pad_const` {pad_const} cannot be safely cast to the data ' + + f'type {out.dtype} of the output array') + if (not is_castable_to(pad_const, out.dtype) + and any(n_new > n_orig + for n_orig, n_new in zip(arr.shape, out.shape))): + raise incompatible_const_error # Handle direction direction, direction_in = str(direction).lower(), direction @@ -476,10 +491,11 @@ def resize_array(arr, newshp, offset=None, pad_mode='constant', pad_const=0, raise ValueError("`pad_const` must be 0 for 'adjoint' direction, " "got {}".format(pad_const)) + fill_with = out.fill_ if impl=='pytorch' else out.fill if direction == 'forward' and pad_mode == 'constant' and pad_const != 0: - out.fill(pad_const) + fill_with(pad_const) else: - out.fill(0) + fill_with(0) # Perform the resizing if direction == 'forward': From 60e0b847865670643cb719d4156624a88a2771fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 7 Oct 2024 14:03:51 +0200 Subject: [PATCH 33/52] One more example using PyTorch storage. --- examples/solvers/pdhg_denoising_pytorch.py | 97 ++++++++++++++++++++++ 1 file changed, 97 insertions(+) create mode 100644 examples/solvers/pdhg_denoising_pytorch.py diff --git a/examples/solvers/pdhg_denoising_pytorch.py b/examples/solvers/pdhg_denoising_pytorch.py new file mode 100644 index 00000000000..a985ca504ab --- /dev/null +++ b/examples/solvers/pdhg_denoising_pytorch.py @@ -0,0 +1,97 @@ +"""Total variation denoising using PDHG. + +Solves the optimization problem + + min_{x >= 0} 1/2 ||x - g||_2^2 + lam || |grad(x)| ||_1 + +Where ``grad`` the spatial gradient and ``g`` is given noisy data. + +For further details and a description of the solution method used, see +https://odlgroup.github.io/odl/guide/pdhg_guide.html in the ODL documentation. +""" + +import numpy as np +import torch +import scipy.misc +import odl +import cProfile + +# Read test image: use only every second pixel, convert integer to float, +# and rotate to get the image upright +image = np.rot90(scipy.datasets.ascent()[::2, ::2], 3).astype('float') +shape = image.shape + +# Rescale max to 1 +image /= image.max() + +# Discretized spaces +space = odl.uniform_discr([0, 0], shape, shape, impl='pytorch') + +# Original image +orig = space.element(image.copy()) + +orig.data.requires_grad = False + +# Add noise +noisy = space.element(image) + 0.1 * odl.phantom.white_noise(orig.space) + +noisy.data.requires_grad = False + +# Gradient operator +gradient = odl.Gradient(space) + +# grad_xmp = gradient(orig) +# grad_xmp.show(title = "Grad-op applied to original") + +# Matrix of operators +op = odl.BroadcastOperator(odl.IdentityOperator(space), gradient) + +# Set up the functionals + +# l2-squared data matching +l2_norm = odl.solvers.L2NormSquared(space).translated(noisy) + +# Isotropic TV-regularization: l1-norm of grad(x) +l1_norm = 0.15 * odl.solvers.L1Norm(gradient.range) + +# Make separable sum of functionals, order must correspond to the operator K +g = odl.solvers.SeparableSum(l2_norm, l1_norm) + +# Non-negativity constraint +f = odl.solvers.IndicatorNonnegativity(op.domain) + +# --- Select solver parameters and solve using PDHG --- # + +# Estimated operator norm, add 10 percent to ensure ||K||_2^2 * sigma * tau < 1 +op_norm = 1.1 * odl.power_method_opnorm(op, xstart=noisy, maxiter=10) + # 3.2833764101732785 +print(f"{op_norm=}") + +niter = 200 # Number of iterations +tau = 1.0 / op_norm # Step size for the primal variable +sigma = 1.0 / op_norm # Step size for the dual variable + +# Optional: pass callback objects to solver +callback = (odl.solvers.CallbackPrintIteration() & + odl.solvers.CallbackShow(step=5)) + +# Starting point +x = op.domain.zero() + +x.data.requires_grad = False + +print("Go solve...") + +# Run algorithm (and display intermediates) +def do_running(): + with torch.no_grad(): + odl.solvers.pdhg(x, f, g, op, niter=niter, tau=tau, sigma=sigma, + callback=callback) + +do_running() +# cProfile.run('do_running()') + +# Display images +orig.show(title='Original Image') +noisy.show(title='Noisy Image') +x.show(title='Reconstruction', force_show=True) From aa1d5b646ae2303dc2def08b66a420da11aa3b3e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 7 Oct 2024 19:21:36 +0200 Subject: [PATCH 34/52] Support PyTorch in the dtype-categorization utils. --- odl/util/utility.py | 12 ++++++++++++ 1 file changed, 12 insertions(+) diff --git a/odl/util/utility.py b/odl/util/utility.py index 9a83ac72ae9..0574206911f 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -366,6 +366,13 @@ def cache_arguments(function): @cache_arguments def is_numeric_dtype(dtype): """Return ``True`` if ``dtype`` is a numeric type.""" + if isinstance(dtype, torch.dtype): + try: + assert(dtype in [torch.float32, torch.float64]) + return True + except AssertionError: + assert(dtype in [torch.complex64, torch.complex128]) + return True dtype = np.dtype(dtype) return np.issubdtype(getattr(dtype, 'base', None), np.number) @@ -399,6 +406,11 @@ def is_real_floating_dtype(dtype): @cache_arguments def is_complex_floating_dtype(dtype): """Return ``True`` if ``dtype`` is a complex floating point type.""" + if isinstance(dtype, torch.dtype): + if(dtype in [torch.float32, torch.float64]): + return False + assert(dtype in [torch.complex64, torch.complex128]) + return True dtype = np.dtype(dtype) return np.issubdtype(getattr(dtype, 'base', None), np.complexfloating) From c97c0ac5e9b358879c4ae3e331a8226fa48eca00 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 7 Oct 2024 19:22:53 +0200 Subject: [PATCH 35/52] Make `as_writable_array` handle and reinstore PyTorch-based elements. --- odl/util/utility.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/odl/util/utility.py b/odl/util/utility.py index 0574206911f..11871bafb37 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -711,7 +711,17 @@ def writable_array(obj, **kwargs): torch_impl = uses_pytorch(obj) try: if torch_impl: - arr = torch.tensor(obj, **kwargs) + if isinstance(obj, torch.Tensor): + arr = obj + elif hasattr(obj, 'data') and isinstance(obj.data, torch.Tensor): + arr = obj.data + else: + if hasattr(obj, 'data'): + if 'dtype' not in kwargs: + kwargs['dtype'] = obj.data.dtype + arr = torch.tensor(obj.data, **kwargs) + else: + arr = torch.tensor(obj, **kwargs) else: arr = np.asarray(obj, **kwargs) yield arr From f5a45de0df5a515fc7c5ee2da6e29d6ca29552f4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 8 Oct 2024 12:48:22 +0200 Subject: [PATCH 36/52] Default Fourier implementation should be based on the space. It does not make much sense to use PyTorch storage but NumPy-based Fourier transform. --- odl/trafos/util/ft_utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/odl/trafos/util/ft_utils.py b/odl/trafos/util/ft_utils.py index d4d4f65dcbf..2f7c439db55 100644 --- a/odl/trafos/util/ft_utils.py +++ b/odl/trafos/util/ft_utils.py @@ -618,7 +618,7 @@ def reciprocal_space(space, axes=None, halfcomplex=False, shift=True, raise ValueError('{} is not a complex data type' ''.format(dtype_repr(dtype))) - impl = kwargs.pop('impl', 'numpy') + impl = kwargs.pop('impl', space.impl) # Calculate range recip_grid = reciprocal_grid(space.grid, shift=shift, From 7a176ae5e68e6d2f91f237a995e271bf2e572405 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 9 Oct 2024 17:14:11 +0200 Subject: [PATCH 37/52] Helpers for generating / converting arrays on NumPy or PyTorch as appropriate. --- odl/util/utility.py | 52 +++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 52 insertions(+) diff --git a/odl/util/utility.py b/odl/util/utility.py index 11871bafb37..6c5b57e2995 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -15,6 +15,7 @@ from collections import OrderedDict from contextlib import contextmanager from itertools import product +from abc import ABC import numpy as np import torch @@ -43,6 +44,8 @@ 'nd_iterator', 'conj_exponent', 'nullcontext', + 'ArrayOnBackendManager', + 'compatible_array_manager', 'writable_array', 'signature_string', 'signature_string_parts', @@ -587,6 +590,55 @@ def is_string(obj): else: return True +class ArrayOnBackendManager(ABC): + def __init__(self): + raise NotImplementedError() + def as_compatible_array(self, arr, **kwargs): + raise NotImplementedError() + def compatible_zeros(self, shape, **kwargs): + raise NotImplementedError() + def compatible_ones(self, shape, **kwargs): + raise NotImplementedError() + def select_dtype(self, arr, dtype): + raise NotImplementedError() + def make_copy(self, arr): + raise NotImplementedError() + +class ArrayOnPytorchManager(ABC): + def __init__(self, device): + self._device = device + def as_compatible_array(self, arr, **kwargs): + return torch.tensor(arr, device = self._device, **kwargs) + def compatible_zeros(self, shape, **kwargs): + return torch.zeros(shape, device = self._device, **kwargs) + def compatible_ones(self, shape, **kwargs): + return torch.ones(shape, device = self._device, **kwargs) + def select_dtype(self, arr, dtype): + if dtype in _CORRESPONDING_PYTORCH_DTYPES: + dtype = _CORRESPONDING_PYTORCH_DTYPES[dtype] + return arr.type(dtype) + def make_copy(self, arr): + return arr.clone().detach() + +class ArrayOnNumPyManager(ABC): + def __init__(self): + pass + def as_compatible_array(self, arr, **kwargs): + return np.array(arr, **kwargs) + def compatible_zeros(self, shape, **kwargs): + return np.zeros(shape, **kwargs) + def compatible_ones(self, shape, **kwargs): + return np.ones(shape, **kwargs) + def select_dtype(self, arr, dtype): + return arr.astype(dtype, copy=False) + def make_copy(self, arr): + return arr.copy() + +def compatible_array_manager(arr): + if uses_pytorch(arr): + return ArrayOnPytorchManager(arr.device) + else: + return ArrayOnNumPyManager() def nd_iterator(shape): """Iterator over n-d cube with shape. From bf3eb16eba67819e55d1b5a19a51665f0c41503a Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 9 Oct 2024 17:34:07 +0200 Subject: [PATCH 38/52] More dtype categorisation with PyTorch. --- odl/util/utility.py | 8 +++++++- 1 file changed, 7 insertions(+), 1 deletion(-) diff --git a/odl/util/utility.py b/odl/util/utility.py index 6c5b57e2995..9dace49b21e 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -402,7 +402,13 @@ def is_real_dtype(dtype): @cache_arguments def is_real_floating_dtype(dtype): """Return ``True`` if ``dtype`` is a real floating point type.""" - dtype = np.dtype(dtype) + if isinstance(dtype, torch.dtype): + if dtype in [torch.complex64, torch.complex128]: + return False + else: + assert(dtype in [torch.float32, torch.float64]) + return True + dtype = np.dtype( dtype) return np.issubdtype(getattr(dtype, 'base', None), np.floating) From 298f7675d32901be1f5bc8aa4c57aa7fb776228d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 9 Oct 2024 18:16:49 +0200 Subject: [PATCH 39/52] Make `fast_1d_tensor_mult` PyTorch-compatible. This is an important auxiliary function for ODL's Fourier transforms (or rather, to their pre- and post-processing). --- odl/util/numerics.py | 27 ++++++++++++++++++++------- 1 file changed, 20 insertions(+), 7 deletions(-) diff --git a/odl/util/numerics.py b/odl/util/numerics.py index 54beeb11d79..b2cdb24fe85 100644 --- a/odl/util/numerics.py +++ b/odl/util/numerics.py @@ -12,7 +12,7 @@ import numpy as np import torch -from odl.util.utility import is_castable_to +from odl.util.utility import is_castable_to, uses_pytorch, compatible_array_manager from odl.util.normalize import normalized_scalar_param_list, safe_int_conv __all__ = ( @@ -210,6 +210,8 @@ def fast_1d_tensor_mult(ndarr, onedim_arrs, axes=None, out=None): The advantage of this approach is that it is memory-friendly and loops over the big array only twice. + TODO update documentation WRT PyTorch + Parameters ---------- ndarr : `array-like` @@ -229,10 +231,18 @@ def fast_1d_tensor_mult(ndarr, onedim_arrs, axes=None, out=None): Result of the modification. If ``out`` was given, the returned object is a reference to it. """ + use_pytorch = uses_pytorch(ndarr) or uses_pytorch(out) + array_mgr = compatible_array_manager(ndarr) + if out is None: - out = np.array(ndarr, copy=True) - else: + if use_pytorch: + out = torch.Tensor(ndarr, copy=True) + else: + out = np.array(ndarr, copy=True) + elif type(out)==type(ndarr): out[:] = ndarr # Self-assignment is free if out is ndarr + else: + raise TypeError(f"{type(ndarr)=} should be the same as {type(out)=}") if not onedim_arrs: raise ValueError('no 1d arrays given') @@ -253,14 +263,17 @@ def fast_1d_tensor_mult(ndarr, onedim_arrs, axes=None, out=None): raise ValueError('`axes` {} out of bounds for {} dimensions' ''.format(axes_in, out.ndim)) + atleast_1d = torch.atleast_1d if use_pytorch else np.atleast_1d + + # Make scalars 1d arrays and squeezable arrays 1d + alist = [atleast_1d(array_mgr.as_compatible_array(a).squeeze()) for a in onedim_arrs] # Make scalars 1d arrays and squeezable arrays 1d - alist = [np.atleast_1d(np.asarray(a).squeeze()) for a in onedim_arrs] if any(a.ndim != 1 for a in alist): raise ValueError('only 1d arrays allowed') if len(axes) < out.ndim: # Make big factor array (start with 0d) - factor = np.array(1.0) + factor = array_mgr.as_compatible_array([1.0]) for ax, arr in zip(axes, alist): # Meshgrid-style slice slc = [None] * out.ndim @@ -274,11 +287,11 @@ def fast_1d_tensor_mult(ndarr, onedim_arrs, axes=None, out=None): # Get the axis to spare for the final multiplication, the one # with the largest stride. - last_ax = np.argmax(out.strides) + last_ax = out.ndim-1 if use_pytorch else np.argmax(out.strides) last_arr = alist[axes.index(last_ax)] # Build the semi-big array and multiply - factor = np.array(1.0) + factor = array_mgr.as_compatible_array([1.0]) for ax, arr in zip(axes, alist): if ax == last_ax: continue From ec19e0f899a89057c8603ed796561b0e96ce16c6 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 9 Oct 2024 18:33:17 +0200 Subject: [PATCH 40/52] Make some Fourier utils PyTorch-compatible. Using the new array-manager classes. --- odl/trafos/util/ft_utils.py | 48 ++++++++++++++++++++++++++++++------- 1 file changed, 40 insertions(+), 8 deletions(-) diff --git a/odl/trafos/util/ft_utils.py b/odl/trafos/util/ft_utils.py index 2f7c439db55..956288eacf6 100644 --- a/odl/trafos/util/ft_utils.py +++ b/odl/trafos/util/ft_utils.py @@ -11,6 +11,7 @@ from __future__ import absolute_import, division, print_function import numpy as np +import torch from odl.util.npy_compat import AVOID_UNNECESSARY_COPY @@ -18,8 +19,10 @@ DiscretizedSpace, uniform_discr_frompartition, uniform_grid, uniform_partition_fromgrid) from odl.set import RealNumbers +from odl.space.base_tensors import Tensor from odl.util import ( complex_dtype, conj_exponent, dtype_repr, fast_1d_tensor_mult, + uses_pytorch, compatible_array_manager, is_complex_floating_dtype, is_numeric_dtype, is_real_dtype, is_real_floating_dtype, is_string, normalized_axes_tuple, normalized_scalar_param_list) @@ -296,7 +299,20 @@ def dft_preprocess_data(arr, shift=True, axes=None, sign='-', out=None): type and ``shift`` is not ``True``. In this case, the return type is the complex counterpart of ``arr.dtype``. """ - arr = np.asarray(arr) + + use_pytorch = uses_pytorch(arr) + array_mgr = compatible_array_manager(arr) + + if use_pytorch: + assert(out is None or isinstance(out, torch.Tensor)), f"{type(out)=}" + else: + if hasattr(arr, 'impl'): + assert(arr.impl=='numpy'), f"{arr.impl=}" + else: + assert(isinstance(arr, np.ndarray)), f"{type(arr)=}" + assert(out is None or isinstance(out, np.ndarray)), f"{type(out)=}" + + arr = array_mgr.as_compatible_array(arr) if not is_numeric_dtype(arr.dtype): raise ValueError('array has non-numeric data type {}' ''.format(dtype_repr(arr.dtype))) @@ -320,7 +336,7 @@ def dft_preprocess_data(arr, shift=True, axes=None, sign='-', out=None): if is_real_dtype(arr.dtype) and not all(shift_list): out = np.array(arr, dtype=complex_dtype(arr.dtype), copy=True) else: - out = arr.copy() + out = array_mgr.make_copy(arr) else: out[:] = arr @@ -338,13 +354,13 @@ def dft_preprocess_data(arr, shift=True, axes=None, sign='-', out=None): def _onedim_arr(length, shift): if shift: # (-1)^indices - factor = np.ones(length, dtype=out.dtype) + factor = array_mgr.compatible_ones(length, dtype=out.dtype) factor[1::2] = -1 else: - factor = np.arange(length, dtype=out.dtype) + factor = array_mgr.as_compatible_array(np.arange(length), dtype=out.dtype) factor *= -imag * np.pi * (1 - 1.0 / length) np.exp(factor, out=factor) - return factor.astype(out.dtype, copy=AVOID_UNNECESSARY_COPY) + return array_mgr.select_dtype(factor, out.dtype, copy=AVOID_UNNECESSARY_COPY) onedim_arrs = [] for axis, shift in zip(axes, shift_list): @@ -460,7 +476,23 @@ def dft_postprocess_data(arr, real_grid, recip_grid, shift, axes, *Numerical Recipes in C - The Art of Scientific Computing* (Volume 3). Cambridge University Press, 2007. """ - arr = np.asarray(arr) + + use_pytorch = uses_pytorch(arr) + array_mgr = compatible_array_manager(arr) + + if use_pytorch: + assert(out is None or isinstance(out, torch.Tensor)), f"{type(out)=}" + assert(arr.dtype in [torch.float32, torch.float64, torch.complex64, torch.complex128]) + else: + if hasattr(arr, 'impl'): + assert(arr.impl=='numpy'), f"{arr.impl=}" + else: + assert(isinstance(arr, np.ndarray)), f"{type(arr)=}" + assert(out is None or isinstance(out, np.ndarray)), f"{type(out)=}" + assert(arr.dtype in map(np.dtype, ['float32', 'float64', 'float128', + 'complex64', 'complex128', 'complex256'])) + + arr = array_mgr.as_compatible_array(arr) if is_real_floating_dtype(arr.dtype): arr = arr.astype(complex_dtype(arr.dtype)) elif not is_complex_floating_dtype(arr.dtype): @@ -468,7 +500,7 @@ def dft_postprocess_data(arr, real_grid, recip_grid, shift, axes, 'data type'.format(dtype_repr(arr.dtype))) if out is None: - out = arr.copy() + out = array_mgr.make_copy(arr) elif out is not arr: out[:] = arr @@ -542,7 +574,7 @@ def dft_postprocess_data(arr, real_grid, recip_grid, shift, axes, else: onedim_arr /= interp_kernel - onedim_arrs.append(onedim_arr.astype(out.dtype, copy=AVOID_UNNECESSARY_COPY)) + onedim_arrs.append(array_mgr.as_compatible_array(onedim_arr, dtype=out.dtype)) fast_1d_tensor_mult(out, onedim_arrs, axes=axes, out=out) return out From c21ef4d5c97ede430f826e4d80a614280cf7eb22 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 9 Oct 2024 18:52:11 +0200 Subject: [PATCH 41/52] Make Fourier transforms robust towards non-NumPy array storage. --- odl/trafos/fourier.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/odl/trafos/fourier.py b/odl/trafos/fourier.py index a4640655e0c..289c8b1e425 100644 --- a/odl/trafos/fourier.py +++ b/odl/trafos/fourier.py @@ -50,7 +50,7 @@ def _select_fft_impl(impl_suggestion: Optional[str], domain_impl: str): impl, impl_in = str(impl).lower(), impl if impl not in _SUPPORTED_FOURIER_IMPLS.get(domain_impl): raise ValueError(f"`impl` '{impl_in}' not supported for" - + " tensors with {domain_impl} implementation.") + + f" tensors with {domain_impl} implementation.") return impl @@ -108,7 +108,7 @@ def __init__(self, inverse, domain, range=None, axes=None, sign='-', ''.format(range)) # Implementation - self.__impl = _select_fft_impl(impl) + self.__impl = _select_fft_impl(impl, domain.impl) # Axes if axes is None: @@ -141,8 +141,9 @@ def __init__(self, inverse, domain, range=None, axes=None, sign='-', shape = np.atleast_1d(ran_shape) range = uniform_discr( - [0] * len(shape), shape - 1, shape, ran_dtype, domain_impl, - nodes_on_bdry=True, exponent=conj_exponent(domain.exponent)) + [0] * len(shape), shape - 1, shape, ran_dtype, + nodes_on_bdry=True, exponent=conj_exponent(domain.exponent), + impl=domain.impl) else: if range.shape != ran_shape: From 5f4b76b290818c08657ba5d6ca349d12b4be86a0 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 11:25:15 +0200 Subject: [PATCH 42/52] Correct bug in Fourier post-processing. It makes no sense to pass the `out` argument as the input parameter. The `x` argument of the `_postprocess` method was completely unused. Upon investigating why this never caused any problems, I found that in all the unit tests `x is out` held true. In that case both are interchangeable, but this cannot in general be assumed. I can imagine no setting where it would actually be necessary to double-pass `out` this way. As for why the old version used `out` as the input argument: this originated in 4e4e928, where the call to `dft_postprocess_data` was refactored: it had previously been in the `_call_pyfftw` method, where indeed the data was stored in `out` (having come out of `pyfftw_call`). It appears that this call was copy&pasted into the then-new `_postprocess` method, but forgotten to change the argument to `x`. --- odl/trafos/fourier.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/odl/trafos/fourier.py b/odl/trafos/fourier.py index 289c8b1e425..eeebf4b2390 100644 --- a/odl/trafos/fourier.py +++ b/odl/trafos/fourier.py @@ -1307,7 +1307,7 @@ def _postprocess(self, x, out=None): # TODO(kohr-h): Add `interp` to operator or simplify it by not # performing interpolation filter return dft_postprocess_data( - out, real_grid=self.domain.grid, recip_grid=self.range.grid, + x, real_grid=self.domain.grid, recip_grid=self.range.grid, shift=self.shifts, axes=self.axes, sign=self.sign, interp='nearest', op='multiply', out=out) From 1ed2d6b51ebbeb1ba5526bef3ad8302b3e1f8955 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 13:53:47 +0200 Subject: [PATCH 43/52] Add Fourier methods using PyTorch. --- odl/trafos/fourier.py | 153 +++++++++++++++++++++++++++++++++++++++++- 1 file changed, 152 insertions(+), 1 deletion(-) diff --git a/odl/trafos/fourier.py b/odl/trafos/fourier.py index eeebf4b2390..2a1356c1112 100644 --- a/odl/trafos/fourier.py +++ b/odl/trafos/fourier.py @@ -11,6 +11,7 @@ from __future__ import absolute_import, division, print_function import numpy as np +import torch from odl.util.npy_compat import AVOID_UNNECESSARY_COPY @@ -189,6 +190,8 @@ def _call(self, x, out, **kwargs): out[:] = self._call_numpy(x.asarray()) case 'pyfftw': out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + case 'pytorch': + out[:] = self._call_pytorch(x.data) case _: raise NotImplementedError(self.impl) @@ -250,6 +253,21 @@ def _call_numpy(self, x): """ raise NotImplementedError('abstract method') + def _call_pytorch(self, x): + """Return ``self(x)`` for PyTorch back-end. + + Parameters + ---------- + x : `torch.Tensor` + Array representing the function to be transformed + + Returns + ------- + out : `torch.Tensor` + Result of the transform + """ + raise NotImplementedError(f'abstract method, not implemented on {type(self)}.') + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` using pyfftw. @@ -483,6 +501,25 @@ def _call_numpy(self, x): return (np.prod(np.take(self.domain.shape, self.axes)) * np.fft.ifftn(x, axes=self.axes)) + def _call_pytorch(self, x): + """Return ``self(x)`` using PyTorch. + + See Also + -------- + DiscreteFourierTransformBase._call_pytorch + """ + assert isinstance(x, torch.Tensor) + + if self.halfcomplex: + return torch.fft.rfftn(x, dim=self.axes) + else: + if self.sign == '-': + return torch.fft.fftn(x, dim=self.axes) + else: + # Need to undo IFFT scaling + return (np.prod(np.take(self.domain.shape, self.axes)) + * torch.fft.ifftn(x, dim=self.axes)) + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` using pyfftw. @@ -639,6 +676,28 @@ def _call_numpy(self, x): return (np.fft.fftn(x, axes=self.axes) / np.prod(np.take(self.domain.shape, self.axes))) + def _call_pytorch(self, x): + """Return ``self(x)`` using PyTorch. + + Parameters + ---------- + x : `torch.Tensor` + Input array to be transformed + + Returns + ------- + out : `torch.Tensor` + Result of the transform + """ + if self.halfcomplex: + return torch.fft.irfftn(x, dim=self.axes) + else: + if self.sign == '+': + return torch.fft.ifftn(x, dim=self.axes) + else: + return (torch.fft.fftn(x, dim=self.axes) + / np.prod(np.take(self.domain.shape, self.axes))) + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` using pyfftw. @@ -918,10 +977,11 @@ def _call(self, x, out, **kwargs): case 'pyfftw': # 0-overhead assignment if asarray() does not copy out[:] = self._call_pyfftw(x.asarray(), out.asarray(), **kwargs) + case 'pytorch': + out[:] = self._call_pytorch(x.asarray(), **kwargs) case _: raise NotImplementedError(self.impl) - def _call_numpy(self, x): """Return ``self(x)`` for numpy back-end. @@ -937,6 +997,21 @@ def _call_numpy(self, x): """ raise NotImplementedError('abstract method') + def _call_pytorch(self, x): + """Return ``self(x)`` for PyTorch back-end. + + Parameters + ---------- + x : `torch.Tensor` + Array representing the function to be transformed + + Returns + ------- + out : `torch.Tensor` + Result of the transform + """ + raise NotImplementedError(f'abstract method, not implemented on {type(self)}.') + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` for pyfftw back-end. @@ -1349,6 +1424,42 @@ def _call_numpy(self, x): self._postprocess(out, out=out) return out + def _call_pytorch(self, x): + """Return ``self(x)`` for PyTorch back-end. + + Parameters + ---------- + x : `torch.Tensor` + Array representing the function to be transformed + + Returns + ------- + out : `torch.Tensor` + Result of the transform + """ + + preproc = self._preprocess(x) + + # The actual call to the FFT library + if self.halfcomplex: + out = torch.fft.rfftn(preproc, dim=self.axes) + else: + if self.sign == '-': + out = torch.fft.fftn(preproc, dim=self.axes) + else: + out = torch.fft.ifftn(preproc, dim=self.axes) + # Numpy's FFT normalizes by 1 / prod(shape[axes]), we + # need to undo that + # TODO(Justus) select PyTorch normalization mode so this + # is unnecessary + out *= float(np.prod(np.take(self.domain.shape, self.axes))) + + # Post-processing accounting for shift, scaling and interpolation + assert(isinstance(out, torch.Tensor)) + out = self._postprocess(out, out=out) + assert(isinstance(out, torch.Tensor)) + return out + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` for pyfftw back-end. @@ -1594,6 +1705,46 @@ def _call_numpy(self, x): else: return out + def _call_pytorch(self, x): + """Return ``self(x)`` for numpy back-end. + + Parameters + ---------- + x : `torch.Tensor` + Array representing the function to be transformed + + Returns + ------- + out : `torch.Tensor` + Result of the transform + """ + # Pre-processing before calculating the DFT + preproc = self._preprocess(x) + + # The actual call to the FFT library + # Normalization by 1 / prod(shape[axes]) is done by Numpy's FFT if + # one of the "i" functions is used. For sign='-' we need to do it + # ourselves. + if self.halfcomplex: + s = tuple(np.asarray(self.range.shape)[list(self.axes)]) + out = torch.fft.irfftn(preproc, dim=self.axes, s=s) + else: + if self.sign == '-': + out = torch.fft.fftn(preproc, dim=self.axes) + out /= np.prod(np.take(self.domain.shape, self.axes)) + else: + out = torch.fft.ifftn(preproc, dim=self.axes) + + # Post-processing in IFT = pre-processing in FT (in-place) + out = self._postprocess(out) + if self.halfcomplex: + assert is_real_dtype(out.dtype) + + if self.range.field == RealNumbers(): + return out.real + else: + return out + def _call_pyfftw(self, x, out, **kwargs): """Implement ``self(x[, out, **kwargs])`` for pyfftw back-end. From 068dbc7c20ab7c3beb1d3fc93d4a768577fe4d52 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 14:08:04 +0200 Subject: [PATCH 44/52] Support the PyTorch-based Fourier transforms. --- odl/trafos/fourier.py | 15 +++++++++------ 1 file changed, 9 insertions(+), 6 deletions(-) diff --git a/odl/trafos/fourier.py b/odl/trafos/fourier.py index 2a1356c1112..36a77ad0a39 100644 --- a/odl/trafos/fourier.py +++ b/odl/trafos/fourier.py @@ -33,8 +33,8 @@ 'FourierTransform', 'FourierTransformInverse') -_SUPPORTED_FOURIER_IMPLS = {'numpy': ('numpy',)} -_DEFAULT_FOURIER_IMPL = {'numpy': 'numpy'} +_SUPPORTED_FOURIER_IMPLS = {'numpy': ('numpy',), 'pytorch': ('pytorch',)} +_DEFAULT_FOURIER_IMPL = {'numpy': 'numpy', 'pytorch': 'pytorch'} if PYFFTW_AVAILABLE: _SUPPORTED_FOURIER_IMPLS['numpy'] += ('pyfftw',) _DEFAULT_FOURIER_IMPL['numpy'] = 'pyfftw' @@ -814,9 +814,12 @@ def __init__(self, inverse, domain, range=None, impl=None, **kwargs): is determined from ``domain`` and the other parameters. The exponent is chosen to be the conjugate ``p / (p - 1)``, which reads as 'inf' for p=1 and 1 for p='inf'. - impl : {'numpy', 'pyfftw'}, optional - Backend for the FFT implementation. The 'pyfftw' backend - is faster but requires the ``pyfftw`` package. + impl : {'numpy', 'pyfftw', 'pytorch'}, optional + Backend for the FFT implementation. NumPy is slow but always + supported. The 'pyfftw' backend is faster but requires the + ``pyfftw`` package. + 'pytorch' requires ``domain`` to be based on PyTorch tensors, + in which case this is the fastest option particularly on GPU. ``None`` selects the fastest available backend. axes : int or sequence of ints, optional Dimensions along which to take the transform. @@ -1655,7 +1658,7 @@ def _postprocess(self, x, out=None): The result is stored in ``out`` if given, otherwise in a temporary or a new array. """ - if out is None: + if out is None and self.impl!='pytorch': if self.range.field == ComplexNumbers(): out = self._tmp_r if self._tmp_r is not None else self._tmp_f elif self.range.field == RealNumbers() and not self.halfcomplex: From c260a19598a19477a12d5a646a83229bacdaeb49 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 14:18:27 +0200 Subject: [PATCH 45/52] Generalize Fourier tests to support arrays other than NumPy. --- odl/test/trafos/fourier_test.py | 45 ++++++++++++++++++--------------- 1 file changed, 24 insertions(+), 21 deletions(-) diff --git a/odl/test/trafos/fourier_test.py b/odl/test/trafos/fourier_test.py index 45bf60c5993..bd02cfc4317 100644 --- a/odl/test/trafos/fourier_test.py +++ b/odl/test/trafos/fourier_test.py @@ -45,8 +45,10 @@ def _params_from_dtype(dtype): halfcomplex = False return halfcomplex, complex_dtype(dtype) +def _dft_domain_impl(impl): + 'numpy' -def _dft_space(shape, dtype='float64'): +def _dft_space(shape, dtype='float64', impl='numpy'): try: ndim = len(shape) except TypeError: @@ -57,6 +59,7 @@ def _dft_space(shape, dtype='float64'): shape, dtype=dtype, nodes_on_bdry=True, + impl = _dft_domain_impl(impl) ) @@ -71,12 +74,12 @@ def sinc(x): def test_dft_init(impl): # Just check if the code runs at all shape = (4, 5) - dom = _dft_space(shape) - dom_nonseq = odl.uniform_discr([0, 0], [1, 1], shape) + dom = _dft_space(shape, impl=impl) + dom_nonseq = odl.uniform_discr([0, 0], [1, 1], shape, impl=impl) dom_f32 = dom.astype('float32') - ran = _dft_space(shape, dtype='complex128') + ran = _dft_space(shape, dtype='complex128', impl=impl) ran_c64 = ran.astype('complex64') - ran_hc = _dft_space((3, 5), dtype='complex128') + ran_hc = _dft_space((3, 5), dtype='complex128', impl=impl) # Implicit range DiscreteFourierTransform(dom, impl=impl) @@ -191,10 +194,10 @@ def test_idft_init(impl): # Just check if the code runs at all; this uses the init function of # DiscreteFourierTransform, so we don't need exhaustive tests here shape = (4, 5) - ran = _dft_space(shape, dtype='complex128') - ran_hc = _dft_space(shape, dtype='float64') - dom = _dft_space(shape, dtype='complex128') - dom_hc = _dft_space((3, 5), dtype='complex128') + ran = _dft_space(shape, dtype='complex128', impl=impl) + ran_hc = _dft_space(shape, dtype='float64', impl=impl) + dom = _dft_space(shape, dtype='complex128', impl=impl) + dom_hc = _dft_space((3, 5), dtype='complex128', impl=impl) # Implicit range DiscreteFourierTransformInverse(dom, impl=impl) @@ -209,7 +212,7 @@ def test_dft_call(impl): # 2d, complex, all ones and random back & forth shape = (4, 5) - dft_dom = _dft_space(shape, dtype='complex64') + dft_dom = _dft_space(shape, dtype='complex64', impl=impl) dft = DiscreteFourierTransform(domain=dft_dom, impl=impl) idft = DiscreteFourierTransformInverse(range=dft_dom, impl=impl) @@ -243,7 +246,7 @@ def test_dft_call(impl): # 2d, halfcomplex, first axis shape = (4, 5) axes = 0 - dft_dom = _dft_space(shape, dtype='float32') + dft_dom = _dft_space(shape, dtype='float32', impl=impl) dft = DiscreteFourierTransform(domain=dft_dom, impl=impl, halfcomplex=True, axes=axes) idft = DiscreteFourierTransformInverse(range=dft_dom, impl=impl, @@ -276,7 +279,7 @@ def test_dft_sign(impl): # 2d, complex, all ones and random back & forth shape = (4, 5) - dft_dom = _dft_space(shape, dtype='complex64') + dft_dom = _dft_space(shape, dtype='complex64', impl=impl) dft_minus = DiscreteFourierTransform(domain=dft_dom, impl=impl, sign='-') dft_plus = DiscreteFourierTransform(domain=dft_dom, impl=impl, sign='+') @@ -297,7 +300,7 @@ def test_dft_sign(impl): # 2d, halfcomplex, first axis shape = (4, 5) axes = (0,) - dft_dom = _dft_space(shape, dtype='float32') + dft_dom = _dft_space(shape, dtype='float32', impl=impl) arr = dft_dom.element([[0, 0, 0, 0, 0], [0, 0, 1, 1, 0], [0, 0, 1, 1, 0], @@ -321,7 +324,7 @@ def test_dft_init_plan(impl): # 2d, halfcomplex, first axis shape = (4, 5) axes = 0 - dft_dom = _dft_space(shape, dtype='float32') + dft_dom = _dft_space(shape, dtype='float32', impl=impl) dft = DiscreteFourierTransform(dft_dom, impl=impl, axes=axes, halfcomplex=True) @@ -399,7 +402,7 @@ def test_fourier_trafo_init_plan(impl, odl_floating_dtype): shape = 10 halfcomplex, _ = _params_from_dtype(dtype) - space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype) + space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype, impl=_dft_domain_impl(impl)) ft = FourierTransform(space_discr, impl=impl, halfcomplex=halfcomplex) if impl != 'pyfftw': @@ -477,7 +480,7 @@ def test_fourier_trafo_call(impl, odl_floating_dtype): shape = 10 halfcomplex, _ = _params_from_dtype(dtype) - space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype) + space_discr = odl.uniform_discr(0, 1, shape, dtype=dtype, impl=_dft_domain_impl(impl)) ft = FourierTransform(space_discr, impl=impl, halfcomplex=halfcomplex) ift = ft.inverse @@ -548,7 +551,7 @@ def test_fourier_trafo_sign(impl, odl_real_floating_dtype): def char_interval(x): return (x >= 0) & (x <= 1) - discr = odl.uniform_discr(-2, 2, 40, impl='numpy', dtype=discrspace_dtype) + discr = odl.uniform_discr(-2, 2, 40, impl=_dft_domain_impl(impl), dtype=discrspace_dtype) ft_minus = FourierTransform(discr, sign='-', impl=impl) ft_plus = FourierTransform(discr, sign='+', impl=impl) @@ -593,7 +596,7 @@ def char_interval(x): return (x >= 0) & (x <= 1) # Complex-to-complex - discr = odl.uniform_discr(-2, 2, 40, impl='numpy', dtype='complex64') + discr = odl.uniform_discr(-2, 2, 40, impl=_dft_domain_impl(impl), dtype='complex64') discr_char = discr.element(char_interval) ft = FourierTransform(discr, sign=sign, impl=impl) @@ -601,7 +604,7 @@ def char_interval(x): assert all_almost_equal(ft.adjoint(ft(char_interval)), discr_char) # Half-complex - discr = odl.uniform_discr(-2, 2, 40, impl='numpy', dtype='float32') + discr = odl.uniform_discr(-2, 2, 40, impl=_dft_domain_impl(impl), dtype='float32') ft = FourierTransform(discr, impl=impl, halfcomplex=True) assert all_almost_equal(ft.inverse(ft(char_interval)), discr_char) @@ -609,7 +612,7 @@ def char_rect(x): return (x[0] >= 0) & (x[0] <= 1) & (x[1] >= 0) & (x[1] <= 1) # 2D with axes, C2C - discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl='numpy', + discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl=_dft_domain_impl(impl), dtype='complex64') discr_rect = discr.element(char_rect) @@ -619,7 +622,7 @@ def char_rect(x): assert all_almost_equal(ft.adjoint(ft(char_rect)), discr_rect) # 2D with axes, halfcomplex - discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl='numpy', + discr = odl.uniform_discr([-2, -2], [2, 2], (20, 10), impl=_dft_domain_impl(impl), dtype='float32') discr_rect = discr.element(char_rect) From 37dbb761490b9ac099f7455b2eb107b2c4ffeddd Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 14:27:47 +0200 Subject: [PATCH 46/52] Add PyTorch to Fourier unit tests. --- odl/test/trafos/fourier_test.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/odl/test/trafos/fourier_test.py b/odl/test/trafos/fourier_test.py index bd02cfc4317..9bdec0f175a 100644 --- a/odl/test/trafos/fourier_test.py +++ b/odl/test/trafos/fourier_test.py @@ -29,6 +29,7 @@ impl = simple_fixture( 'impl', [pytest.param('numpy'), + pytest.param('pytorch'), pytest.param('pyfftw', marks=skip_if_no_pyfftw)] ) exponent = simple_fixture('exponent', [2.0, 1.0, float('inf'), 1.5]) @@ -46,7 +47,7 @@ def _params_from_dtype(dtype): return halfcomplex, complex_dtype(dtype) def _dft_domain_impl(impl): - 'numpy' + return 'pytorch' if impl=='pytorch' else 'numpy' def _dft_space(shape, dtype='float64', impl='numpy'): try: @@ -398,6 +399,9 @@ def test_fourier_trafo_init_plan(impl, odl_floating_dtype): # Not supported, skip if dtype == np.dtype('float16') and impl == 'pyfftw': return + elif (dtype in [np.dtype('float128'), np.dtype('complex256')] + and impl == 'pytorch'): + return shape = 10 halfcomplex, _ = _params_from_dtype(dtype) @@ -477,6 +481,9 @@ def test_fourier_trafo_call(impl, odl_floating_dtype): # Not supported, skip if dtype == np.dtype('float16') and impl == 'pyfftw': return + elif (dtype in [np.dtype('float16'), np.dtype('float128'), np.dtype('complex256')] + and impl == 'pytorch'): + return shape = 10 halfcomplex, _ = _params_from_dtype(dtype) From 74c75a55fe23b38d29562257eaba449c56633fc4 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 14 Oct 2024 14:33:24 +0200 Subject: [PATCH 47/52] Add half-precision dtypes for PyTorch. --- odl/space/pytorch_tensors.py | 3 ++- odl/util/utility.py | 4 +++- 2 files changed, 5 insertions(+), 2 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index b159fabbad7..48740e1c8b4 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -488,7 +488,8 @@ def available_dtypes(): Currently only a conservative selection of the types supported by Pytorch. """ - return [np.float32, np.float64, np.complex64, np.complex128] + return [np.float16, np.float32, np.float64, + np.complex64, np.complex128] @staticmethod def default_dtype(field=None): diff --git a/odl/util/utility.py b/odl/util/utility.py index 9dace49b21e..355b00e99c6 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -543,7 +543,9 @@ def complex_dtype(dtype, default=None): else: return np.dtype((complex_base_dtype, dtype.shape)) -_CORRESPONDING_PYTORCH_DTYPES = {np.dtype('float32'): torch.float32, +_CORRESPONDING_PYTORCH_DTYPES = { + np.dtype('float16'): torch.float16, + np.dtype('float32'): torch.float32, np.dtype('float64'): torch.float64, np.dtype('complex64'): torch.complex64, np.dtype('complex128'): torch.complex128} From 8c0dad6958f5502088aa73afb939d1295ad90419 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 21 Oct 2024 15:13:43 +0200 Subject: [PATCH 48/52] Avoid Torch warning/error messages when managing arrays that are already on Torch. The old `torch.tensor` converter did just the right thing: construct a new tensor when given lists or NumPy arrays, and simply clone the input if it was already a PyTorch tensor. For some reason, Torch has deprecated this behaviour, so it is now necessary to hard-code the different possibilities. --- odl/util/utility.py | 12 +++++++++++- 1 file changed, 11 insertions(+), 1 deletion(-) diff --git a/odl/util/utility.py b/odl/util/utility.py index 355b00e99c6..96c126865bf 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -616,7 +616,17 @@ class ArrayOnPytorchManager(ABC): def __init__(self, device): self._device = device def as_compatible_array(self, arr, **kwargs): - return torch.tensor(arr, device = self._device, **kwargs) + dtype = kwargs.get('dtype', None) + if isinstance(arr, torch.Tensor): + arr = arr.detach() + if dtype is not None and arr.dtype!=kwargs['dtype']: + arr = arr.type(dtype) + if self._device is not None and arr.device!=self._device: + return arr.to(self._device) + else: + return arr + else: + return torch.tensor(arr, device = self._device, **kwargs) def compatible_zeros(self, shape, **kwargs): return torch.zeros(shape, device = self._device, **kwargs) def compatible_ones(self, shape, **kwargs): From 2a0af48fbf39f0fee16a307ebdf971bec2d5af15 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Mon, 21 Oct 2024 16:11:32 +0200 Subject: [PATCH 49/52] Use the `ArrayOnBackendManager` classes for generating PyTorch-based TensorSpace objects. This removes the necessity for some redundant checks/conversions. Backend-specific arrays are used in several different places; TensorSpace-element construction is only one of them. --- odl/space/pytorch_tensors.py | 24 +++++++++++++----------- 1 file changed, 13 insertions(+), 11 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 48740e1c8b4..6b31a747f89 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -24,7 +24,7 @@ from odl.space.weighting import ( ArrayWeighting, ConstWeighting, CustomDist, CustomInner, CustomNorm, Weighting) -from odl.util.utility import _CORRESPONDING_PYTORCH_DTYPES +from odl.util.utility import ArrayOnPytorchManager, _CORRESPONDING_PYTORCH_DTYPES from odl.util import ( dtype_str, is_floating_dtype, is_numeric_dtype, is_real_dtype, nullcontext, signature_string, writable_array) @@ -422,11 +422,16 @@ def element(self, inp=None, data_ptr=None, order=None): if order is not None and str(order).upper() not in ('C'): raise ValueError(f"Only row-major order supported ('C'), not '{order}'.") - if inp is None and data_ptr is None: - arr = torch.empty(self.shape, dtype=self._torch_dtype, device=self._torch_device) - + def wrapped_array(arr): + if arr.shape != self.shape: + raise ValueError('shape of `inp` not equal to space shape: ' + '{} != {}'.format(arr.shape, self.shape)) return self.element_type(self, arr) + if inp is None and data_ptr is None: + return wrapped_array(torch.empty( + self.shape, dtype=self._torch_dtype, device=self._torch_device)) + elif inp is None and data_ptr is not None: if order is None: raise ValueError('`order` cannot be None for element ' @@ -437,7 +442,8 @@ def element(self, inp=None, data_ptr=None, order=None): as_numpy_array = np.ctypeslib.as_array(as_ctype_array) arr = as_numpy_array.view(dtype=self._torch_dtype) arr = arr.reshape(self.shape, order=order) - return self.element_type(self, torch.Tensor(arr)) + return wrapped_array(torch.tensor( + arr, dtype=self._torch_dtype, device=self._torch_device)) elif inp is not None and data_ptr is None: if inp in self and order is None: @@ -445,12 +451,8 @@ def element(self, inp=None, data_ptr=None, order=None): return inp # TODO avoid copy when it's not necessary - arr = torch.tensor(inp, dtype=self._torch_dtype, device=self._torch_device) - - if arr.shape != self.shape: - raise ValueError('shape of `inp` not equal to space shape: ' - '{} != {}'.format(arr.shape, self.shape)) - return self.element_type(self, arr) + return wrapped_array(ArrayOnPytorchManager(device=self._torch_device) + .as_compatible_array(inp, dtype=self._torch_dtype)) else: raise TypeError('cannot provide both `inp` and `data_ptr`') From 42b07bcd43f9b7316709f9e3c1b789fb2dc2f58d Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Tue, 3 Dec 2024 19:07:21 +0100 Subject: [PATCH 50/52] Implemented the methods related to in- vs out-of-place selection for PyTorch spaces. --- odl/space/pytorch_tensors.py | 35 +++++++++++++++++++++++++++++++++-- 1 file changed, 33 insertions(+), 2 deletions(-) diff --git a/odl/space/pytorch_tensors.py b/odl/space/pytorch_tensors.py index 6b31a747f89..4aa1f3f8d08 100644 --- a/odl/space/pytorch_tensors.py +++ b/odl/space/pytorch_tensors.py @@ -19,7 +19,8 @@ import torch from odl.set.sets import ComplexNumbers, RealNumbers -from odl.set.space import LinearSpaceTypeError +from odl.set.space import (LinearSpaceTypeError, + NumOperationParadigmSupport, SupportedNumOperationParadigms) from odl.space.base_tensors import Tensor, TensorSpace from odl.space.weighting import ( ArrayWeighting, ConstWeighting, CustomDist, CustomInner, CustomNorm, @@ -285,6 +286,8 @@ def __init__(self, shape, dtype=None, **kwargs): # No weighting, i.e., weighting with constant 1.0 self.__weighting = PytorchTensorSpaceConstWeighting(1.0, exponent) + self._use_in_place_ops = kwargs.pop('use_in_place_ops', True) + # Make sure there are no leftover kwargs if kwargs: raise TypeError('got unknown keyword arguments {}'.format(kwargs)) @@ -294,6 +297,22 @@ def impl(self): """Name of the implementation back-end: ``'pytorch'``.""" return 'pytorch' + @property + def supported_num_operation_paradigms(self) -> NumOperationParadigmSupport: + """PyTorch supports both in-place and out of place operations, but the + former are problematic especially when automatic differentiation is + used: PyTorch needs to ensure the modification does not interfere with + the backwards pass. This makes the performance much worse than for the + out-of-place style.""" + if self._use_in_place_ops: + return SupportedNumOperationParadigms( + in_place = NumOperationParadigmSupport.SUPPORTED, + out_of_place = NumOperationParadigmSupport.PREFERRED) + else: + return SupportedNumOperationParadigms( + in_place = NumOperationParadigmSupport.NOT_SUPPORTED, + out_of_place = NumOperationParadigmSupport.PREFERRED) + @property def default_order(self): """Default (and only) storage order for new elements in this space: ``'C'``.""" @@ -554,7 +573,11 @@ def _lincomb(self, a, x1, b, x2, out): >>> result is out True """ - torch.add(input=a*x1.data, other=x2.data, alpha=b, out=out.data) + if self._use_in_place_ops and out is not None: + torch.add(input=a*x1.data, other=x2.data, alpha=b, out=out.data) + else: + assert(out is None) + return self.element(a * x1.data + b * x2.data) def _dist(self, x1, x2): """Return the distance between ``x1`` and ``x2``. @@ -873,6 +896,14 @@ def data(self): """The `torch.Tensor` representing the data of ``self``.""" return self.__data + def _assign(self, other, avoid_deep_copy): + """Assign the values of ``other``, which is assumed to be in the + same space, to ``self``.""" + if avoid_deep_copy or not self.space._use_in_place_ops: + self.__data = other.__data + else: + self.__data[:] = other.__data + def asarray(self, out=None): """Extract the data of this array as a ``torch.Tensor``. From a6ace7e6fd62f01c4747f79e93053dbbe100f4f3 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 30 Apr 2025 15:22:17 +0200 Subject: [PATCH 51/52] Add a `copy` argument to the array-manager's dtype-switching method. --- odl/util/utility.py | 16 +++++++++++----- 1 file changed, 11 insertions(+), 5 deletions(-) diff --git a/odl/util/utility.py b/odl/util/utility.py index 96c126865bf..4ff17d55d14 100644 --- a/odl/util/utility.py +++ b/odl/util/utility.py @@ -16,6 +16,7 @@ from contextlib import contextmanager from itertools import product from abc import ABC +from typing import Optional import numpy as np import torch @@ -607,7 +608,7 @@ def compatible_zeros(self, shape, **kwargs): raise NotImplementedError() def compatible_ones(self, shape, **kwargs): raise NotImplementedError() - def select_dtype(self, arr, dtype): + def select_dtype(self, arr, dtype, copy: Optional[bool]): raise NotImplementedError() def make_copy(self, arr): raise NotImplementedError() @@ -631,10 +632,15 @@ def compatible_zeros(self, shape, **kwargs): return torch.zeros(shape, device = self._device, **kwargs) def compatible_ones(self, shape, **kwargs): return torch.ones(shape, device = self._device, **kwargs) - def select_dtype(self, arr, dtype): + def select_dtype(self, arr, dtype, copy): if dtype in _CORRESPONDING_PYTORCH_DTYPES: dtype = _CORRESPONDING_PYTORCH_DTYPES[dtype] - return arr.type(dtype) + # PyTorch (as of version 2.7) only supports the values False and True + # for the `copy` argument, the former being a non-binding request to + # avoid a copy if it is not necessary. + if copy==AVOID_UNNECESSARY_COPY: + copy = False + return arr.type(dtype, copy=copy) def make_copy(self, arr): return arr.clone().detach() @@ -647,8 +653,8 @@ def compatible_zeros(self, shape, **kwargs): return np.zeros(shape, **kwargs) def compatible_ones(self, shape, **kwargs): return np.ones(shape, **kwargs) - def select_dtype(self, arr, dtype): - return arr.astype(dtype, copy=False) + def select_dtype(self, arr, dtype, copy): + return arr.astype(dtype, copy=copy) def make_copy(self, arr): return arr.copy() From 58355f001357c787b8b380b61a47fbe82131e5de Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Justus=20Sagem=C3=BCller?= Date: Wed, 30 Apr 2025 15:42:51 +0200 Subject: [PATCH 52/52] Start adapting the tensor space tests for PyTorch. --- odl/test/space/tensors_test.py | 17 +++++++++++++++++ 1 file changed, 17 insertions(+) diff --git a/odl/test/space/tensors_test.py b/odl/test/space/tensors_test.py index e722d29303e..645f0fdb58d 100644 --- a/odl/test/space/tensors_test.py +++ b/odl/test/space/tensors_test.py @@ -22,6 +22,10 @@ NumpyTensor, NumpyTensorSpace, NumpyTensorSpaceArrayWeighting, NumpyTensorSpaceConstWeighting, NumpyTensorSpaceCustomDist, NumpyTensorSpaceCustomInner, NumpyTensorSpaceCustomNorm) +from odl.space.pytorch_tensors import ( + PytorchTensor, PytorchTensorSpace, PytorchTensorSpaceArrayWeighting, + PytorchTensorSpaceConstWeighting, PytorchTensorSpaceCustomDist, + PytorchTensorSpaceCustomInner, PytorchTensorSpaceCustomNorm) from odl.util.testutils import ( all_almost_equal, all_equal, noise_array, noise_element, noise_elements, simple_fixture) @@ -72,6 +76,19 @@ def _weighting_cls(impl, kind): return NumpyTensorSpaceCustomDist else: assert False + elif impl == 'pytorch': + if kind == 'array': + return PytorchTensorSpaceArrayWeighting + elif kind == 'const': + return PytorchTensorSpaceConstWeighting + elif kind == 'inner': + return PytorchTensorSpaceCustomInner + elif kind == 'norm': + return PytorchTensorSpaceCustomNorm + elif kind == 'dist': + return PytorchTensorSpaceCustomDist + else: + assert False else: assert False