examples: Update superstepping script to use the examples.model

Author: JDBetteridge

examples/timestepping/acoustic_superstep_2d.py

@@ -1,81 +1,22 @@
 """ Script that demonstrates the functionality of the superstep in 2D
 Acoustic wave equation with source injection
 """
-import shutil
-import urllib.request
+import os
 from argparse import ArgumentParser
-from collections.abc import Callable
-from dataclasses import dataclass
-from pathlib import Path
 
 import matplotlib.pyplot as plt
 import numpy as np
-from scipy.interpolate import interpn
 
 from devito import (
     ConditionalDimension,
     Eq,
-    Function,
-    Grid,
     Operator,
     SparseTimeFunction,
     TimeFunction,
     solve,
 )
 from devito.timestepping.superstep import superstep_generator
-
-
-@dataclass
-class model:
-    name: str
-    velocity: Callable
-    # Spatial Domain
-    shape: tuple[int]
-    origin: tuple[float]
-    extent: tuple[float]
-    # Souce Location
-    source: tuple[float]
-    # Time Domain
-    t0: float
-    t1: float
-    t2: float
-    critical_dt: float
-    # Plotting
-    zlim: float
-
-
-def layered_velocity(grid, step=0):
-    velocity = Function(name="layered", grid=grid, space_order=(2, step, step))
-    q = grid.shape[1]//4
-    velocity.data[:] = 1500
-    velocity.data[:, q:2*q] = 2000
-    velocity.data[:, 2*q:3*q] = 2500
-    velocity.data[:, 3*q:] = 3000
-    return velocity
-
-
-def marmousi(grid, step=0):
-    # Grab dataset from pwd or download
-    url = 'https://github.com/devitocodes/data/raw/refs/heads/master/Simple2D/vp_marmousi_bi'  # noqa: E501
-    filename = Path('marmousi.np')
-    shape = (1601, 401)
-    if not filename.exists():
-        with urllib.request.urlopen(url) as response, open(filename, 'wb') as fh:
-            shutil.copyfileobj(response, fh)
-    data = np.fromfile(filename, dtype=np.float32, sep='').reshape(*shape)
-    cropped = data[650:1051, 35:]
-
-    xs = np.linspace(0, 1, cropped.shape[0])
-    ys = np.linspace(0, 1, cropped.shape[1])
-
-    xd = np.linspace(0, 1, grid.shape[0])
-    yd = np.linspace(0, 1, grid.shape[1])
-
-    velocity = Function(name="marmousi", grid=grid, space_order=(2, step, step))
-    velocity.data[:] = 1000*interpn(
-        (xs, ys), cropped, np.meshgrid(xd, yd), method='nearest'
-    ).T
-    return velocity
+from examples.seismic import demo_model, SeismicModel
 
 
 def ricker(t, f=10, A=1):
@@ -88,31 +29,35 @@ def ricker(t, f=10, A=1):
     return A * (1 - 2 * trm) * np.exp(-trm)
 
 
-def acoustic_model(model, step=1, snapshots=1):
+def acoustic_model(model, t0, t1, t2, critical_dt, source, step=1, snapshots=1):
     # Construct 2D Grid
-    grid = Grid(shape=model.shape, extent=model.extent)
-    x, y = grid.dimensions
-
-    t0, t1, t2 = model.t0, model.t1, model.t2
-
-    velocity = model.velocity(grid, step)
-    u = TimeFunction(name="u", grid=grid, time_order=2, space_order=2)
+    x, y = model.grid.dimensions
+    velocity = model.vp
+    u = TimeFunction(name="u", grid=model.grid, time_order=2, space_order=2)
 
     pde = (1/velocity**2)*u.dt2 - u.laplace
     stencil = Eq(u.forward, solve(pde, u.forward))
 
-    nt1 = int(np.ceil((t1 - t0)/model.critical_dt))
+    nt1 = int(np.ceil((t1 - t0)/critical_dt))
     dt = (t1 - t0)/nt1
 
     # Source
     t = np.linspace(t0, t1, nt1)
     rick = ricker(t)
     source = SparseTimeFunction(
-        name="ricker", npoint=1, coordinates=[model.source], nt=nt1, grid=grid,
-        time_order=2, space_order=4
+        name="ricker",
+        npoint=1,
+        coordinates=[source],
+        nt=nt1,
+        grid=model.grid,
+        time_order=2,
+        space_order=4
     )
     source.data[:, 0] = rick
-    src_term = source.inject(field=u.forward, expr=source*velocity*velocity*dt*dt)
+    src_term = source.inject(
+        field=u.forward,
+        expr=source*velocity**2*dt**2
+    )
 
     op1 = Operator([stencil] + src_term)
     op1(time=nt1 - 1, dt=dt)
@@ -121,7 +66,12 @@ def acoustic_model(model, step=1, snapshots=1):
     idx = nt1 % 3
     if step == 1:
         # Non-superstep case
-        new_u = TimeFunction(name="new_u", grid=grid, time_order=2, space_order=2)
+        new_u = TimeFunction(
+            name="new_u",
+            grid=model.grid,
+            time_order=2,
+            space_order=2
+        )
         stencil = [stencil.subs(
             {u.forward: new_u.forward, u: new_u, u.backward: new_u.backward}
         )]
@@ -131,16 +81,23 @@ def acoustic_model(model, step=1, snapshots=1):
     else:
         new_u, new_u_p, *stencil = superstep_generator(u, stencil.rhs, step, nt=nt1)
 
-    nt2 = int(np.ceil((t2 - t1)/model.critical_dt))
+    nt2 = int(np.ceil((t2 - t1)/critical_dt))
     dt = (t2 - t1)/nt2
 
     # Snapshot the solution
     factor = int(np.ceil(nt2/(snapshots + 1)))
-    t_sub = ConditionalDimension('t_sub', parent=grid.time_dim, factor=factor)
+    t_sub = ConditionalDimension(
+        't_sub',
+        parent=model.grid.time_dim,
+        factor=factor
+    )
     u_save = TimeFunction(
-        name='usave', grid=grid,
-        time_order=0, space_order=2,
-        save=snapshots//step + 1, time_dim=t_sub
+        name='usave',
+        grid=model.grid,
+        time_order=0,
+        space_order=2,
+        save=snapshots//step + 1,
+        time_dim=t_sub
     )
     save = Eq(u_save, new_u)
 
@@ -153,72 +110,82 @@ def acoustic_model(model, step=1, snapshots=1):
     return u_save.data
 
 
-layered_model = model(
-    name='layered',
-    velocity=layered_velocity,
-    shape=(101, 101),
-    origin=(0., 0.),
-    extent=(1000, 1000),  # 1kmx1km
-    source=(500, 20),
-    t0=0,
-    t1=0.2,
-    t2=0.65,
-    critical_dt=0.002357,
-    zlim=30
-)
-
-marmousi_model = model(
-    name='marmousi',
-    velocity=marmousi,
-    shape=(274, 301),
-    origin=(4875., 262.5),
-    extent=(3000, 2737.5),  # 3kmx2.7km
-    source=(1000, 1200),
-    t0=0,
-    t1=0.2,
-    t2=0.5,
-    critical_dt=0.0013728,
-    zlim=20
-)
-
-
 if __name__ == '__main__':
     parser = ArgumentParser()
     parser.add_argument('--model', default='layered', choices=['layered', 'marmousi'])
     args = parser.parse_args()
 
-    model = layered_model if args.model == 'layered' else marmousi_model
+    t0 = 0
+    t1 = 0.2
+    if args.model == 'layered':
+        source = (500, 20)
+        t2 = 0.65
+        critical_dt=0.002357
+        zlim = 30
+    else:  # Marmousi
+        # This requires the `devitocodes/data` repository, which we
+        # assume to be checked out at `$VIRTUAL_ENV/src/data`.
+        source = (1500, 1500)
+        t2 = 0.5
+        critical_dt=0.0013728
+        zlim = 20
+        tmp_model = demo_model(
+            'marmousi-isotropic',
+            space_order=2,
+            data_path=f'{os.environ["VIRTUAL_ENV"]}/src/data',
+            nbl=0
+        )
+        cropped = tmp_model.vp.data[400:701, -321:-20]
 
     # Supersteps
     k = [1, 4]
     # Snapshots
     m = 13
     fig, axes = plt.subplots(len(k), m)
 
-    # Velocity model
-    grid = Grid(shape=model.shape, extent=model.extent)
-    v = model.velocity(grid)
-    plot_extent = [
-        model.origin[0],
-        model.origin[0] + model.extent[0],
-        model.origin[1] + model.extent[1],
-        model.origin[1]
-    ]
-
     for step, ax_row in zip(k, axes, strict=True):
-        data = acoustic_model(model, step=step, snapshots=m)
-        time = np.linspace(model.t1, model.t2, (m - 1)//step + 1)
+        # Redefine the model every iteration because we need to adjust
+        # the space order
+        if args.model == 'layered':
+            model = demo_model(
+                'layers-isotropic',
+                space_order=(2, step, step),
+                nlayers=4,
+                vp_top=1500,
+                vp_bottom=3000,
+                nbl=0
+            )
+        else:  # Marmousi
+            model = SeismicModel(
+                space_order=(2, step, step),
+                vp=1000*cropped,
+                nbl=0,
+                origin=(0, 0),
+                shape=cropped.shape,
+                spacing=(10, 10)
+            )
+
+        plot_extent = [
+            model.origin[0],
+            model.origin[0] + model.grid.extent[0],
+            model.origin[1] + model.grid.extent[1],
+            model.origin[1]
+        ]
+        data = acoustic_model(
+            model, t0, t1, t2, critical_dt, source, step=step, snapshots=m
+        )
+        time = np.linspace(t1, t2, (m - 1)//step + 1)
         idx = 0
         for ii, ax in enumerate(ax_row):
             if ii % step == 0:
                 ax.imshow(
                     data[idx, :, :].T,
                     extent=plot_extent,
-                    vmin=-model.zlim, vmax=model.zlim,
+                    vmin=-zlim, vmax=zlim,
                     cmap='seismic'
                 )
-                ax.imshow(v.data.T, cmap='grey', extent=plot_extent, alpha=0.2)
-                ax.set_title(f't = {time[idx]:0.3f}')
+                ax.imshow(model.vp.data.T, cmap='grey', extent=plot_extent, alpha=0.2)
+                ax.set_title(f't={time[idx]:0.3f}')
                 idx += 1
                 if ii > 0:
                     ax.set_xticklabels([])
@@ -227,13 +194,21 @@ def acoustic_model(model, step=1, snapshots=1):
                     xticks = ax.get_xticks()
                     ax.set_xticks(np.array((
                         model.origin[0],
-                        round(model.origin[0] + model.extent[0]/2, -3),
-                        model.origin[0] + model.extent[0]
+                        model.origin[0] + model.grid.extent[0]
                     )))
-                    ax.set_xlim(model.origin[0], model.origin[0] + model.extent[0])
+                    ax.set_xlim(
+                        model.origin[0],
+                        model.origin[0] + model.grid.extent[0]
+                    )
                     yticks = ax.get_yticks()
-                    ax.set_yticks(np.concat(((model.origin[1],), yticks[2::2])))
-                    ax.set_ylim(model.origin[1] + model.extent[1], model.origin[1])
+                    ax.set_yticks(np.array((
+                        model.origin[1],
+                        model.origin[1] + model.grid.extent[1]
+                    )))
+                    ax.set_ylim(
+                        model.origin[1] + model.grid.extent[1],
+                        model.origin[1]
+                    )
             else:
                 ax.remove()
 
@@ -246,4 +221,4 @@ def acoustic_model(model, step=1, snapshots=1):
         wspace=0.06,
         hspace=0.06
     )
-    fig.savefig(f'{model.name}.png', dpi=300)
+    fig.savefig(f'{args.model}.png', dpi=300)