.. _numpy_in_depth:

######################################
Using ODL with NumPy, SciPy or PyTorch
######################################

`NumPy <http://www.numpy.org/>`_ is the traditional library for array computations in Python, and is still used by most major numerical packages.
It provides optimized `Array objects <http://docs.scipy.org/doc/numpy/reference/arrays.html>`_ that allow efficient storage of large arrays.
It also provides several optimized algorithms for many of the functions used in numerical programming, such as taking the cosine or adding two arrays.

`SciPy <http://www.scipy.org/>`_ is a library built on top of NumPy providing more advanced algorithms such as linear solvers, statistics, signal and image processing etc.

`PyTorch <https://pytorch.org/>`_ is best known as a deep learning framework, but also useful as a general-purpose, GPU-accelerated array library.

Many operations are more naturally performed using one of those libraries than with ODL, and with that in mind ODL has been designed such that interfacing with them is as easy and fast as possible.

Casting vectors to and from arrays
==================================
ODL vectors are stored in an abstract way, enabling storage on the CPU, GPU, using different backends which can be switched using the ``impl`` argument when declaring the space.
This allows algorithms to be written in a generalized and storage-agnostic manner.
Still, it is often convenient to be able to access the raw data either for inspection or manipulation, perhaps to initialize a vector, or to call an external function.

To cast a NumPy array to an element of an ODL vector space, one can simply call the `LinearSpace.element` method in an appropriate space::

   >>> import odl
   >>> import numpy as np
   >>> r3 = odl.rn(3)
   >>> arr = np.array([1, 2, 3])
   >>> x = r3.element(arr)
   >>> x
   rn(3).element([ 1., 2., 3. ])

``element`` works not only for NumPy arrays, but also for raw arrays of any library supporting the DLPack standard.

   >>> import torch
   >>> x_t = r3.element(torch.tensor([4, 5, 6]))
   >>> x_t
   rn(3).element([ 4., 5., 6. ])

This element will still internally be stored using NumPy: storage is determined by the space.

   >>> type(x_t.element)
   <class 'numpy.ndarray'>

To store in PyTorch instead, only the space declaration needs to be modified, by the ``impl`` argument (whose default is ``'numpy'``). Again, it is then possible to generate elements from any source:

   >>> r3_t = odl.rn(3, impl='pytorch')
   >>> type(r3_t.element(arr).data)
   <class 'torch.Tensor'>

.. note::
  Relying on the automatic copying of the `LinearSpace.element` method is not necessarily a good idea: for one thing, DLPack support is still somewhat inconsistent in PyTorch as of 2025; for another, it circumvents the device-preserving policy of ODL (i.e. it will in general incur copying of data between different devices, which can take considerable time).
  As a rule of thumb, you should only at the start of a computation declare spaces and call `element` on them. Inside of your algorithms' loops, you should instead use already defined spaces and elements, and modify them with ODL operators.

The other way around, casting ODL vector space elements to NumPy arrays can be done through the member function `Tensor.asarray`. This returns a view if possible::

   >>> x.asarray()
   array([ 1.,  2.,  3.])

`Tensor.asarray` only yields a NumPy array if the space has ``impl='numpy'``.
If for example ``impl='pytorch'``, it gives a `torch.Tensor` instead.

   >>> r3_t.element(arr).asarray()
   tensor([1., 2., 3.], dtype=torch.float64)

.. note::
  For simple ℝⁿ spaces, instead of `asarray` one can also access the `data` attribute directly. That is not recommended for user code, though.

These methods work with any ODL object represented by an array.
For example, in discretizations, a two-dimensional array can be used::

   >>> space = odl.uniform_discr([0, 0], [1, 1], shape=(3, 3))
   >>> arr = np.array([[1, 2, 3],
   ...                 [4, 5, 6],
   ...                 [7, 8, 9]])
   >>> x = space.element(arr)
   >>> x.asarray()
   array([[ 1.,  2.,  3.],
          [ 4.,  5.,  6.],
          [ 7.,  8.,  9.]])

Using ODL objects with array-based functions
============================================
Although ODL offers its own interface to formulate mathematical algorithms (which we recommend using), there are situations where one needs to manipulate objects on the raw array level.

.. note::
  ODL versions 0.7 and 0.8 allowed directly applying NumPy ufuncs to ODL objects.
   This is not allowed anymore in ODL 1.x, since the ufunc compatibility mechanism interfered with high-performance support for other backends.

..
  TODO link to migration guide.

Apart from unwrapping the contained arrays and `.element`-wrapping their modified versions again (see above), there is also the option to modify as space element in-place using some NumPy function (or any function defined on backend-specific arrays). For this purpose we have the `writable_array` context manager that exposes a raw array which gets automatically assigned back to the ODL object::

    >>> x = odl.rn(3).element([1,2,3])
    >>> with odl.util.writable_array(x) as x_arr:
    ...     np.cumsum(x_arr, out=x_arr)
    >>> x
    rn(3).element([ 1.,  3.,  6.])

.. note::
    The re-assignment is a no-op if ``x`` has a single array as its data container, hence the operation will be as fast as manipulating ``x`` directly.
    The same syntax also works with other data containers, but in this case, copies are usually necessary.


NumPy functions as Operators
============================
It is often useful to write an `Operator` wrapping NumPy or other low-level functions, thus allowing full access to the ODL ecosystem.
The convolution operation, written as ODL operator, could look like this::

   >>> class MyConvolution(odl.Operator):
   ...     """Operator for convolving with a given kernel."""
   ...
   ...     def __init__(self, kernel):
   ...         """Initialize the convolution."""
   ...         self.kernel = kernel
   ...
   ...         # Initialize operator base class.
   ...         # This operator maps from the space of vector to the same space and is linear
   ...         super(MyConvolution, self).__init__(
   ...             domain=kernel.space, range=kernel.space, linear=True)
   ...
   ...     def _call(self, x):
   ...         # The output of an Operator is automatically cast to an ODL object
   ...         return self.range.element(np.convolve(x.asarray(), self.kernel.asarray(), mode='same'))

This operator can then be called on its domain elements::

   >>> kernel = odl.rn(3).element([1, 2, 1])
   >>> conv_op = MyConvolution(kernel)
   >>> conv_op([1, 2, 3])
   rn(3).element([ 4.,  8.,  8.])

N.B. the input list `[1,2,3]` is automatically wrapped into `conv_op.domain.element` by the `Operator` base class before the low-level call; in production code it is recommended to do this explicitly for better control.

Such operators can also be used with any of the ODL operator functionalities such as multiplication with scalar, composition, etc::

   >>> scaled_op = 2 * conv_op  # scale output by 2
   >>> scaled_op([1, 2, 3])
   rn(3).element([  8.,  16.,  16.])
   >>> y = odl.rn(3).element([1, 1, 1])
   >>> inner_product_op = odl.InnerProductOperator(y)
   >>> # Create composition with inner product operator with [1, 1, 1].
   >>> # When called on a vector, the result should be the sum of the
   >>> # convolved vector.
   >>> composed_op = inner_product_op @ conv_op
   >>> composed_op([1, 2, 3])
   20.0

For more information on ODL Operators, how to implement them and their features, see the guide on `operators_in_depth`.

Using ODL with SciPy linear solvers
===================================
SciPy includes `a series of very competent solvers <http://docs.scipy.org/doc/scipy/reference/sparse.linalg.html>`_ that may be useful in solving some linear problems.
If you have invested some effort into writing an ODL operator, or perhaps wish to use a pre-existing operator, then the function `as_scipy_operator` creates a Python object that can be used in SciPy's linear solvers.
Here is a simple example of solving Poisson's equation :math:`- \Delta u = f` on the interval :math:`[0, 1]`::

   >>> space = odl.uniform_discr(0, 1, 5)
   >>> op = -odl.Laplacian(space)
   >>> f = space.element(lambda x: (x > 0.4) & (x < 0.6))  # indicator function on [0.4, 0.6]
   >>> u, status = scipy.sparse.linalg.cg(odl.as_scipy_operator(op), f.asarray())
   >>> u
   array([ 0.02,  0.04,  0.06,  0.04,  0.02])

Of course, this also could (and should!) be done with ODL's own version of the solver:

   >>> x = op.domain.element()
   >>> odl.solvers.conjugate_gradient(op=op, x=x, rhs=f, niter=100)
   >>> x
   uniform_discr(0.0, 1.0, 5).element([ 0.02,  0.04,  0.06,  0.04,  0.02])

.. _adapteroperatorguide:

Switching between array backends
===================================

Some operators require low-level functions that are only available on CPU.
Others are implemented in e.g. PyTorch and may only be practical to evaluate on GPU.
Although there is usually a way to convert between these representations, the required data-copying takes time and ODL does deliberately not do this automatically.
Instead, composing two operators that act on different devices or -backends gives an immediate error message.

   >>> op_np = odl.MultiplyOperator(odl.rn(4, impl='numpy').element([1,2,3,4]))
   >>> op_pt = odl.MultiplyOperator(odl.rn(4, impl='pytorch').element([10,20,30,40]))
   >>> op_np @ op_pt
   Traceback (most recent call last):
     File "<stdin>", line 1, in <module>
     File "/home/justussa/progwrit/python/odl/odl/core/operator/operator.py", line 861, in __matmul__
       return self.__mul__(other)
              ^^^^^^^^^^^^^^^^^^^
     File "/home/justussa/progwrit/python/odl/odl/core/operator/operator.py", line 781, in __mul__
       return OperatorComp(self, other)
              ^^^^^^^^^^^^^^^^^^^^^^^^^
     File "/home/justussa/progwrit/python/odl/odl/core/operator/operator.py", line 1312, in __init__
       raise OpTypeError(
   odl.OpTypeError: `range` rn(4, 'float64', 'pytorch') of the right operator x * [ 10.,  20.,  30.,  40.] not equal to the domain rn(4) of the left operator x * [ 1.,  2.,  3.,  4.]
   
In situations where you are sure you need a copying step you can use adapter-operators, which wrap backend-specific device and implementation conversion in operator syntax.
Adapters are conceptually identity operators but with different backends / devices for their domain and range.
For example, to change the device from CPU to the GPU, use `DeviceChange` with ``domain_device='cpu'`` and ``range_device='cuda:0'``.
Note that the domain-device is the left argument, range the right argument. (This can look somewhat confusing because the dataflow in a chain of `@`-composed operators is right-to-left!)

A device change requires an array backend that supports both of the devices.
The backend can also be changed with an adapter:

   >>> op_np @ odl.ArrayBackendChange('pytorch', 'numpy') @ op_pt
   OperatorComp(OperatorComp(MultiplyOperator([ 1.,  2.,  3.,  4.]), _ImplChangeOperator(domain=rn(4, 'float64', 'pytorch'), range_impl='numpy')), MultiplyOperator([ 10.,  20.,  30.,  40.]))

Notice that the `ArrayBackendChange` adapter has been implicitly cast to `_ImplChangeOperator`, which is a proper ODL operator.
`ArrayBackendChange` is itself a subclass of `AdapterOperator`, not of `Operator`.
The difference is that an adapter does not have a fixed domain and range, but can act on spaces of any shape / discretization etc., so long as they support the desired device change.

Adapters can be composed with operators as above, or directly applied to elements:

   >>> odl.ArrayBackendChange('numpy', 'pytorch')(odl.rn(2).element([8,9]))
   rn(2, 'float64', 'pytorch').element([ 8.,  9.])

If all you need is changing the backend of a single element, this could also be done with the `TensorSpace.to_impl` / `TensorSpace.to_device` methods:

   >>> odl.rn(2).element([8,9]).to_impl('pytorch')
   rn(2, 'float64', 'pytorch').element([ 8.,  9.])

These methods are not standalone objects though, so for conversion within operator pipelines it is necessary to use `AdapterOperator`s.

Product spaces can have multiple different backends and/or devices. To change only one of them, use `subspace_index`:

   >>> p = odl.ProductSpace(odl.rn(2), odl.rn(3)).element(([0,1], [10,11,12]))
   >>> odl.ArrayBackendChange('numpy', 'pytorch', subspace_index=1)(p)
   rn(2) x rn(3, 'float64', 'pytorch').element([
       [ 0.,  1.],
       [ 10.,  11.,  12.]
   ])