Merge pull request #14 from mrava87/master

MDD method

Author: mrava87

docs/source/api/index.rst

@@ -61,6 +61,7 @@ Wave-Equation processing
    :toctree: generated/
 
    MDC
+   MDD
    Marchenko
    Demigration
 

pylops_distributed/LinearOperator.py

@@ -209,7 +209,7 @@ def _adjoint(self):
                                               self.todask[0]),
                                      dtype=self.dtype)
 
-    def div1(self, y, niter=100):
+    def div(self, y, niter=100):
         r"""Solve the linear problem :math:`\mathbf{y}=\mathbf{A}\mathbf{x}`.
 
         Overloading of operator ``/`` to improve expressivity of

pylops_distributed/waveeqprocessing/__init__.py

@@ -1,4 +1,4 @@
 from .mdd import MDC
-#from .mdd import MDD
+from .mdd import MDD
 from .marchenko import Marchenko
 from .lsm import Demigration, LSM

pylops_distributed/waveeqprocessing/marchenko.py

@@ -79,7 +79,7 @@ class Marchenko():
 
     """
     def __init__(self, R, nt, dt=0.004, dr=1., wav=None, toff=0.0,
-                 nsmooth=10, saveRt=True, prescaled=False, dtype='float32'):
+                 nsmooth=10, saveRt=False, prescaled=False, dtype='float32'):
         # Save inputs into class
         self.nt = nt
         self.dt = dt
@@ -105,7 +105,6 @@ def __init__(self, R, nt, dt=0.004, dr=1., wav=None, toff=0.0,
         # Save R
         self.Rtwosided_fft = R
 
-
     def apply_onepoint(self, trav, dist=None, G0=None, nfft=None,
                        rtm=False, greens=False, dottest=False,
                        **kwargs_cgls):

pylops_distributed/waveeqprocessing/mdd.py

@@ -1,14 +1,18 @@
 import logging
+import numpy as np
+import dask.array as da
 from pylops.waveeqprocessing.mdd import _MDC
 
+from pylops_distributed.utils import dottest as Dottest
 from pylops_distributed import Identity, Transpose
 from pylops_distributed.signalprocessing import FFT, Fredholm1
+from pylops_distributed.optimization.cg import cgls
 
 logging.basicConfig(format='%(levelname)s: %(message)s', level=logging.WARNING)
 
 
 def MDC(G, nt, nv, dt=1., dr=1., twosided=True,
-        saveGt=True, conj=False, prescaled=False,
+        saveGt=False, conj=False, prescaled=False,
         compute=(False, False), todask=(False, False)):
     r"""Multi-dimensional convolution.
 
@@ -73,3 +77,126 @@ def MDC(G, nt, nv, dt=1., dr=1., twosided=True,
                                       (nt, G.shape[1], nv)),
                            'todask': (todask[1], False),
                            'compute':(False, compute[0])})
+
+
+
+def MDD(G, d, dt=0.004, dr=1., nfmax=None, wav=None,
+        twosided=True, adjoint=False, dottest=False,
+        saveGt=False, add_negative=True, **kwargs_cgls):
+    r"""Multi-dimensional deconvolution.
+
+    Solve multi-dimensional deconvolution problem using
+    :py:func:`scipy.sparse.linalg.lsqr` iterative solver.
+
+    Parameters
+    ----------
+    G : :obj:`dask.array.ndarray`
+        Multi-dimensional convolution kernel in frequency domain of size
+        :math:`[n_{f,max} \times n_s \times n_r]`
+    d : :obj:`dask.array.ndarray`
+        Data in time domain :math:`[n_t \times n_s (\times n_vs)]` if
+        ``twosided=False`` or ``twosided=True`` and ``add_negative=True``
+        (with only positive times) or size
+        :math:`[2*n_t-1 \times n_s (\times n_vs)]` if ``twosided=True``
+    dt : :obj:`float`, optional
+        Sampling of time integration axis
+    dr : :obj:`float`, optional
+        Sampling of receiver integration axis
+    nfmax : :obj:`int`, optional
+        Index of max frequency to include in deconvolution process
+    wav : :obj:`numpy.ndarray`, optional
+        Wavelet to convolve to the inverted model and psf
+        (must be centered around its index in the middle of the array).
+        If ``None``, the outputs of the inversion are returned directly.
+    twosided : :obj:`bool`, optional
+        MDC operator and data both negative and positive time (``True``)
+        or only positive (``False``)
+    add_negative : :obj:`bool`, optional
+        Add negative side to MDC operator and data (``True``) or not
+        (``False``)- operator and data are already provided with both positive
+        and negative sides. To be used only with ``twosided=True``.
+    adjoint : :obj:`bool`, optional
+        Compute and return adjoint(s)
+    dottest : :obj:`bool`, optional
+        Apply dot-test
+    saveGt : :obj:`bool`, optional
+        Save ``G`` and ``G^H`` to speed up the computation of adjoint of
+        :class:`pylops_distributed.signalprocessing.Fredholm1` (``True``) or
+        create ``G^H`` on-the-fly (``False``) Note that ``saveGt=True`` will be
+        faster but double the amount of required memory
+    **kwargs_cgls
+        Arbitrary keyword arguments for
+        :py:func:`pylops_distributed.optimization.cg.cgls` solver
+
+    Returns
+    -------
+    minv : :obj:`dask.array.ndarray`
+        Inverted model of size :math:`[n_t \times n_r (\times n_{vs})]`
+        for ``twosided=False`` or
+        :math:`[2*n_t-1 \times n_r (\times n_vs)]` for ``twosided=True``
+    madj : :obj:`dask.array.ndarray`
+        Adjoint model of size :math:`[n_t \times n_r (\times n_{vs})]`
+        for ``twosided=False`` or
+        :math:`[2*n_t-1 \times n_r (\times n_r) ]` for ``twosided=True``
+
+    See Also
+    --------
+    MDC : Multi-dimensional convolution
+
+    Notes
+    -----
+    Refer to :class:`pylops.waveeqprocessing.MDD` for implementation
+    details. Note that this implementation is currently missing the ``wav``
+    and ``causality_precond=False`` options.
+
+    """
+    nf, ns, nr = G.shape
+    nt = d.shape[0]
+    if len(d.shape) == 2:
+        nv = 1
+    else:
+        nv = d.shape[2]
+    if twosided:
+        if add_negative:
+            nt2 = 2 * nt - 1
+        else:
+            nt2 = nt
+            nt = (nt2 + 1) // 2
+        nfmax_allowed = int(np.ceil((nt2+1)/2))
+    else:
+        nt2 = nt
+        nfmax_allowed = nt
+
+    # Fix nfmax to be at maximum equal to half of the size of fft samples
+    if nfmax is None or nfmax > nfmax_allowed:
+        nfmax = nfmax_allowed
+        logging.warning('nfmax set equal to ceil[(nt+1)/2=%d]' % nfmax)
+
+    # Add negative part to data and model
+    if twosided and add_negative:
+        if nv == 1:
+            d = da.concatenate((da.zeros((nt - 1, ns)), d), axis=0)
+        else:
+            d = da.concatenate((da.zeros((nt - 1, ns, nv)), d), axis=0)
+
+    # Define MDC linear operator
+    MDCop = MDC(G, nt2, nv=nv, dt=dt, dr=dr,
+                twosided=twosided, saveGt=saveGt)
+    if dottest:
+        Dottest(MDCop, nt2 * ns * nv, nt2 * nr * nv, verb=True)
+
+    # Adjoint
+    if adjoint:
+        madj = MDCop.H * d.flatten()
+        madj = da.squeeze(madj.reshape(nt2, nr, nv))
+
+    # Inverse
+    minv = cgls(MDCop, d.flatten(), **kwargs_cgls)[0]
+    minv = da.squeeze(minv.reshape(nt2, nr, nv))
+    #if wav is not None:
+    #    minv = sp_convolve1d(minv, wav, axis=-1)
+
+    if adjoint:
+        return minv, madj
+    else:
+        return minv

pytests/test_waveeqprocessing.py

@@ -8,6 +8,7 @@
 from pylops.waveeqprocessing.mdd import MDC
 
 from pylops_distributed.waveeqprocessing.mdd import MDC as dMDC
+from pylops_distributed.waveeqprocessing.mdd import MDD
 from pylops_distributed.utils import dottest
 
 PAR = {'ox': 0, 'dx': 2, 'nx': 10,
@@ -144,20 +145,32 @@ def test_MDC_1virtualsource(par):
     # Define MDC linear operator
     Gwav_fft = np.fft.fft(Gwav, par['nt2'], axis=-1)
     Gwav_fft = Gwav_fft[..., :par['nfmax']]
+    dGwav_fft = da.from_array(Gwav_fft.transpose(2, 0, 1))
 
-    dMDCop = dMDC(da.from_array(Gwav_fft.transpose(2, 0, 1)),
-                  nt=par['nt2'], nv=1, dt=par['dt'], dr=par['dx'],
+    dMDCop = dMDC(dGwav_fft, nt=par['nt2'], nv=1, dt=par['dt'], dr=par['dx'],
                   twosided=par['twosided'])
     MDCop = MDC(Gwav_fft.transpose(2, 0, 1), nt=par['nt2'], nv=1,
                 dt=par['dt'], dr=par['dx'], twosided=par['twosided'],
                 transpose=False, dtype='float32')
     dottest(dMDCop, par['nt2'] * par['ny'], par['nt2'] * par['nx'],
             chunks=((par['nt2'] * par['ny'], par['nt2'] * par['nx'])))
 
+    # Compare results with pylops implementation
     mwav = mwav.T
-    dy = (dMDCop * da.from_array(mwav.flatten())).compute()
+    dy = dMDCop * da.from_array(mwav.flatten())
     y = MDCop * mwav.flatten()
-    assert_array_almost_equal(dy, y, decimal=5)
+    assert_array_almost_equal(dy.compute(), y, decimal=5)
+
+    # Apply mdd function
+    dy = dy.reshape(par['nt2'], par['ny'])
+    print(dy)
+    minv = MDD(dGwav_fft,
+               dy[par['nt'] - 1:] if par['twosided'] else dy,
+               dt=par['dt'], dr=par['dx'], nfmax=par['nfmax'],
+               twosided=par['twosided'], adjoint=False, dottest=False,
+               **dict(niter=50))
+    print('minv', minv)
+    assert_array_almost_equal(mwav, minv.compute(), decimal=2)
 
 
 @pytest.mark.parametrize("par", [(par1), (par2), (par3), (par4),
@@ -204,8 +217,9 @@ def test_MDC_Nvirtualsources(par):
     # Define MDC linear operator
     Gwav_fft = np.fft.fft(Gwav, par['nt2'], axis=-1)
     Gwav_fft = Gwav_fft[..., :par['nfmax']]
+    dGwav_fft = da.from_array(Gwav_fft.transpose(2, 0, 1))
 
-    dMDCop = dMDC(da.from_array(Gwav_fft.transpose(2, 0, 1)), nt=par['nt2'],
+    dMDCop = dMDC(dGwav_fft, nt=par['nt2'],
                   nv=par['nx'], dt=par['dt'], dr=par['dx'],
                   twosided=par['twosided'])
     MDCop = MDC(Gwav_fft.transpose(2, 0, 1), nt=par['nt2'], nv=par['nx'],
@@ -217,7 +231,20 @@ def test_MDC_Nvirtualsources(par):
             chunks=((par['nt2'] * par['ny'] * par['nx'],
                      par['nt2'] * par['nx'] * par['nx'])))
 
+    # Compare results with pylops implementation
     mwav = mwav.T
-    dy = (dMDCop * da.from_array(mwav.flatten())).compute()
+    dy = (dMDCop * da.from_array(mwav.flatten()))
     y = MDCop * mwav.flatten()
-    assert_array_almost_equal(dy, y, decimal=5)
+    assert_array_almost_equal(dy.compute(), y, decimal=5)
+
+    # Apply mdd function
+    dy = dy.reshape(par['nt2'], par['ny'], par['nx'])
+
+    print(dy.shape)
+    minv = MDD(dGwav_fft,
+               dy[par['nt'] - 1:] if par['twosided'] else dy,
+               dt=par['dt'], dr=par['dx'], nfmax=par['nfmax'],
+               twosided=par['twosided'], adjoint=False,
+               dottest=False, **dict(niter=50))
+    print(mwav.shape, minv)
+    assert_array_almost_equal(mwav, minv.compute(), decimal=2)

tutorials/marchenko.py

@@ -10,8 +10,14 @@
 This allows lazy loading using a `Dask <https://dask.org>`_
 array and distributing over frequencies the computation of the various
 Fredholm integrals involved in the forward model.
+
+**NOTE:** do not expect this code to run any fast than its `pylops equivalent
+<https://pylops.readthedocs.io/en/latest/tutorials/marchenko.html#sphx-glr-tutorials-marchenko-py>`_.
+for small datasets. The pylops-distributed framework should only be used when
+dealing with large datasets that do not fit in memory and benefit from
+distributed computing.
+
 """
-# sphinx_gallery_thumbnail_number = 3
 import warnings
 import numpy as np
 import dask.array as da
@@ -23,6 +29,8 @@
 warnings.filterwarnings('ignore')
 plt.close('all')
 
+# sphinx_gallery_thumbnail_number = 4
+
 ###############################################################################
 # Let's start by defining some input parameters and loading the test data