Merge pull request #2670 from hyperspy/cupy_support

Cupy support

.github/workflows/tests.yml

@@ -63,16 +63,16 @@ jobs:
           python --version
           pip --version
 
-      - name: Install
-        shell: bash
-        run: |
-          pip install ${{ env.PIP_ARGS }} .'${{ matrix.PIP_SELECTOR }}'
-
       - name: Install oldest supported version
         if: ${{ matrix.OLDEST_SUPPORTED_VERSION }}
         run: |
           pip install ${{ matrix.DEPENDENCIES }}
 
+      - name: Install
+        shell: bash
+        run: |
+          pip install ${{ env.PIP_ARGS }} .'${{ matrix.PIP_SELECTOR }}'
+
       - name: Pip list
         run: |
           pip list

doc/conf.py

@@ -245,7 +245,8 @@
 
 
 # Add the hyperspy website to the intersphinx domains
-intersphinx_mapping = {'python': ('https://docs.python.org/3', None),
+intersphinx_mapping = {'cupy': ('https://docs.cupy.dev/en/stable', None),
+                       'python': ('https://docs.python.org/3', None),
                        'h5py': ('https://docs.h5py.org/en/stable', None),
                        'hyperspyweb': ('https://hyperspy.org/', None),
                        'matplotlib': ('https://matplotlib.org', None),

doc/user_guide/big_data.rst

@@ -264,6 +264,51 @@ instead:
     <LazySignal2D, title: , dimensions: (200, 200|512, 512)>
     >>> s.plot(navigator='slider')
 
+.. versionadded:: 1.7
+
+.. _big_data.gpu:
+
+GPU support
+-----------
+
+Lazy data processing on GPUs requires explicitly transferring the data to the
+GPU.
+
+On linux, it is recommended to use the
+`dask_cuda <https://docs.rapids.ai/api/dask-cuda/stable/index.html>`_ library 
+(not supported on windows) to manage the dask scheduler. As for CPU lazy
+processing, if the dask scheduler is not specified, the default scheduler
+will be used.
+
+.. code-block:: python
+
+    >>> from dask_cuda import LocalCUDACluster
+    >>> from dask.distributed import Client
+    >>> cluster = LocalCUDACluster()
+    >>> client = Client(cluster)
+
+.. code-block:: python
+
+    >>> import hyperspy.api as hs
+    >>> import cupy as cp
+    >>> import dask.array as da
+    >>> # Create a dask array
+    >>> data = da.random.random(size=(20, 20, 100, 100))
+    >>> print(data)
+    ... dask.array<random_sample, shape=(20, 20, 100, 100), dtype=float64,
+    ... chunksize=(20, 20, 100, 100), chunktype=numpy.ndarray>
+    >>> # convert the dask chunks from numpy array to cupy array
+    >>> data = data.map_blocks(cp.asarray)
+    >>> print(data)
+    ... dask.array<random_sample, shape=(20, 20, 100, 100), dtype=float64,
+    ... chunksize=(20, 20, 100, 100), chunktype=cupy.ndarray>
+    >>> # Create the signal
+    >>> s = hs.signals.Signal2D(data).as_lazy()
+
+.. note::
+    See the dask blog on `Richardson Lucy (RL) deconvolution <https://blog.dask.org/2020/11/12/deconvolution>`_
+    for an example of lazy processing on GPUs using dask and cupy
+
 
 .. _FitBigData-label:
 

doc/user_guide/signal.rst

@@ -84,7 +84,7 @@ This also applies to the :py:attr:`~.signal.BaseSignal.metadata`.
     └── Signal
 	└── signal_type =
 
-Instead of using a list of *axes dictionaries* ``[dict0, dict1]`` during signal 
+Instead of using a list of *axes dictionaries* ``[dict0, dict1]`` during signal
 initialization, you can also pass a list of *axes objects*: ``[axis0, axis1]``.
 
 
@@ -292,7 +292,7 @@ For example, a numpy ragged array can be created as follow:
        >>> arr = np.array([[1, 2, 3], [1]], dtype=object)
        >>> arr
        array([list([1, 2, 3]), list([1])], dtype=object)
-       
+
 Note that the array shape is (2, ):
 
    .. code-block:: python
@@ -329,10 +329,10 @@ To create a hyperspy signal of a numpy ragged array:
 
        >>> s.axes_manager
        <Axes manager, axes: (2|ragged)>
-                   Name |   size |  index |  offset |   scale |  units 
-       ================ | ====== | ====== | ======= | ======= | ====== 
-            <undefined> |      2 |      0 |       0 |       1 | <undefined> 
-       ---------------- | ------ | ------ | ------- | ------- | ------ 
+                   Name |   size |  index |  offset |   scale |  units
+       ================ | ====== | ====== | ======= | ======= | ======
+            <undefined> |      2 |      0 |       0 |       1 | <undefined>
+       ---------------- | ------ | ------ | ------- | ------- | ------
             Ragged axis |               Variable length
 
 .. note::
@@ -345,7 +345,7 @@ To create a hyperspy signal of a numpy ragged array:
            array([[1, 2],
                   [1, 2]], dtype=object)
 
-       
+
     Unlike in the previous example, here the array is not ragged, because
     the length of the nested sequences are equal (2) and numpy will create
     an array of shape (2, 2) instead of (2, ) as in the previous example of
@@ -374,11 +374,11 @@ shape of the ragged array:
        >>> s = hs.signals.BaseSignal(arr)
        >>> s
        <BaseSignal, title: , dimensions: (|2)>
-       
+
        >>> s.ragged = True
        >>> s
-       <BaseSignal, title: , dimensions: (2|ragged)>       
-       
+       <BaseSignal, title: , dimensions: (2|ragged)>
+
 
 .. _signal.binned:
 
@@ -1211,7 +1211,7 @@ signal:
     uint8
 
     # By default `dtype=None`, the dtype is determined by the behaviour of
-    # numpy.sum, in this case, unsigned integer of the same precision as the 
+    # numpy.sum, in this case, unsigned integer of the same precision as the
     # platform interger
     >>> s4 = s.rebin(scale=(5, 2, 1))
     >>> print(s4.data.dtype)
@@ -1864,3 +1864,38 @@ and `y` direction, while the offset is determined by the ``offset`` parameter.
 The fulcrum of the linear ramp is at the origin and the slopes are given in
 units of the axis with the according scale taken into account. Both are
 available via the :py:class:`~.axes.AxesManager` of the signal.
+
+.. versionadded:: 1.7
+
+.. _gpu_processing:
+
+GPU support
+-----------
+
+GPU processing is supported thanks to the numpy dispatch mechanism of array functions
+- read `NEP-18 <https://numpy.org/neps/nep-0018-array-function-protocol.html>`_
+and `NEP-35 <https://numpy.org/neps/nep-0035-array-creation-dispatch-with-array-function.html>`_
+for more information. It means that most HyperSpy functions will work on a GPU
+if the data is a :py:class:`cupy.ndarray` and the required functions are
+implemented in ``cupy``.
+
+.. note::
+    GPU processing with hyperspy requires numpy>=1.20 and dask>=2021.3.0, to be
+    able to use NEP-18 and NEP-35.
+
+.. code-block:: python
+
+    >>> import cupy as cp
+    >>> # Create a cupy array (on GPU device)
+    >>> data = cp.random.random(size=(20, 20, 100, 100))
+    >>> s = hs.signals.Signal2D(data)
+    >>> type(s.data)
+    ... cupy._core.core.ndarray
+
+Two convenience methods are available to transfer data between the host and
+the (GPU) device memory:
+
+- :py:meth:`~.signal.BaseSignal.to_host`
+- :py:meth:`~.signal.BaseSignal.to_device`
+
+For lazy processing, see the :ref:`corresponding section<big_data.gpu>`.

hyperspy/_components/doniach.py

@@ -27,7 +27,7 @@
 sqrt2pi = math.sqrt(2 * math.pi)
 
 
-tiny = np.finfo(np.float64).eps
+tiny = np.finfo(float).eps
 
 
 class Doniach(Expression):

hyperspy/_components/gaussian.py

@@ -50,7 +50,7 @@ def _estimate_gaussian_parameters(signal, x1, x2, only_current):
         centre_shape[i] = 1
 
     centre = np.sum(X.reshape(X_shape) * data, i) / np.sum(data, i)
-    sigma = np.sqrt(np.abs(np.sum((X.reshape(X_shape) - centre.reshape(
+    sigma = np.sqrt(abs(np.sum((X.reshape(X_shape) - centre.reshape(
         centre_shape)) ** 2 * data, i) / np.sum(data, i)))
     height = data.max(i)
 

hyperspy/_components/lorentzian.py

@@ -44,9 +44,9 @@ def _estimate_lorentzian_parameters(signal, x1, x2, only_current):
     cdf = np.cumsum(data,i)
     cdfnorm = cdf/np.max(cdf, i).reshape(centre_shape)
 
-    icentre = np.argmin(np.abs(0.5 - cdfnorm), i)
-    igamma1 = np.argmin(np.abs(0.75 - cdfnorm), i)
-    igamma2 = np.argmin(np.abs(0.25 - cdfnorm), i)
+    icentre = np.argmin(abs(0.5 - cdfnorm), i)
+    igamma1 = np.argmin(abs(0.75 - cdfnorm), i)
+    igamma2 = np.argmin(abs(0.25 - cdfnorm), i)
 
     if isinstance(data, da.Array):
         icentre, igamma1, igamma2 = da.compute(icentre, igamma1, igamma2)

hyperspy/_components/power_law.py

@@ -22,6 +22,7 @@
 
 from hyperspy._components.expression import Expression
 
+from hyperspy.misc.utils import get_numpy_kwargs
 
 _logger = logging.getLogger(__name__)
 
@@ -132,24 +133,22 @@ def estimate_parameters(self, signal, x1, x2, only_current=False,
         else:
             s = signal
         if s._lazy:
-            import dask.array as da
-            log = da.log
             I1 = s.isig[i1:i3].integrate1D(2j).data
             I2 = s.isig[i3:i2].integrate1D(2j).data
         else:
             from hyperspy.signal import BaseSignal
             shape = s.data.shape[:-1]
-            I1_s = BaseSignal(np.empty(shape, dtype='float'))
-            I2_s = BaseSignal(np.empty(shape, dtype='float'))
+            kw = get_numpy_kwargs(s.data)
+            I1_s = BaseSignal(np.empty(shape, dtype='float', **kw))
+            I2_s = BaseSignal(np.empty(shape, dtype='float', **kw))
             # Use the `out` parameters to avoid doing the deepcopy
             s.isig[i1:i3].integrate1D(2j, out=I1_s)
             s.isig[i3:i2].integrate1D(2j, out=I2_s)
             I1 = I1_s.data
             I2 = I2_s.data
-            log = np.log
         with np.errstate(divide='raise'):
             try:
-                r = 2 * (log(I1) - log(I2)) / (log(x2) - log(x1))
+                r = 2 * (np.log(I1) - np.log(I2)) / (np.log(x2) - np.log(x1))
                 k = 1 - r
                 A = k * I2 / (x2 ** k - x3 ** k)
                 if s._lazy:
@@ -162,15 +161,15 @@ def estimate_parameters(self, signal, x1, x2, only_current=False,
                 _logger.warning('Power-law parameter estimation failed '
                                 'because of a "divide-by-zero" error.')
                 return False
+
         if only_current is True:
             self.r.value = r
             self.A.value = A
             return True
+
         if out:
             return A, r
         else:
-            if self.A.map is None:
-                self._create_arrays()
             self.A.map['values'][:] = A
             self.A.map['is_set'][:] = True
             self.r.map['values'][:] = r

hyperspy/_components/skew_normal.py

@@ -51,17 +51,17 @@ def _estimate_skewnormal_parameters(signal, x1, x2, only_current):
     a1 = np.sqrt(2 / np.pi)
     b1 = (4 / np.pi - 1) * a1
     m1 = np.sum(X.reshape(X_shape) * data, i) / np.sum(data, i)
-    m2 = np.abs(np.sum((X.reshape(X_shape) - m1.reshape(x0_shape)) ** 2 * data, i)
+    m2 = abs(np.sum((X.reshape(X_shape) - m1.reshape(x0_shape)) ** 2 * data, i)
               / np.sum(data, i))
-    m3 = np.abs(np.sum((X.reshape(X_shape) - m1.reshape(x0_shape)) ** 3 * data, i)
+    m3 = abs(np.sum((X.reshape(X_shape) - m1.reshape(x0_shape)) ** 3 * data, i)
               / np.sum(data, i))
 
     x0 = m1 - a1 * (m3 / b1) ** (1 / 3)
     scale = np.sqrt(m2 + a1 ** 2 * (m3 / b1) ** (2 / 3))
     delta = np.sqrt(1 / (a1**2 + m2 * (b1 / m3) ** (2 / 3)))
     shape = delta / np.sqrt(1 - delta**2)
 
-    iheight = np.argmin(np.abs(X.reshape(X_shape) - x0.reshape(x0_shape)), i)
+    iheight = np.argmin(abs(X.reshape(X_shape) - x0.reshape(x0_shape)), i)
     # height is the value of the function at x0, shich has to be computed
     # differently for dask array (lazy) and depending on the dimension
     if isinstance(data, da.Array):
@@ -275,5 +275,5 @@ def mode(self):
             return self.x0.value
         else:
             m0 = muz - self.skewness * sigmaz / 2 - np.sign(self.shape.value) \
-                / 2 * np.exp(- 2 * np.pi / np.abs(self.shape.value))
+                / 2 * np.exp(- 2 * np.pi / abs(self.shape.value))
             return self.x0.value + self.scale.value * m0

hyperspy/_signals/complex_signal.py

@@ -203,11 +203,11 @@ def unwrapped_phase(self, wrap_around=False, seed=None,
     unwrapped_phase.__doc__ %= (SHOW_PROGRESSBAR_ARG, PARALLEL_ARG, MAX_WORKERS_ARG)
 
     def __call__(self, axes_manager=None, power_spectrum=False,
-                 fft_shift=False):
+                 fft_shift=False, as_numpy=None):
         value = super().__call__(axes_manager=axes_manager,
-                                 fft_shift=fft_shift)
+                                 fft_shift=fft_shift, as_numpy=as_numpy)
         if power_spectrum:
-            value = np.abs(value)**2
+            value = abs(value)**2
         return value
 
     def plot(self,
@@ -345,7 +345,7 @@ def argand_diagram(self, size=[256, 256], range=None):
             argand_diagram.axes_manager.signal_axes[0].name = 'Real'
             units_real = None
         argand_diagram.axes_manager.signal_axes[0].offset = real_edges[0]
-        argand_diagram.axes_manager.signal_axes[0].scale = np.abs(real_edges[0] - real_edges[1])
+        argand_diagram.axes_manager.signal_axes[0].scale = abs(real_edges[0] - real_edges[1])
 
         if self.imag.metadata.Signal.has_item('quantity'):
             quantity_imag, units_imag = parse_quantity(self.imag.metadata.Signal.quantity)
@@ -354,7 +354,7 @@ def argand_diagram(self, size=[256, 256], range=None):
             argand_diagram.axes_manager.signal_axes[1].name = 'Imaginary'
             units_imag = None
         argand_diagram.axes_manager.signal_axes[1].offset = imag_edges[0]
-        argand_diagram.axes_manager.signal_axes[1].scale = np.abs(imag_edges[0] - imag_edges[1])
+        argand_diagram.axes_manager.signal_axes[1].scale = abs(imag_edges[0] - imag_edges[1])
         if units_real:
             argand_diagram.axes_manager.signal_axes[0].units = units_real
         if units_imag: