Merge pull request #9 from mrava87/master

Avoid repeating matvec/rmatvec in compuond operators (sum, prod, etc.)

Author: mrava87

pylops_distributed/LinearOperator.py

@@ -201,8 +201,12 @@ def adjoint(self):
     def _adjoint(self):
         """Default implementation of _adjoint; defers to rmatvec."""
         shape = (self.shape[1], self.shape[0])
-        return _CustomLinearOperator(shape, matvec=self.rmatvec,
-                                     rmatvec=self.matvec,
+        return _CustomLinearOperator(shape, matvec=self._rmatvec,
+                                     rmatvec=self._matvec,
+                                     compute=(self.compute[1],
+                                              self.compute[0]),
+                                     todask=(self.todask[1],
+                                              self.todask[0]),
                                      dtype=self.dtype)
 
     def div1(self, y, niter=100):
@@ -286,7 +290,7 @@ def _adjoint(self):
                                      matvec=self.__rmatvec_impl,
                                      rmatvec=self.__matvec_impl,
                                      dtype=self.dtype,
-                                     compute=self.compute,
+                                     compute=(self.compute[1], self.compute[0]),
                                      todask=(self.todask[1], self.todask[0]))
 
 
@@ -319,13 +323,13 @@ def __init__(self, A, B):
         self.args = (Ac, Bc)
 
     def _matvec(self, x):
-        return self.args[0].matvec(x) + self.args[1].matvec(x)
+        return self.args[0]._matvec(x) + self.args[1]._matvec(x)
 
     def _rmatvec(self, x):
-        return self.args[0].rmatvec(x) + self.args[1].rmatvec(x)
+        return self.args[0]._rmatvec(x) + self.args[1]._rmatvec(x)
 
     def _matmat(self, x):
-        return self.args[0].matmat(x) + self.args[1].matmat(x)
+        return self.args[0]._matmat(x) + self.args[1]._matmat(x)
 
     def _adjoint(self):
         A, B = self.args
@@ -345,8 +349,8 @@ def __init__(self, A, B):
                                                      dtype=A.dtype, Op=None,
                                                      explicit=A.explicit and
                                                               B.explicit,
-                                                     compute=(B.compute[0],
-                                                              A.compute[1]),
+                                                     compute=(A.compute[0],
+                                                              B.compute[1]),
                                                      todask=(B.todask[0],
                                                              A.todask[1]))
         # Force compute and todask not to be applied to individual operators
@@ -359,13 +363,13 @@ def __init__(self, A, B):
         self.args = (Ac, Bc)
 
     def _matvec(self, x):
-        return self.args[0].matvec(self.args[1].matvec(x))
+        return self.args[0]._matvec(self.args[1]._matvec(x))
 
     def _rmatvec(self, x):
-        return self.args[1].rmatvec(self.args[0].rmatvec(x))
+        return self.args[1]._rmatvec(self.args[0]._rmatvec(x))
 
     def _matmat(self, x):
-        return self.args[0].matmat(self.args[1].matmat(x))
+        return self.args[0]._matmat(self.args[1]._matmat(x))
 
     def _adjoint(self):
         A, B = self.args
@@ -390,13 +394,13 @@ def __init__(self, A, alpha):
         self.args = (Ac, alpha)
 
     def _matvec(self, x):
-        return self.args[1] * self.args[0].matvec(x)
+        return self.args[1] * self.args[0]._matvec(x)
 
     def _rmatvec(self, x):
-        return np.conj(self.args[1]) * self.args[0].rmatvec(x)
+        return np.conj(self.args[1]) * self.args[0]._rmatvec(x)
 
     def _matmat(self, x):
-        return self.args[1] * self.args[0].matmat(x)
+        return self.args[1] * self.args[0]._matmat(x)
 
     def _adjoint(self):
         A, alpha = self.args
@@ -428,13 +432,13 @@ def _power(self, fun, x, compute):
         return res
 
     def _matvec(self, x):
-        return self._power(self.args[0].matvec, x, self.compute[0])
+        return self._power(self.args[0]._matvec, x, self.compute[0])
 
     def _rmatvec(self, x):
-        return self._power(self.args[0].rmatvec, x, self.compute[1])
+        return self._power(self.args[0]._rmatvec, x, self.compute[1])
 
     def _matmat(self, x):
-        return self._power(self.args[0].matmat, x, self.compute[0])
+        return self._power(self.args[0]._matmat, x, self.compute[0])
 
     def _adjoint(self):
         A, p = self.args
@@ -489,4 +493,4 @@ def aslinearoperator(Op):
         return Op
     else:
         return LinearOperator(Op.shape, Op.dtype, Op, explicit=Op.explicit,
-                              compute=Op.compute, todask=Op.todask)
+                              compute=Op.compute, todask=Op.todask)

pylops_distributed/waveeqprocessing/mdd.py

@@ -74,5 +74,4 @@ def MDC(G, nt, nv, dt=1., dr=1., twosided=True,
                 args_FFT1={'chunks': ((nt, G.shape[1], nv),
                                       (nt, G.shape[1], nv)),
                            'todask': (todask[1], False),
-                           'compute':(compute[0], False)})
-
+                           'compute':(False, compute[0])})

pytests/test_waveeqprocessing.py

@@ -60,6 +60,48 @@
 par8['nfmax'] = int(np.ceil((PAR['nt']+1.)/2))-30
 
 
+@pytest.mark.parametrize("par", [(par1)])
+def test_MDC_compute(par):
+    """Ensure that forward and adjoint of MDC return numpy array when
+    compute=True
+    """
+    par['nt2'] = par['nt']
+    v = 1500
+    it0_m = 25
+    t0_m = it0_m * par['dt']
+    theta_m = 0
+    amp_m = 1.
+
+    it0_G = np.array([25, 50, 75])
+    t0_G = it0_G * par['dt']
+    theta_G = (0, 0, 0)
+    phi_G = (0, 0, 0)
+    amp_G = (1., 0.6, 2.)
+
+    # Create axis
+    t, _, x, y = makeaxis(par)
+
+    # Create wavelet
+    wav = ricker(t[:41], f0=par['f0'])[0]
+
+    # Generate model
+    _, mwav = linear2d(x, t, v, t0_m, theta_m, amp_m, wav)
+    # Generate operator
+    _, Gwav = linear3d(x, y, t, v, t0_G, theta_G, phi_G, amp_G, wav)
+
+    # Define MDC linear operator
+    Gwav_fft = np.fft.fft(Gwav, par['nt2'], axis=-1)
+    Gwav_fft = Gwav_fft[..., :par['nfmax']]
+
+    dMDCop = dMDC(da.from_array(Gwav_fft.transpose(2, 0, 1)),
+                  nt=par['nt2'], nv=1, dt=par['dt'], dr=par['dx'],
+                  twosided=par['twosided'], todask=(True, True),
+                  compute=(True, True))
+
+    assert isinstance(dMDCop.matvec(np.ones(dMDCop.shape[1])), np.ndarray)
+    assert isinstance(dMDCop.rmatvec(np.ones(dMDCop.shape[0])), np.ndarray)
+
+
 @pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4),
                                  (par5), (par6), (par7), (par8)])
 def test_MDC_1virtualsource(par):