repr: thesis scripts

examples/seismic/acoustic/bench_acoustic_thesis.py

@@ -0,0 +1,97 @@
+# Based on the implementation of the Devito acoustic example implementation
+import numpy as np
+from devito import TimeFunction, Eq, Operator, solve, norm
+from examples.seismic import Model, TimeAxis
+
+import matplotlib.pyplot as plt
+from mpl_toolkits.mplot3d import Axes3D
+
+import argparse
+
+parser = argparse.ArgumentParser(description='Process arguments.')
+
+parser.add_argument("-d", "--shape", default=(11, 11, 11), type=int, nargs="+",
+                    help="Number of grid points along each axis")
+parser.add_argument("-so", "--space_order", default=4,
+                    type=int, help="Space order of the simulation")
+parser.add_argument("-to", "--time_order", default=2,
+                    type=int, help="Time order of the simulation")
+parser.add_argument("-tn", "--tn", default=40,
+                    type=int, help="Simulation time in millisecond")
+parser.add_argument("-bls", "--blevels", default=2, type=int, nargs="+",
+                    help="Block levels")
+args = parser.parse_args()
+
+
+# Define a physical size
+nx, ny, nz = args.shape
+nt = args.tn
+
+shape = (nx, ny, nz)  # Number of grid point (nx, ny, nz)
+spacing = (10., 10., 10.)  # Grid spacing in m. The domain size is now 1km by 1km
+origin = (0., 0., 0.)  # What is the location of the top left corner. This is necessary to define
+# the absolute location of the source and receivers
+
+# Define a velocity profile. The velocity is in km/s
+v = np.empty(shape, dtype=np.float32)
+v[:, :, :51] = 1.5
+v[:, :, 51:] = 2.5
+
+# With the velocity and model size defined, we can create the seismic model that
+# encapsulates this properties. We also define the size of the absorbing layer as 10 grid points
+so = args.space_order
+to = args.time_order
+
+model = Model(vp=v, origin=origin, shape=shape, spacing=spacing,
+              space_order=so, nbl=10, bcs="damp")
+
+# plot_velocity(model)
+
+t0 = 0.  # Simulation starts a t=0
+tn = nt  # Simulation last 1 second (1000 ms)
+dt = model.critical_dt  # Time step from model grid spacing
+
+time_range = TimeAxis(start=t0, stop=tn, step=dt)
+
+# Define the wavefield with the size of the model and the time dimension
+u = TimeFunction(name="u", grid=model.grid, time_order=to, space_order=so)
+
+u.data[0, int(nx/2), int(ny/3), 10] = 5
+u.data[0, int(nx/3), int(ny/2), 10] = 1
+
+# We can now write the PDE
+pde = model.m * u.dt2 - u.laplace + model.damp * u.dt
+
+# The PDE representation is as on paper
+pde
+
+stencil = Eq(u.forward, solve(pde, u.forward))
+stencil
+
+op = Operator([stencil], subs=model.spacing_map)
+op.apply(time=time_range.num-1, dt=model.critical_dt)
+# op.apply(time=time_range.num-1, dt=model.critical_dt, **{'x0_blk0_size': 16, 'y0_blk0_size': 8})
+print(norm(u))
+
+# Initialize data
+u.data[:] = 0
+u.data[0, int(nx/2), int(ny/3), 10] = 5
+u.data[0, int(nx/3), int(ny/2), 10] = 1
+
+op1 = Operator([stencil], subs=model.spacing_map,
+               opt=('advanced', {'skewing': True, 'blocklevels': 2}))
+# op1.apply(time=time_range.num-2, dt=model.critical_dt, **{'time0_blk0_size': 282, 'x0_blk0_size': 32, 'x0_blk1_size': 4, 'y0_blk0_size': 32, 'y0_blk1_size': 4})
+op1.apply(time=time_range.num-2, dt=model.critical_dt)
+print(norm(u))
+
+# import pdb;pdb.set_trace()
+
+import matplotlib.pyplot as plt
+from examples.cfd import plot_field
+from examples.seismic import plot_image
+
+# r_square = (u.data[0])
+# plot_3D_array_slices(r_square)
+
+plot_image(u.data[0,:,:,10], cmap="viridis")
+plt.show()

examples/seismic/tti/bench_tti_thesis.py

@@ -0,0 +1,89 @@
+import numpy as np
+import pytest
+
+from devito import Constant, Function, smooth, norm, info
+
+from examples.seismic import demo_model, setup_geometry, seismic_args
+from examples.seismic.tti import AnisotropicWaveSolver
+
+
+def tti_setup(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=250.0,
+              kernel='centered', space_order=4, nbl=10, preset='layers-tti',
+              **kwargs):
+
+    # Two layer model for true velocity
+    model = demo_model(preset, shape=shape, spacing=spacing,
+                       space_order=space_order, nbl=nbl, **kwargs)
+
+    # Source and receiver geometries
+    geometry = setup_geometry(model, tn)
+
+    return AnisotropicWaveSolver(model, geometry=geometry, space_order=space_order,
+                                 kernel=kernel, **kwargs)
+
+
+def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=250.0,
+        autotune=False, space_order=4, nbl=10, preset='layers-tti',
+        kernel='centered', full_run=False, checkpointing=False, **kwargs):
+
+    solver = tti_setup(shape=shape, spacing=spacing, tn=tn, space_order=space_order,
+                       nbl=nbl, kernel=kernel, 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, y`` just below surface)
+    u, v, summary = solver.forward_plain(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, [u, v]
+
+    # 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.data**(-2))
+
+    return summary.gflopss, summary.oi, summary.timings, [u, v]
+
+
+@pytest.mark.parametrize('shape', [(51, 51), (16, 16, 16)])
+@pytest.mark.parametrize('kernel', ['centered', 'staggered'])
+def test_tti_stability(shape, kernel):
+    spacing = tuple([20]*len(shape))
+    _, _, _, [rec, _, _] = run(shape=shape, spacing=spacing, kernel=kernel,
+                               tn=16000.0, nbl=0)
+    assert np.isfinite(norm(rec))
+
+
+if __name__ == "__main__":
+    description = ("Example script to execute a TTI forward operator.")
+    parser = seismic_args(description)
+    parser.add_argument('--noazimuth', dest='azi', default=False, action='store_true',
+                        help="Whether or not to use an azimuth angle")
+    parser.add_argument("-k", dest="kernel", default='centered',
+                        choices=['centered', 'staggered'],
+                        help="Choice of finite-difference kernel")
+    args = parser.parse_args()
+
+    # Switch to TTI kernel if input is acoustic kernel
+    preset = 'layers-tti-noazimuth' if args.azi else 'layers-tti'
+
+    # Preset parameters
+    ndim = args.ndim
+    shape = args.shape[:args.ndim]
+    spacing = tuple(ndim * [20.0])
+    tn = args.tn if args.tn > 0 else (750. if ndim < 3 else 1250.)
+
+    run(shape=shape, spacing=spacing, nbl=args.nbl, tn=tn,
+        space_order=args.space_order, autotune=args.autotune, dtype=args.dtype,
+        opt=args.opt, kernel=args.kernel, preset=preset,
+        checkpointing=args.checkpointing, full_run=args.full)