Added tutorial on 2D LS-RTM through devito

.gitignore

@@ -9,3 +9,4 @@ dask-worker-space
 /devel
 experiments
 publishing.md
+tutorials/devito

tutorials/born_devito.py

@@ -0,0 +1,140 @@
+import numpy as np
+import occamypy as o
+from typing import Tuple, List
+from devitoseismic import AcquisitionGeometry, demo_model, SeismicModel
+from devitoseismic.acoustic import AcousticWaveSolver
+import devito
+
+devito.configuration['log-level'] = 'ERROR'
+
+
+def create_models(args: dict) -> Tuple[SeismicModel, SeismicModel, SeismicModel]:
+    hard = demo_model('layers-isotropic', origin=(0., 0.),
+                      shape=args["shape"], spacing=args["spacing"],
+                      nbl=args["nbl"], grid=None, nlayers=2)
+    smooth = demo_model('layers-isotropic', origin=(0., 0.),
+                        shape=args["shape"], spacing=args["spacing"],
+                        nbl=args["nbl"], grid=hard.grid, nlayers=2)
+
+    devito.gaussian_smooth(smooth.vp, sigma=args["filter_sigma"])
+
+    water = demo_model('layers-isotropic', origin=(0., 0.),
+                       shape=args["shape"], spacing=args["spacing"],
+                       nbl=args["nbl"], grid=hard.grid, nlayers=1)
+
+    return hard, smooth, water
+
+
+def build_src_coordinates(x: float, z: float) -> np.ndarray:
+    src = np.empty((1, 2), dtype=np.float32)
+    src[0, :] = x
+    src[0, -1] = z
+    return src
+
+
+def build_rec_coordinates(model: SeismicModel, args: dict) -> np.ndarray:
+    """Receivers equispaced on the whole domain"""
+    rec = np.empty((args["nreceivers"], 2))
+    rec[:, 0] = np.linspace(0, model.domain_size[0], num=args["nreceivers"])
+    rec[:, 1] = args["rec_depth"]
+
+    return rec
+
+
+def direct_arrival_mask(data: o.Vector, rec_pos: np.ndarray, src_pos: np.ndarray,
+                        vel_sep: float = 1500., offset: float = 0.) -> o.Vector:
+    dt = data.ax_info[0].d
+
+    direct = np.sqrt(np.sum((src_pos - rec_pos) ** 2, axis=1)) / vel_sep
+    direct += offset
+
+    mask = data.clone().zero()
+
+    iwin = np.round(direct / dt).astype(int)
+    for i in range(rec_pos.shape[0]):
+        mask[iwin[i]:, i] = 1.
+
+    return mask
+
+
+def _propagate_shot(model: SeismicModel, rec_pos: np.ndarray, src_pos: np.ndarray, param: dict) -> o.VectorNumpy:
+    geometry = AcquisitionGeometry(model, rec_pos, src_pos, **param)
+    solver = AcousticWaveSolver(model, geometry, **param)
+
+    devito.clear_cache()
+
+    # propagate (source -> receiver data)
+    data = o.VectorNumpy(solver.forward()[0].data.__array__())
+
+    data.ax_info = [o.AxInfo(geometry.nt, geometry.t0, geometry.dt / 1000, "time [s]"),
+                    o.AxInfo(geometry.nrec, float(rec_pos[0][0]), float(rec_pos[1][0] - rec_pos[0][0]), "rec pos x [m]")]
+
+    devito.clear_cache()
+    return data
+
+
+def propagate_shots(model: SeismicModel, rec_pos: np.ndarray, src_pos: List[np.ndarray], param: dict):
+    if len(src_pos) == 1:
+        return _propagate_shot(model=model, rec_pos=rec_pos, src_pos=src_pos[0], param=param)
+    else:
+        return o.superVector([_propagate_shot(model=model, rec_pos=rec_pos, src_pos=s, param=param) for s in src_pos])
+
+
+class BornSingleSource(o.Operator):
+
+    def __init__(self, velocity: SeismicModel, src_pos: np.ndarray, rec_pos: np.ndarray, args: dict):
+
+        # store params
+        self.src_pos = src_pos
+        self.rec_pos = rec_pos
+        self.nbl = args["nbl"]
+
+        # build geometry and acoustic solver
+        self.geometry = AcquisitionGeometry(velocity, rec_pos, src_pos, **args)
+        self.solver = AcousticWaveSolver(velocity, self.geometry, **args)
+
+        # allocate vectors
+        self.velocity = o.VectorNumpy(velocity.vp.data.__array__())
+        self.velocity.ax_info = [
+            o.AxInfo(velocity.vp.shape[0], velocity.origin[0] - self.nbl * velocity.spacing[0], velocity.spacing[0],
+                     "x [m]"),
+            o.AxInfo(velocity.vp.shape[1], velocity.origin[1] - self.nbl * velocity.spacing[1], velocity.spacing[1],
+                     "z [m]")]
+
+        csg = o.VectorNumpy((self.geometry.nt, self.geometry.nrec))
+        csg.ax_info = [o.AxInfo(self.geometry.nt, self.geometry.t0, self.geometry.dt / 1000, "time [s]"),
+                       o.AxInfo(self.geometry.nrec, float(rec_pos[0][0]), float(rec_pos[1][0] - rec_pos[0][0]),
+                                "rec pos x [m]")]
+
+        super(BornSingleSource, self).__init__(self.velocity, csg)
+
+        # store source wavefield
+        self.src_wfld = self.solver.forward(save=True)[1]
+
+    def __str__(self):
+        return "DeviBorn"
+
+    def forward(self, add, model, data):
+        """Modeling function: image -> residual data"""
+        self.checkDomainRange(model, data)
+        if not add:
+            data.zero()
+
+        recs = self.solver.jacobian(dmin=model[:])[0]
+        data[:] += recs.data.__array__()
+
+        return
+
+    def adjoint(self, add, model, data):
+        """Adjoint function: data -> image"""
+        self.checkDomainRange(model, data)
+        if not add:
+            model.zero()
+
+        recs = self.geometry.rec.copy()
+        recs.data[:] = data[:]
+
+        img = self.solver.gradient(rec=recs, u=self.src_wfld)[0]
+        model[:] += img.data.__array__()
+
+        return

tutorials/devitoseismic/__init__.py

@@ -0,0 +1,5 @@
+from .model import *  # noqa
+from .source import *  # noqa
+from .plotting import *  # noqa
+from .preset_models import *  # noqa
+from .utils import *  # noqa

tutorials/devitoseismic/acoustic/__init__.py

@@ -0,0 +1,3 @@
+from .operators import *  # noqa
+from .wavesolver import *  # noqa
+from .acoustic_example import *  # noqa

tutorials/devitoseismic/acoustic/acoustic_example.py

@@ -0,0 +1,101 @@
+import numpy as np
+import pytest
+
+from devito.logger import info
+from devito import Constant, Function, smooth, norm
+from . import AcousticWaveSolver
+from .. import demo_model, setup_geometry, seismic_args
+
+
+def acoustic_setup(shape=(50, 50, 50), spacing=(15.0, 15.0, 15.0),
+                   tn=500., kernel='OT2', space_order=4, nbl=10,
+                   preset='layers-isotropic', fs=False, **kwargs):
+    model = demo_model(preset, space_order=space_order, shape=shape, nbl=nbl,
+                       dtype=kwargs.pop('dtype', np.float32), spacing=spacing,
+                       fs=fs, **kwargs)
+
+    # Source and receiver geometries
+    geometry = setup_geometry(model, tn)
+
+    # Create solver object to provide relevant operators
+    solver = AcousticWaveSolver(model, geometry, kernel=kernel,
+                                space_order=space_order, **kwargs)
+    return solver
+
+
+def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=1000.0,
+        space_order=4, kernel='OT2', nbl=40, full_run=False, fs=False,
+        autotune=False, preset='layers-isotropic', checkpointing=False, **kwargs):
+
+    solver = acoustic_setup(shape=shape, spacing=spacing, nbl=nbl, tn=tn,
+                            space_order=space_order, kernel=kernel, fs=fs,
+                            preset=preset, **kwargs)
+
+    info("Applying Forward")
+    # Whether or not we save the whole time history. We only need the full wavefield
+    # with 'save=True' if we compute the gradient without checkpointing, if we use
+    # checkpointing, PyRevolve will take care of the time history
+    save = full_run and not checkpointing
+    # Define receiver geometry (spread across x, just below surface)
+    rec, u, summary = solver.forward(save=save, autotune=autotune)
+
+    if preset == 'constant':
+        # With  a new m as Constant
+        v0 = Constant(name="v", value=2.0, dtype=np.float32)
+        solver.forward(save=save, vp=v0)
+        # With a new vp as a scalar value
+        solver.forward(save=save, vp=2.0)
+
+    if not full_run:
+        return summary.gflopss, summary.oi, summary.timings, [rec, u.csg_nonlinear]
+
+    # Smooth velocity
+    initial_vp = Function(name='v0', grid=solver.model.grid, space_order=space_order)
+    smooth(initial_vp, solver.model.vp)
+    dm = np.float32(initial_vp.data ** (-2) - solver.model.vp.csg_nonlinear ** (-2))
+
+    info("Applying Adjoint")
+    solver.adjoint(rec, autotune=autotune)
+    info("Applying Born")
+    solver.jacobian(dm, autotune=autotune)
+    info("Applying Gradient")
+    solver.jacobian_adjoint(rec, u, autotune=autotune, checkpointing=checkpointing)
+    return summary.gflopss, summary.oi, summary.timings, [rec, u.csg_nonlinear]
+
+
+@pytest.mark.parametrize('shape', [(101,), (51, 51), (16, 16, 16)])
+@pytest.mark.parametrize('k', ['OT2', 'OT4'])
+def test_isoacoustic_stability(shape, k):
+    spacing = tuple([20]*len(shape))
+    _, _, _, [rec, _] = run(shape=shape, spacing=spacing, tn=20000.0, nbl=0, kernel=k)
+    assert np.isfinite(norm(rec))
+
+
+@pytest.mark.parametrize('fs, normrec, dtype', [(True, 369.955, np.float32),
+                                                (False, 459.1678, np.float64)])
+def test_isoacoustic(fs, normrec, dtype):
+    _, _, _, [rec, _] = run(fs=fs, dtype=dtype)
+    assert np.isclose(norm(rec), normrec, rtol=1e-3, atol=0)
+
+
+if __name__ == "__main__":
+    description = ("Example script for a set of acoustic operators.")
+    parser = seismic_args(description)
+    parser.add_argument('--fs', dest='fs', default=False, action='store_true',
+                        help="Whether or not to use a freesurface")
+    parser.add_argument("-k", dest="kernel", default='OT2',
+                        choices=['OT2', 'OT4'],
+                        help="Choice of finite-difference kernel")
+    args = parser.parse_args()
+
+    # 3D preset parameters
+    ndim = args.ndim
+    shape = args.shape[:args.ndim]
+    spacing = tuple(ndim * [15.0])
+    tn = args.tn if args.tn > 0 else (750. if ndim < 3 else 1250.)
+
+    preset = 'constant-isotropic' if args.constant else 'layers-isotropic'
+    run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn, fs=args.fs,
+        space_order=args.space_order, preset=preset, kernel=args.kernel,
+        autotune=args.autotune, opt=args.opt, full_run=args.full,
+        checkpointing=args.checkpointing, dtype=args.dtype)