Allow passing numpy types to operators

Author: mrava87

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/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))

pylops_gpu/utils/torch2numpy.py

@@ -2,6 +2,25 @@
 import numpy as np
 
 
+
+def numpytype_from_strtype(strtype):
+    """Convert str into equivalent numpy type
+
+    Parameters
+    ----------
+    strtype : :obj:`str`
+        String type
+
+    Returns
+    -------
+    numpytype : :obj:`numpy.dtype`
+        Numpy equivalent type
+
+    """
+    numpytype = np.dtype(strtype)
+    return numpytype
+
+
 def numpytype_from_torchtype(torchtype):
     """Convert torch type into equivalent numpy type
 
@@ -12,11 +31,15 @@ def numpytype_from_torchtype(torchtype):
 
     Returns
     -------
-    numpytype : :obj:`torch.dtype`
+    numpytype : :obj:`numpy.dtype`
         Numpy equivalent type
 
     """
-    numpytype = torch.scalar_tensor(1, dtype=torchtype).numpy().dtype
+    if isinstance(torchtype, torch.dtype):
+        numpytype = torch.scalar_tensor(1, dtype=torchtype).numpy().dtype
+    else:
+        # in case it is already a numpy dtype
+        numpytype = torchtype
     return numpytype
 
 
@@ -25,7 +48,7 @@ def torchtype_from_numpytype(numpytype):
 
     Parameters
     ----------
-    numpytype : :obj:`torch.dtype`
+    numpytype : :obj:`numpy.dtype`
         Numpy type
 
     Returns
@@ -40,5 +63,10 @@ def torchtype_from_numpytype(numpytype):
     returned.
 
     """
-    torchtype = torch.from_numpy(np.real(np.ones(1, dtype=numpytype))).dtype
+    if isinstance(numpytype, torch.dtype):
+        # in case it is already a torch dtype
+        torchtype = numpytype
+    else:
+        torchtype = \
+            torch.from_numpy(np.real(np.ones(1, dtype=numpytype_from_strtype(numpytype)))).dtype
     return torchtype