Merge pull request #1602 from ofmla/tti_1st_order_operators

examples: TTI 1st order operators

examples/seismic/tti/operators.py

@@ -341,7 +341,7 @@ def particle_velocity_fields(model, space_order):
     return vx, vz, vy
 
 
-def kernel_staggered_2d(model, u, v, space_order):
+def kernel_staggered_2d(model, u, v, space_order, **kwargs):
     """
     TTI finite difference. The equation solved is:
 
@@ -350,35 +350,60 @@ def kernel_staggered_2d(model, u, v, space_order):
     m * v.dt = - sqrt(1 + 2 delta) vx.dx - vz.dz + Fh
     m * u.dt = - (1 + 2 epsilon) vx.dx - sqrt(1 + 2 delta) vz.dz + Fv
     """
+    # Forward or backward
+    forward = kwargs.get('forward', True)
+
     dampl = 1 - model.damp
     m, epsilon, delta = model.m, model.epsilon, model.delta
     costheta, sintheta = trig_func(model)
     epsilon = 1 + 2 * epsilon
     delta = sqrt(1 + 2 * delta)
     s = model.grid.stepping_dim.spacing
     x, z = model.grid.dimensions
+
+    # Get source
+    qu = kwargs.get('qu', 0)
+    qv = kwargs.get('qv', 0)
+
     # Staggered setup
     vx, vz, _ = particle_velocity_fields(model, space_order)
 
-    # Stencils
-    phdx = costheta * u.dx - sintheta * u.dy
-    u_vx = Eq(vx.forward, dampl * vx - dampl * s * phdx)
+    if forward:
+        # Stencils
+        phdx = costheta * u.dx - sintheta * u.dy
+        u_vx = Eq(vx.forward, dampl * vx - dampl * s * phdx)
+
+        pvdz = sintheta * v.dx + costheta * v.dy
+        u_vz = Eq(vz.forward, dampl * vz - dampl * s * pvdz)
 
-    pvdz = sintheta * v.dx + costheta * v.dy
-    u_vz = Eq(vz.forward, dampl * vz - dampl * s * pvdz)
+        dvx = costheta * vx.forward.dx - sintheta * vx.forward.dy
+        dvz = sintheta * vz.forward.dx + costheta * vz.forward.dy
 
-    dvx = costheta * vx.forward.dx - sintheta * vx.forward.dy
-    dvz = sintheta * vz.forward.dx + costheta * vz.forward.dy
+        # u and v equations
+        pv_eq = Eq(v.forward, dampl * (v - s / m * (delta * dvx + dvz)) + s / m * qv)
+        ph_eq = Eq(u.forward, dampl * (u - s / m * (epsilon * dvx + delta * dvz)) +
+                   s / m * qu)
+    else:
+        # Stencils
+        phdx = ((costheta*epsilon*u).dx - (sintheta*epsilon*u).dy +
+                (costheta*delta*v).dx - (sintheta*delta*v).dy)
+        u_vx = Eq(vx.backward, dampl * vx + dampl * s * phdx)
 
-    # u and v equations
-    pv_eq = Eq(v.forward, dampl * (v - s / m * (delta * dvx + dvz)))
+        pvdz = ((sintheta*delta*u).dx + (costheta*delta*u).dy +
+                (sintheta*v).dx + (costheta*v).dy)
+        u_vz = Eq(vz.backward, dampl * vz + dampl * s * pvdz)
 
-    ph_eq = Eq(u.forward, dampl * (u - s / m * (epsilon * dvx + delta * dvz)))
+        dvx = (costheta * vx.backward).dx - (sintheta * vx.backward).dy
+        dvz = (sintheta * vz.backward).dx + (costheta * vz.backward).dy
+
+        # u and v equations
+        pv_eq = Eq(v.backward, dampl * (v + s / m * dvz))
+        ph_eq = Eq(u.backward, dampl * (u + s / m * dvx))
 
     return [u_vx, u_vz] + [pv_eq, ph_eq]
 
 
-def kernel_staggered_3d(model, u, v, space_order):
+def kernel_staggered_3d(model, u, v, space_order, **kwargs):
     """
     TTI finite difference. The equation solved is:
 
@@ -388,41 +413,83 @@ def kernel_staggered_3d(model, u, v, space_order):
     m * v.dt = - sqrt(1 + 2 delta) (vx.dx + vy.dy) - vz.dz + Fh
     m * u.dt = - (1 + 2 epsilon) (vx.dx + vy.dy) - sqrt(1 + 2 delta) vz.dz + Fv
     """
+    # Forward or backward
+    forward = kwargs.get('forward', True)
+
     dampl = 1 - model.damp
     m, epsilon, delta = model.m, model.epsilon, model.delta
     costheta, sintheta, cosphi, sinphi = trig_func(model)
     epsilon = 1 + 2 * epsilon
     delta = sqrt(1 + 2 * delta)
     s = model.grid.stepping_dim.spacing
     x, y, z = model.grid.dimensions
+
+    # Get source
+    qu = kwargs.get('qu', 0)
+    qv = kwargs.get('qv', 0)
+
     # Staggered setup
     vx, vz, vy = particle_velocity_fields(model, space_order)
-    # Stencils
-    phdx = (costheta * cosphi * u.dx +
-            costheta * sinphi * u.dyc -
-            sintheta * u.dzc)
-    u_vx = Eq(vx.forward, dampl * vx - dampl * s * phdx)
-
-    phdy = -sinphi * u.dxc + cosphi * u.dy
-    u_vy = Eq(vy.forward, dampl * vy - dampl * s * phdy)
-
-    pvdz = (sintheta * cosphi * v.dxc +
-            sintheta * sinphi * v.dyc +
-            costheta * v.dz)
-    u_vz = Eq(vz.forward, dampl * vz - dampl * s * pvdz)
-
-    dvx = (costheta * cosphi * vx.forward.dx +
-           costheta * sinphi * vx.forward.dyc -
-           sintheta * vx.forward.dzc)
-    dvy = -sinphi * vy.forward.dxc + cosphi * vy.forward.dy
-    dvz = (sintheta * cosphi * vz.forward.dxc +
-           sintheta * sinphi * vz.forward.dyc +
-           costheta * vz.forward.dz)
-    # u and v equations
-    pv_eq = Eq(v.forward, dampl * (v - s / m * (delta * (dvx + dvy) + dvz)))
-
-    ph_eq = Eq(u.forward, dampl * (u - s / m * (epsilon * (dvx + dvy) +
-                                                delta * dvz)))
+
+    if forward:
+        # Stencils
+        phdx = (costheta * cosphi * u.dx +
+                costheta * sinphi * u.dyc -
+                sintheta * u.dzc)
+        u_vx = Eq(vx.forward, dampl * vx - dampl * s * phdx)
+
+        phdy = -sinphi * u.dxc + cosphi * u.dy
+        u_vy = Eq(vy.forward, dampl * vy - dampl * s * phdy)
+
+        pvdz = (sintheta * cosphi * v.dxc +
+                sintheta * sinphi * v.dyc +
+                costheta * v.dz)
+        u_vz = Eq(vz.forward, dampl * vz - dampl * s * pvdz)
+
+        dvx = (costheta * cosphi * vx.forward.dx +
+               costheta * sinphi * vx.forward.dyc -
+               sintheta * vx.forward.dzc)
+        dvy = -sinphi * vy.forward.dxc + cosphi * vy.forward.dy
+        dvz = (sintheta * cosphi * vz.forward.dxc +
+               sintheta * sinphi * vz.forward.dyc +
+               costheta * vz.forward.dz)
+        # u and v equations
+        pv_eq = Eq(v.forward, dampl * (v - s / m * (delta * (dvx + dvy) + dvz)) +
+                   s / m * qv)
+
+        ph_eq = Eq(u.forward, dampl * (u - s / m * (epsilon * (dvx + dvy) +
+                                                    delta * dvz)) + s / m * qu)
+    else:
+        # Stencils
+        phdx = ((costheta * cosphi * epsilon*u).dx +
+                (costheta * sinphi * epsilon*u).dyc -
+                (sintheta * epsilon*u).dzc + (costheta * cosphi * delta*v).dx +
+                                             (costheta * sinphi * delta*v).dyc -
+                                             (sintheta * delta*v).dzc)
+        u_vx = Eq(vx.backward, dampl * vx + dampl * s * phdx)
+
+        phdy = (-(sinphi * epsilon*u).dxc + (cosphi * epsilon*u).dy -
+                (sinphi * delta*v).dxc + (cosphi * delta*v).dy)
+        u_vy = Eq(vy.backward, dampl * vy + dampl * s * phdy)
+
+        pvdz = ((sintheta * cosphi * delta*u).dxc +
+                (sintheta * sinphi * delta*u).dyc +
+                (costheta * delta*u).dz + (sintheta * cosphi * v).dxc +
+                                          (sintheta * sinphi * v).dyc +
+                                          (costheta * v).dz)
+        u_vz = Eq(vz.backward, dampl * vz + dampl * s * pvdz)
+
+        dvx = ((costheta * cosphi * vx.backward).dx +
+               (costheta * sinphi * vx.backward).dyc -
+               (sintheta * vx.backward).dzc)
+        dvy = (-sinphi * vy.backward).dxc + (cosphi * vy.backward).dy
+        dvz = ((sintheta * cosphi * vz.backward).dxc +
+               (sintheta * sinphi * vz.backward).dyc +
+               (costheta * vz.backward).dz)
+        # u and v equations
+        pv_eq = Eq(v.backward, dampl * (v + s / m * dvz))
+
+        ph_eq = Eq(u.backward, dampl * (u + s / m * (dvx + dvy)))
 
     return [u_vx, u_vy, u_vz] + [pv_eq, ph_eq]
 
@@ -473,16 +540,17 @@ def ForwardOperator(model, geometry, space_order=4,
     stencils = FD_kernel(model, u, v, space_order)
 
     # Source and receivers
-    stencils += src.inject(field=u.forward, expr=src * dt**2 / m)
-    stencils += src.inject(field=v.forward, expr=src * dt**2 / m)
+    expr = src * dt / m if kernel == 'staggered' else src * dt**2 / m
+    stencils += src.inject(field=u.forward, expr=expr)
+    stencils += src.inject(field=v.forward, expr=expr)
     stencils += rec.interpolate(expr=u + v)
 
     # Substitute spacing terms to reduce flops
     return Operator(stencils, subs=model.spacing_map, name='ForwardTTI', **kwargs)
 
 
 def AdjointOperator(model, geometry, space_order=4,
-                    **kwargs):
+                    kernel='centered', **kwargs):
     """
     Construct an adjoint modelling operator in an tti media.
 
@@ -495,29 +563,36 @@ def AdjointOperator(model, geometry, space_order=4,
         receivers (SparseTimeFunction) and their position.
     space_order : int, optional
         Space discretization order.
+    kernel : str, optional
+        Type of discretization, centered or shifted
     """
 
     dt = model.grid.time_dim.spacing
     m = model.m
-    time_order = 2
+    time_order = 1 if kernel == 'staggered' else 2
+    if kernel == 'staggered':
+        stagg_p = stagg_r = NODE
+    else:
+        stagg_p = stagg_r = None
 
     # Create symbols for forward wavefield, source and receivers
-    p = TimeFunction(name='p', grid=model.grid, save=None, time_order=time_order,
-                     space_order=space_order)
-    r = TimeFunction(name='r', grid=model.grid, save=None, time_order=time_order,
-                     space_order=space_order)
+    p = TimeFunction(name='p', grid=model.grid, staggered=stagg_p,
+                     time_order=time_order, space_order=space_order)
+    r = TimeFunction(name='r', grid=model.grid, staggered=stagg_r,
+                     time_order=time_order, space_order=space_order)
     srca = PointSource(name='srca', grid=model.grid, time_range=geometry.time_axis,
                        npoint=geometry.nsrc)
     rec = Receiver(name='rec', grid=model.grid, time_range=geometry.time_axis,
                    npoint=geometry.nrec)
 
     # FD kernels of the PDE
-    FD_kernel = kernels[('centered', len(model.shape))]
+    FD_kernel = kernels[(kernel, len(model.shape))]
     stencils = FD_kernel(model, p, r, space_order, forward=False)
 
     # Construct expression to inject receiver values
-    stencils += rec.inject(field=p.backward, expr=rec * dt**2 / m)
-    stencils += rec.inject(field=r.backward, expr=rec * dt**2 / m)
+    expr = rec * dt / m if kernel == 'staggered' else rec * dt**2 / m
+    stencils += rec.inject(field=p.backward, expr=expr)
+    stencils += rec.inject(field=r.backward, expr=expr)
 
     # Create interpolation expression for the adjoint-source
     stencils += srca.interpolate(expr=p + r)

examples/seismic/tti/tti_example.py

@@ -18,7 +18,8 @@ def tti_setup(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=250.0,
     # Source and receiver geometries
     geometry = setup_geometry(model, tn)
 
-    return AnisotropicWaveSolver(model, geometry, space_order=space_order, **kwargs)
+    return AnisotropicWaveSolver(model, geometry, space_order=space_order,
+                                 kernel=kernel, **kwargs)
 
 
 def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=250.0,
@@ -29,7 +30,7 @@ def run(shape=(50, 50, 50), spacing=(20.0, 20.0, 20.0), tn=250.0,
                        nbl=nbl, kernel=kernel, **kwargs)
     info("Applying Forward")
 
-    rec, u, v, summary = solver.forward(autotune=autotune, kernel=kernel)
+    rec, u, v, summary = solver.forward(autotune=autotune)
 
     if not full_run:
         return summary.gflopss, summary.oi, summary.timings, [rec, u, v]