feature: enabled use of cupy arrays

Author: mrava87

pylops_mpi/DistributedArray.py

@@ -5,6 +5,7 @@
 from enum import Enum
 
 from pylops.utils import DTypeLike, NDArray
+from pylops.utils.backend import get_module
 
 
 class Partition(Enum):
@@ -78,6 +79,8 @@ class DistributedArray:
         Axis along which distribution occurs. Defaults to ``0``.
     local_shapes : :obj:`list`, optional
         List of tuples representing local shapes at each rank.
+    engine : :obj:`str`, optional
+        Engine used to store array (``numpy`` or ``cupy``)
     dtype : :obj:`str`, optional
         Type of elements in input array. Defaults to ``numpy.float64``.
     """
@@ -86,6 +89,7 @@ def __init__(self, global_shape: Union[Tuple, Integral],
                  base_comm: Optional[MPI.Comm] = MPI.COMM_WORLD,
                  partition: Partition = Partition.SCATTER, axis: int = 0,
                  local_shapes: Optional[List[Tuple]] = None,
+                 engine: Optional[str] = "numpy",
                  dtype: Optional[DTypeLike] = np.float64):
         if isinstance(global_shape, Integral):
             global_shape = (global_shape,)
@@ -103,7 +107,8 @@ def __init__(self, global_shape: Union[Tuple, Integral],
         self._check_local_shapes(local_shapes)
         self._local_shape = local_shapes[base_comm.rank] if local_shapes else local_split(global_shape, base_comm,
                                                                                           partition, axis)
-        self._local_array = np.empty(shape=self.local_shape, dtype=self.dtype)
+        self._engine = engine
+        self._local_array = get_module(engine).empty(shape=self.local_shape, dtype=self.dtype)
 
     def __getitem__(self, index):
         return self.local_array[index]
@@ -160,6 +165,16 @@ def local_shape(self):
         """
         return self._local_shape
 
+    @property
+    def engine(self):
+        """Engine of the Distributed array
+
+        Returns
+        -------
+        engine : :obj:`str`
+        """
+        return self._engine
+
     @property
     def local_array(self):
         """View of the Local Array
@@ -334,6 +349,7 @@ def __neg__(self):
                                partition=self.partition,
                                axis=self.axis,
                                local_shapes=self.local_shapes,
+                               engine=self.engine,
                                dtype=self.dtype)
         arr[:] = -self.local_array
         return arr
@@ -365,6 +381,7 @@ def add(self, dist_array):
                                     dtype=self.dtype,
                                     partition=self.partition,
                                     local_shapes=self.local_shapes,
+                                    engine=self.engine,
                                     axis=self.axis)
         SumArray[:] = self.local_array + dist_array.local_array
         return SumArray
@@ -387,6 +404,7 @@ def multiply(self, dist_array):
                                         dtype=self.dtype,
                                         partition=self.partition,
                                         local_shapes=self.local_shapes,
+                                        engine=self.engine,
                                         axis=self.axis)
         if isinstance(dist_array, DistributedArray):
             # multiply two DistributedArray
@@ -480,6 +498,7 @@ def conj(self):
                                 partition=self.partition,
                                 axis=self.axis,
                                 local_shapes=self.local_shapes,
+                                engine=self.engine,
                                 dtype=self.dtype)
         conj[:] = self.local_array.conj()
         return conj
@@ -492,6 +511,7 @@ def copy(self):
                                partition=self.partition,
                                axis=self.axis,
                                local_shapes=self.local_shapes,
+                               engine=self.engine,
                                dtype=self.dtype)
         arr[:] = self.local_array
         return arr

pylops_mpi/basicoperators/BlockDiag.py

@@ -3,8 +3,9 @@
 from mpi4py import MPI
 from typing import Optional, Sequence
 
-from pylops.utils import DTypeLike
 from pylops import LinearOperator
+from pylops.utils import DTypeLike
+from pylops.utils.backend import get_module
 
 from pylops_mpi import MPILinearOperator, MPIStackedLinearOperator
 from pylops_mpi import DistributedArray, StackedDistributedArray
@@ -113,22 +114,26 @@ def __init__(self, ops: Sequence[LinearOperator],
 
     @reshaped(forward=True, stacking=True)
     def _matvec(self, x: DistributedArray) -> DistributedArray:
-        y = DistributedArray(global_shape=self.shape[0], local_shapes=self.local_shapes_n, dtype=self.dtype)
+        ncp = get_module(x.engine)
+        y = DistributedArray(global_shape=self.shape[0], local_shapes=self.local_shapes_n, 
+                             engine=x.engine, dtype=self.dtype)
         y1 = []
         for iop, oper in enumerate(self.ops):
             y1.append(oper.matvec(x.local_array[self.mmops[iop]:
                                                 self.mmops[iop + 1]]))
-        y[:] = np.concatenate(y1)
+        y[:] = ncp.concatenate(ncp.asarray(y1))
         return y
 
     @reshaped(forward=False, stacking=True)
     def _rmatvec(self, x: DistributedArray) -> DistributedArray:
-        y = DistributedArray(global_shape=self.shape[1], local_shapes=self.local_shapes_m, dtype=self.dtype)
+        ncp = get_module(x.engine)
+        y = DistributedArray(global_shape=self.shape[1], local_shapes=self.local_shapes_m,
+                             engine=x.engine, dtype=self.dtype)
         y1 = []
         for iop, oper in enumerate(self.ops):
             y1.append(oper.rmatvec(x.local_array[self.nnops[iop]:
                                                  self.nnops[iop + 1]]))
-        y[:] = np.concatenate(y1)
+        y[:] = ncp.concatenate(ncp.asarray(y1))
         return y
 
 

pylops_mpi/basicoperators/VStack.py

@@ -5,6 +5,7 @@
 
 from pylops import LinearOperator
 from pylops.utils import DTypeLike
+from pylops.utils.backend import get_module
 
 from pylops_mpi import (
     MPILinearOperator,
@@ -116,22 +117,26 @@ def __init__(self, ops: Sequence[LinearOperator],
         super().__init__(shape=shape, dtype=dtype, base_comm=base_comm)
 
     def _matvec(self, x: DistributedArray) -> DistributedArray:
+        ncp = get_module(x.engine)
         if x.partition is not Partition.BROADCAST:
             raise ValueError(f"x should have partition={Partition.BROADCAST}, {x.partition} != {Partition.BROADCAST}")
-        y = DistributedArray(global_shape=self.shape[0], local_shapes=self.local_shapes_n, dtype=self.dtype)
+        y = DistributedArray(global_shape=self.shape[0], local_shapes=self.local_shapes_n, 
+                             engine=x.engine, dtype=self.dtype)
         y1 = []
         for iop, oper in enumerate(self.ops):
             y1.append(oper.matvec(x.local_array))
-        y[:] = np.concatenate(y1)
+        y[:] = ncp.concatenate(y1)
         return y
 
     @reshaped(forward=False, stacking=True)
     def _rmatvec(self, x: DistributedArray) -> DistributedArray:
-        y = DistributedArray(global_shape=self.shape[1], partition=Partition.BROADCAST, dtype=self.dtype)
+        ncp = get_module(x.engine)
+        y = DistributedArray(global_shape=self.shape[1], partition=Partition.BROADCAST, 
+                             engine=x.engine, dtype=self.dtype)
         y1 = []
         for iop, oper in enumerate(self.ops):
             y1.append(oper.rmatvec(x.local_array[self.nnops[iop]: self.nnops[iop + 1]]))
-        y1 = np.sum(y1, axis=0)
+        y1 = ncp.sum(ncp.asarray(y1), axis=0)
         y[:] = self.base_comm.allreduce(y1, op=MPI.SUM)
         return y
 

pylops_mpi/optimization/cls_basic.py

@@ -337,7 +337,7 @@ def setup(self,
         self.rank = x.rank
         self.c = r.copy()
         self.q = self.Op.matvec(self.c)
-        self.kold = np.abs(r.dot(r.conj()))
+        self.kold = float(np.abs(r.dot(r.conj())))
 
         # create variables to track the residual norm and iterations
         self.cost = []
@@ -373,13 +373,13 @@ def step(self, x: Union[DistributedArray, StackedDistributedArray],
 
         """
 
-        a = self.kold / (self.q.dot(self.q.conj()) + self.damp * self.c.dot(self.c.conj()))
+        a = float(self.kold / (self.q.dot(self.q.conj()) + self.damp * self.c.dot(self.c.conj())))
         x += a * self.c
         self.s -= a * self.q
         damped_x = self.damp * x
         r = self.Op.rmatvec(self.s) - damped_x
-        k = np.abs(r.dot(r.conj()))
-        b = k / self.kold
+        k = float(np.abs(r.dot(r.conj())))
+        b = float(k / self.kold)
         self.c = r + b * self.c
         self.q = self.Op.matvec(self.c)
         self.kold = k

pylops_mpi/utils/decorators.py

@@ -54,7 +54,8 @@ def wrapper(self, x: DistributedArray):
                 local_shapes = None
                 global_shape = getattr(self, "dims")
             arr = DistributedArray(global_shape=global_shape,
-                                   local_shapes=local_shapes, axis=0, dtype=x.dtype)
+                                   local_shapes=local_shapes, axis=0, 
+                                   engine=x.engine, dtype=x.dtype)
             arr_local_shapes = np.asarray(arr.base_comm.allgather(np.prod(arr.local_shape)))
             x_local_shapes = np.asarray(x.base_comm.allgather(np.prod(x.local_shape)))
             # Calculate num_ghost_cells required for each rank