def born(model, src_coords, rcv_coords, wavelet, space_order=8,
save=False, q=None, free_surface=False, isic=False, ws=None):
"""
Low level propagator, to be used through `interface.py`
Compute adjoint wavefield v = adjoint(F(m))*y
and related quantities (||v||_w, v(xsrc))
"""
# Setting adjoint wavefield
u = wavefield(model, space_order, save=save, nt=wavelet.shape[0])
ul = wavefield(model, space_order, name="l")
# Extended source
q = q or wf_as_src(u, w=0)
q = extented_src(model, ws, wavelet, q=q)
# Set up PDE expression and rearrange
pde, fsu = wave_kernel(model, u, fs=free_surface, q=q)
pdel, fsul = wave_kernel(model, ul, q=lin_src(model, u, isic=isic), fs=free_surface)
# Setup source and receiver
geom_expr, _, _ = src_rec(model, u, src_coords=src_coords, wavelet=wavelet)
geom_exprl, _, rcvl = src_rec(model, ul, rec_coords=rcv_coords, nt=wavelet.shape[0])
# Create operator and run
subs = model.spacing_map
op = Operator(pde + geom_expr + pdel + geom_exprl + fsu + fsul,
subs=subs, name="born"+name(model))
op(**op_kwargs(model, fs=free_surface))
# Output
return rcvl.data, u
def forward(model,
src_coords,
rcv_coords,
wavelet,
space_order=8,
save=False,
q=None,
return_op=False,
freq_list=None,
dft_sub=None,
ws=None,
t_sub=1,
**kwargs):
"""
Low level propagator, to be used through `interface.py`
Compute forward wavefield u = A(m)^{-1}*f and related quantities (u(xrcv))
"""
# Number of time steps
nt = as_tuple(q)[0].shape[0] if wavelet is None else wavelet.shape[0]
# Setting forward wavefield
u = wavefield(model, space_order, save=save, nt=nt, t_sub=t_sub)
# Expression for saving wavefield if time subsampling is used
u_save, eq_save = wavefield_subsampled(model, u, nt, t_sub)
# Add extended source
q = q or wf_as_src(u, w=0)
q = extented_src(model, ws, wavelet, q=q)
# Set up PDE expression and rearrange
pde = wave_kernel(model, u, q=q)
# Setup source and receiver
geom_expr, _, rcv = src_rec(model,
u,
src_coords=src_coords,
nt=nt,
rec_coords=rcv_coords,
wavelet=wavelet)
# On-the-fly Fourier
dft, dft_modes = otf_dft(u, freq_list, model.critical_dt, factor=dft_sub)
# Create operator and run
subs = model.spacing_map
op = Operator(pde + dft + geom_expr + eq_save,
subs=subs,
name="forward" + name(model),
opt=opt_op(model))
op.cfunction
if return_op:
return op, u, rcv
summary = op()
# Output
return rcv, dft_modes or (u_save if t_sub > 1 else u), summary
def wri_func(model, src_coords, wavelet, rec_coords, recin, yin, space_order=8,
isic=False, ws=None, t_sub=1, grad="m", grad_corr=False,
alpha_op=False, w_fun=None, eps=0):
"""
Time domain wavefield reconstruction inversion wrapper
"""
# F(m0) * q if y is not an input and compute y = r(m0)
if yin is None or grad_corr:
y, u0, _ = forward(model, src_coords, rec_coords, wavelet, save=grad_corr,
space_order=space_order, ws=ws)
ydat = recin[:] - y.data[:]
else:
ydat = yin
# Compute wavefield vy = adjoint(F(m0))*y and norm on the fly
srca, v, norm_v, _ = adjoint(model, ydat, src_coords, rec_coords,
norm_v=True, w_fun=w_fun,
save=grad is not None)
c1 = 1 / (recin.shape[1])
c2 = np.log(np.prod(model.shape))
# <PTy, d-F(m)*f> = <PTy, d>-<adjoint(F(m))*PTy, f>
ndt = np.sqrt(model.critical_dt)
PTy_dot_r = ndt**2 * (np.dot(ydat.reshape(-1), recin.reshape(-1)) -
np.dot(srca.data.reshape(-1), wavelet.reshape(-1)))
norm_y = ndt * np.linalg.norm(ydat)
# alpha
α = compute_optalpha(c2*norm_y, c1*norm_v, eps, comp_alpha=alpha_op)
# Lagrangian evaluation
fun = -.5 * c1 * α**2 * norm_v + c2 * α * PTy_dot_r - eps * np.abs(α) * norm_y
gradm = grady = None
if grad is not None:
w = weight_fun(w_fun, model, src_coords)
w = c1*α/w**2 if w is not None else c1*α
Q = wf_as_src(v, w=w)
rcv, gradm, _ = forward_grad(model, src_coords, rec_coords, c2*wavelet, q=Q, v=v)
# Compute gradient wrt y
if grad_corr or grad in ["all", "y"]:
grady = c2 * recin - rcv.data[:]
if norm_y != 0:
grady -= np.abs(eps) * ydat / norm_y
# Correcting for reduced gradient
if not grad_corr:
gradm = gradm.data
else:
gradm_corr, _ = gradient(model, grady, rec_coords, u0)
# Reduced gradient post-processing
gradm = gradm.data + gradm_corr.data
return fun, α * gradm, grady
def forward_grad(model,
src_coords,
rcv_coords,
wavelet,
v,
space_order=8,
q=None,
ws=None,
isic=False,
w=None,
freq=None,
**kwargs):
"""
Low level propagator, to be used through `interface.py`
Compute forward wavefield u = A(m)^{-1}*f and related quantities (u(xrcv))
"""
# Number of time steps
nt = as_tuple(q)[0].shape[0] if wavelet is None else wavelet.shape[0]
# Setting forward wavefield
u = wavefield(model, space_order, save=False)
# Add extended source
q = q or wf_as_src(u, w=0)
q = extented_src(model, ws, wavelet, q=q)
# Set up PDE expression and rearrange
pde = wave_kernel(model, u, q=q)
# Setup source and receiver
geom_expr, _, rcv = src_rec(model,
u,
src_coords=src_coords,
nt=nt,
rec_coords=rcv_coords,
wavelet=wavelet)
# Setup gradient wrt m
gradm = Function(name="gradm", grid=model.grid)
g_expr = grad_expr(gradm, v, u, model, w=w, isic=isic, freq=freq)
# Create operator and run
subs = model.spacing_map
op = Operator(pde + geom_expr + g_expr,
subs=subs,
name="forward_grad" + name(model),
opt=opt_op(model))
summary = op()
# Output
return rcv, gradm, summary
def born(model,
src_coords,
rcv_coords,
wavelet,
space_order=8,
save=False,
q=None,
isic=False,
ws=None,
t_sub=1):
"""
Low level propagator, to be used through `interface.py`
Compute adjoint wavefield v = adjoint(F(m))*y
and related quantities (||v||_w, v(xsrc))
"""
nt = wavelet.shape[0]
# Setting wavefield
u = wavefield(model, space_order, save=save, nt=nt, t_sub=t_sub)
ul = wavefield(model, space_order, name="l")
# Expression for saving wavefield if time subsampling is used
u_save, eq_save = wavefield_subsampled(model, u, nt, t_sub)
# Extended source
q = q or wf_as_src(u, w=0)
q = extented_src(model, ws, wavelet, q=q)
# Set up PDE expression and rearrange
pde, tmpu = wave_kernel(model, u, q=q)
pdel, tmpul = wave_kernel(model, ul, q=lin_src(model, u, isic=isic))
# Setup source and receiver
geom_expr, _, _ = src_rec(model, u, src_coords=src_coords, wavelet=wavelet)
geom_exprl, _, rcvl = src_rec(model,
ul,
rec_coords=rcv_coords,
nt=wavelet.shape[0])
# Create operator and run
subs = model.spacing_map
op = Operator(tmpu + tmpul + pde + geom_expr + geom_exprl + pdel + eq_save,
subs=subs,
name="born" + name(model),
opt=opt_op(model, no_ms=ws is not None))
summary = op()
# Output
return rcvl, (u_save if t_sub > 1 else u), summary
def objTWRIdual_devito(model, y, src_coords, rcv_coords, wav,
dat, Filter, eps, mode="eval", objfact=np.float32(1),
comp_alpha=True, grad_corr=False, weight_fun_pars=None,
dt=None, space_order=8):
"""
Evaluate TWRI objective functional/gradients for current (m, y)
"""
# Setting time sampling
if dt is None:
dt = model.critical_dt
# Computing y in reduced mode (= residual) if not provided
u0 = None
y_was_None = y is None
if y_was_None:
u0rcv, u0 = forward(model, src_coords, rcv_coords, wav, dt=dt,
space_order=space_order, save=(mode == "grad") and grad_corr)
y = applyfilt(dat-u0rcv, Filter)
PTy = applyfilt_transp(y, Filter)
else:
PTy = y
# Normalization constants
nx = np.float32(model.vp.size)
nt, nr = np.float32(y.shape)
etaf = npla.norm(wav.reshape(-1)) / np.sqrt(nt * nx)
etad = npla.norm(applyfilt(dat, Filter).reshape(-1)) / np.sqrt(nt * nr)
# Compute wavefield vy = adjoint(F(m))*Py
norm_vPTy2, vPTy_src, vPTy = adjoint_y(model, PTy, src_coords, rcv_coords,
weight_fun_pars=weight_fun_pars, dt=dt,
space_order=space_order, save=(mode == "grad"))
# <PTy, d-F(m)*f> = <PTy, d>-<adjoint(F(m))*PTy, f>
PTy_dot_r = (np.dot(PTy.reshape(-1), dat.reshape(-1)) -
np.dot(vPTy_src.reshape(-1), wav.reshape(-1)))
# ||y||
norm_y = npla.norm(y.reshape(-1))
# Optimal alpha
c1 = etaf**np.float32(2) / (np.float32(4) * etad**np.float32(2) * nx * nt)
c2 = np.float32(1) / (etad * nr * nt)
c3 = eps / np.sqrt(nr * nt)
alpha = compute_optalpha(c1*norm_vPTy2, c2*PTy_dot_r, c3*norm_y,
comp_alpha=comp_alpha)
# Lagrangian evaluation
fun = (alpha * (-alpha * c1 * norm_vPTy2 + c2 * PTy_dot_r) -
np.abs(alpha) * c3 * norm_y)
# Gradient computation
if mode == "grad":
# Set up extebded source
w = 2.0 * c1 / c2 * alpha
if weight_fun_pars is not None:
w /= weight_fun(weight_fun_pars, model, src_coords)**2
Q = wf_as_src(vPTy, w=w)
# Setup gradient wrt m
u = wavefield(model, space_order)
gradm = Function(name="gradm", grid=model.grid)
g_exp = grad_expr(gradm, u, vPTy, w=alpha * c2)
rcv, _ = forward(model, src_coords, rcv_coords, wav, dt=dt,
space_order=space_order, q=Q, extra_expr=g_exp, u=u)
# Compute gradient wrt y
if not y_was_None or grad_corr:
norm_y = npla.norm(y)
grady_data = alpha * c2 * applyfilt(dat - rcv.data, Filter)
if norm_y != 0:
grady_data -= np.abs(alpha) * c3 * y / norm_y
# Correcting for reduced gradient
if not y_was_None or (y_was_None and not grad_corr):
gradm_data = gradm.data
else:
gradm_corr = gradient(model, applyfilt_transp(grady_data, Filter), rcv_coords,
u0, dt=dt, space_order=space_order, w=1)
# Reduced gradient post-processing
gradm_data = gradm.data + gradm_corr.data
# Return output
if mode == "eval":
return fun / objfact
elif mode == "grad" and y_was_None:
return fun / objfact, gradm_data / objfact
elif mode == "grad" and not y_was_None:
return fun / objfact, gradm_data / objfact, grady_data / objfact
def born(model,
src_coords,
rcv_coords,
wavelet,
space_order=8,
save=False,
q=None,
return_op=False,
isic=False,
freq_list=None,
dft_sub=None,
ws=None,
t_sub=1,
nlind=False):
"""
Low level propagator, to be used through `interface.py`
Compute linearized wavefield U = J(m)* δ m
and related quantities.
"""
nt = wavelet.shape[0]
# Setting wavefield
u = wavefield(model, space_order, save=save, nt=nt, t_sub=t_sub)
ul = wavefield(model, space_order, name="l")
# Expression for saving wavefield if time subsampling is used
u_save, eq_save = wavefield_subsampled(model, u, nt, t_sub)
# Extended source
q = q or wf_as_src(u, w=0)
q = extented_src(model, ws, wavelet, q=q)
# Set up PDE expression and rearrange
pde = wave_kernel(model, u, q=q)
if model.dm == 0:
pdel = []
else:
pdel = wave_kernel(model, ul, q=lin_src(model, u, isic=isic))
# Setup source and receiver
geom_expr, _, rcvnl = src_rec(model,
u,
rec_coords=rcv_coords if nlind else None,
src_coords=src_coords,
wavelet=wavelet)
geom_exprl, _, rcvl = src_rec(model, ul, rec_coords=rcv_coords, nt=nt)
# On-the-fly Fourier
dft, dft_modes = otf_dft(u, freq_list, model.critical_dt, factor=dft_sub)
# Create operator and run
subs = model.spacing_map
op = Operator(pde + geom_expr + geom_exprl + pdel + dft + eq_save,
subs=subs,
name="born" + name(model),
opt=opt_op(model))
op.cfunction
outrec = (rcvl, rcvnl) if nlind else rcvl
if return_op:
return op, u, outrec
summary = op()
# Output
return outrec, dft_modes or (u_save if t_sub > 1 else u), summary
def twri_fun(self,
y,
src_coords,
rcv_coords,
wav,
dat,
Filter,
eps,
mode="eval",
objfact=np.float32(1),
comp_alpha=True,
grad_corr=False,
weight_fun_pars=None):
"""
Evaluate TWRI objective functional/gradients for current (m, y)
"""
clear_cache()
dt = self.model.critical_dt
# Computing y in reduced mode (= residual) if not provided
u0 = None
y_was_None = y is None
if y_was_None:
u0rcv, u0 = self.forward_run(wav,
src_coords,
rcv_coords,
save=(mode == "grad") and grad_corr)
y = applyfilt(dat - u0rcv, Filter)
PTy = applyfilt_transp(y, Filter)
else:
PTy = y
# Normalization constants
nx = np.float32(self.model.vp.size)
nt, nr = np.float32(y.shape)
etaf = npla.norm(wav.reshape(-1)) / np.sqrt((nt * dt) * nx)
etad = npla.norm(applyfilt(dat, Filter).reshape(-1)) / np.sqrt(
(nt * dt) * nr)
# Compute wavefield vy = adjoint(F(m))*Py
norm_vPTy2, vPTy_src, vPTy = self.adjoint_y_run(
PTy,
src_coords,
rcv_coords,
weight_fun_pars=weight_fun_pars,
save=(mode == "grad"))
# <PTy, d-F(m)*f> = <PTy, d>-<adjoint(F(m))*PTy, f>
PTy_dot_r = (np.dot(PTy.reshape(-1), dat.reshape(-1)) -
np.dot(vPTy_src.reshape(-1), wav.reshape(-1)))
# ||y||
norm_y = np.sqrt(dt) * npla.norm(y.reshape(-1))
# Optimal alpha
c1 = etaf**np.float32(2) / (np.float32(4) * etad**np.float32(2) * nx *
(nt * dt))
c2 = np.float32(1) / (etad * nr * (nt * dt))
c3 = eps / np.sqrt(nr * (nt * dt))
alpha = compute_optalpha(c1 * norm_vPTy2,
c2 * PTy_dot_r,
c3 * norm_y,
comp_alpha=comp_alpha)
# Lagrangian evaluation
fun = (alpha * (-alpha * c1 * norm_vPTy2 + c2 * PTy_dot_r) -
np.abs(alpha) * c3 * norm_y)
# Gradient computation
if mode == "grad":
# Set up extebded source
w = 2.0 * c1 / c2 * alpha
if weight_fun_pars is not None:
w /= weight_fun(weight_fun_pars, self.model, src_coords)**2
Q = wf_as_src(vPTy, w=w)
# Setup gradient wrt m
rcv, _, gradm = self.forward_run(wav,
src_coords,
rcv_coords,
q=Q,
grad=True,
w=alpha * c2,
v=vPTy)
# Compute gradient wrt y
if not y_was_None or grad_corr:
norm_y = npla.norm(y)
grady_data = alpha * c2 * applyfilt(dat - rcv.data, Filter)
if norm_y != 0:
grady_data -= np.abs(alpha) * c3 * y / norm_y
# Correcting for reduced gradient
if not y_was_None or (y_was_None and not grad_corr):
gradm_data = gradm.data
else:
gradm_corr = self.grad_run(
rcv_coords, applyfilt_transp(grady_data, Filter), u0)
# Reduced gradient post-processing
gradm_data = gradm.data + gradm_corr.data
# Return output
if mode == "eval":
return fun / objfact
elif mode == "grad" and y_was_None:
return fun / objfact, -gradm_data / objfact
elif mode == "grad" and not y_was_None:
return fun / objfact, -gradm_data / objfact, grady_data / objfact
