def forward_rec_w(model, weight, wavelet, rec_coords, space_order=8):
"""
Forward modeling of an extended source with receivers Pr*F*Pw^T*w
Parameters
----------
model: Model
Physical model
weights: Array
Spatial distribution of the extended source.
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Shot record
"""
rec, _, _ = forward(model,
None,
rec_coords,
wavelet,
save=False,
ws=weight,
space_order=space_order)
return rec.data
def forward_rec(model, src_coords, wavelet, rec_coords, space_order=8):
"""
Forward modeling of a point source with receivers Pr*F*Ps^T*q.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Shot record
"""
rec, _, _ = forward(model,
src_coords,
rec_coords,
wavelet,
save=False,
space_order=space_order)
return rec.data
def forward_wf_src_norec(model, u, space_order=8):
"""
Forward modeling of a full wavefield source without receiver F*u.
Parameters
----------
model: Model
Physical model
u: TimeFunction or Array
Time-space dependent wavefield
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Wavefield
"""
wf_src = TimeFunction(name='wf_src',
grid=model.grid,
time_order=2,
space_order=space_order,
save=u.shape[0])
if isinstance(u, TimeFunction):
wf_src._data = u._data
else:
wf_src.data[:] = u[:]
_, u, _ = forward(model,
None,
None,
None,
space_order=space_order,
save=True,
q=wf_src)
return u.data
def forward_wf_src(model, u, rec_coords, space_order=8, free_surface=False):
"""
Forward modeling of a full wavefield source Pr*F*u.
Parameters
----------
model: Model
Physical model
u: TimeFunction or Array
Time-space dependent wavefield
rec_coords: Array
Coordiantes of the receiver(s)
space_order: Int (optional)
Spatial discretization order, defaults to 8
free_surface: Bool (optional)
Whether or not to use a free surface
Returns
----------
Array
Shot record
"""
wf_src = TimeFunction(name='wf_src', grid=model.grid, time_order=2,
space_order=space_order, save=u.shape[0])
if isinstance(u, TimeFunction):
wf_src._data = u._data
else:
wf_src.data[:] = u[:]
rec, _ = forward(model, None, rec_coords, None, space_order=space_order,
free_surface=free_surface, q=wf_src)
return rec.data
def forward_no_rec(model, src_coords, wavelet, space_order=8):
"""
Forward modeling of a point source without receiver.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
space_order: Int (optional)
Spatial discretization order, defaults to 8
Returns
----------
Array
Wavefield
"""
_, u, _ = forward(model,
src_coords,
None,
wavelet,
space_order=space_order,
save=True)
return u.data
def forward_rec_wf(model, src_coords, wavelet, rec_coords, t_sub=1,
space_order=8, free_surface=False):
"""
Forward modeling of a point source Pr*F*Ps^T*q and return wavefield.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
space_order: Int (optional)
Spatial discretization order, defaults to 8
free_surface: Bool (optional)
Whether or not to use a free surface
Returns
----------
Array
Shot record
TimeFunction
Wavefield
"""
rec, u = forward(model, src_coords, rec_coords, wavelet, save=True, t_sub=t_sub,
space_order=space_order, free_surface=free_surface)
return rec.data, u
def forward_no_rec(model, src_coords, wavelet, space_order=8, free_surface=False):
"""
Forward modeling of a point source without receiver.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
space_order: Int (optional)
Spatial discretization order, defaults to 8
free_surface: Bool (optional)
Whether or not to use a free surface
Returns
----------
Array
Wavefield
"""
_, u = forward(model, src_coords, None, wavelet, space_order=space_order,
save=True, free_surface=free_surface)
return u.data
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 J_adjoint_freq(model, src_coords, wavelet, rec_coords, recin, space_order=8,
free_surface=False, freq_list=[], is_residual=False, return_obj=False,
dft_sub=None, isic=False, ws=None, t_sub=1):
"""
Jacobian (adjoint fo born modeling operator) operator on a shot record
as a source (i.e data residual). Outputs the gradient with Frequency
compression (on-the-fly DFT).
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
recin: Array
Receiver data
space_order: Int (optional)
Spatial discretization order, defaults to 8
free_surface: Bool (optional)
Whether or not to use a free surface
freq_list: List
List of frequencies for on-the-fly DFT
dft_sub: Int
Subsampling factor for on-the-fly DFT
isic : Bool
Whether or not to use ISIC imaging condition
ws : Array
Extended source spatial distribution
is_residual: Bool
Whether to treat the input as the residual or as the observed data
Returns
----------
Array
Adjoint jacobian on the input data (gradient)
"""
rec, u = forward(model, src_coords, rec_coords, wavelet, save=False,
space_order=space_order, free_surface=free_surface,
freq_list=freq_list, dft_sub=dft_sub, ws=ws)
# Residual and gradient
if not is_residual:
recin[:] = rec.data[:] - recin[:] # input is observed data
g = gradient(model, recin, rec_coords, u, space_order=space_order, isic=isic,
free_surface=free_surface, freq=freq_list, dft_sub=dft_sub)
if return_obj:
return .5*model.critical_dt*np.linalg.norm(recin)**2, g.data
return g.data
def J_adjoint_standard(model, src_coords, wavelet, rec_coords, recin, space_order=8,
free_surface=False, is_residual=False, return_obj=False,
isic=False, ws=None, t_sub=1):
"""
Adjoint Jacobian (adjoint fo born modeling operator) operator on a shot record
as a source (i.e data residual). Outputs the gradient with standard
zero lag cross correlation over time.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
recin: Array
Receiver data
space_order: Int (optional)
Spatial discretization order, defaults to 8
free_surface: Bool (optional)
Whether or not to use a free surface
isic : Bool
Whether or not to use ISIC imaging condition
ws : Array
Extended source spatial distribution
is_residual: Bool
Whether to treat the input as the residual or as the observed data
Returns
----------
Array
Adjoint jacobian on the input data (gradient)
"""
rec, u = forward(model, src_coords, rec_coords, wavelet, save=True, ws=ws,
space_order=space_order, free_surface=free_surface, t_sub=t_sub)
# Residual and gradient
if not is_residual:
recin[:] = rec.data[:] - recin[:] # input is observed data
g = gradient(model, recin, rec_coords, u, space_order=space_order,
free_surface=free_surface, isic=isic)
if return_obj:
return .5*model.critical_dt*np.linalg.norm(recin)**2, g.data
return g.data
size = comm.size
#########################################################################################
# Source wavelet
tn = 2000.
dt_shot = model0.critical_dt
nt = int(tn/dt_shot)
time_s = np.linspace(0, tn, nt)
wavelet = Ricker(0.015, time_s)
#########################################################################################
# Devito operator
# d_obs = forward(model0, src_coords, rec_coords, wavelet, save=False, t_sub=1)[0]
# #d_lin = born(model0, src_coords, rec_coords, wavelet, isic=True)[0]
# # Gradient
# u0 = forward(model0, src_coords, rec_coords, wavelet, save=True, t_sub=4)[1]
# g = gradient(model0, d_obs, d_obs.coordinates, u0, isic=True)
# Devito operator
d_obs = forward(model0, src_coords, rec_coords, wavelet, save=False, t_sub=4)[0]
d_gather, rec_gather = collect_shot(comm, d_obs)
if rank == 0:
segy_write(d_gather, src_coords[:, 0], src_coords[:, 1], rec_coords[:, 0], rec_coords[:, 1], dt_shot, 'shot_n_2.segy')
new_time_range = TimeAxis(start=start, stop=stop, num=num)
dt = new_time_range.step
to_interp = np.asarray(rec.data)
data = np.zeros((num, to_interp.shape[1]))
for i in range(to_interp.shape[1]):
tck = interpolate.splrep(time, to_interp[:, i], k=3)
data[:, i] = interpolate.splev(new_time_range.time_values, tck)
coords_loc = np.asarray(rec.coordinates.data)
# Return new object
return data, coords_loc
# Devito operator
d_obs, u0, summary1 = forward(model,
src.coordinates.data,
rec_coords,
src.data,
save=True,
t_sub=12)
grad, summary2 = gradient(model, d_obs, rec_coords, u0, isic=True)
grad.data[:, 0:66] = 0 # mute water column
# Remove pml and pad
rtm = grad.data[model.nbl:-model.nbl, model.nbl:-model.nbl] # remove padding
rtm = extent_gradient(shape_full, origin_full, shape, origin, spacing, rtm)
plt.figure()
plt.imshow(d_obs.data, vmin=-1e-1, vmax=1e-1, cmap='gray', aspect='auto')
plt.figure()
plt.imshow(np.transpose(rtm), vmin=-2e0, vmax=2e0, cmap='gray', aspect='auto')
plt.show()
t0 = 0.
tn = 1300.
dt = model.critical_dt
nt = int(1 + (tn - t0) / dt)
time_axis = np.linspace(t0, tn, nt)
# Source
f1 = 0.008
src1 = RickerSource(name='src', grid=model.grid, f0=f1, time=time_axis)
src1.coordinates.data[0, :] = np.array(model.domain_size) * 0.5
src1.coordinates.data[0, -1] = 20.
# Receiver for observed data
rec_t = Receiver(name='rec_t', grid=model.grid, npoint=301, ntime=nt)
rec_t.coordinates.data[:, 0] = np.linspace(0., 3000., num=301)
rec_t.coordinates.data[:, 1] = 20.
# Test data and source
d_hat, u1, _ = forward(model, src1.coordinates.data, rec_t.coordinates.data,
src1.data)
# Adjoint
q0, _, _ = adjoint(model, d_hat, src1.coordinates.data, rec_t.coordinates.data)
# Adjoint test
a = inner(d_hat, d_hat)
b = inner(q0, src1)
print("Adjoint test F")
print("a = %2.2e, b = %2.2e, diff = %2.2e: " % (a, b, a - b))
print("Relative error: ", a / b - 1)
def J_adjoint_checkpointing(model,
src_coords,
wavelet,
rec_coords,
recin,
space_order=8,
is_residual=False,
n_checkpoints=None,
maxmem=None,
return_obj=False,
isic=False,
ws=None,
t_sub=1):
"""
Jacobian (adjoint fo born modeling operator) operator on a shot record
as a source (i.e data residual). Outputs the gradient with Checkpointing.
Parameters
----------
model: Model
Physical model
src_coords: Array
Coordiantes of the source(s)
wavelet: Array
Source signature
rec_coords: Array
Coordiantes of the receiver(s)
recin: Array
Receiver data
space_order: Int (optional)
Spatial discretization order, defaults to 8
checkpointing: Bool
Whether or not to use checkpointing
n_checkpoints: Int
Number of checkpoints for checkpointing
maxmem: Float
Maximum memory to use for checkpointing
isic : Bool
Whether or not to use ISIC imaging condition
ws : Array
Extended source spatial distribution
is_residual: Bool
Whether to treat the input as the residual or as the observed data
Returns
----------
Array
Adjoint jacobian on the input data (gradient)
"""
# Optimal checkpointing
op_f, u, rec_g = forward(model,
src_coords,
rec_coords,
wavelet,
space_order=space_order,
return_op=True,
ws=ws)
op, g, v = gradient(model,
recin,
rec_coords,
u,
space_order=space_order,
return_op=True,
isic=isic)
nt = wavelet.shape[0]
rec = Receiver(name='rec',
grid=model.grid,
ntime=nt,
coordinates=rec_coords)
cp = DevitoCheckpoint([uu for uu in as_tuple(u)])
if maxmem is not None:
memsize = (cp.size * u.data.itemsize)
n_checkpoints = int(np.floor(maxmem * 10**6 / memsize))
# Op arguments
uk = {uu.name: uu for uu in as_tuple(u)}
vk = {**uk, **{vv.name: vv for vv in as_tuple(v)}}
uk.update({'rcv%s' % as_tuple(u)[0].name: rec_g})
vk.update({'src%s' % as_tuple(v)[0].name: rec})
# Wrapped ops
wrap_fw = CheckpointOperator(op_f, vp=model.vp, **uk)
wrap_rev = CheckpointOperator(op, vp=model.vp, **vk)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt - 2)
wrp.apply_forward()
# Residual and gradient
if is_residual is True: # input data is already the residual
rec.data[:] = recin[:]
else:
rec.data[:] = rec.data[:] - recin[:] # input is observed data
wrp.apply_reverse()
if return_obj:
return .5 * model.critical_dt * norm(rec)**2, g.data
return g.data
rec_t.coordinates.data[:, 0] = np.linspace(0., 3000., num=301)
rec_t.coordinates.data[:, 1] = 20.
# Linearized data
print("Forward J")
dD_hat, u0l, _ = born(model,
src.coordinates.data,
rec_t.coordinates.data,
src.data,
save=True)
# Forward
print("Forward")
_, u0, _ = forward(model,
src.coordinates.data,
rec_t.coordinates.data,
src.data,
save=True)
# gradient
print("Adjoint J")
dm_hat, _ = gradient(model, dD_hat, rec_t.coordinates.data, u0)
# Adjoint test
a = model.critical_dt * inner(dD_hat, dD_hat)
b = inner(dm_hat, model.dm)
if is_tti:
c = np.linalg.norm(u0[0].data.flatten() - u0l[0].data.flatten(), np.inf)
else:
c = np.linalg.norm(u0.data.flatten() - u0l.data.flatten(), np.inf)
#########################################################################################
# Source wavelet
tn = 1000.
dt_shot = model.critical_dt
nt = int(tn / dt_shot)
time_s = np.linspace(0, tn, nt)
wavelet = Ricker(0.015, time_s)
#########################################################################################
# Devito operator
d_obs = born(model, src_coords, rec_coords, wavelet, save=False)[0]
u0 = forward(model, src_coords, rec_coords, wavelet, save=True, t_sub=8)[1]
grad_dist = gradient(model, d_obs, d_obs.coordinates, u0, isic=False)[0]
if rank > 0:
# Send result to master
comm.send(model.m.local_indices, dest=0, tag=10)
comm.send(grad_dist.data, dest=0, tag=11)
else: # Master
# Initialize full array
grad = np.empty(shape=model.m.shape_global, dtype='float32')
grad[model.m.local_indices] = grad_dist.data
# Collect gradients
for j in range(1, size):
local_indices = comm.recv(source=j, tag=10)
