Merge pull request #10 from mrava87/master

TorchOperator cupy integration

Author: mrava87

README.md

@@ -83,7 +83,22 @@ This repository is organized as follows:
 You need **Python 3.5 or greater**.
 
 #### From PyPi
-Coming soon...
+
+If you want to use PyLops-gpu within your codes,
+install it in your Python-gpu environment by typing the following command in your terminal:
+
+```
+pip install pylops-gpu
+```
+
+Open a python terminal and type:
+
+```
+import pylops_gpu
+```
+
+If you do not see any error, you should be good to go, enjoy!
+
 
 #### From Github
 

examples/plot_convolve.py

@@ -22,7 +22,7 @@
 plt.close('all')
 
 ###############################################################################
-# We will start by creating a zero signal of lenght :math:`nt` and we will
+# We will start by creating a zero signal of length :math:`nt` and we will
 # place a unitary spike at its center. We also create our filter to be
 # applied by means of :py:class:`pylops_gpu.signalprocessing.Convolve1D`
 # operator.

pylops_gpu/TorchOperator.py

@@ -1,5 +1,13 @@
 import torch
 
+from torch.utils.dlpack import from_dlpack, to_dlpack
+from pylops.utils import deps
+
+if deps.cupy_enabled:
+    import cupy as cp
+else:
+    cp = None
+
 
 class _TorchOperator(torch.autograd.Function):
     """Wrapper class for PyLops operators into Torch functions
@@ -16,20 +24,47 @@ def forward(ctx, x, forw, adj, pylops, device):
         ctx.pylops = pylops
         ctx.device = device
 
+        # prepare input
         if ctx.pylops:
-            x = x.cpu().detach().numpy()
+            if ctx.device == 'cpu':
+                # bring x to cpu and numpy
+                x = x.cpu().detach().numpy()
+            else:
+                # pass x to cupy using DLPack
+                x = cp.fromDlpack(to_dlpack(x))
+
+        # apply forward operator
         y = ctx.forw(x)
+
+        # prepare output
         if ctx.pylops:
-            y = torch.from_numpy(y).to(ctx.device)
+            if ctx.device == 'cpu':
+                # move y to torch and device
+                y = torch.from_numpy(y).to(ctx.device)
+            else:
+                # move y to torch and device
+                y = from_dlpack(y.toDlpack())
         return y
 
     @staticmethod
     def backward(ctx, y):
+        # prepare input
         if ctx.pylops:
-            y = y.cpu().detach().numpy()
+            if ctx.device == 'cpu':
+                y = y.cpu().detach().numpy()
+            else:
+                # pass x to cupy using DLPack
+                y = cp.fromDlpack(to_dlpack(y))
+
+        # apply adjoint operator
         x = ctx.adj(y)
+
+        # prepare output
         if ctx.pylops:
-            x = torch.from_numpy(x).to(ctx.device)
+            if ctx.device == 'cpu':
+                x = torch.from_numpy(x).to(ctx.device)
+            else:
+                x = from_dlpack(x.toDlpack())
         return  x, None, None, None, None
 
 
@@ -75,6 +110,7 @@ def __init__(self, Op, batch=False, pylops=False, device='cpu'):
         else:
             self.matvec = lambda x: Op.matmat(x, kfirst=True)
             self.rmatvec = lambda x: Op.rmatmat(x, kfirst=True)
+        self.Top = _TorchOperator.apply
 
     def apply(self, x):
         """Apply forward pass to input vector
@@ -90,5 +126,5 @@ def apply(self, x):
             Output array resulting from the application of the operator to ``x``.
 
         """
-        return _TorchOperator.apply(x, self.matvec, self.rmatvec,
-                                    self.pylops, self.device)
+        return self.Top(x, self.matvec, self.rmatvec,
+                        self.pylops, self.device)

pylops_gpu/avo/poststack.py

@@ -6,9 +6,10 @@
 from pylops import MatrixMult, FirstDerivative
 from pylops.utils.signalprocessing import convmtx, nonstationary_convmtx
 from pylops.signalprocessing import Convolve1D
-#from pylops.avo.poststack import _PoststackLinearModelling
+from pylops.avo.poststack import _PoststackLinearModelling
 
 from pylops_gpu.utils import dottest as Dottest
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 from pylops_gpu import MatrixMult as gMatrixMult
 from pylops_gpu import FirstDerivative as gFirstDerivative
 from pylops_gpu import SecondDerivative as gSecondDerivative
@@ -19,24 +20,16 @@
 
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
 
-
+"""
 def _PoststackLinearModelling(wav, nt0, spatdims=None, explicit=False,
                               sparse=False, _MatrixMult=MatrixMult,
                               _Convolve1D=Convolve1D,
                               _FirstDerivative=FirstDerivative,
                               args_MatrixMult={}, args_Convolve1D={},
                               args_FirstDerivative={}):
+    # define dtype to be used (ensure wav.dtype rules that of operator)
+    dtype = torchtype_from_numpytype(wav.dtype)
 
-
-    """Post-stack linearized seismic modelling operator.
-
-    Used to be able to provide operators from different libraries to
-    PoststackLinearModelling. It operates in the same way as public method
-    (PoststackLinearModelling) but has additional input parameters allowing
-    passing a different operator and additional arguments to be passed to such
-    operator.
-
-    """
     if len(wav.shape) == 2 and wav.shape[0] != nt0:
         raise ValueError('Provide 1d wavelet or 2d wavelet composed of nt0 '
                          'wavelets')
@@ -67,23 +60,27 @@ def _PoststackLinearModelling(wav, nt0, spatdims=None, explicit=False,
         M = np.dot(C, D)
         if sparse:
             M = csc_matrix(M)
-        Pop = _MatrixMult(M, dims=spatdims, **args_MatrixMult)
+        Pop = _MatrixMult(M, dims=spatdims, dtype=dtype, **args_MatrixMult)
     else:
         # Create wavelet operator
         if len(wav.shape) == 1:
             Cop = _Convolve1D(np.prod(np.array(dims)), h=wav,
                               offset=len(wav) // 2, dir=0, dims=dims,
-                              **args_Convolve1D)
+                              dtype=dtype, **args_Convolve1D)
         else:
             Cop = _MatrixMult(nonstationary_convmtx(wav, nt0,
                                                     hc=wav.shape[1] // 2,
                                                     pad=(nt0, nt0)),
-                              dims=spatdims, **args_MatrixMult)
+                              dims=spatdims, dtype=dtype,
+                              **args_MatrixMult)
         # Create derivative operator
         Dop = _FirstDerivative(np.prod(np.array(dims)), dims=dims,
-                               dir=0, sampling=1., **args_FirstDerivative)
+                               dir=0, sampling=1., dtype=dtype,
+                               **args_FirstDerivative)
+
         Pop = Cop * Dop
     return Pop
+"""
 
 
 def PoststackLinearModelling(wav, nt0, spatdims=None, explicit=False,
@@ -129,8 +126,9 @@ def PoststackLinearModelling(wav, nt0, spatdims=None, explicit=False,
     implementation details.
 
     """
-    if not isinstance(wav, torch.Tensor) and not explicit:
-        wav = torch.from_numpy(wav).to(device)
+    # ensure wav is always numpy, it will be converted later back to torch
+    if isinstance(wav, torch.Tensor):
+        wav = wav.cpu().numpy()
     return _PoststackLinearModelling(wav, nt0, spatdims=spatdims,
                                      explicit=explicit, sparse=False,
                                      _MatrixMult=gMatrixMult,

pylops_gpu/basicoperators/Diagonal.py

@@ -5,6 +5,7 @@
 from pylops_gpu import LinearOperator
 from pylops_gpu.utils.complex import conj, flatten, reshape, \
     complextorch_fromnumpy
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 
 
 class Diagonal(LinearOperator):
@@ -78,7 +79,7 @@ def __init__(self, diag, dims=None, dir=0, device='cpu',
         self.device = device
         self.togpu = togpu
         self.tocpu = tocpu
-        self.dtype = dtype
+        self.dtype = torchtype_from_numpytype(dtype)
         self.explicit = False
         self.Op = None
 

pylops_gpu/basicoperators/FirstDerivative.py

@@ -2,6 +2,7 @@
 
 from pylops_gpu import LinearOperator
 from pylops_gpu.signalprocessing import Convolve1D
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 
 
 """
@@ -39,7 +40,7 @@ class FirstDerivative(LinearOperator):
     tocpu : :obj:`tuple`, optional
         Move data and model from gpu to cpu after applying ``matvec`` and
         ``rmatvec``, respectively (only when ``device='gpu'``)
-    dtype : :obj:`torch.dtype`, optional
+    dtype : :obj:`torch.dtype` or :obj:`np.dtype`, optional
         Type of elements in input array.
 
     Attributes
@@ -63,6 +64,9 @@ class FirstDerivative(LinearOperator):
     def __init__(self, N, dims=None, dir=0, sampling=1., device='cpu',
                  togpu=(False, False), tocpu=(False, False),
                  dtype=torch.float32):
+        # convert dtype to torch.dtype
+        dtype = torchtype_from_numpytype(dtype)
+
         h = torch.torch.tensor([0.5, 0, -0.5],
                                dtype=dtype).to(device) / sampling
         self.device = device
@@ -73,4 +77,4 @@ def __init__(self, N, dims=None, dir=0, sampling=1., device='cpu',
         self.explicit = False
         self.Op = Convolve1D(N, h, offset=1, dims=dims, dir=dir,
                              zero_edges=True, device=device,
-                             togpu=togpu, tocpu=tocpu, dtype=dtype)
+                             togpu=togpu, tocpu=tocpu, dtype=self.dtype)

pylops_gpu/basicoperators/MatrixMult.py

@@ -4,7 +4,8 @@
 from pytorch_complex_tensor import ComplexTensor
 from pylops_gpu.LinearOperator import LinearOperator
 from pylops_gpu.utils.complex import conj, reshape, flatten
-from pylops_gpu.utils.torch2numpy import numpytype_from_torchtype
+from pylops_gpu.utils.torch2numpy import numpytype_from_torchtype, \
+    torchtype_from_numpytype
 
 
 class MatrixMult(LinearOperator):
@@ -29,7 +30,7 @@ class MatrixMult(LinearOperator):
     tocpu : :obj:`tuple`, optional
         Move data and model from gpu to cpu after applying ``matvec`` and
         ``rmatvec``, respectively (only when ``device='gpu'``)
-    dtype : :obj:`torch.dtype`, optional
+    dtype : :obj:`torch.dtype` or :obj:`np.dtype`, optional
         Type of elements in input array.
 
     Attributes
@@ -49,10 +50,12 @@ class MatrixMult(LinearOperator):
     def __init__(self, A, dims=None, device='cpu',
                  togpu=(False, False), tocpu=(False, False),
                  dtype=torch.float32):
+        # convert A to torch tensor if provided as numpy array numpy
         if not isinstance(A, (torch.Tensor, ComplexTensor)):
-            self.complex = True if np.iscomplexobj(A) else False
+            dtype = numpytype_from_torchtype(dtype)
             self.A = \
                 torch.from_numpy(A.astype(numpytype_from_torchtype(dtype))).to(device)
+            self.complex = True if np.iscomplexobj(A) else False
         else:
             self.complex = True if isinstance(A, ComplexTensor) else False
             self.A = A
@@ -76,7 +79,7 @@ def __init__(self, A, dims=None, device='cpu',
         self.device = device
         self.togpu = togpu
         self.tocpu = tocpu
-        self.dtype = dtype
+        self.dtype = torchtype_from_numpytype(dtype)
         self.explicit = True
         self.Op = None
 

pylops_gpu/basicoperators/Restriction.py

@@ -2,6 +2,7 @@
 import numpy as np
 
 from pylops_gpu.LinearOperator import LinearOperator
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 
 
 class Restriction(LinearOperator):
@@ -78,7 +79,7 @@ def __init__(self, M, iava, dims=None, dir=0, inplace=True,
         self.device = device
         self.togpu = togpu
         self.tocpu = tocpu
-        self.dtype = dtype
+        self.dtype = torchtype_from_numpytype(dtype)
         self.explicit = True
         self.Op = None
 

pylops_gpu/basicoperators/SecondDerivative.py

@@ -2,6 +2,7 @@
 
 from pylops_gpu import LinearOperator
 from pylops_gpu.signalprocessing import Convolve1D
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 
 
 class SecondDerivative(LinearOperator):
@@ -28,7 +29,7 @@ class SecondDerivative(LinearOperator):
     tocpu : :obj:`tuple`, optional
         Move data and model from gpu to cpu after applying ``matvec`` and
         ``rmatvec``, respectively (only when ``device='gpu'``)
-    dtype : :obj:`torch.dtype`, optional
+    dtype : :obj:`torch.dtype` or :obj:`np.dtype`, optional
         Type of elements in input array.
 
     Attributes
@@ -52,6 +53,9 @@ class SecondDerivative(LinearOperator):
     def __init__(self, N, dims=None, dir=0, sampling=1., device='cpu',
                  togpu=(False, False), tocpu=(False, False),
                  dtype=torch.float32):
+        # convert dtype to torch.dtype
+        dtype = torchtype_from_numpytype(dtype)
+
         h = torch.torch.tensor([1., -2, 1.],
                                dtype=dtype).to(device) / sampling**2
         self.device = device

pylops_gpu/signalprocessing/Convolve1D.py

@@ -3,6 +3,7 @@
 
 from torch.nn.functional import pad
 from pylops_gpu import LinearOperator
+from pylops_gpu.utils.torch2numpy import torchtype_from_numpytype
 
 
 class Convolve1D(LinearOperator):
@@ -16,7 +17,7 @@ class Convolve1D(LinearOperator):
     ----------
     N : :obj:`int`
         Number of samples in model.
-    h : :obj:`torch.Tensor`
+    h : :obj:`torch.Tensor` or :obj:`numpy.ndarray`
         1d compact filter to be convolved to input signal
     offset : :obj:`int`
         Index of the center of the compact filter
@@ -55,6 +56,13 @@ class Convolve1D(LinearOperator):
     def __init__(self, N, h, offset=0, dims=None, dir=0, zero_edges=False,
                  device='cpu', togpu=(False, False), tocpu=(False, False),
                  dtype=torch.float32):
+        # convert dtype to torch.dtype
+        if not isinstance(dtype, torch.dtype):
+            dtype = torchtype_from_numpytype(dtype)
+
+        # convert h to torch if numpy
+        if not isinstance(h, torch.Tensor):
+            h = torch.from_numpy(h).to(device)
         self.nh = h.size()[0]
         self.h = h.reshape(1, 1, self.nh)
         self.offset = 2*(self.nh // 2 - int(offset))