import matplotlib.pyplot as plt

from devito import *
from devitopro import *
from recipes import AcousticIsotropic, layered_model
from devitopro.transforms import fft, ifft
from devitopro.types.dense import FFTFunction

class UpDownAcousticIsotropic(AcousticIsotropic):

    def sensitivity(self, params=('m',)):
        u = self.saved_state['u_save']
        v = self.state['u']
        vdt = v.dt if self.options['time_slots'] == 2 else -v.dt
        tcond = u.dimensions[0]

        cdtype = np.complex64

        # RTM imaging condition with up down decomposition
        # https://library.seg.org/doi/epub/10.1190/1.3533914
        # 1. Computes fourier transform of u and v
        # 2. Set negative wavenumbers to zero
        # 3. Computes inverse fourier transform u* and v*
        # 4. Computes imaging condition as 2 * re(u* * v*)

        # Currently, using FFTFunction is required to avoid issues with padding and
        # get best performance.
        u_f = FFTFunction(name='u_f', grid=self.grid, dtype=cdtype, space_order=0)
        v_f = FFTFunction(name='v_f', grid=self.grid, dtype=cdtype, space_order=0)

        z = self.grid.dimensions[-1]
        nz = self.grid.shape[-1]
    
        mask = Function(name='mask', grid=self.grid, space_order=0, dimensions=(z,), shape=(nz,))
        mask.data[:] = 1
        mask.data[nz//2+1:] = 0

        # Forward fft
        eqs_fwd = [Eq(u_f, u.dt), Eq(v_f, vdt, implicit_dims=(tcond,)),
                   fft(u_f, u_f, dims=z, implicit_dims=(tcond,)),
                   fft(v_f, v_f, dims=z, implicit_dims=(tcond,))]
        # Set negative wavenumbers to zero
        filter = [Eq(u_f, mask * u_f, implicit_dims=(tcond,)),
                  Eq(v_f, mask * v_f, implicit_dims=(tcond,))]
        # Inverse fft
        eqs_inv = [ifft(u_f, u_f, dims=z, implicit_dims=(tcond,)),
                   ifft(v_f, v_f, dims=z, implicit_dims=(tcond,))]
        # Imaging condition
        dm = self.perturbation(param=params[0])
        dt = u.indices[0].spacing
        imaging = Eq(dm, dm - 2 * dt * Real(u_f * v_f), implicit_dims=(tcond,))
        return eqs_fwd + filter + eqs_inv + [imaging]
so = 8
model = layered_model('iso-acoustic', nlayers=2, shape=(601, 151), spacing=(20, 20), nbl=80, space_order=so)
model0 = layered_model('iso-acoustic', nlayers=2, shape=(601, 151), spacing=(20, 20), space_order=so, dt=model.critical_dt, nbl=80)
model0.smooth(('vp',), sigma=5)
options = {'nt': 2001, 'f0': 0.005, 'space_order': so}
#NBVAL_IGNORE_OUTPUT

slicevp = model.vp.data[80:-80, 80:-80]
slicevp0 = model0.vp.data[80:-80, 80:-80]
plt.figure(figsize=(12, 6))
plt.subplot(121)
plt.imshow(slicevp.T, aspect="auto")
plt.title("P-wave velocity")
plt.colorbar()
plt.subplot(122)
plt.imshow(slicevp0.T, aspect="auto")
plt.title("P-wave velocity")
plt.colorbar()
plt.show()

solver = AcousticIsotropic(model, options)
solver0 = AcousticIsotropic(model0, options)
#NBVAL_IGNORE_OUTPUT

# Forward in true model
solver.forward()
# Forward in background model
solver0.forward(save=True)
Operator `dampvp` ran in 0.01 s
Operator `initdamp` ran in 0.01 s
Operator `ForwardAcousticIsotropic` ran in 0.41 s
Operator `dampvp` ran in 0.01 s
Operator `initdamp` ran in 0.01 s
Operator `ForwardAcousticIsotropic` ran in 0.65 s
PerformanceSummary([(PerfKey(name='section0', rank=None),
                     PerfEntry(time=0.00014199999999999998, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section1', rank=None),
                     PerfEntry(time=0.6231060000000042, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section2', rank=None),
                     PerfEntry(time=0.0006620000000000083, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section3', rank=None),
                     PerfEntry(time=0.011138000000000014, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section4', rank=None),
                     PerfEntry(time=1.4000000000000003e-05, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section5', rank=None),
                     PerfEntry(time=0.005716000000000093, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[]))])
#NBVAL_IGNORE_OUTPUT
plt.subplot(2,2,1)
plt.imshow(solver.rec.data, aspect="auto", vmin=-.1, vmax=.1, cmap="seismic")
plt.title("True data")
plt.subplot(2,2,2)
plt.imshow(solver0.rec.data, aspect="auto", vmin=-.1, vmax=.1, cmap="seismic")
plt.title("Background data")
plt.subplot(2,2,3)
plt.imshow(solver.rec.data-solver0.rec.data, aspect="auto", vmin=-.1, vmax=.1, cmap="seismic")
plt.title("Residual")
# Imaging
solver0.rec.data[:] -= solver.rec.data[:]
solver0.jacobian_adjoint()
slicedm = np.array(solver0.perturbation(param='m').data[80:-80, 80:-80])
clip = np.quantile(np.abs(slicedm.reshape(-1)), 0.9)
plt.subplot(2,2,4)
plt.imshow(slicedm.T, aspect="auto", cmap="Greys", vmin=-clip, vmax=clip)
plt.imshow(slicevp.T, aspect="auto", cmap="terrain", alpha=.5)
plt.title("Gradient")
plt.tight_layout()
plt.show()
Operator `JacobianAdjAcousticIsotropic` ran in 0.71 s

nz = model.grid.shape[-1]
mask = Function(name='mask', grid=model.grid, space_order=0,
                dimensions=(model.grid.dimensions[-1],),
                shape=(nz,))
mask.data[:] = 1
mask.data[nz//2:] = 0
#NBVAL_IGNORE_OUTPUT

solver0updown = UpDownAcousticIsotropic(model0, options)

configuration['log-level'] = 'DEBUG'
solver0updown.jacobian_adjoint(rec=solver0.rec, u_save=solver0.saved_state['u_save'])
Allocating host memory for src(2001, 1) [8 KB]
Allocating host memory for src_coords(1, 2) [8 B]
Allocating host memory for rec_coords(599, 2) [5 KB]
Allocating host memory for mask(320,) [1 KB]
Hint: avoid non-uniform halo along `x`. Detected `[(0, 0), (4, 4)]`
Hint: avoid non-uniform halo along `y`. Detected `[(0, 0), (4, 4)]`
Operator `JacobianAdjUpDownAcousticIsotropic` generated in 1.19 s
  * lowering.IET: 0.56 s (47.1 %)
     * specializing.IET: 0.41 s (34.5 %)
  * lowering.Clusters: 0.53 s (44.6 %)
     * specializing.Clusters: 0.42 s (35.4 %)
Flops reduction after symbolic optimization: [84 --> 62]
Operator `JacobianAdjUpDownAcousticIsotropic` fetched `/var/folders/5k/328zp6912nd1h1bqnylt0wph0000gp/T/devito-jitcache-uid502/bae8262decb5fc1144c8510d1eeef96b0baa3c5c.c` in 0.11 s from jit-cache
Allocating host memory for dm(769, 320) [961 KB]
Allocating host memory for u(3, 769, 320) [3 MB]
Allocating host memory for u_f(761, 315) [2 MB]
Allocating host memory for v_f(761, 315) [2 MB]
Operator `JacobianAdjUpDownAcousticIsotropic` arguments preprocess: 0.01 s
Operator `JacobianAdjUpDownAcousticIsotropic` arguments postprocess: 0.01 s
Operator `JacobianAdjUpDownAcousticIsotropic` ran in 2.96 s
Global performance: [8.90 GFlops/s, 0.16 GPts/s]
Global performance <w/o setup>: [1.10 s, 0.44 GPts/s]
Local performance:
  * section0 ran in 0.01 s
  * section1 ran in 0.63 s [OI=1.35, 33.46 GFlops/s, 0.77 GPts/s]
  * section2 ran in 0.04 s [OI=3.20, 7.84 GFlops/s, 0.13 GPts/s]
  * section3 ran in 0.11 s [OI=0.25, 9.05 GFlops/s]
  * section4 ran in 0.11 s [OI=0.13, 9.07 GFlops/s]
  * section5 ran in 0.01 s
  * section6 ran in 0.01 s
  * section7 ran in 0.16 s [OI=0.05, 6.14 GFlops/s]
  * section8 ran in 0.06 s [OI=0.16, 39.79 GFlops/s]
  * section9 ran in 0.01 s
Performance[mode=advanced] arguments: {'nthreads': 1, 'nthreads_nonaffine': 1, 'pthreads': 1}
PerformanceSummary([(PerfKey(name='section0', rank=None),
                     PerfEntry(time=0.000261, gflopss=1.5325670498084293e-05, gpointss=0.0, oi=5.281571119524067e-07, ops=2, itershapes=((2, 1), (2, 1)))),
                    (PerfKey(name='section1', rank=None),
                     PerfEntry(time=0.6222709999999975, gflopss=33.45268295646123, gpointss=0.7602882490104825, oi=1.3496645037041715, ops=44, itershapes=((1999, 761, 311),))),
                    (PerfKey(name='section2', rank=None),
                     PerfEntry(time=0.03912900000000003, gflopss=7.833950675969224, gpointss=0.12240547931201913, oi=3.2, ops=64, itershapes=((1999, 599, 2, 2), (1999, 599, 2, 2)))),
                    (PerfKey(name='section3', rank=None),
                     PerfEntry(time=0.10455500000000006, gflopss=9.049884347950835, gpointss=0.0, oi=0.24975012493753124, ops=2, itershapes=((1999, 761, 311),))),
                    (PerfKey(name='section4', rank=None),
                     PerfEntry(time=0.10432899999999995, gflopss=9.06948842603687, gpointss=0.0, oi=0.12496874218554639, ops=2, itershapes=((1999, 761, 311),))),
                    (PerfKey(name='section5', rank=None),
                     PerfEntry(time=0.004535999999999988, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section6', rank=None),
                     PerfEntry(time=0.004764000000000012, gflopss=0.0, gpointss=0.0, oi=0.0, ops=0, itershapes=[])),
                    (PerfKey(name='section7', rank=None),
                     PerfEntry(time=0.15419000000000016, gflopss=6.136653855632655, gpointss=0.0, oi=0.05, ops=2, itershapes=((1999, 761, 311),))),
                    (PerfKey(name='section8', rank=None),
                     PerfEntry(time=0.05946199999999996, gflopss=39.78215742827355, gpointss=0.0, oi=0.15625, ops=5, itershapes=((1999, 761, 311),))),
                    (PerfKey(name='section9', rank=None),
                     PerfEntry(time=4.699999999999999e-05, gflopss=0.08506382978723405, gpointss=0.0, oi=1.055258437574074e-06, ops=2, itershapes=((1999, 1), (1999, 1))))])
#NBVAL_IGNORE_OUTPUT

slicedm1 = np.array(solver0updown.perturbation(param='m').data[80:-80, 80:-80])
clip = np.max(np.abs(slicedm.reshape(-1))) / 4
clip1 = np.max(np.abs(slicedm1.reshape(-1))) / 4
plt.figure(figsize=(12, 4))
plt.subplot(121)
plt.imshow(slicedm.T, aspect="auto", cmap="seismic", vmin=-clip, vmax=clip)
plt.subplot(122)
plt.imshow(slicedm1.T, aspect="auto", cmap="seismic", vmin=-clip1, vmax=clip1)
plt.title("Gradient")
plt.tight_layout()
plt.show()

#NBVAL_IGNORE_OUTPUT

import numpy as np

assert np.isclose(np.linalg.norm(slicedm1), 3.499*1e10, rtol=1e-3)
assert np.isclose(np.linalg.norm(slicedm), 5.497*1e5, rtol=1e-3)
Back to top