Add device parameter to ToCupy operator and to_cupy utility

pylops/basicoperators/tocupy.py

@@ -22,6 +22,9 @@ class ToCupy(LinearOperator):
         Number of samples for each dimension
     dtype : :obj:`str`, optional
         Type of elements in input array.
+    device : :obj:`int`, optional
+       GPU device where the array will be transferred to.
+       By default, device 0 is used.
     name : :obj:`str`, optional
         Name of operator (to be used by :func:`pylops.utils.describe.describe`)
 
@@ -52,13 +55,15 @@ def __init__(
         self,
         dims: Union[int, InputDimsLike],
         dtype: DTypeLike = "float64",
+        device: int = 0,
         name: str = "C",
     ) -> None:
+        self.device = device
         dims = _value_or_sized_to_tuple(dims)
         super().__init__(dtype=np.dtype(dtype), dims=dims, dimsd=dims, name=name)
 
     def _matvec(self, x: NDArray) -> NDArray:
-        return to_cupy(x)
+        return to_cupy(x, device=self.device)
 
     def _rmatvec(self, x: NDArray) -> NDArray:
         return to_numpy(x)

pylops/utils/backend.py

@@ -485,13 +485,16 @@ def get_real_dtype(dtype: DTypeLike) -> DTypeLike:
     return np.real(np.ones(1, dtype)).dtype
 
 
-def to_cupy(x: ArrayLike) -> ArrayLike:
+def to_cupy(x: ArrayLike, device: int = 0) -> ArrayLike:
     """Convert x to cupy array if cupy is available
 
     Parameters
     ----------
     x : :obj:`numpy.ndarray`, :obj:`cupy.ndarray` or :obj:`jax.Array`
         Array to evaluate
+    device : :obj:`int`, optional
+        GPU device where the array will be transferred to.
+        by default, device 0 is used.
 
     Returns
     -------
@@ -501,7 +504,8 @@ def to_cupy(x: ArrayLike) -> ArrayLike:
     """
     if deps.cupy_enabled:
         if cp.get_array_module(x) != cp:
-            x = cp.asarray(x)
+            with cp.cuda.Device(device):
+                x = cp.asarray(x)
     return x
 
 

pytests/test_basicoperators.py

@@ -718,3 +718,19 @@ def test_ToCupy(par):
     xadj = Top.H * y
     assert_array_equal(x, xadj)
     assert_array_equal(y, np.asarray(x))
+
+
+@pytest.mark.parametrize("par", [(par1), (par2), (par1j), (par2j), (par3)])
+def test_ToCupy_device(par):
+    """Forward and adjoint for ToCupy operator with explicit device parameter
+    (checking device is correctly stored and used)
+    """
+    Top = ToCupy(par["nx"], dtype=par["dtype"], device=0)
+
+    np.random.seed(10)
+    x = npp.random.randn(par["nx"]) + par["imag"] * npp.random.randn(par["nx"])
+    y = Top * x
+    xadj = Top.H * y
+    assert_array_equal(x, xadj)
+    assert_array_equal(y, np.asarray(x))
+    assert Top.device == 0