devitocodes / devito

DSL and compiler framework for automated finite-differences and stencil computation
http://www.devitoproject.org
MIT License
565 stars 229 forks source link

Customized coefficients for the staggered grid are not correctly substituted. #2485

Open HaoHu-seis opened 1 week ago

HaoHu-seis commented 1 week ago

The customized coefficients for the staggered grid FD are not correctly substituted, at least for the items, tau_zz.dz, vx.dx, vy.dy, vz.dz, for isotropic elastic case. For VTI/TTI, we should double-check again.

The MFE:

from devito import Eq, Operator, VectorTimeFunction, TensorTimeFunction, Grid
from devito import div, grad, diag, solve
import numpy as np
from devito import Coefficient, Substitutions

# set up the model
space_order = 4
shape = (51, 51, 51)
spacing = (10., 10., 10.)
extent = (500, 500, 500)

grid = Grid(shape=shape, extent= extent)

# first-order particle-velocity/strain elastic wave equation
v = VectorTimeFunction(name='v', grid=grid,
                       save=None,
                       space_order=space_order, time_order=1,
                       coefficients='symbolic')  # the customized coefficients
tau = TensorTimeFunction(name='tau', grid=grid,
                         save=None,
                         space_order=space_order, time_order=1,
                         coefficients='symbolic')  # the customized coefficients

# customized FD coefficients
# weights_opt = np.array([0.4304542E-1, -0.1129136E+1, 0.1129136E+1, -0.4304542E-12])
# weights_opt = np.array([0.1, -1, 1, -0.1])
weights_opt = np.array([0, 0, 0, 0])

x, y, z = grid.dimensions

list_subs_coef = []
for u in v.flat():
    for idim, idim_symbol in enumerate(grid.dimensions):
            list_subs_coef.append(Coefficient(
                1, u, idim_symbol, weights_opt/idim_symbol.spacing))
coeffs_v = Substitutions(*list_subs_coef)

list_subs_coef = []
for u in tau.flat():
    for idim, idim_symbol in enumerate(grid.dimensions):
            list_subs_coef.append(Coefficient(
                1, u, idim_symbol, weights_opt/idim_symbol.spacing))
coeffs_t = Substitutions(*list_subs_coef)

lam = 1.
mu = 1.
b = 1.

# Particle velocity
eq_v = v.dt - b * div(tau)
# Stress
e = (grad(v.forward) + grad(v.forward).T)
eq_tau = tau.dt - lam * diag(div(v.forward)) - mu * e

u_v = Eq(v.forward, solve(eq_v, v.forward),
         coefficients=coeffs_v)
print(u_v.evaluate)

u_tau = Eq(tau.forward, solve(eq_tau, tau.forward),
           coefficients=coeffs_t)
print(u_tau.evaluate)

OP = Operator([u_v, u_tau])
print(OP)

The coefficients for items, tau_zz.dz, vx.dx, vy.dy, vz.dz are not correctly substituted from the generated C code:

              v_x[t1][x + 4][y + 4][z + 4] = v_x[t0][x + 4][y + 4][z + 4];
              v_y[t1][x + 4][y + 4][z + 4] = v_y[t0][x + 4][y + 4][z + 4];
              v_z[t1][x + 4][y + 4][z + 4] = dt*(1.0F*r0*(4.16666667e-2F*(tau_zz[t0][x + 4][y + 4][z + 3] - tau_zz[t0][x + 4][y + 4][z + 6]) + 1.1250F*(-tau_zz[t0][x + 4][y + 4][z + 4] + tau_zz[t0][x + 4][y + 4][z + 5])) + r1*v_z[t0][x + 4][y + 4][z + 4]);
              tau_xx[t1][x + 4][y + 4][z + 4] = tau_xx[t0][x + 4][y + 4][z + 4];
              tau_xy[t1][x + 4][y + 4][z + 4] = tau_xy[t0][x + 4][y + 4][z + 4];
              tau_xz[t1][x + 4][y + 4][z + 4] = tau_xz[t0][x + 4][y + 4][z + 4];

              float r4 = -v_z[t1][x + 4][y + 4][z + 3] + v_z[t1][x + 4][y + 4][z + 4];
              float r5 = v_z[t1][x + 4][y + 4][z + 2] - v_z[t1][x + 4][y + 4][z + 5];
              float r6 = -v_x[t1][x + 3][y + 4][z + 4] + v_x[t1][x + 4][y + 4][z + 4];
              float r7 = v_x[t1][x + 2][y + 4][z + 4] - v_x[t1][x + 5][y + 4][z + 4];
              float r8 = -v_y[t1][x + 4][y + 3][z + 4] + v_y[t1][x + 4][y + 4][z + 4];
              float r9 = v_y[t1][x + 4][y + 2][z + 4] - v_y[t1][x + 4][y + 5][z + 4];
              tau_yy[t1][x + 4][y + 4][z + 4] = dt*(r1*tau_yy[t0][x + 4][y + 4][z + 4] + 2.0F*r3*(1.1250F*r8 + 4.16666667e-2F*r9) + 1.0F*(r0*(1.1250F*r4 + 4.16666667e-2F*r5) + r2*(1.1250F*r6 + 4.16666667e-2F*r7)));
              tau_yz[t1][x + 4][y + 4][z + 4] = tau_yz[t0][x + 4][y + 4][z + 4];
              tau_zz[t1][x + 4][y + 4][z + 4] = dt*(r0*(1.1250F*r4 + 4.16666666642413e-2F*r5) + r1*tau_zz[t0][x + 4][y + 4][z + 4] + r2*(1.1250F*r6 + 4.16666666642413e-2F*r7) + r3*(1.1250F*r8 + 4.16666666642413e-2F*r9));