minor: added dims and dimsd to MatrixMult

Changed dims in input parameters to otherdims to be able to have
dims and dimsd consistent with the other operators.

Author: mrava87

pylops/avo/poststack.py

@@ -87,7 +87,7 @@ def _PoststackLinearModelling(
         M = ncp.dot(C, D)
         if sparse:
             M = get_csc_matrix(wav)(M)
-        Pop = _MatrixMult(M, dims=spatdims, dtype=dtype, **args_MatrixMult)
+        Pop = _MatrixMult(M, otherdims=spatdims, dtype=dtype, **args_MatrixMult)
     else:
         # Create wavelet operator
         if len(wav.shape) == 1:
@@ -102,7 +102,7 @@ def _PoststackLinearModelling(
         else:
             Cop = _MatrixMult(
                 nonstationary_convmtx(wav, nt0, hc=wav.shape[1] // 2, pad=(nt0, nt0)),
-                dims=spatdims,
+                otherdims=spatdims,
                 dtype=dtype,
                 **args_MatrixMult
             )
@@ -350,7 +350,7 @@ def PoststackInversion(
                     minv = get_lstsq(data)(PP, datarn, **kwargs_solver)[0]
                 else:
                     # solve regularized normal equations simultaneously
-                    PPop_reg = MatrixMult(PP, dims=nspatprod)
+                    PPop_reg = MatrixMult(PP, otherdims=nspatprod)
                     if ncp == np:
                         minv = lsqr(PPop_reg, datar.ravel(), **kwargs_solver)[0]
                     else:
@@ -364,7 +364,7 @@ def PoststackInversion(
                 # create regularized normal eqs. and solve them simultaneously
                 PP = ncp.dot(PPop.A.T, PPop.A) + epsI * ncp.eye(nt0, dtype=PPop.A.dtype)
                 datarn = PPop.A.T * datar.reshape(nt0, nspatprod)
-                PPop_reg = MatrixMult(PP, dims=nspatprod)
+                PPop_reg = MatrixMult(PP, otherdims=nspatprod)
                 minv = get_lstsq(data)(PPop_reg.A, datarn.ravel(), **kwargs_solver)[0]
         else:
             # solve unregularized normal equations simultaneously with lop

pylops/avo/prestack.py

@@ -182,7 +182,7 @@ def PrestackLinearModelling(
 
         # Combine operators
         M = ncp.dot(C, ncp.dot(G, D))
-        Preop = MatrixMult(M, dims=spatdims, dtype=dtype)
+        Preop = MatrixMult(M, otherdims=spatdims, dtype=dtype)
 
     else:
         # Create wavelet operator
@@ -504,7 +504,7 @@ def PrestackInversion(
         ]
         if explicit:
             PPop = MatrixMult(
-                np.vstack([Op.A for Op in PPop]), dims=nspat, dtype=PPop[0].A.dtype
+                np.vstack([Op.A for Op in PPop]), otherdims=nspat, dtype=PPop[0].A.dtype
             )
         else:
             PPop = VStack(PPop)
@@ -560,7 +560,7 @@ def PrestackInversion(
                     minv = get_lstsq(data)(PP, datarn, **kwargs_solver)[0]
                 else:
                     # solve regularized normal equations simultaneously
-                    PPop_reg = MatrixMult(PP, dims=nspatprod)
+                    PPop_reg = MatrixMult(PP, otherdims=nspatprod)
                     if ncp == np:
                         minv = lsqr(PPop_reg, datarn.ravel(), **kwargs_solver)[0]
                     else:
@@ -574,7 +574,7 @@ def PrestackInversion(
             #    # create regularized normal eqs. and solve them simultaneously
             #    PP = np.dot(PPop.A.T, PPop.A) + epsI * np.eye(nt0*nm)
             #    datarn = PPop.A.T * datar.reshape(nt0*ntheta, nspatprod)
-            #    PPop_reg = MatrixMult(PP, dims=ntheta*nspatprod)
+            #    PPop_reg = MatrixMult(PP, otherdims=ntheta*nspatprod)
             #    minv = lstsq(PPop_reg, datarn.ravel(), **kwargs_solver)[0]
         else:
             # solve unregularized normal equations simultaneously with lop

pylops/basicoperators/MatrixMult.py

@@ -5,6 +5,7 @@
 from scipy.sparse.linalg import inv
 
 from pylops import LinearOperator
+from pylops.utils._internal import _value_or_list_like_to_array
 from pylops.utils.backend import get_array_module
 
 logging.basicConfig(format="%(levelname)s: %(message)s", level=logging.WARNING)
@@ -20,7 +21,7 @@ class MatrixMult(LinearOperator):
     ----------
     A : :obj:`numpy.ndarray` or :obj:`scipy.sparse` matrix
         Matrix.
-    dims : :obj:`tuple`, optional
+    otherdims : :obj:`tuple`, optional
         Number of samples for each other dimension of model
         (model/data will be reshaped and ``A`` applied multiple times
         to each column of the model/data).
@@ -29,6 +30,10 @@ class MatrixMult(LinearOperator):
 
     Attributes
     ----------
+    dimsd : :obj:`tuple`
+        Shape of the array after the forward, but before linearization.
+
+        For example, ``y_reshaped = (Op * x.ravel()).reshape(Op.dimsd)``.
     shape : :obj:`tuple`
         Operator shape
     explicit : :obj:`bool`
@@ -39,30 +44,29 @@ class MatrixMult(LinearOperator):
 
     """
 
-    def __init__(self, A, dims=None, dtype="float64"):
+    def __init__(self, A, otherdims=None, dtype="float64"):
         ncp = get_array_module(A)
         self.A = A
         if isinstance(A, ncp.ndarray):
             self.complex = np.iscomplexobj(A)
         else:
             self.complex = np.iscomplexobj(A.data)
-        if dims is None:
+        if otherdims is None:
+            self.dims, self.dimsd = A.shape[1], A.shape[0]
             self.reshape = False
             self.shape = A.shape
             self.explicit = True
         else:
-            if isinstance(dims, int):
-                dims = (dims,)
+            otherdims = _value_or_list_like_to_array(otherdims)
+            self.otherdims = np.array(otherdims, dtype=int)
+            self.dims, self.dimsd = np.insert(
+                self.otherdims, 0, self.A.shape[1]
+            ), np.insert(self.otherdims, 0, self.A.shape[0])
+            self.dimsflatten, self.dimsdflatten = np.insert(
+                [np.prod(self.otherdims)], 0, self.A.shape[1]
+            ), np.insert([np.prod(self.otherdims)], 0, self.A.shape[0])
             self.reshape = True
-            self.dims = np.array(dims, dtype=int)
-            self.reshapedims = [
-                np.insert([np.prod(self.dims)], 0, self.A.shape[1]),
-                np.insert([np.prod(self.dims)], 0, self.A.shape[0]),
-            ]
-            self.shape = (
-                A.shape[0] * np.prod(self.dims),
-                A.shape[1] * np.prod(self.dims),
-            )
+            self.shape = (np.prod(self.dimsd), np.prod(self.dims))
             self.explicit = False
         self.dtype = np.dtype(dtype)
         # Check dtype for correctness (upcast to complex when A is complex)
@@ -75,7 +79,7 @@ def __init__(self, A, dims=None, dtype="float64"):
     def _matvec(self, x):
         ncp = get_array_module(x)
         if self.reshape:
-            x = ncp.reshape(x, self.reshapedims[0])
+            x = ncp.reshape(x, self.dimsflatten)
         y = self.A.dot(x)
         if self.reshape:
             return y.ravel()
@@ -85,7 +89,7 @@ def _matvec(self, x):
     def _rmatvec(self, x):
         ncp = get_array_module(x)
         if self.reshape:
-            x = ncp.reshape(x, self.reshapedims[1])
+            x = ncp.reshape(x, self.dimsdflatten)
         if self.complex:
             y = (self.A.T.dot(x.conj())).conj()
         else:

pytests/test_basicoperators.py

@@ -127,7 +127,7 @@ def test_MatrixMult_repeated(par):
     G = np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32") + par[
         "imag"
     ] * np.random.normal(0, 10, (par["ny"], par["nx"])).astype("float32")
-    Gop = MatrixMult(G, dims=5, dtype=par["dtype"])
+    Gop = MatrixMult(G, otherdims=5, dtype=par["dtype"])
     assert dottest(
         Gop, par["ny"] * 5, par["nx"] * 5, complexflag=0 if par["imag"] == 1 else 3
     )