feat: dims and dimsd for relevant basicoperators

Author: cako

pylops/basicoperators/CausalIntegration.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from pylops import LinearOperator
-from pylops.utils._internal import _value_or_list_like_to_array
+from pylops.utils._internal import _value_or_list_like_to_tuple
 
 
 class CausalIntegration(LinearOperator):
@@ -97,16 +97,18 @@ def __init__(
         kind="full",
         removefirst=False,
     ):
-        self.dims = _value_or_list_like_to_array(dims)
+        self.dims = _value_or_list_like_to_tuple(dims)
         self.axis = axis
         self.sampling = sampling
         self.kind = kind
         if kind == "full" and halfcurrent:  # ensure backcompatibility
             self.kind = "half"
         self.removefirst = removefirst
-        self.dimsd = self.dims.copy()
+        dimsd = list(self.dims)
         if self.removefirst:
-            self.dimsd[self.axis] -= 1
+            dimsd[self.axis] -= 1
+        self.dimsd = tuple(dimsd)
+
         self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False

pylops/basicoperators/Conj.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_tuple
 from pylops.utils.backend import get_array_module
 
 
@@ -41,7 +42,9 @@ class Conj(LinearOperator):
     """
 
     def __init__(self, dims, dtype="complex128"):
-        self.shape = (np.prod(np.array(dims)), np.prod(np.array(dims)))
+        self.dims = self.dimsd = _value_or_list_like_to_tuple(dims)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False
         self.clinear = False

pylops/basicoperators/Diagonal.py

@@ -3,6 +3,7 @@
 import numpy as np
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_tuple
 from pylops.utils.backend import get_array_module, to_cupy_conditional
 
 
@@ -64,28 +65,26 @@ def __init__(self, diag, dims=None, axis=-1, dtype="float64"):
         self.diag = diag.ravel()
         self.complex = True if ncp.iscomplexobj(self.diag) else False
 
-        if dims is None:
-            self.shape = (len(self.diag), len(self.diag))
-            self.dims = None
-            self.reshape = False
-        else:
-            diagdims = [1] * len(dims)
-            diagdims[axis] = dims[axis]
-            self.diag = self.diag.reshape(diagdims)
-            self.shape = (np.prod(dims), np.prod(dims))
-            self.dims = dims
-            self.reshape = True
+        ncp = get_array_module(diag)
+        self.diag = diag.ravel()
+        self.complex = True if ncp.iscomplexobj(self.diag) else False
+        self.dims = self.dimsd = (
+            (len(self.diag),) if dims is None else _value_or_list_like_to_tuple(dims)
+        )
+
+        diagdims = np.ones_like(self.dims)
+        diagdims[axis] = self.dims[axis]
+        self.diag = self.diag.reshape(diagdims)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False
 
     def _matvec(self, x):
         if type(self.diag) != type(x):
             self.diag = to_cupy_conditional(x, self.diag)
-        if not self.reshape:
-            y = self.diag * x.ravel()
-        else:
-            x = x.reshape(self.dims)
-            y = self.diag * x
+        x = x.reshape(self.dims)
+        y = self.diag * x
         return y.ravel()
 
     def _rmatvec(self, x):
@@ -95,11 +94,8 @@ def _rmatvec(self, x):
             diagadj = self.diag.conj()
         else:
             diagadj = self.diag
-        if not self.reshape:
-            y = diagadj * x.ravel()
-        else:
-            x = x.reshape(self.dims)
-            y = diagadj * x
+        x = x.reshape(self.dims)
+        y = diagadj * x
         return y.ravel()
 
     def matrix(self):

pylops/basicoperators/DirectionalDerivative.py

@@ -67,8 +67,10 @@ def FirstDirectionalDerivative(
     else:
         Dop = Diagonal(v.ravel(), dtype=dtype)
     Sop = Sum(dims=[len(dims)] + list(dims), axis=0, dtype=dtype)
-    ddop = Sop * Dop * Gop
-    return LinearOperator(ddop)
+    ddop = LinearOperator(Sop * Dop * Gop)
+    ddop.dims = ddop.dimsd = dims
+    ddop.sampling = sampling
+    return ddop
 
 
 def SecondDirectionalDerivative(dims, v, sampling=1, edge=False, dtype="float64"):
@@ -118,5 +120,7 @@ def SecondDirectionalDerivative(dims, v, sampling=1, edge=False, dtype="float64"
     in the literature.
     """
     Dop = FirstDirectionalDerivative(dims, v, sampling=sampling, edge=edge, dtype=dtype)
-    ddop = -Dop.H * Dop
-    return LinearOperator(ddop)
+    ddop = LinearOperator(-Dop.H * Dop)
+    ddop.dims = ddop.dimsd = dims
+    ddop.sampling = sampling
+    return ddop

pylops/basicoperators/FirstDerivative.py

@@ -4,7 +4,7 @@
 from numpy.core.multiarray import normalize_axis_index
 
 from pylops import LinearOperator
-from pylops.utils._internal import _value_or_list_like_to_array
+from pylops.utils._internal import _value_or_list_like_to_tuple
 from pylops.utils.backend import get_array_module
 
 
@@ -73,13 +73,13 @@ def __init__(
         dtype="float64",
         kind="centered",
     ):
-        self.dims = _value_or_list_like_to_array(dims)
+        self.dims = self.dimsd = _value_or_list_like_to_tuple(dims)
         self.axis = normalize_axis_index(axis, len(self.dims))
         self.sampling = sampling
         self.edge = edge
         self.kind = kind
-        N = np.prod(self.dims)
-        self.shape = (N, N)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False
 

pylops/basicoperators/Flip.py

@@ -3,7 +3,7 @@
 import numpy as np
 
 from pylops import LinearOperator
-from pylops.utils._internal import _value_or_list_like_to_array
+from pylops.utils._internal import _value_or_list_like_to_tuple
 
 
 class Flip(LinearOperator):
@@ -46,10 +46,10 @@ class Flip(LinearOperator):
     """
 
     def __init__(self, dims, axis=-1, dtype="float64"):
-        self.dims = _value_or_list_like_to_array(dims)
+        self.dims = self.dimsd = _value_or_list_like_to_tuple(dims)
         self.axis = axis
-        N = np.prod(self.dims)
-        self.shape = (N, N)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False
 

pylops/basicoperators/Gradient.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pylops.basicoperators import FirstDerivative, VStack
+from pylops.utils._internal import _value_or_list_like_to_tuple
 
 
 def Gradient(dims, sampling=1, edge=False, dtype="float64", kind="centered"):
@@ -58,9 +59,9 @@ def Gradient(dims, sampling=1, edge=False, dtype="float64", kind="centered"):
     axes are instead summed together.
 
     """
+    dims = _value_or_list_like_to_tuple(dims)
     ndims = len(dims)
-    if isinstance(sampling, (int, float)):
-        sampling = [sampling] * ndims
+    sampling = _value_or_list_like_to_tuple(sampling, repeat=ndims)
 
     gop = VStack(
         [
@@ -75,4 +76,9 @@ def Gradient(dims, sampling=1, edge=False, dtype="float64", kind="centered"):
             for iax in range(ndims)
         ]
     )
+    gop.dims = dims
+    gop.dimsd = (ndims, *gop.dims)
+    gop.sampling = sampling
+    gop.edge = edge
+    gop.kind = kind
     return gop

pylops/basicoperators/Imag.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_tuple
 from pylops.utils.backend import get_array_module
 
 
@@ -43,7 +44,9 @@ class Imag(LinearOperator):
     """
 
     def __init__(self, dims, dtype="complex128"):
-        self.shape = (np.prod(np.array(dims)), np.prod(np.array(dims)))
+        self.dims = self.dimsd = _value_or_list_like_to_tuple(dims)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.rdtype = np.real(np.ones(1, self.dtype)).dtype
         self.explicit = False

pylops/basicoperators/Pad.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_tuple
 
 
 class Pad(LinearOperator):
@@ -60,14 +61,19 @@ class Pad(LinearOperator):
     def __init__(self, dims, pad, dtype="float64"):
         if np.any(np.array(pad) < 0):
             raise ValueError("Padding must be positive or zero")
-        self.dims = dims
+        self.reshape = False if isinstance(dims, int) else True
+        self.dims = _value_or_list_like_to_tuple(dims)
         self.pad = pad
-        self.reshape = False if isinstance(self.dims, int) else True
         if self.reshape:
-            self.dimsd = [dim + p[0] + p[1] for dim, p in zip(dims, pad)]
+            dimsd = [
+                dim + before + after
+                for dim, (before, after) in zip(self.dims, self.pad)
+            ]
         else:
-            self.dimsd = dims + pad[0] + pad[1]
-        self.shape = (np.prod(np.array(self.dimsd)), np.prod(np.array(self.dims)))
+            dimsd = [self.dims[0] + pad[0] + pad[1]]
+        self.dimsd = tuple(dimsd)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.explicit = False
 
@@ -82,8 +88,8 @@ def _matvec(self, x):
     def _rmatvec(self, x):
         if self.reshape:
             y = x.reshape(self.dimsd)
-            for ax, pad in enumerate(self.pad):
-                y = np.take(y, np.arange(pad[0], pad[0] + self.dims[ax]), axis=ax)
+            for ax, (before, _) in enumerate(self.pad):
+                y = np.take(y, np.arange(before, before + self.dims[ax]), axis=ax)
         else:
-            y = x[self.pad[0] : self.pad[0] + self.dims]
+            y = x[self.pad[0] : self.pad[0] + self.dims[0]]
         return y.ravel()

pylops/basicoperators/Real.py

@@ -1,6 +1,7 @@
 import numpy as np
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_tuple
 from pylops.utils.backend import get_array_module
 
 
@@ -43,7 +44,9 @@ class Real(LinearOperator):
     """
 
     def __init__(self, dims, dtype="complex128"):
-        self.shape = (np.prod(np.array(dims)), np.prod(np.array(dims)))
+        self.dims = self.dimsd = _value_or_list_like_to_tuple(dims)
+
+        self.shape = (np.prod(self.dimsd), np.prod(self.dims))
         self.dtype = np.dtype(dtype)
         self.rdtype = np.real(np.ones(1, self.dtype)).dtype
         self.explicit = False