def test_forward():
tn = 1000.
shape = (101, 101)
nbl = 40
model = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
model0 = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
gaussian_smooth(model0.vp, sigma=(1, 1))
geometry0 = setup_geometry(model0, tn)
geometry = setup_geometry(model, tn)
# Pure Devito
solver = AcousticWaveSolver(model, geometry, space_order=8)
solver0 = AcousticWaveSolver(model0, geometry0, space_order=8)
d = solver.forward()[0]
d0, u0 = solver0.forward(save=True, vp=model0.vp)[:2]
residual = d.data - d0.data
rec = geometry0.rec
rec.data[:] = -residual[:]
grad_devito = solver0.jacobian_adjoint(rec, u0)[0].data
grad_devito = np.array(grad_devito)[nbl:-nbl, nbl:-nbl]
# Devito4PyTorch
d = torch.from_numpy(np.array(d.data)).to(device)
forward_modeling = ForwardModelingLayer(model0, geometry0, device)
m0 = np.array(model0.vp.data**(-2))[nbl:-nbl, nbl:-nbl]
m0 = torch.Tensor(m0).unsqueeze(0).unsqueeze(0).to(device)
m0.requires_grad = True
loss = 0.5 * torch.norm(forward_modeling(m0) - d)**2
grad = torch.autograd.grad(loss, m0, create_graph=False)[0]
# Test
rel_err = np.linalg.norm(grad.cpu().numpy() -
grad_devito) / np.linalg.norm(grad_devito)
assert np.isclose(rel_err, 0., atol=1.e-6)
def test_forward_born():
tn = 1000.
shape = (101, 101)
nbl = 40
model = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
model0 = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
gaussian_smooth(model0.vp, sigma=(1, 1))
geometry0 = setup_geometry(model0, tn)
geometry = setup_geometry(model, tn)
# Pure Devito
solver = AcousticWaveSolver(model, geometry, space_order=8)
solver0 = AcousticWaveSolver(model0, geometry0, space_order=8)
d = solver.forward(vp=model.vp)[0]
d0, u0 = solver0.forward(save=True, vp=model0.vp)[:2]
d_lin = d.data - d0.data
rec = geometry0.rec
rec.data[:] = -d_lin[:]
grad_devito = solver0.jacobian_adjoint(rec, u0)[0].data
grad_devito = np.array(grad_devito)[nbl:-nbl, nbl:-nbl]
# Devito4PyTorch
d_lin = torch.from_numpy(np.array(d_lin)).to(device)
forward_born = ForwardBornLayer(model0, geometry0, device)
dm_est = torch.zeros([1, 1, shape[0], shape[1]],
requires_grad=True,
device=device)
loss = 0.5 * torch.norm(forward_born(dm_est) - d_lin)**2
grad = torch.autograd.grad(loss, dm_est, create_graph=False)[0]
# Test
rel_err = np.linalg.norm(grad.cpu().numpy() -
grad_devito) / np.linalg.norm(grad_devito)
assert np.isclose(rel_err, 0., atol=1.e-6)
def viscoelastic_setup(shape=(50, 50),
spacing=(15.0, 15.0),
tn=500.,
space_order=4,
nbl=10,
constant=True,
**kwargs):
preset = 'constant-viscoelastic' if constant else 'layers-viscoelastic'
model = demo_model(preset,
space_order=space_order,
shape=shape,
nbl=nbl,
dtype=kwargs.pop('dtype', np.float32),
spacing=spacing)
# Source and receiver geometries
geometry = setup_geometry(model, tn)
# Create solver object to provide relevant operators
solver = ViscoelasticWaveSolver(model,
geometry,
space_order=space_order,
**kwargs)
return solver
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,
space_order=space_order,
kernel=kernel,
**kwargs)
def acoustic_setup(shape=(50, 50, 50),
spacing=(15.0, 15.0, 15.0),
tn=500.,
kernel='OT2',
space_order=4,
nbl=10,
preset='layers-isotropic',
fs=False,
**kwargs):
model = demo_model(preset,
space_order=space_order,
shape=shape,
nbl=nbl,
dtype=kwargs.pop('dtype', np.float32),
spacing=spacing,
fs=fs,
**kwargs)
# Source and receiver geometries
geometry = setup_geometry(model, tn)
# Create solver object to provide relevant operators
solver = AcousticWaveSolver(model,
geometry,
kernel=kernel,
space_order=space_order,
**kwargs)
return solver
def viscoacoustic_setup(shape=(50, 50),
spacing=(15.0, 15.0),
tn=500.,
space_order=4,
nbl=40,
preset='layers-viscoacoustic',
kernel='sls',
time_order=2,
**kwargs):
model = demo_model(preset,
space_order=space_order,
shape=shape,
nbl=nbl,
dtype=kwargs.pop('dtype', np.float32),
spacing=spacing)
# Source and receiver geometries
geometry = setup_geometry(model, tn)
# Create solver object to provide relevant operators
solver = ViscoacousticWaveSolver(model,
geometry,
space_order=space_order,
kernel=kernel,
time_order=time_order,
**kwargs)
return solver
def test_default_geom(shape):
vp = np.ones(shape)
o = tuple([0] * len(shape))
d = tuple([10] * len(shape))
model = Model(o, d, shape, 4, vp, nbl=20, dt=1)
assert model.critical_dt == 1
geometry = setup_geometry(model, 250)
nrec = shape[0] * (shape[1] if len(shape) > 2 else 1)
assert geometry.grid == model.grid
assert geometry.nrec == nrec
assert geometry.nsrc == 1
assert geometry.src_type == "Ricker"
assert geometry.rec.shape == (251, nrec)
assert norm(geometry.rec) == 0
assert geometry.src.shape == (251, 1)
assert norm(geometry.new_src(src_type=None)) == 0
rec2 = geometry.rec.resample(num=501)
assert rec2.shape == (501, nrec)
assert rec2.grid == model.grid
assert geometry.new_rec(name="bonjour").name == "bonjour"
assert geometry.new_src(name="bonjour").name == "bonjour"
def test_adjoint_born():
tn = 1000.
shape = (101, 101)
nbl = 40
model = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
model0 = demo_model('layers-isotropic',
origin=(0., 0.),
shape=shape,
spacing=(10., 10.),
nbl=nbl,
nlayers=2,
space_order=8)
gaussian_smooth(model0.vp, sigma=(1, 1))
geometry0 = setup_geometry(model0, tn)
# Pure Devito
solver0 = AcousticWaveSolver(model0, geometry0, space_order=8)
dm = model.vp.data**(-2) - model0.vp.data**(-2)
grad_devito = np.array(solver0.jacobian(dm)[0].data)
# Devito4PyTorch
dm = torch.from_numpy(np.array(dm[nbl:-nbl, nbl:-nbl])).to(device)
d_est = torch.zeros(geometry0.rec.data.shape,
requires_grad=True,
device=device)
adjoint_born = AdjointBornLayer(model0, geometry0, device)
loss = 0.5 * torch.norm(adjoint_born(d_est) - dm)**2
# Negative gradient to be compared with linearized data computed above
grad = -torch.autograd.grad(loss, d_est, create_graph=False)[0]
# Test
rel_err = np.linalg.norm(grad.cpu().numpy() -
grad_devito) / np.linalg.norm(grad_devito)
assert np.isclose(rel_err, 0., atol=1.e-6)
def acoustic_sa_setup(shape=(50, 50, 50),
spacing=(10.0, 10.0, 10.0),
tn=500.,
space_order=8,
nbl=10,
**kwargs):
# SA parameters
if space_order < 8:
info("Low space order not supported, running space_order=8")
space_order = 8
qmin = 0.1
qmax = 1000.0
fpeak = 0.010
omega = 2.0 * np.pi * fpeak
vp = 1.5 * np.ones(shape)
b = 1.0 * np.ones(shape)
init_damp = lambda func, nbl: setup_w_over_q(
func, omega, qmin, qmax, nbl, sigma=0)
o = tuple([0] * len(shape))
spacing = spacing[:len(shape)]
model = Model(origin=o,
shape=shape,
vp=vp,
b=b,
spacing=spacing,
nbl=nbl,
space_order=space_order,
bcs=init_damp,
dtype=kwargs.pop('dtype', np.float32),
**kwargs)
# Source and receiver geometries
geometry = setup_geometry(model, tn)
# Create solver object to provide relevant operators
solver = SaIsoAcousticWaveSolver(model,
geometry,
space_order=space_order,
**kwargs)
return solver
def elastic_VTI_setup(origin=(0., 0., 0.),
spacing=(15.0, 15.0, 15.0),
shape=(50, 50, 50),
space_order=4,
vp=2.0,
vs=1.0,
rho=1.8,
epsilon=0.25,
delta=0.10,
gamma=0.05,
nbl=10,
tn=500.,
constant=False,
**kwargs):
model = ModelElasticVTI(origin=origin,
spacing=spacing,
shape=shape,
space_order=space_order,
vp=vp,
vs=vs,
rho=rho,
epsilon=epsilon,
delta=delta,
gamma=gamma,
nbl=nbl,
dtype=kwargs.pop('dtype', np.float32),
**kwargs)
# Source and receiver geometries
geometry = setup_geometry(model, tn)
# Create solver object to provide relevant operators
solver = ElasticVTIWaveSolver(model,
geometry,
space_order=space_order,
**kwargs)
return solver
def test_gradient_equivalence(self, shape, kernel, space_order, preset,
nbl, dtype, tolerance, spacing, tn):
""" This test asserts that the gradient calculated through the following three
expressions should match within floating-point precision:
- grad = sum(-u.dt2 * v)
- grad = sum(-u * v.dt2)
- grad = sum(-u.dt * v.dt)
The computation has the following number of operations:
u.dt2 (5 ops) * v = 6ops * 500 (nt) ~ 3000 ops ~ 1e4 ops
Hence tolerances are eps * ops = 1e-4 (sp) and 1e-13 (dp)
"""
model = demo_model(preset,
space_order=space_order,
shape=shape,
nbl=nbl,
dtype=dtype,
spacing=spacing)
m = model.m
v_true = model.vp
geometry = setup_geometry(model, tn)
dt = model.critical_dt
src = geometry.src
rec = geometry.rec
rec_true = geometry.rec
rec0 = geometry.rec
s = model.grid.stepping_dim.spacing
u = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order,
save=geometry.nt)
eqn_fwd = iso_stencil(u, model, kernel)
src_term = src.inject(field=u.forward, expr=src * s**2 / m)
rec_term = rec.interpolate(expr=u)
fwd_op = Operator(eqn_fwd + src_term + rec_term,
subs=model.spacing_map,
name='Forward')
v0 = Function(name='v0',
grid=model.grid,
space_order=space_order,
dtype=dtype)
smooth(v0, model.vp)
grad_u = Function(name='gradu', grid=model.grid)
grad_v = Function(name='gradv', grid=model.grid)
grad_uv = Function(name='graduv', grid=model.grid)
v = TimeFunction(name='v',
grid=model.grid,
save=None,
time_order=2,
space_order=space_order)
s = model.grid.stepping_dim.spacing
eqn_adj = iso_stencil(v, model, kernel, forward=False)
receivers = rec.inject(field=v.backward, expr=rec * s**2 / m)
gradient_update_v = Eq(grad_v, grad_v - u * v.dt2)
grad_op_v = Operator(eqn_adj + receivers + [gradient_update_v],
subs=model.spacing_map,
name='GradientV')
gradient_update_u = Eq(grad_u, grad_u - u.dt2 * v)
grad_op_u = Operator(eqn_adj + receivers + [gradient_update_u],
subs=model.spacing_map,
name='GradientU')
gradient_update_uv = Eq(grad_uv, grad_uv + u.dt * v.dt)
grad_op_uv = Operator(eqn_adj + receivers + [gradient_update_uv],
subs=model.spacing_map,
name='GradientUV')
fwd_op.apply(dt=dt, vp=v_true, rec=rec_true)
fwd_op.apply(dt=dt, vp=v0, rec=rec0)
residual = Receiver(name='rec',
grid=model.grid,
data=(rec0.data - rec_true.data),
time_range=geometry.time_axis,
coordinates=geometry.rec_positions,
dtype=dtype)
grad_op_u.apply(dt=dt, vp=v0, rec=residual)
# Reset v before calling the second operator since the object is shared
v.data[:] = 0.
grad_op_v.apply(dt=dt, vp=v0, rec=residual)
v.data[:] = 0.
grad_op_uv.apply(dt=dt, vp=v0, rec=residual)
assert (np.allclose(grad_u.data,
grad_v.data,
rtol=tolerance,
atol=tolerance))
assert (np.allclose(grad_u.data,
grad_uv.data,
rtol=tolerance,
atol=tolerance))