def adjoint_born(model,
rec_coords,
rec_data,
u=None,
op_forward=None,
is_residual=False,
space_order=8,
nb=40,
isic=False,
dt=None,
n_checkpoints=None,
maxmem=None):
clear_cache()
# Parameters
nt = rec_data.shape[0]
if dt is None:
dt = model.critical_dt
m, rho, damp = model.m, model.rho, model.damp
# Create adjoint wavefield and gradient
v = TimeFunction(name='v',
grid=model.grid,
time_order=2,
space_order=space_order)
gradient = Function(name='gradient', grid=model.grid)
# Set up PDE and rearrange
vlaplace, rho = acoustic_laplacian(v, rho)
H = symbols('H')
eqn = m / rho * v.dt2 - H - damp * v.dt
stencil = solve(eqn, v.backward, simplify=False, rational=False)[0]
expression = [Eq(v.backward, stencil.subs({H: vlaplace}))]
# Data at receiver locations as adjoint source
rec_g = Receiver(name='rec_g',
grid=model.grid,
ntime=nt,
coordinates=rec_coords)
if op_forward is None:
rec_g.data[:] = rec_data[:]
adj_src = rec_g.inject(field=v.backward,
offset=model.nbpml,
expr=rec_g * rho * dt**2 / m)
# Gradient update
if u is None:
u = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order)
if isic is not True:
gradient_update = [Eq(gradient, gradient - dt * u.dt2 / rho * v)]
else:
# sum u.dx * v.dx fo x in dimensions.
# space_order//2
diff_u_v = sum([
first_derivative(u, dim=d, order=space_order // 2) *
first_derivative(v, dim=d, order=space_order // 2)
for d in u.space_dimensions
])
gradient_update = [
Eq(gradient, gradient - dt * (u * v.dt2 * m + diff_u_v) / rho)
]
# Create operator and run
set_log_level('ERROR')
expression += adj_src + gradient_update
subs = model.spacing_map
subs[u.grid.time_dim.spacing] = dt
op = Operator(expression,
subs=subs,
dse='advanced',
dle='advanced',
name="Gradient%s" % randint(1e5))
# Optimal checkpointing
if op_forward is not None:
rec = Receiver(name='rec',
grid=model.grid,
ntime=nt,
coordinates=rec_coords)
cp = DevitoCheckpoint([u])
if maxmem is not None:
n_checkpoints = int(
np.floor(maxmem * 10**6 / (cp.size * u.data.itemsize)))
wrap_fw = CheckpointOperator(op_forward, u=u, m=model.m, rec=rec)
wrap_rev = CheckpointOperator(op, u=u, v=v, m=model.m, rec_g=rec_g)
# 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_g.data[:] = rec_data[:]
else:
rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data
fval = .5 * np.dot(rec_g.data[:].flatten(),
rec_g.data[:].flatten()) * dt
wrp.apply_reverse()
else:
op()
clear_cache()
if op_forward is not None and is_residual is not True:
return fval, gradient.data
else:
return gradient.data
def adjoint_born(model,
rec_coords,
rec_data,
u=None,
op_forward=None,
is_residual=False,
space_order=8,
nb=40,
isic=False,
dt=None):
clear_cache()
# Parameters
nt = rec_data.shape[0]
if dt is None:
dt = model.critical_dt
m, damp = model.m, model.damp
# Create adjoint wavefield and gradient
v = TimeFunction(name='v',
grid=model.grid,
time_order=2,
space_order=space_order)
gradient = Function(name='gradient', grid=model.grid)
# Set up PDE and rearrange
eqn = m * v.dt2 - v.laplace - damp * v.dt
stencil = solve(eqn, v.backward)[0]
expression = [Eq(v.backward, stencil)]
# Data at receiver locations as adjoint source
rec_g = Receiver(name='rec_g',
grid=model.grid,
ntime=nt,
coordinates=rec_coords)
if op_forward is None:
rec_g.data[:] = rec_data[:]
adj_src = rec_g.inject(field=v.backward,
offset=model.nbpml,
expr=rec_g * dt**2 / m)
# Gradient update
if u is None:
u = TimeFunction(name='u',
grid=model.grid,
time_order=2,
space_order=space_order)
if isic is not True:
gradient_update = [Eq(gradient, gradient - u * v.dt2)
] # zero-lag cross-correlation imaging condition
else:
# linearized inverse scattering imaging condition (Op't Root et al. 2010; Whitmore and Crawley 2012)
if len(model.shape) == 2:
gradient_update = [
Eq(gradient,
gradient - (u * v.dt2 * m + u.dx * v.dx + u.dy * v.dy))
]
else:
gradient_update = [
Eq(
gradient, gradient -
(u * v.dt2 * m + u.dx * v.dx + u.dy * v.dy + u.dz * v.dz))
]
# Create operator and run
set_log_level('ERROR')
expression += adj_src + gradient_update
op = Operator(expression,
subs=model.spacing_map,
dse='advanced',
dle='advanced',
name="Gradient%s" % randint(1e5))
# Optimal checkpointing
if op_forward is not None:
rec = Receiver(name='rec',
grid=model.grid,
ntime=nt,
coordinates=rec_coords)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(op_forward,
u=u,
m=model.m.data,
rec=rec,
dt=dt)
wrap_rev = CheckpointOperator(op,
u=u,
v=v,
m=model.m.data,
rec_g=rec_g,
dt=dt)
# 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_g.data[:] = rec_data[:]
else:
rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data
fval = .5 * np.linalg.norm(rec_g.data[:])**2
wrp.apply_reverse()
else:
op(dt=dt)
clear_cache()
if op_forward is not None and is_residual is not True:
return fval, gradient.data
else:
return gradient.data
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
def adjoint_born(model, rec_coords, rec_data, u=None, op_forward=None, is_residual=False,
space_order=8, isic=False, dt=None, n_checkpoints=None, maxmem=None,
free_surface=False, tsub_factor=1, checkpointing=False):
clear_cache()
# Parameters
nt = rec_data.shape[0]
if dt is None:
dt = model.critical_dt
m, rho, damp = model.m, model.rho, model.damp
# Create adjoint wavefield and gradient
v = TimeFunction(name='v', grid=model.grid, time_order=2, space_order=space_order)
gradient = Function(name='gradient', grid=model.grid)
# Set up PDE and rearrange
vlaplace, rho = acoustic_laplacian(v, rho)
stencil = damp * (2.0 * v - damp * v.forward + dt**2 * rho / m * vlaplace)
expression = [Eq(v.backward, stencil)]
# Data at receiver locations as adjoint source
rec_g = Receiver(name='rec_g', grid=model.grid, ntime=nt, coordinates=rec_coords)
if op_forward is None:
rec_g.data[:] = rec_data[:]
adj_src = rec_g.inject(field=v.backward, expr=rec_g * rho * dt**2 / m)
# Gradient update
if u is None:
u = TimeFunction(name='u', grid=model.grid, time_order=2, space_order=space_order)
if isic is not True:
gradient_update = [Inc(gradient, - dt * u.dt2 / rho * v)]
else:
# sum u.dx * v.dx fo x in dimensions.
# space_order//2
diff_u_v = sum([first_derivative(u, dim=d, fd_order=space_order//2)*
first_derivative(v, dim=d, fd_order=space_order//2)
for d in u.space_dimensions])
gradient_update = [Inc(gradient, - tsub_factor * dt * (u * v.dt2 * m + diff_u_v) / rho)]
# Create operator and run
# Free surface
if free_surface is True:
expression += freesurface(v, space_order//2, model.nbpml, forward=False)
expression += adj_src + gradient_update
subs = model.spacing_map
subs[u.grid.time_dim.spacing] = dt
op = Operator(expression, subs=subs, dse='advanced', dle='advanced')
# Optimal checkpointing
summary1 = None
summary2 = None
if op_forward is not None and checkpointing is True:
rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords)
cp = DevitoCheckpoint([u])
if maxmem is not None:
n_checkpoints = int(np.floor(maxmem * 10**6 / (cp.size * u.data.itemsize)))
wrap_fw = CheckpointOperator(op_forward, u=u, m=model.m, rec=rec)
wrap_rev = CheckpointOperator(op, u=u, v=v, m=model.m, rec_g=rec_g)
# 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_g.data[:] = rec_data[:]
else:
rec_g.data[:] = rec.data[:] - rec_data[:] # input is observed data
fval = .5*np.dot(rec_g.data[:].flatten(), rec_g.data[:].flatten()) * dt
wrp.apply_reverse()
elif op_forward is not None and checkpointing is False:
# Compile first
cf1 = op_forward.cfunction
cf2 = op.cfunction
# Run forward and adjoint
summary1 = op_forward.apply()
if is_residual is True:
rec_g.data[:] = rec_data[:]
else:
rec = Receiver(name='rec', grid=model.grid, ntime=nt, coordinates=rec_coords)
rec_g.data[:] = rec.data[:] - rec_data[:]
fval = .5*np.dot(rec_g.data[:].flatten(), rec_g.data[:].flatten()) * dt
summary2 = op.apply()
else:
cf = op.cfunction
summary1 = op.apply()
clear_cache()
if op_forward is not None and is_residual is not True:
if summary2 is not None:
return fval, gradient.data, summary1, summary2
elif summary1 is not None:
return fval, gradient.data, summary1
else:
return fval, gradient.data
else:
if summary2 is not None:
return gradient.data, summary1, summary2
elif summary1 is not None:
return gradient.data, summary1
else:
return gradient.data