Which Idea is Better for Achieving PML Boundary?

[zhangxiaoshuotttt]

@mloubout Hi,

Currently, the absorbing boundary is not efficient, especially when aiming for elastic or 3D FWI and LSRTM, as it requires a significant amount of grid space for the absorbing area.

Based on your tutorial, I would like to implement the PML boundary. While I find your method for achieving the absorbing boundary elegant, modifying the forward and adjoint formulas will be necessary for successfully implementing PML or HABC.

I have two ideas, and I would appreciate your guidance on them. I want to ensure that the PML can be utilized conveniently if I need to change my formulas again (please correct me if I'm wrong).

1. I could define a pml_stencil in the operators.py file and incorporate it within the ForwardOperator function, similar to the examples you've provided. However, the drawback of this approach is that if I wish to implement FWI or RTM, I will also need to modify other functions in the Operators module. Also, if I want to achieve another formula, I need to rewrite this.

2. Alternatively, can we achieve the PML functionality solely within the model.py? I'm considering adding the wavefield u as a parameter inside the Model function, allowing us to implement the PML based on this parameter. However, is this feasible ?I would like your insightful advice on this matter.

I am open to any other suggestions or comments you may have. Thank you for your time and patience.

Sincerely,
Shuo

[mloubout]

I would probably look at the first option. The Model is not really "propagation" related but more a plain representation of the subsurface irrespective of waves. The examples are more of a set of tutorial-ish propagators that can be used for research but you can use it as a skeleton for your own PDEs or go a separate ways if there is a design and interface that would fit your needs better.

PS:

I'm in and out for a few weeks, so apologies for the short answers

[zhangxiaoshuotttt]

Thank you for your insightful response. I also prefer the first option, as the second one is difficult for me to achieve.
I appreciate your time and assistance.

[zhangxiaoshuotttt]

Hi,

I’m trying to learn from your (example) and modify the related files so we can use the solver to achieve FWI and LSRTM. I believe the same methods can be extended to the elastic example as well.

I have attached my files and would appreciate your feedback. I’m particularly thinking about improving the implementation of the PML. An efficient absorbing boundary is crucial, especially if we want to extend this to 3D or field data.

Here’s what I’ve done so far and where I need help:
1. In model.py, I modified it to generate subdomains and parameters for PML.
2. In wavesolver.py, I disabled self.model._initialize_bcs(bcs="damp") since I want to test only the PML in this case.
3. In operators.py, I updated iso_stencil to apply the damping equation in the absorbing area.

Right now, the forward simulation works fine, but when I try to calculate the gradient, I get some errors. I think this might be related to the acquisition setup, but I’m not sure.

If you have time to review my files or share your ideas, I would greatly appreciate it. If I’ve made any mistakes or there’s a better way to do this, please let me know.

Thank you so much for your time and help!

I will show my codes as follows.

[zhangxiaoshuotttt]

Here are my codes for testing:

import numpy as np
%matplotlib inline

from devito import configuration
configuration['log-level'] = 'WARNING'

from model_new import Model
#NBVAL_IGNORE_OUTPUT
from examples.seismic import plot_velocity, plot_perturbation
# Compute synthetic data with forward operator 
from wavesolver_new import AcousticWaveSolver

# Create the computational grid
nx, nz = 201, 201       # Number of grid points
dx, dz = 10.0, 10.0     # Grid spacing

shape = (nx, nz)
spacing = (dx, dz)
origin = (0., 0.)

space_order = 8

# Define velocity model v(x)
# Define a velocity profile. The velocity is in km/s
v = np.empty(shape, dtype=np.float32)
v[:] = 1.5  # Background velocity

v[:,100:] = 2.5 

nbpml = 80
model = Model(origin=origin, shape=shape,spacing=spacing, 
              vp=v, space_order=space_order, nbl=nbpml,bcs="pml")


# Define a new model for calculating kernnel

v0 = np.empty(shape, dtype=np.float32)
v0[:] = 1.5  # Background velocity

model0 = Model(origin=origin, shape=shape,spacing=spacing, 
              vp=v0,space_order=space_order, nbl=nbpml,bcs="pml")

#NBVAL_IGNORE_OUTPUT
# Define acquisition geometry: source
from examples.seismic import AcquisitionGeometry

t0 = 0.
tn = 4000. 
f0 = 0.010
# First, position source centrally in all dimensions, then set depth
src_coordinates = np.empty((1, 2))
src_coordinates[0, :] = 200.  # Depth is 20m
src_coordinates[0, 0] = 1000

nshots = 21  # Need good covergae in shots, one every two grid points
nreceivers = 101  # One recevier every grid point
# Define acquisition geometry: receivers

# Initialize receivers for synthetic and imaging data
rec_coordinates = np.empty((nreceivers, 2))
rec_coordinates[:, 1] = 200.
rec_coordinates[:, 0] = np.linspace(0, model.domain_size[0], num=nreceivers)

# Geometry

geometry = AcquisitionGeometry(model, rec_coordinates, src_coordinates, t0, tn, f0=f0, src_type='Ricker')
# We can plot the time signature to see the wavelet
geometry.src.show()

#NBVAL_IGNORE_OUTPUT
# Plot acquisition geometry
plot_velocity(model, source=geometry.src_positions,
              receiver=geometry.rec_positions[::4, :])
              
 
solver = AcousticWaveSolver(model, geometry, space_order=space_order)
true_d , _, _ = solver.forward(vp=model.vp)

# Compute initial data with forward operator 
smooth_d, _, _ = solver.forward(vp=model0.vp)

#NBVAL_IGNORE_OUTPUT
# Plot shot record for true and smooth velocity model and the difference
from examples.seismic import plot_shotrecord

plot_shotrecord(true_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data, model, t0, tn)
plot_shotrecord(smooth_d.data - true_d.data, model, t0, tn)

#NBVAL_IGNORE_OUTPUT

# Prepare the varying source locations sources
source_locations = np.empty((nshots, 2), dtype=np.float32)
source_locations[:, 0] = np.linspace(200., 1800, num=nshots)
source_locations[:, 1] = 200

plot_velocity(model, source=source_locations)

from devito import Eq, Operator

# Computes the residual between observed and synthetic data into the residual
def compute_residual(residual, dobs, dsyn):
    if residual.grid.distributor.is_parallel:
        # If we run with MPI, we have to compute the residual via an operator
        # First make sure we can take the difference and that receivers are at the 
        # same position
        assert np.allclose(dobs.coordinates.data[:], dsyn.coordinates.data)
        assert np.allclose(residual.coordinates.data[:], dsyn.coordinates.data)
        # Create a difference operator
        diff_eq = Eq(residual, dsyn.subs({dsyn.dimensions[-1]: residual.dimensions[-1]}) -
                               dobs.subs({dobs.dimensions[-1]: residual.dimensions[-1]}))
        Operator(diff_eq)()
    else:
        # A simple data difference is enough in serial
        residual.data[:] = dsyn.data[:] - dobs.data[:]
    
    return residual
    
    # Create FWI gradient kernel 
from devito import Function, TimeFunction, norm
from examples.seismic import Receiver

import scipy
def fwi_gradient(vp_in):    
    # Create symbols to hold the gradient
    grad = Function(name="grad", grid=model.grid)
    # Create placeholders for the data residual and data
    residual = Receiver(name='residual', grid=model.grid,
                        time_range=geometry.time_axis, 
                        coordinates=geometry.rec_positions)
    d_obs = Receiver(name='d_obs', grid=model.grid,
                     time_range=geometry.time_axis, 
                     coordinates=geometry.rec_positions)
    d_syn = Receiver(name='d_syn', grid=model.grid,
                     time_range=geometry.time_axis, 
                     coordinates=geometry.rec_positions)
    objective = 0.
    for i in range(nshots):
        # Update source location
        geometry.src_positions[0, :] = source_locations[i, :]
        
        # Generate synthetic data from true model
        _, _, _ = solver.forward(vp=model.vp, rec=d_obs)
        #print("Source:",i,"Location:",geometry.src_positions[0, :])
        # Compute smooth data and full forward wavefield u0
        _, u0, _ = solver.forward(vp=vp_in, save=True, rec=d_syn)
        
        # Compute gradient from data residual and update objective function 
        compute_residual(residual, d_obs, d_syn)
        #print("Residual shape:",residual.data.shape)
        objective += .5*norm(residual)**2
        #print("Objective value:",objective)
        solver.gradient(rec=residual, u=u0, vp=vp_in, grad=grad)
    
    return objective, grad
    
    # Compute gradient of initial model
ff, update = fwi_gradient(model0.vp)
assert np.isclose(ff, 57283, rtol=1e0)

#NBVAL_IGNORE_OUTPUT
from devito import mmax
from examples.seismic import plot_image

# Plot the FWI gradient
plot_image(-update.data, vmin=-1e4, vmax=1e4, cmap="jet")

# Plot the difference between the true and initial model.
# This is not known in practice as only the initial model is provided.
plot_image(model0.vp.data - model.vp.data, vmin=-1e-1, vmax=1e-1, cmap="jet")

# Show what the update does to the model
alpha = .5 / mmax(update)
plot_image(model0.vp.data + alpha*update.data, vmin=2.5, vmax=3.0, cmap="jet")

[zhangxiaoshuotttt]

Here are the errors:

/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c: In function ‘Forward’:
/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c:222:23: error: ‘posx’ undeclared (first use in this function)
  222 |           if (rsrcx + posx >= x_m - 1 && rsrcy + posy >= y_m - 1 && rsrcx + posx <= x_M + 1 && rsrcy + posy <= y_M + 1)
      |                       ^~~~
/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c:222:23: note: each undeclared identifier is reported only once for each function it appears in
/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c:222:50: error: ‘posy’ undeclared (first use in this function)
  222 |           if (rsrcx + posx >= x_m - 1 && rsrcy + posy >= y_m - 1 && rsrcx + posx <= x_M + 1 && rsrcy + posy <= y_M + 1)
      |                                                  ^~~~
/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c:224:87: error: ‘px’ undeclared (first use in this function)
  224 |             float r0 = (dt*dt)*(vp[posx + 8][posy + 8]*vp[posx + 8][posy + 8])*(rsrcx*px + (1 - rsrcx)*(1 - px))*(rsrcy*py + (1 - rsrcy)*(1 - py))*src[time][p_src];
      |                                                                                       ^~
/tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c:224:121: error: ‘py’ undeclared (first use in this function)
  224 |             float r0 = (dt*dt)*(vp[posx + 8][posy + 8]*vp[posx + 8][posy + 8])*(rsrcx*px + (1 - rsrcx)*(1 - px))*(rsrcy*py + (1 - rsrcy)*(1 - py))*src[time][p_src];
      |                                                                                                                         ^~
FAILED compiler invocation: gcc -march=native -O3 -g -fPIC -Wall -std=c99 -Wno-unused-result -Wno-unused-variable -Wno-unused-but-set-variable -ffast-math -mprefer-vector-width=512 -shared -fopenmp /tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.c -lm -o /tmp/devito-jitcache-uid212010/66e8eefb40bf267ff6fd0da42ae2bd623f1ddb45.so

---------------------------------------------------------------------------
CompileError                              Traceback (most recent call last)
Cell In[27], line 2
      1 # Compute gradient of initial model
----> 2 ff, update = fwi_gradient(model0.vp)
      3 assert np.isclose(ff, 57283, rtol=1e0)
      5 #NBVAL_IGNORE_OUTPUT

Cell In[26], line 28, in fwi_gradient(vp_in)
     25 _, _, _ = solver.forward(vp=model.vp, rec=d_obs)
     26 #print("Source:",i,"Location:",geometry.src_positions[0, :])
     27 # Compute smooth data and full forward wavefield u0
---> 28 _, u0, _ = solver.forward(vp=vp_in, save=True, rec=d_syn)
     30 # Compute gradient from data residual and update objective function 
     31 compute_residual(residual, d_obs, d_syn)

File ~/devito/examples/seismic/acoustic/wavesolver_new.py:113, in AcousticWaveSolver.forward(self, src, rec, u, model, save, **kwargs)
    110 kwargs.update(model.physical_params(**kwargs))
    112 # Execute operator and return wavefield and receiver data
--> 113 summary = self.op_fwd(save).apply(src=src, rec=rec, u=u,
    114                                   dt=kwargs.pop('dt', self.dt), **kwargs)
    116 return rec, u, summary

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/devito/operator/operator.py:870, in Operator.apply(self, **kwargs)
    868 arg_values = [args[p.name] for p in self.parameters]
    869 try:
--> 870     cfunction = self.cfunction
    871     with self._profiler.timer_on('apply', comm=args.comm):
    872         retval = cfunction(*arg_values)

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/devito/operator/operator.py:752, in Operator.cfunction(self)
    750 """The JIT-compiled C function as a ctypes.FuncPtr object."""
    751 if self._lib is None:
--> 752     self._jit_compile()
    753     self._lib = self._compiler.load(self._soname)
    754     self._lib.name = self._soname

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/devito/operator/operator.py:738, in Operator._jit_compile(self)
    736 if self._lib is None:
    737     with self._profiler.timer_on('jit-compile'):
--> 738         recompiled, src_file = self._compiler.jit_compile(self._soname, str(self))
    740     elapsed = self._profiler.py_timers['jit-compile']
    741     if recompiled:

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/devito/arch/compiler.py:376, in Compiler.jit_compile(self, soname, code)
    374 with warnings.catch_warnings():
    375     warnings.simplefilter('ignore')
--> 376     _, _, _, recompiled = compile_from_string(self, target, code, src_file,
    377                                               cache_dir=cache_dir, debug=debug,
    378                                               sleep_delay=sleep_delay)
    380 return recompiled, src_file

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/codepy/jit.py:439, in compile_from_string(toolchain, name, source_string, source_name, cache_dir, debug, wait_on_error, debug_recompile, object, source_is_binary, sleep_delay)
    437     toolchain.build_object(ext_file, source_paths, debug=debug)
    438 else:
--> 439     toolchain.build_extension(ext_file, source_paths, debug=debug)
    441 if info_path is not None:
    442     import pickle

File ~/miniconda3/envs/Devito/lib/python3.9/site-packages/codepy/toolchain.py:211, in GCCLikeToolchain.build_extension(self, ext_file, source_files, debug)
    208 import sys
    209 print("FAILED compiler invocation: {}".format(" ".join(cc_cmdline)),
    210       file=sys.stderr)
--> 211 raise CompileError("module compilation failed")

CompileError: module compilation failed

[zhangxiaoshuotttt]

Here are codes in model.py:

import numpy as np
from sympy import finite_diff_weights as fd_w
try:
    import pytest
except:
    pass
import sys
from devito import (Grid, SubDomain, Function, Constant, warning, TimeFunction,
                    SubDimension, Eq, Inc, Operator, div, sin, Abs,NODE)
from devito.builtins import initialize_function, gaussian_smooth, mmax, mmin
from devito.tools import as_tuple

__all__ = ['SeismicModel', 'Model', 'ModelElastic',
           'ModelViscoelastic', 'ModelViscoacoustic']


def initialize_damp(damp, padsizes, spacing, abc_type="damp", fs=False):
    """
    Initialize damping field with an absorbing boundary layer.

    Parameters
    ----------
    damp : Function
        The damping field for absorbing boundary condition.
    nbl : int
        Number of points in the damping layer.
    spacing :
        Grid spacing coefficient.
    mask : bool, optional
        whether the dampening is a mask or layer.
        mask => 1 inside the domain and decreases in the layer
        not mask => 0 inside the domain and increase in the layer
    """

    eqs = [Eq(damp, 1.0 if abc_type == "mask" else 0.0)]
    for (nbl, nbr), d in zip(padsizes, damp.dimensions):
        if not fs or d is not damp.dimensions[-1]:
            dampcoeff = 1.5 * np.log(1.0 / 0.001) / (nbl)
            # left
            dim_l = SubDimension.left(name='abc_%s_l' % d.name, parent=d,
                                      thickness=nbl)
            pos = Abs((nbl - (dim_l - d.symbolic_min) + 1) / float(nbl))
            val = dampcoeff * (pos - sin(2*np.pi*pos)/(2*np.pi))
            val = -val if abc_type == "mask" else val
            eqs += [Inc(damp.subs({d: dim_l}), val/d.spacing)]
        # right
        dampcoeff = 1.5 * np.log(1.0 / 0.001) / (nbr)
        dim_r = SubDimension.right(name='abc_%s_r' % d.name, parent=d,
                                   thickness=nbr)
        pos = Abs((nbr - (d.symbolic_max - dim_r) + 1) / float(nbr))
        val = dampcoeff * (pos - sin(2*np.pi*pos)/(2*np.pi))
        val = -val if abc_type == "mask" else val
        eqs += [Inc(damp.subs({d: dim_r}), val/d.spacing)]

    Operator(eqs, name='initdamp')()

def fdamp(x,z,i,x0pml,x1pml,z0pml,z1pml,lx,lz):
    
    quibar  = 0.05
          
    if(i==1):
        a = np.where(x<=x0pml,(np.abs(x-x0pml)/lx),np.where(x>=x1pml,(np.abs(x-x1pml)/lx),0.))
        fdamp = quibar*(a-(1./(2.*np.pi))*np.sin(2.*np.pi*a))
    if(i==2):
        a = np.where(z<=z0pml,(np.abs(z-z0pml)/lz),np.where(z>=z1pml,(np.abs(z-z1pml)/lz),0.))
        fdamp = quibar*(a-(1./(2.*np.pi))*np.sin(2.*np.pi*a))
      
    return fdamp

def generatemdamp(x0,x1,z0,z1,nptx,nptz,hxv,hzv,x0pml,x1pml,z0pml,z1pml,lx,lz):
    
    X0     = np.linspace(x0,x1,nptx)    
    Z0     = np.linspace(z0,z1,nptz)
    X0grid,Z0grid = np.meshgrid(X0,Z0)
    X1   = np.linspace((x0+0.5*hxv),(x1-0.5*hxv),nptx-1)
    Z1   = np.linspace((z0+0.5*hzv),(z1-0.5*hzv),nptz-1)
    X1grid,Z1grid = np.meshgrid(X1,Z1)
   
    D01 = np.zeros((nptx,nptz))
    D02 = np.zeros((nptx,nptz))
    D11 = np.zeros((nptx,nptz))
    D12 = np.zeros((nptx,nptz))
    
    D01 = np.transpose(fdamp(X0grid,Z0grid,1,x0pml,x1pml,z0pml,z1pml,lx,lz))
    D02 = np.transpose(fdamp(X0grid,Z0grid,2,x0pml,x1pml,z0pml,z1pml,lx,lz))
  
    D11 = np.transpose(fdamp(X1grid,Z1grid,1,x0pml,x1pml,z0pml,z1pml,lx,lz))
    D12 = np.transpose(fdamp(X1grid,Z1grid,2,x0pml,x1pml,z0pml,z1pml,lx,lz))
    
    return D01, D02, D11, D12


class PhysicalDomain(SubDomain):
    name = 'physdomain'
    def __init__(self, so, fs=False, npml=20):
        super().__init__()
        self.so = so
        self.fs = fs
        self.npml = npml

    def define(self, dimensions):
        x, z = dimensions
        if self.fs:
            # With free surface: PML on left, right, bottom only
            return {
                x: ('middle', self.npml, self.npml),
                z: ('middle', self.npml, 0)  # No top PML
            }
        else:
            # Without free surface: PML on all sides
            return {
                x: ('middle', self.npml, self.npml),
                z: ('middle', self.npml, self.npml)
            }

class PMLLeft(SubDomain):
    name = 'pml_left'
    def __init__(self,so, npml=20):
        super().__init__()
        self.so = so
        self.npml = npml
    def define(self, dimensions):
        x, z = dimensions
        return {x: ('left', self.npml), z: z}

class PMLRight(SubDomain):
    name = 'pml_right'
    def __init__(self,so, npml=20):
        super().__init__()
        self.so = so
        self.npml = npml
    def define(self, dimensions):
        x, z = dimensions
        return {x: ('right', self.npml), z: z}

class PMLBottom(SubDomain):
    name = 'pml_bottom'
    def __init__(self,so, npml=20):
        super().__init__()
        self.so = so
        self.npml = npml
    def define(self, dimensions):
        x, z = dimensions
        return {x: ('middle', self.npml, self.npml),
                z: ('left', self.npml)}

class PMLTop(SubDomain):
    name = 'pml_top'
    def __init__(self,so, fs=False, npml=20):
        super().__init__()
        self.so = so
        self.fs = fs
        self.npml = npml
    def define(self, dimensions):
        if self.fs:
            return {}  # No top PML when free surface is enabled
        x, z = dimensions
        return {x: ('middle', self.npml, self.npml),
                z: ('right', self.npml)}

class FSDomain(SubDomain):

    name = 'fsdomain'

    def __init__(self, so):
        super().__init__()
        self.size = so

    def define(self, dimensions):
        """
        Definition of the upper section of the domain for wrapped indices FS.
        """

        return {d: (d if not d == dimensions[-1] else ('left', self.size))
                for d in dimensions}


class GenericModel:
    """
    General model class with common properties
    """
    def __init__(self, origin, spacing, shape, space_order, nbl=20,
                 dtype=np.float32, subdomains=(), bcs="damp", grid=None,
                 fs=False,pml=True):
        self.shape = shape
        self.space_order = space_order
        self.nbl = int(nbl)
        self.origin = tuple([dtype(o) for o in origin])
        self.fs = fs
        self.pml = pml
        # Default setup
        origin_pml = [dtype(o - s*nbl) for o, s in zip(origin, spacing)]
        shape_pml = np.array(shape) + 2 * self.nbl

        # Model size depending on freesurface
        physdomain = PhysicalDomain(space_order, fs=fs)
        subdomains = subdomains + (physdomain,)
        # For PML, only setup PML subdomains
        if bcs == "pml":
            # Only setup PML subdomains
            PMLLeftdomain = PMLLeft(space_order,npml=self.nbl)
            PMLRightdomain = PMLRight(space_order,npml=self.nbl)
            PMLBottomdomain = PMLBottom(space_order,npml=self.nbl)
            subdomains = subdomains + (PMLLeftdomain, PMLRightdomain, PMLBottomdomain,)
            if not self.fs:
                PMLTopdomain = PMLTop(space_order,fs=self.fs, npml=self.nbl)
                subdomains = subdomains + (PMLTopdomain,)

        if fs:
            fsdomain = FSDomain(space_order)
            subdomains = subdomains + (fsdomain,)
            origin_pml[-1] = origin[-1]
            shape_pml[-1] -= self.nbl

        # Origin of the computational domain with boundary to inject/interpolate
        # at the correct index
        if grid is None:
            # Physical extent is calculated per cell, so shape - 1
            extent = tuple(np.array(spacing) * (shape_pml - 1))
            self.grid = Grid(extent=extent, shape=shape_pml, origin=origin_pml,
                             dtype=dtype, subdomains=subdomains)

        else:
            self.grid = grid

        self._physical_parameters = set()
        self.damp = None
        if not pml:
            print("Now using simple damping")
            self._initialize_bcs(bcs=bcs)

    def _initialize_bcs(self, bcs="damp"):
        # Create dampening field as symbol `damp`
        if self.nbl == 0:
            self.damp = 1 if bcs == "mask" else 0
            return
    
        # First initialization
        init = self.damp is None
        # Get current Function if already initialized
        self.damp = self.damp or Function(name="damp", grid=self.grid,
                                          space_order=self.space_order)
        if callable(bcs):
            bcs(self.damp, self.nbl)
        else:
            re_init = ((bcs == "mask" and mmin(self.damp) == 0) or
                       (bcs == "damp" and mmax(self.damp) == 1))
            if init or re_init:
                if re_init and not init:
                    bcs_o = "damp" if bcs == "mask" else "mask"
                    warning("Re-initializing damp profile from %s to %s" % (bcs_o, bcs))
                    warning("Model has to be created with `bcs=\"%s\"`"
                            "for this WaveSolver" % bcs)
                initialize_damp(self.damp, self.padsizes, self.spacing,
                                abc_type=bcs, fs=self.fs)
        self._physical_parameters.update(['damp'])

    @property
    def padsizes(self):
        """
        Padding size for each dimension.
        """
        padsizes = [(self.nbl, self.nbl) for _ in range(self.dim-1)]
        padsizes.append((0 if self.fs else self.nbl, self.nbl))
        return padsizes

    def physical_params(self, **kwargs):
        """
        Return all set physical parameters and update to input values if provided
        """
        known = [getattr(self, i) for i in self.physical_parameters]
        return {i.name: kwargs.get(i.name, i) or i for i in known}

    def _gen_phys_param(self, field, name, space_order, is_param=True,
                        default_value=0, avg_mode='arithmetic', **kwargs):
        if field is None:
            return default_value
        if isinstance(field, np.ndarray):
            function = Function(name=name, grid=self.grid, space_order=space_order,
                                parameter=is_param, avg_mode=avg_mode)
            initialize_function(function, field, self.padsizes)
        else:
            function = Constant(name=name, value=field, dtype=self.grid.dtype)
        self._physical_parameters.update([name])
        return function

    @property
    def physical_parameters(self):
        return as_tuple(self._physical_parameters)

    @property
    def dim(self):
        """
        Spatial dimension of the problem and model domain.
        """
        return self.grid.dim

    @property
    def spacing(self):
        """
        Grid spacing for all fields in the physical model.
        """
        return self.grid.spacing

    @property
    def space_dimensions(self):
        """
        Spatial dimensions of the grid
        """
        return self.grid.dimensions

    @property
    def spacing_map(self):
        """
        Map between spacing symbols and their values for each `SpaceDimension`.
        """
        return self.grid.spacing_map

    @property
    def dtype(self):
        """
        Data type for all assocaited data objects.
        """
        return self.grid.dtype

    @property
    def domain_size(self):
        """
        Physical size of the domain as determined by shape and spacing
        """
        return tuple((d-1) * s for d, s in zip(self.shape, self.spacing))


class SeismicModel(GenericModel):
    """
    The physical model used in seismic inversion processes.

    Parameters
    ----------
    origin : tuple of floats
        Origin of the model in m as a tuple in (x,y,z) order.
    spacing : tuple of floats
        Grid size in m as a Tuple in (x,y,z) order.
    shape : tuple of int
        Number of grid points size in (x,y,z) order.
    space_order : int
        Order of the spatial stencil discretisation.
    vp : array_like or float
        Velocity in km/s.
    nbl : int, optional
        The number of absorbin layers for boundary damping.
    bcs: str or callable
        Absorbing boundary type ("damp" or "mask") or initializer.
    dtype : np.float32 or np.float64
        Defaults to np.float32.
    epsilon : array_like or float, optional
        Thomsen epsilon parameter (0<epsilon<1).
    delta : array_like or float
        Thomsen delta parameter (0<delta<1), delta<epsilon.
    theta : array_like or float
        Tilt angle in radian.
    phi : array_like or float
        Asymuth angle in radian.
    b : array_like or float
        Buoyancy.
    vs : array_like or float
        S-wave velocity.
    qp : array_like or float
        P-wave attenuation.
    qs : array_like or float
        S-wave attenuation.
    """
    _known_parameters = ['vp', 'damp', 'vs', 'b', 'epsilon', 'delta',
                         'theta', 'phi', 'qp', 'qs', 'lam', 'mu']

    def __init__(self, origin, spacing, shape, space_order, vp, nbl=20, fs=False,pml=True,
                 dtype=np.float32, subdomains=(), bcs="mask", grid=None, **kwargs):
        super().__init__(origin, spacing, shape, space_order, nbl,
                         dtype, subdomains, grid=grid, bcs=bcs, fs=fs,pml=pml)
        
        if pml:
            self._create_pml_arrays()
        # Initialize physics
        self._initialize_physics(vp, space_order, **kwargs)
        sys.stdout.flush()  
        # User provided dt
        self._dt = kwargs.get('dt')
        # Some wave equation need a rescaled dt that can't be infered from the model
        # parameters, such as isoacoustic OT4 that can use a dt sqrt(3) larger than
        # isoacoustic OT2. This property should be set from a wavesolver or after model
        # instanciation only via model.dt_scale = value.
        self._dt_scale = 1

    def _initialize_physics(self, vp, space_order, **kwargs):
        """
        Initialize physical parameters and type of physics from inputs.
        The types of physics supported are:
        - acoustic: [vp, b]
        - elastic: [vp, vs, b] represented through Lame parameters [lam, mu, b]
        - visco-acoustic: [vp, b, qp]
        - visco-elastic: [vp, vs, b, qs]
        - vti: [vp, epsilon, delta]
        - tti: [epsilon, delta, theta, phi]
        """
        params = []
        # Buoyancy
        b = kwargs.get('b', 1)


        # Initialize elastic with Lame parametrization
        if 'vs' in kwargs:
            vs = kwargs.pop('vs')
            self.lam = self._gen_phys_param((vp**2 - 2. * vs**2)/b, 'lam', space_order,
                                            is_param=True)
            # Need to add small value to avoid division by zero
            if isinstance(vs, np.ndarray):
                vs = vs + 1e-12
            self.mu = self._gen_phys_param(vs**2 / b, 'mu', space_order, is_param=True,
                                           avg_mode='harmonic')
        else:
            # All other seismic models have at least a velocity
            self.vp = self._gen_phys_param(vp, 'vp', space_order)
        # Initialize rest of the input physical parameters
        for name in self._known_parameters:
            if kwargs.get(name) is not None:
                field = self._gen_phys_param(kwargs.get(name), name, space_order)
                setattr(self, name, field)
                params.append(name)

    def _create_pml_arrays(self):
        """
        Creates the PML arrays just once, stored in self.
        Here is where you define:
          vel1, phi1, phi2, dampx0, dampz0, dampx1, dampz1, ...
        Then iso_stencil can reference model.vel1, model.phi1, ...
        """
        grid = self.grid
        x, z = grid.dimensions
        so = self.space_order
        nptx, nptz = grid.shape

        # 1) velocity in PML
        self.vel1 = Function(name="vel1", grid=grid, space_order=so, staggered=(x,z))
        self.vel1.data[0:nptx-1, 0:nptz-1] = 0.0

        # 2) time functions for boundary
        self.phi1 = TimeFunction(name="phi1", grid=grid, space_order=so,
                                 time_order=2, staggered=(x,z))
        self.phi2 = TimeFunction(name="phi2", grid=grid, space_order=so,
                                 time_order=2, staggered=(x,z))

        # 3) compute your PML damping arrays only once
        #    e.g. from generatemdamp(...) or the user-coded approach
        x0, z0 = grid.origin
        x1, z1 = grid.extent
        hx, hz = grid.spacing
        npml = self.nbl
        x0pml = x0 + npml*hx
        x1pml = x1 - npml*hx
        z0pml = z0 + npml*hz
        z1pml = z1 - npml*hz
        lx = npml*hx
        lz = npml*hz
        D01, D02, D11, D12 = generatemdamp(x0, x1, z0, z1, nptx, nptz,
                                          hx, hz, x0pml, x1pml, z0pml, z1pml,
                                          lx, lz)

        # 4) store them in Functions
        self.dampx0 = Function(name="dampx0", grid=grid, space_order=so, staggered=NODE)
        self.dampz0 = Function(name="dampz0", grid=grid, space_order=so, staggered=NODE)
        self.dampx0.data[:] = D01
        self.dampz0.data[:] = D02

        self.dampx1 = Function(name="dampx1", grid=grid, space_order=so, staggered=(x,z))
        self.dampz1 = Function(name="dampz1", grid=grid, space_order=so, staggered=(x,z))
        self.dampx1.data[0:nptx-1,0:nptz-1] = D11
        self.dampz1.data[0:nptx-1,0:nptz-1] = D12

        # apply boundary fill for dampx1, dampz1
        self.dampx1.data[nptx-1,0:nptz-1] = self.dampx1.data[nptx-2,0:nptz-1]
        self.dampx1.data[0:nptx,nptz-1]   = self.dampx1.data[0:nptx,nptz-2]
        self.dampz1.data[nptx-1,0:nptz-1] = self.dampz1.data[nptx-2,0:nptz-1]
        self.dampz1.data[0:nptx,nptz-1]   = self.dampz1.data[0:nptx,nptz-2]

        # => done. Now iso_stencil can do: model.vel1, model.phi1, etc.
        # And crucially, we do NOT do self._physical_parameters.update(...) for any
        # of these. So it won't pass them as arguments to the operator.

        return  # Done creating PML arrays

    @property
    def _max_vp(self):
        if 'vp' in self._physical_parameters:
            return mmax(self.vp)
        else:
            return np.sqrt(mmin(self.b) * (mmax(self.lam) + 2 * mmax(self.mu)))

    @property
    def _thomsen_scale(self):
        # Update scale for tti
        if 'epsilon' in self._physical_parameters:
            return np.sqrt(1 + 2 * mmax(self.epsilon))
        return 1

    @property
    def dt_scale(self):
        return self._dt_scale

    @dt_scale.setter
    def dt_scale(self, val):
        self._dt_scale = val

    @property
    def _cfl_coeff(self):
        """
        Courant number from the physics and spatial discretization order.
        The CFL coefficients are described in:
        - https://doi.org/10.1137/0916052 for the elastic case
        - https://library.seg.org/doi/pdf/10.1190/1.1444605 for the acoustic case
        """
        # Elasic coefficient (see e.g )
        if 'lam' in self._physical_parameters or 'vs' in self._physical_parameters:
            coeffs = fd_w(1, range(-self.space_order//2+1, self.space_order//2+1), .5)
            c_fd = sum(np.abs(coeffs[-1][-1])) / 2
            return np.sqrt(self.dim) / self.dim / c_fd
        a1 = 4  # 2nd order in time
        coeffs = fd_w(2, range(-self.space_order, self.space_order+1), 0)[-1][-1]
        return np.sqrt(a1/float(self.grid.dim * sum(np.abs(coeffs))))

    @property
    def critical_dt(self):
        """
        Critical computational time step value from the CFL condition.
        """
        # For a fixed time order this number decreases as the space order increases.
        #
        # The CFL condtion is then given by
        # dt <= coeff * h / (max(velocity))
        dt = self._cfl_coeff * np.min(self.spacing) / (self._thomsen_scale*self._max_vp)
        dt = self.dtype("%.3e" % (self.dt_scale * dt))
        if self._dt:
            return self._dt
        return dt

    def update(self, name, value):
        """
        Update the physical parameter param.
        """
        try:
            param = getattr(self, name)
        except AttributeError:
            # No physical parameter with tha name, create it
            setattr(self, name, self._gen_phys_param(value, name, self.space_order))
            return
        # Update the physical parameter according to new value
        if isinstance(value, np.ndarray):
            if value.shape == param.shape:
                param.data[:] = value[:]
            elif value.shape == self.shape:
                initialize_function(param, value, self.nbl)
            else:
                raise ValueError("Incorrect input size %s for model" % value.shape +
                                 " %s without or %s with padding" % (self.shape,
                                                                     param.shape))
        else:
            param.data = value

    @property
    def m(self):
        """
        Squared slowness.
        """
        return 1 / (self.vp * self.vp)

    @property
    def dm(self):
        """
        Create a simple model perturbation from the velocity as `dm = div(vp)`.
        """
        dm = Function(name="dm", grid=self.grid, space_order=self.space_order)
        Operator(Eq(dm, div(self.vp)), subs=self.spacing_map)()
        return dm

    def smooth(self, physical_parameters, sigma=5.0):
        """
        Apply devito.gaussian_smooth to model physical parameters.

        Parameters
        ----------
        physical_parameters : string or tuple of string
            Names of the fields to be smoothed.
        sigma : float
            Standard deviation of the smoothing operator.
        """
        model_parameters = self.physical_params()
        for i in physical_parameters:
            gaussian_smooth(model_parameters[i], sigma=sigma)
        return


@pytest.mark.parametrize('shape', [(51, 51), (16, 16, 16)])
def test_model_update(shape):

    # Set physical properties as numpy arrays
    vp = np.full(shape, 2.5, dtype=np.float32)  # vp constant of 2.5 km/s
    qp = 3.516*((vp[:]*1000.)**2.2)*10**(-6)  # Li's empirical formula
    b = 1 / (0.31*(vp[:]*1000.)**0.25)  # Gardner's relation

    # Define a simple visco-acoustic model from the physical parameters above
    va_model = SeismicModel(space_order=4, vp=vp, qp=qp, b=b, nbl=10,
                            origin=tuple([0. for _ in shape]), shape=shape,
                            spacing=[20. for _ in shape],)

    # Define a simple acoustic model with vp=1.5 km/s
    model = SeismicModel(space_order=4, vp=np.full_like(vp, 1.5), nbl=10,
                         origin=tuple([0. for _ in shape]), shape=shape,
                         spacing=[20. for _ in shape],)

    # Define a velocity Function
    vp_fcn = Function(name='vp0', grid=model.grid, space_order=4)
    vp_fcn.data[:] = 2.5

    # Test 1. Update vp of acoustic model from array
    model.update('vp', vp)
    assert np.array_equal(va_model.vp.data, model.vp.data)

    # Test 2. Update vp of acoustic model from Function
    model.update('vp', vp_fcn.data)
    assert np.array_equal(va_model.vp.data, model.vp.data)

    # Test 3. Create a new physical parameter in the acoustic model from array
    model.update('qp', qp)
    assert np.array_equal(va_model.qp.data, model.qp.data)

    # Make a set of physical parameters from each model
    tpl1_set = set(va_model.physical_parameters)
    tpl2_set = set(model.physical_parameters)

    # Physical parameters in either set but not in the intersection.
    diff_phys_par = tuple(tpl1_set ^ tpl2_set)

    # Turn acoustic model (it is just lacking 'b') into a visco-acoustic model
    slices = tuple(slice(model.nbl, -model.nbl) for _ in range(model.dim))
    for i in diff_phys_par:
        # Test 4. Create a new physical parameter in the acoustic model from function
        model.update(i, getattr(va_model, i).data[slices])
        assert np.array_equal(getattr(model, i).data, getattr(va_model, i).data)


# For backward compatibility
Model = SeismicModel
ModelElastic = SeismicModel
ModelViscoelastic = SeismicModel
ModelViscoacoustic = SeismicModel

[zhangxiaoshuotttt]

Here are codes in wavesolver.py:

from devito import Function, TimeFunction, DevitoCheckpoint, CheckpointOperator, Revolver
from devito.tools import memoized_meth
from operators_new import (
    ForwardOperator, AdjointOperator, GradientOperator, BornOperator
)


class AcousticWaveSolver:
    """
    Solver object that provides operators for seismic inversion problems
    and encapsulates the time and space discretization for a given problem
    setup.

    Parameters
    ----------
    model : Model
        Physical model with domain parameters.
    geometry : AcquisitionGeometry
        Geometry object that contains the source (SparseTimeFunction) and
        receivers (SparseTimeFunction) and their position.
    kernel : str, optional
        Type of discretization, centered or shifted.
    space_order: int, optional
        Order of the spatial stencil discretisation. Defaults to 4.
    """
    def __init__(self, model, geometry, kernel='OT2', space_order=4, **kwargs):
        self.model = model
        #self.model._initialize_bcs(bcs="damp")
        self.geometry = geometry

        assert self.model.grid == geometry.grid

        self.space_order = space_order
        self.kernel = kernel

        # Cache compiler options
        self._kwargs = kwargs

    @property
    def dt(self):
        # Time step can be \sqrt{3}=1.73 bigger with 4th order
        if self.kernel == 'OT4':
            return self.model.dtype(1.73 * self.model.critical_dt)
        return self.model.critical_dt

    @memoized_meth
    def op_fwd(self, save=None):
        """Cached operator for forward runs with buffered wavefield"""
        return ForwardOperator(self.model, save=save, geometry=self.geometry,
                               kernel=self.kernel, space_order=self.space_order,
                               **self._kwargs)

    @memoized_meth
    def op_adj(self):
        """Cached operator for adjoint runs"""
        return AdjointOperator(self.model, save=None, geometry=self.geometry,
                               kernel=self.kernel, space_order=self.space_order,
                               **self._kwargs)

    @memoized_meth
    def op_grad(self, save=True):
        """Cached operator for gradient runs"""
        return GradientOperator(self.model, save=save, geometry=self.geometry,
                                kernel=self.kernel, space_order=self.space_order,
                                **self._kwargs)

    @memoized_meth
    def op_born(self):
        """Cached operator for born runs"""
        return BornOperator(self.model, save=None, geometry=self.geometry,
                            kernel=self.kernel, space_order=self.space_order,
                            **self._kwargs)

    def forward(self, src=None, rec=None, u=None, model=None, save=None, **kwargs):
        """
        Forward modelling function that creates the necessary
        data objects for running a forward modelling operator.

        Parameters
        ----------
        src : SparseTimeFunction or array_like, optional
            Time series data for the injected source term.
        rec : SparseTimeFunction or array_like, optional
            The interpolated receiver data.
        u : TimeFunction, optional
            Stores the computed wavefield.
        model : Model, optional
            Object containing the physical parameters.
        vp : Function or float, optional
            The time-constant velocity.
        save : bool, optional
            Whether or not to save the entire (unrolled) wavefield.

        Returns
        -------
        Receiver, wavefield and performance summary
        """
        # Source term is read-only, so re-use the default
        src = src or self.geometry.src
        # Create a new receiver object to store the result
        rec = rec or self.geometry.rec

        # Create the forward wavefield if not provided
        u = u or TimeFunction(name='u', grid=self.model.grid,
                              save=self.geometry.nt if save else None,
                              time_order=2, space_order=self.space_order)

        model = model or self.model
        # Pick vp from model unless explicitly provided
        kwargs.update(model.physical_params(**kwargs))

        # Execute operator and return wavefield and receiver data
        summary = self.op_fwd(save).apply(src=src, rec=rec, u=u,
                                          dt=kwargs.pop('dt', self.dt), **kwargs)

        return rec, u, summary

    def adjoint(self, rec, srca=None, v=None, model=None, **kwargs):
        """
        Adjoint modelling function that creates the necessary
        data objects for running an adjoint modelling operator.

        Parameters
        ----------
        rec : SparseTimeFunction or array-like
            The receiver data. Please note that
            these act as the source term in the adjoint run.
        srca : SparseTimeFunction or array-like
            The resulting data for the interpolated at the
            original source location.
        v: TimeFunction, optional
            The computed wavefield.
        model : Model, optional
            Object containing the physical parameters.
        vp : Function or float, optional
            The time-constant velocity.

        Returns
        -------
        Adjoint source, wavefield and performance summary.
        """
        # Create a new adjoint source and receiver symbol
        srca = srca or self.geometry.new_src(name='srca', src_type=None)

        # Create the adjoint wavefield if not provided
        v = v or TimeFunction(name='v', grid=self.model.grid,
                              time_order=2, space_order=self.space_order)

        model = model or self.model
        # Pick vp from model unless explicitly provided
        kwargs.update(model.physical_params(**kwargs))

        # Execute operator and return wavefield and receiver data
        summary = self.op_adj().apply(srca=srca, rec=rec, v=v,
                                      dt=kwargs.pop('dt', self.dt), **kwargs)
        return srca, v, summary

    def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None,
                         checkpointing=False, **kwargs):
        """
        Gradient modelling function for computing the adjoint of the
        Linearized Born modelling function, ie. the action of the
        Jacobian adjoint on an input data.

        Parameters
        ----------
        rec : SparseTimeFunction
            Receiver data.
        u : TimeFunction
            Full wavefield `u` (created with save=True).
        v : TimeFunction, optional
            Stores the computed wavefield.
        grad : Function, optional
            Stores the gradient field.
        model : Model, optional
            Object containing the physical parameters.
        vp : Function or float, optional
            The time-constant velocity.

        Returns
        -------
        Gradient field and performance summary.
        """
        dt = kwargs.pop('dt', self.dt)
        # Gradient symbol
        grad = grad or Function(name='grad', grid=self.model.grid)

        # Create the forward wavefield
        v = v or TimeFunction(name='v', grid=self.model.grid,
                              time_order=2, space_order=self.space_order)

        model = model or self.model
        # Pick vp from model unless explicitly provided
        kwargs.update(model.physical_params(**kwargs))

        if checkpointing:
            u = TimeFunction(name='u', grid=self.model.grid,
                             time_order=2, space_order=self.space_order)
            cp = DevitoCheckpoint([u])
            n_checkpoints = None
            wrap_fw = CheckpointOperator(self.op_fwd(save=False),
                                         src=src or self.geometry.src,
                                         u=u, dt=dt, **kwargs)
            wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v,
                                          rec=rec, dt=dt, grad=grad, **kwargs)

            # Run forward
            wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
            wrp.apply_forward()
            summary = wrp.apply_reverse()
        else:
            summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt,
                                           **kwargs)
        return grad, summary

    def jacobian(self, dmin, src=None, rec=None, u=None, U=None, model=None, **kwargs):
        """
        Linearized Born modelling function that creates the necessary
        data objects for running an adjoint modelling operator.

        Parameters
        ----------
        src : SparseTimeFunction or array_like, optional
            Time series data for the injected source term.
        rec : SparseTimeFunction or array_like, optional
            The interpolated receiver data.
        u : TimeFunction, optional
            The forward wavefield.
        U : TimeFunction, optional
            The linearized wavefield.
        model : Model, optional
            Object containing the physical parameters.
        vp : Function or float, optional
            The time-constant velocity.
        """
        # Source term is read-only, so re-use the default
        src = src or self.geometry.src
        # Create a new receiver object to store the result
        rec = rec or self.geometry.rec

        # Create the forward wavefields u and U if not provided
        u = u or TimeFunction(name='u', grid=self.model.grid,
                              time_order=2, space_order=self.space_order)
        U = U or TimeFunction(name='U', grid=self.model.grid,
                              time_order=2, space_order=self.space_order)

        model = model or self.model
        # Pick vp from model unless explicitly provided
        kwargs.update(model.physical_params(**kwargs))

        # Execute operator and return wavefield and receiver data
        summary = self.op_born().apply(dm=dmin, u=u, U=U, src=src, rec=rec,
                                       dt=kwargs.pop('dt', self.dt), **kwargs)
        return rec, u, U, summary

    # Backward compatibility
    born = jacobian
    gradient = jacobian_adjoint

[zhangxiaoshuotttt]

Here are codes in operatos.py:

#Shuo Zhang 2025/01/15

from devito import Eq, Operator, Function, TimeFunction, Inc, solve, sign
from devito.symbolics import retrieve_functions, INT, retrieve_derivatives
import numpy as np

def freesurface(model, eq):
    """
    Generate the stencil that mirrors the field as a free surface modeling for
    the acoustic wave equation.

    Parameters
    ----------
    model : Model
        Physical model.
    eq : Eq
        Time-stepping stencil (time update) to mirror at the freesurface.
    """
    lhs, rhs = eq.args
    # Get vertical dimension and corresponding subdimension
    fsdomain = model.grid.subdomains['fsdomain']
    zfs = fsdomain.dimensions[-1]
    z = zfs.parent

    # Retrieve vertical derivatives
    dzs = {d for d in retrieve_derivatives(rhs) if z in d.dims}
    # Remove inner duplicate
    dzs = dzs - {d for D in dzs for d in retrieve_derivatives(D.expr) if z in d.dims}
    dzs = {d: d._eval_at(lhs).evaluate for d in dzs}

    # Finally get functions for evaluated derivatives
    funcs = {f for f in retrieve_functions(dzs.values())}

    mapper = {}
    # Antisymmetric mirror at negative indices
    # TODO: Make a proper "mirror_indices" tool function
    for f in funcs:
        zind = f.indices[-1]
        if (zind - z).as_coeff_Mul()[0] < 0:
            s = sign(zind.subs({z: zfs, z.spacing: 1}))
            mapper.update({f: s * f.subs({zind: INT(abs(zind))})})

    # Mapper for vertical derivatives
    dzmapper = {d: v.subs(mapper) for d, v in dzs.items()}

    fs_eq = [eq.func(lhs, rhs.subs(dzmapper), subdomain=fsdomain)]
    fs_eq.append(eq.func(lhs._subs(z, 0), 0, subdomain=fsdomain))

    return fs_eq


def laplacian(field, model, kernel):
    """
    Spatial discretization for the isotropic acoustic wave equation. For a 4th
    order in time formulation, the 4th order time derivative is replaced by a
    double laplacian:
    H = (laplacian + s**2/12 laplacian(1/m*laplacian))

    Parameters
    ----------
    field : TimeFunction
        The computed solution.
    model : Model
        Physical model.
    """
    if kernel not in ['OT2', 'OT4']:
        raise ValueError("Unrecognized kernel")
    s = model.grid.time_dim.spacing
    biharmonic = field.biharmonic(1/model.m) if kernel == 'OT4' else 0
    return field.laplace + s**2/12 * biharmonic

def fdamp(x,z,i,x0pml,x1pml,z0pml,z1pml,lx,lz):
    
    quibar  = 0.05
          
    if(i==1):
        a = np.where(x<=x0pml,(np.abs(x-x0pml)/lx),np.where(x>=x1pml,(np.abs(x-x1pml)/lx),0.))
        fdamp = quibar*(a-(1./(2.*np.pi))*np.sin(2.*np.pi*a))
    if(i==2):
        a = np.where(z<=z0pml,(np.abs(z-z0pml)/lz),np.where(z>=z1pml,(np.abs(z-z1pml)/lz),0.))
        fdamp = quibar*(a-(1./(2.*np.pi))*np.sin(2.*np.pi*a))
      
    return fdamp

def generatemdamp(x0,x1,z0,z1,nptx,nptz,hxv,hzv,x0pml,x1pml,z0pml,z1pml,lx,lz):
    
    X0     = np.linspace(x0,x1,nptx)    
    Z0     = np.linspace(z0,z1,nptz)
    X0grid,Z0grid = np.meshgrid(X0,Z0)
    X1   = np.linspace((x0+0.5*hxv),(x1-0.5*hxv),nptx-1)
    Z1   = np.linspace((z0+0.5*hzv),(z1-0.5*hzv),nptz-1)
    X1grid,Z1grid = np.meshgrid(X1,Z1)
   
    D01 = np.zeros((nptx,nptz))
    D02 = np.zeros((nptx,nptz))
    D11 = np.zeros((nptx,nptz))
    D12 = np.zeros((nptx,nptz))
    
    D01 = np.transpose(fdamp(X0grid,Z0grid,1,x0pml,x1pml,z0pml,z1pml,lx,lz))
    D02 = np.transpose(fdamp(X0grid,Z0grid,2,x0pml,x1pml,z0pml,z1pml,lx,lz))
  
    D11 = np.transpose(fdamp(X1grid,Z1grid,1,x0pml,x1pml,z0pml,z1pml,lx,lz))
    D12 = np.transpose(fdamp(X1grid,Z1grid,2,x0pml,x1pml,z0pml,z1pml,lx,lz))
    
    return D01, D02, D11, D12

def iso_stencil(field, model, kernel, **kwargs):
    """
    Stencil for the acoustic isotropic wave-equation:
    u.dt2 - H + damp*u.dt = 0.

    Parameters
    ----------
    field : TimeFunction
        The computed solution.
    model : Model
        Physical model.
    kernel : str, optional
        Type of discretization, 'OT2' or 'OT4'.
    q : TimeFunction, Function or float
        Full-space/time source of the wave-equation.
    forward : bool, optional
        Whether to propagate forward (True) or backward (False) in time.
    """
    # Forward or backward
    forward = kwargs.get('forward', True)
    # Define time step to be updated
    unext = field.forward if forward else field.backward
    udt = field.dt if forward else field.dt.T
    # Get the spacial FD
    lap = laplacian(field, model, kernel)
    # Get source
    q = kwargs.get('q', 0)

    # stencil for PML
    bcs = model.pml
    if bcs:
        (x, z) = model.grid.dimensions
        vel1 = model.vel1
        phi1 = model.phi1
        phi2 = model.phi2
        dampx0 = model.dampx0
        dampz0 = model.dampz0
        dampx1 = model.dampx1
        dampz1 = model.dampz1
        
        # In physical domain, the wave equation is:
        lhs_interior = model.m * field.dt2 - lap 
        eq_phys = Eq(unext, solve(lhs_interior - q, unext), subdomain=model.grid.subdomains['physdomain'])

        (hx,hz) = model.grid.spacing_map  
        t       = model.grid.stepping_dim

        pde02a  = field.dt2   + (dampx0+dampz0)*field.dt + (dampx0*dampz0)*field - lap

        pde02b  = - (0.5/hx)*(phi1[t,x,z-1]+phi1[t,x,z]-phi1[t,x-1,z-1]-phi1[t,x-1,z])
        pde02c  = - (0.5/hz)*(phi2[t,x-1,z]+phi2[t,x,z]-phi2[t,x-1,z-1]-phi2[t,x,z-1])

        if model.fs:
            subds = ['pml_left','pml_right','pml_bottom']
        else:
            subds = ['pml_left','pml_right','pml_bottom','pml_top']
        
        #Now in PML Domain
        lhs02  = pde02a + pde02b + pde02c
        stencil02 = [Eq(unext,solve(lhs02 - q, unext), subdomain = model.grid.subdomains[subds[i]]) for i in range(0,len(subds))]

        pde10 = phi1.dt + dampx1*0.5*(phi1.forward+phi1)
        a1    = field[t+1,x+1,z] + field[t+1,x+1,z+1] - field[t+1,x,z] - field[t+1,x,z+1] 
        a2    = field[t,x+1,z]   + field[t,x+1,z+1]   - field[t,x,z]   - field[t,x,z+1] 
        pde11 = -(dampz1-dampx1)*0.5*(0.5/hx)*(a1+a2)*vel1**2

        lhs1  = pde10+pde11
        stencil1 = [Eq(phi1.forward, solve(lhs1,phi1.forward),subdomain = model.grid.subdomains[subds[i]]) for i in range(0,len(subds))]
                                           
        pde20 = phi2.dt + dampz1*0.5*(phi2.forward+phi2) 
        b1    = field[t+1,x,z+1] + field[t+1,x+1,z+1] - field[t+1,x,z] - field[t+1,x+1,z] 
        b2    = field[t,x,z+1]   + field[t,x+1,z+1]   - field[t,x,z]   - field[t,x+1,z] 
        pde21 = -(dampx1-dampz1)*0.5*(0.5/hz)*(b1+b2)*vel1**2
        lhs2  = pde20+pde21
        stencil2 = [Eq(phi2.forward, solve(lhs2,phi2.forward),subdomain = model.grid.subdomains[subds[i]]) for i in range(0,len(subds))]

        eqns = [eq_phys,stencil02,stencil1,stencil2]
                # Add free surface conditions if necessary
        if model.fs:
            eqns += freesurface(model, eq_phys)
        return eqns
    else: 

        # Define PDE and update rule
        eq_time = solve(model.m * field.dt2 - lap - q + model.damp * udt, unext)

        # Time-stepping stencil.
        eqns = [Eq(unext, eq_time, subdomain=model.grid.subdomains['physdomain'])]

    # Add free surface
    if model.fs:
        eqns.append(freesurface(model, Eq(unext, eq_time)))
    return eqns


def ForwardOperator(model, geometry, space_order=4,
                    save=False, kernel='OT2', **kwargs):
    """
    Construct a forward modelling operator in an acoustic medium.

    Parameters
    ----------
    model : Model
        Object containing the physical parameters.
    geometry : AcquisitionGeometry
        Geometry object that contains the source (SparseTimeFunction) and
        receivers (SparseTimeFunction) and their position.
    space_order : int, optional
        Space discretization order.
    save : int or Buffer, optional
        Saving flag, True saves all time steps. False saves three timesteps.
        Defaults to False.
    kernel : str, optional
        Type of discretization, 'OT2' or 'OT4'.
    """
    m = model.m

    # Create symbols for forward wavefield, source and receivers
    u = TimeFunction(name='u', grid=model.grid,
                     save=geometry.nt if save else None,
                     time_order=2, space_order=space_order)
    src = geometry.src
    rec = geometry.rec

    s = model.grid.stepping_dim.spacing
    eqn = iso_stencil(u, model, kernel)

    # Construct expression to inject source values
    src_term = src.inject(field=u.forward, expr=src * s**2 / m)

    # Create interpolation expression for receivers
    rec_term = rec.interpolate(expr=u)

    # Substitute spacing terms to reduce flops
    return Operator(eqn + src_term + rec_term, subs=model.spacing_map,
                    name='Forward', **kwargs)


def AdjointOperator(model, geometry, space_order=4,
                    kernel='OT2', **kwargs):
    """
    Construct an adjoint modelling operator in an acoustic media.

    Parameters
    ----------
    model : Model
        Object containing the physical parameters.
    geometry : AcquisitionGeometry
        Geometry object that contains the source (SparseTimeFunction) and
        receivers (SparseTimeFunction) and their position.
    space_order : int, optional
        Space discretization order.
    kernel : str, optional
        Type of discretization, 'OT2' or 'OT4'.
    """
    m = model.m

    v = TimeFunction(name='v', grid=model.grid, save=None,
                     time_order=2, space_order=space_order)
    srca = geometry.new_src(name='srca', src_type=None)
    rec = geometry.rec

    s = model.grid.stepping_dim.spacing
    eqn = iso_stencil(v, model, kernel, forward=False)

    # Construct expression to inject receiver values
    receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)

    # Create interpolation expression for the adjoint-source
    source_a = srca.interpolate(expr=v)

    # Substitute spacing terms to reduce flops
    return Operator(eqn + receivers + source_a, subs=model.spacing_map,
                    name='Adjoint', **kwargs)


def GradientOperator(model, geometry, space_order=4, save=True,
                     kernel='OT2', **kwargs):
    """
    Construct a gradient operator in an acoustic media.

    Parameters
    ----------
    model : Model
        Object containing the physical parameters.
    geometry : AcquisitionGeometry
        Geometry object that contains the source (SparseTimeFunction) and
        receivers (SparseTimeFunction) and their position.
    space_order : int, optional
        Space discretization order.
    save : int or Buffer, optional
        Option to store the entire (unrolled) wavefield.
    kernel : str, optional
        Type of discretization, centered or shifted.
    """
    m = model.m

    # Gradient symbol and wavefield symbols
    grad = Function(name='grad', grid=model.grid)
    u = TimeFunction(name='u', grid=model.grid, save=geometry.nt if save
                     else None, time_order=2, space_order=space_order)
    v = TimeFunction(name='v', grid=model.grid, save=None,
                     time_order=2, space_order=space_order)
    rec = geometry.rec

    s = model.grid.stepping_dim.spacing
    eqn = iso_stencil(v, model, kernel, forward=False)

    if kernel == 'OT2':
        gradient_update = Inc(grad, - u * v.dt2)
    elif kernel == 'OT4':
        gradient_update = Inc(grad, - u * v.dt2 - s**2 / 12.0 * u.biharmonic(m**(-2)) * v)
    # Add expression for receiver injection
    receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)

    # Substitute spacing terms to reduce flops
    return Operator(eqn + receivers + [gradient_update], subs=model.spacing_map,
                    name='Gradient', **kwargs)


def BornOperator(model, geometry, space_order=4,
                 kernel='OT2', **kwargs):
    """
    Construct an Linearized Born operator in an acoustic media.

    Parameters
    ----------
    model : Model
        Object containing the physical parameters.
    geometry : AcquisitionGeometry
        Geometry object that contains the source (SparseTimeFunction) and
        receivers (SparseTimeFunction) and their position.
    space_order : int, optional
        Space discretization order.
    kernel : str, optional
        Type of discretization, centered or shifted.
    """
    m = model.m

    # Create source and receiver symbols
    src = geometry.src
    rec = geometry.rec

    # Create wavefields and a dm field
    u = TimeFunction(name="u", grid=model.grid, save=None,
                     time_order=2, space_order=space_order)
    U = TimeFunction(name="U", grid=model.grid, save=None,
                     time_order=2, space_order=space_order)
    dm = Function(name="dm", grid=model.grid, space_order=0)

    s = model.grid.stepping_dim.spacing
    eqn1 = iso_stencil(u, model, kernel)
    eqn2 = iso_stencil(U, model, kernel, q=-dm*u.dt2)

    # Add source term expression for u
    source = src.inject(field=u.forward, expr=src * s**2 / m)

    # Create receiver interpolation expression from U
    receivers = rec.interpolate(expr=U)

    # Substitute spacing terms to reduce flops
    return Operator(eqn1 + source + eqn2 + receivers, subs=model.spacing_map,
                    name='Born', **kwargs)

[zhangxiaoshuotttt]

I am sorry for using such a not elegant way to show my codes, once again, thanks for your time and help.