def test_forward_with_breaks(shape, kernel, space_order):
""" Test running forward in one go and "with breaks"
and ensure they produce the same result
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
time_order = 2
nrec = shape[0]
solver = acoustic_setup(shape=shape, spacing=spacing, tn=tn,
space_order=space_order, kernel=kernel)
grid = solver.model.grid
rec = Receiver(name='rec', grid=grid, time_range=solver.geometry.time_axis,
npoint=nrec)
rec.coordinates.data[:, :] = solver.geometry.rec_positions
dt = solver.model.critical_dt
u = TimeFunction(name='u', grid=grid, time_order=2, space_order=space_order)
cp = DevitoCheckpoint([u])
wrap_fw = CheckpointOperator(solver.op_fwd(save=False), rec=rec,
src=solver.geometry.src, u=u, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, dt=dt, rec=rec)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, rec._time_range.num-time_order)
rec1, u1, summary = solver.forward()
wrp.apply_forward()
assert(np.allclose(u1.data, u.data))
assert(np.allclose(rec1.data, rec.data))
def test_forward_with_breaks(shape, time_order, space_order):
""" Test running forward in one go and "with breaks"
and ensure they produce the same result
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
example = CheckpointingExample(shape, spacing, tn, time_order, space_order)
m0, dm = example.initial_estimate()
cp = DevitoCheckpoint([example.forward_field])
wrap_fw = CheckpointOperator(example.forward_operator,
u=example.forward_field,
rec=example.rec,
m=m0,
src=example.src,
dt=example.dt)
wrap_rev = CheckpointOperator(example.gradient_operator,
u=example.forward_field,
v=example.adjoint_field,
m=m0,
rec=example.rec_g,
grad=example.grad,
dt=example.dt)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, example.nt - time_order)
example.forward_operator.apply(u=example.forward_field,
rec=example.rec,
m=m0,
src=example.src,
dt=example.dt)
u_temp = np.copy(example.forward_field.data)
rec_temp = np.copy(example.rec.data)
example.forward_field.data[:] = 0
wrp.apply_forward()
assert (np.allclose(u_temp, example.forward_field.data))
assert (np.allclose(rec_temp, example.rec.data))
def test_forward_with_breaks(shape, kernel, space_order):
""" Test running forward in one go and "with breaks"
and ensure they produce the same result
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
time_order = 2
nrec = shape[0]
solver = acoustic_setup(shape=shape, spacing=spacing, tn=tn,
space_order=space_order, kernel=kernel)
grid = solver.model.grid
rec = Receiver(name='rec', grid=grid, time_range=solver.geometry.time_axis,
npoint=nrec)
rec.coordinates.data[:, :] = solver.geometry.rec_positions
dt = solver.model.critical_dt
u = TimeFunction(name='u', grid=grid, time_order=2, space_order=space_order)
cp = DevitoCheckpoint([u])
wrap_fw = CheckpointOperator(solver.op_fwd(save=False), rec=rec,
src=solver.geometry.src, u=u, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, dt=dt, rec=rec)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, rec._time_range.num-time_order)
rec1, u1, summary = solver.forward()
wrp.apply_forward()
assert(np.allclose(u1.data, u.data))
assert(np.allclose(rec1.data, rec.data))
def test_number_of_saves_in_forward(nt, ncp):
cp = SimpleCheckpoint()
f = SimpleOperator()
b = SimpleOperator()
rev = Revolver(cp, f, b, ncp, nt)
assert(cp.save_counter == 0)
rev.apply_forward()
assert(cp.save_counter == min(ncp, nt-1))
def jacobian_adjoint(self, rec, u, src=None, v=None, grad=None, model=None,
checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u : TimeFunction
Full wavefield `u` (created with save=True).
v : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v', grid=self.model.grid,
time_order=2, space_order=self.space_order)
model = model or self.model
# Pick vp from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if checkpointing:
u = TimeFunction(name='u', grid=self.model.grid,
time_order=2, space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=src or self.geometry.src,
u=u, dt=dt, **kwargs)
wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v,
rec=rec, dt=dt, grad=grad, **kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, dt=dt,
**kwargs)
return grad, summary
def test_forward_nt(nt, ncp):
cp = SimpleCheckpoint()
f = SimpleOperator()
b = SimpleOperator()
rev = Revolver(cp, f, b, ncp, nt)
assert(f.counter == 0)
rev.apply_forward()
assert(f.counter == nt)
def gradient(self, rec, u, v=None, grad=None, m=None, checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
:param recin: Receiver data as a numpy array
:param u: Symbol for full wavefield `u` (created with save=True)
:param v: (Optional) Symbol to store the computed wavefield
:param grad: (Optional) Symbol to store the gradient field
:returns: Gradient field and performance summary
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v', grid=self.model.grid,
time_order=2, space_order=self.space_order)
# Pick m from model unless explicitly provided
m = m or self.model.m
if checkpointing:
u = TimeFunction(name='u', grid=self.model.grid,
time_order=2, space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False), src=self.geometry.src,
u=u, m=m, dt=dt)
wrap_rev = CheckpointOperator(self.op_grad(save=False), u=u, v=v,
m=m, rec=rec, dt=dt, grad=grad)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_grad().apply(rec=rec, grad=grad, v=v, u=u, m=m,
dt=dt, **kwargs)
return grad, summary
def gradient(self, m0, maxmem=None):
cp = DevitoCheckpoint([self.forward_field])
n_checkpoints = None
if maxmem is not None:
n_checkpoints = int(
floor(maxmem * 10**6 /
(cp.size * self.forward_field.data.itemsize)))
wrap_fw = CheckpointOperator(self.forward_operator,
u=self.forward_field,
rec=self.rec,
m=m0,
src=self.src,
dt=self.dt)
wrap_rev = CheckpointOperator(self.gradient_operator,
u=self.forward_field,
v=self.adjoint_field,
m=m0,
rec=self.rec_g,
grad=self.grad,
dt=self.dt)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
self.nt - self.time_order)
wrp.apply_forward()
self.rec_g.data[:] = self.rec.data[:] - self.rec_t.data[:]
wrp.apply_reverse()
# The result is in grad
return self.grad.data, self.rec.data
def test_ptr_loads_and_saves(nt, ncp):
cp = SimpleCheckpoint()
f = SimpleOperator()
b = SimpleOperator()
rev = Revolver(cp, f, b, ncp, nt)
rev.apply_forward()
rev.apply_reverse()
assert(cp.save_pointers == cp.load_pointers)
assert(len(cp.save_pointers) == min(ncp, nt - 1))
def test_num_loads_and_saves(nt, ncp):
cp = SimpleCheckpoint()
f = SimpleOperator()
b = SimpleOperator()
rev = Revolver(cp, f, b, ncp, nt)
rev.apply_forward()
assert(cp.load_counter == 0)
rev.apply_reverse()
assert(cp.load_counter >= cp.save_counter)
def revolve_factors(nt, ncp):
fwd = CounterOperator()
rev = CounterOperator()
cp = CounterCheckpoint()
revolver = Revolver(cp, fwd, rev, ncp,
nt) #, compression_params={'scheme': None})
revolver.apply_forward()
assert (fwd.counter == nt)
revolver.apply_reverse()
assert (rev.counter == nt)
return fwd.counter, cp.savecount, cp.loadcount
def verify(space_order=4, kernel='OT4', nbpml=40, filename='',
compression_params={}, **kwargs):
solver = acoustic_setup(shape=(10, 10), spacing=(10, 10), nbpml=10, tn=50,
space_order=space_order, kernel=kernel, **kwargs)
# solver = overthrust_setup(filename=filename, tn=50, nbpml=nbpml,
# space_order=space_order, kernel=kernel,
# **kwargs)
u = TimeFunction(name='u', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
cp = DevitoCheckpoint([u])
n_checkpoints = None
dt = solver.dt
v = TimeFunction(name='v', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
grad = Function(name='grad', grid=solver.model.grid)
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.geometry.src, u=u, rec=rec, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, v=v,
rec=rec, dt=dt, grad=grad)
nt = rec.data.shape[0] - 2
print("Verifying for %d timesteps" % nt)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
wrp.apply_forward()
wrp.apply_reverse()
print(wrp.profiler.timings)
with Timer([]) as tf:
rec2, u2, _ = solver.forward(save=True)
with Timer([]) as tr:
grad2, _ = solver.gradient(rec=rec2, u=u2)
error = grad.data - grad2.data
# to_hdf5(error, 'zfp_grad_errors.h5')
print("Error norm", np.linalg.norm(error))
# assert(np.allclose(grad.data, grad2.data))
print("Checkpointing implementation is numerically verified")
print("Verification took %d ms for forward and %d ms for reverse" % (tf.elapsed, tr.elapsed))
def test_complete(scheme):
nt = 100
ncp = 10
shape = (10, 10)
a = np.zeros(shape)
fwd = IncrementOperator(1, a)
rev = IncrementOperator(-1, a)
cp = YoCheckpoint(a)
compression_params = {'scheme': scheme}
revolver = Revolver(cp,
fwd,
rev,
ncp,
nt,
compression_params=compression_params)
revolver.apply_forward()
assert (np.all(np.isclose(a, np.zeros(shape) + nt)))
revolver.apply_reverse()
assert (np.all(np.isclose(a, np.zeros(shape))))
def calculate_lossy_gradient(i, solver, vp, grad, path_prefix, to=2, so=4,
n_checkpoints=1000, compression_params=None):
true_d, source_location = load_shot(i, path_prefix)
dt = solver.dt
# Update source location
solver.geometry.src_positions[0, :] = source_location[:]
if compression_params is None:
print("Using default compression params")
compression_params = {}
# Compute smooth data and full forward wavefield u0
u = TimeFunction(name='u', grid=solver.model.grid, time_order=to,
space_order=so, save=None)
v = TimeFunction(name='v', grid=solver.model.grid, time_order=to,
space_order=so)
residual = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
smooth_d = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
fwd_op = solver.op_fwd(save=False)
rev_op = solver.op_grad(save=False)
wrap_fw = CheckpointOperator(fwd_op, src=solver.geometry.src, u=u,
rec=smooth_d, vp=vp, dt=dt)
wrap_rev = CheckpointOperator(rev_op, vp=vp, u=u, v=v, rec=residual,
grad=grad, dt=dt)
cp = DevitoCheckpoint([u])
nt = smooth_d.data.shape[0] - 2
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
wrp.apply_forward()
# Compute gradient from data residual
residual.data[:] = smooth_d.data[:] - true_d[:]
wrp.apply_reverse()
def checkpointed_run(space_order=4, ncp=None, kernel='OT4', nbpml=40,
filename='', compression_params={}, tn=1000, **kwargs):
solver = overthrust_setup(filename=filename, tn=tn, nbpml=nbpml,
space_order=space_order, kernel=kernel, **kwargs)
u = TimeFunction(name='u', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec', grid=solver.model.grid,
time_range=solver.geometry.time_axis,
coordinates=solver.geometry.rec_positions)
cp = DevitoCheckpoint([u])
n_checkpoints = ncp
dt = solver.dt
v = TimeFunction(name='v', grid=solver.model.grid, time_order=2,
space_order=solver.space_order)
grad = Function(name='grad', grid=solver.model.grid)
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.geometry.src, u=u,
rec=rec, dt=dt)
wrap_rev = CheckpointOperator(solver.op_grad(save=False), u=u, v=v,
rec=rec, dt=dt, grad=grad)
fw_timings = []
rev_timings = []
nt = rec.data.shape[0] - 2
print("Running %d timesteps" % (nt))
print(compression_params)
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, nt,
compression_params=compression_params)
with Timer(fw_timings):
wrp.apply_forward()
with Timer(rev_timings):
wrp.apply_reverse()
return grad, wrp, fw_timings, rev_timings
def test_compression_is_used():
nt = 100
ncp = 10
shape = (10, 10)
a = np.zeros(shape)
fwd = IncrementOperator(1, a)
rev = IncrementOperator(-1, a)
cp = YoCheckpoint(a)
counters = [0, 0]
compressor, decompressor = compressors[None], decompressors[None]
def this_compressor(params, data):
counters[0] += 1
return compressor(params, data)
def this_decompressor(params, data):
counters[1] += 1
return decompressor(params, data)
compression_params = {
'scheme': 'custom',
'compressor': this_compressor,
'decompressor': this_decompressor
}
revolver = Revolver(cp,
fwd,
rev,
ncp,
nt,
compression_params=compression_params)
revolver.apply_forward()
assert (counters[0] >= ncp)
revolver.apply_reverse()
assert (counters[1] >= counters[0])
def jacobian_adjoint(self, rec, p, pa=None, grad=None, r=None, model=None,
checkpointing=False, **kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
p : TimeFunction
Full wavefield `p` (created with save=True).
pa : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
r : TimeFunction, optional
The computed attenuation memory variable.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
qp : Function, optional
The P-wave quality factor.
b : Function, optional
The time-constant inverse density.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
pa = pa or TimeFunction(name='pa', grid=self.model.grid,
time_order=self.time_order, space_order=self.space_order,
staggered=NODE)
model = model or self.model
# Pick vp and physical parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if checkpointing:
p = TimeFunction(name='p', grid=self.model.grid,
time_order=self.time_order, space_order=self.space_order,
staggered=NODE)
r = TimeFunction(name="r", grid=self.model.grid, time_order=self.time_order,
space_order=self.space_order, staggered=NODE)
cp = DevitoCheckpoint([p, r])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src, p=p, r=r, dt=dt, **kwargs)
wrap_rev = CheckpointOperator(self.op_grad(save=False), p=p, pa=pa, rec=rec,
dt=dt, grad=grad, **kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, rec.data.shape[0]-2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
# Memory variable:
r = r or TimeFunction(name="r", grid=self.model.grid,
time_order=self.time_order,
space_order=self.space_order, staggered=NODE)
summary = self.op_grad().apply(rec=rec, grad=grad, pa=pa, p=p, r=r, dt=dt,
**kwargs)
return grad, summary
def test_index_alignment(const):
""" A much simpler test meant to ensure that the forward and reverse indices are
correctly aligned (i.e. u * v , where u is the forward field and v the reverse field
corresponds to the correct timesteps in u and v). The forward operator does u = u + 1
which means that the field a will be equal to nt (0 -> 1 -> 2 -> 3), the number of
timesteps this operator is run for. The field at the last time step of the forward is
used to initialise the field v for the reverse pass. The reverse operator does
v = v - 1, which means that if the reverse operator is run for the same number of
timesteps as the forward operator, v should be 0 at the last time step
(3 -> 2 -> 1 -> 0). There is also a grad = grad + u * v accumulator in the reverse
operator. If the alignment is correct, u and v should have the same value at every
time step:
0 -> 1 -> 2 -> 3 u
0 <- 1 <- 2 <- 3 v
and hence grad = 0*0 + 1*1 + 2*2 + 3*3 = sum(n^2) where n -> [0, nt]
If the test fails, the resulting number can tell you how the fields are misaligned
"""
n = 4
grid = Grid(shape=(2, 2))
order_of_eqn = 1
modulo_factor = order_of_eqn + 1
nt = n - order_of_eqn
u = TimeFunction(name='u', grid=grid, save=n)
# Increment one in the forward pass 0 -> 1 -> 2 -> 3
fwd_op = Operator(Eq(u.forward, u + 1. * const))
# Invocation 1
fwd_op(time=nt - 1, constant=1)
last_time_step_v = nt % modulo_factor
# Last time step should be equal to the number of timesteps we ran
assert (np.allclose(u.data[nt, :, :], nt))
v = TimeFunction(name='v', grid=grid, save=None)
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Decrement one in the reverse pass 3 -> 2 -> 1 -> 0
adj_eqn = Eq(v, v.forward - 1. * const)
adj_op = Operator(adj_eqn)
# Invocation 2
adj_op(time=nt - 1, constant=1)
# Last time step should be back to 0
assert (np.allclose(v.data[0, :, :], 0))
# Reset v to run the backward again
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
prod = Function(name="prod", grid=grid)
# Multiply u and v and add them
# = 3*3 + 2*2 + 1*1 + 0*0
prod_eqn = Eq(prod, prod + u * v)
comb_op = Operator([adj_eqn, prod_eqn])
# Invocation 3
comb_op(time=nt - 1, constant=1)
final_value = sum([n**2 for n in range(nt)])
# Final value should be sum of squares of first nt natural numbers
assert (np.allclose(prod.data, final_value))
# Now reset to repeat all the above tests with checkpointing
prod.data[:] = 0
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Checkpointed version doesn't require to save u
u_nosave = TimeFunction(name='u_n', grid=grid)
# change equations to use new symbols
fwd_eqn_2 = Eq(u_nosave.forward, u_nosave + 1. * const)
fwd_op_2 = Operator(fwd_eqn_2)
cp = DevitoCheckpoint([u_nosave])
wrap_fw = CheckpointOperator(fwd_op_2, constant=1)
prod_eqn_2 = Eq(prod, prod + u_nosave * v)
comb_op_2 = Operator([adj_eqn, prod_eqn_2])
wrap_rev = CheckpointOperator(comb_op_2, constant=1)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, nt)
# Invocation 4
wrp.apply_forward()
assert (np.allclose(u_nosave.data[last_time_step_v, :, :], nt))
# Invocation 5
wrp.apply_reverse()
assert (np.allclose(v.data[0, :, :], 0))
assert (np.allclose(prod.data, final_value))
def fwi_gradient_checkpointed(vp_in,
model,
geometry,
nshots,
client,
solver_params,
n_checkpoints=1000,
compression_params=None):
# Create symbols to hold the gradient and residual
grad = Function(name="grad", grid=model.grid)
vp = Function(name="vp", grid=model.grid)
smooth_d = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
residual = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
objective = 0.
time_order = 2
vp_in = vec2mat(vp_in, model.shape)
assert (model.vp.shape == vp_in.shape)
vp.data[:] = vp_in[:]
filename = solver_params['filename']
solver = overthrust_solver_iso(filename, datakey="m0")
dt = solver.dt
nt = smooth_d.data.shape[0] - 2
u = TimeFunction(name='u',
grid=model.grid,
time_order=time_order,
space_order=4)
v = TimeFunction(name='v',
grid=model.grid,
time_order=time_order,
space_order=4)
fwd_op = solver.op_fwd(save=False)
rev_op = solver.op_grad(save=False)
cp = DevitoCheckpoint([u])
for i in range(nshots):
true_d, source_location = load_shot(i)
# Update source location
solver.geometry.src_positions[0, :] = source_location[:]
# Compute smooth data and full forward wavefield u0
u.data[:] = 0.
residual.data[:] = 0.
v.data[:] = 0.
smooth_d.data[:] = 0.
wrap_fw = CheckpointOperator(fwd_op,
src=solver.geometry.src,
u=u,
rec=smooth_d,
vp=vp,
dt=dt)
wrap_rev = CheckpointOperator(rev_op,
vp=vp,
u=u,
v=v,
rec=residual,
grad=grad,
dt=dt)
wrp = Revolver(cp,
wrap_fw,
wrap_rev,
n_checkpoints,
nt,
compression_params=compression_params)
wrp.apply_forward()
# Compute gradient from data residual and update objective function
residual.data[:] = smooth_d.data[:] - true_d[:]
objective += .5 * np.linalg.norm(residual.data.ravel())**2
wrp.apply_reverse()
print(wrp.profiler.summary())
return objective, -np.ravel(grad.data).astype(np.float64)
def jacobian_adjoint(self,
rec,
u0,
v0,
du=None,
dv=None,
dm=None,
model=None,
checkpointing=False,
kernel='centered',
**kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u0 : TimeFunction
The computed background wavefield.
v0 : TimeFunction, optional
The computed background wavefield.
du : Function or float
The computed perturbed wavefield.
dv : Function or float
The computed perturbed wavefield.
dm : Function, optional
Stores the gradient field.
model : Model, optional
Object containing the physical parameters.
vp : Function or float, optional
The time-constant velocity.
epsilon : Function or float, optional
The time-constant first Thomsen parameter.
delta : Function or float, optional
The time-constant second Thomsen parameter.
theta : Function or float, optional
The time-constant Dip angle (radians).
phi : Function or float, optional
The time-constant Azimuth angle (radians).
Returns
-------
Gradient field and performance summary.
"""
if kernel != 'centered':
raise ValueError(
'Only centered kernel is supported for the jacobian_adj')
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
dm = dm or Function(name='dm', grid=self.model.grid)
# Create the perturbation wavefields if not provided
du = du or TimeFunction(name='du',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
dv = dv or TimeFunction(name='dv',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
model = model or self.model
# Pick vp and Thomsen parameters from model unless explicitly provided
kwargs.update(model.physical_params(**kwargs))
if self.model.dim < 3:
kwargs.pop('phi', None)
if checkpointing:
u0 = TimeFunction(name='u0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
v0 = TimeFunction(name='v0',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
cp = DevitoCheckpoint([u0, v0])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src,
u=u0,
v=v0,
dt=dt,
**kwargs)
wrap_rev = CheckpointOperator(self.op_jacadj(save=False),
u0=u0,
v0=v0,
du=du,
dv=dv,
rec=rec,
dm=dm,
dt=dt,
**kwargs)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
rec.data.shape[0] - 2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
summary = self.op_jacadj().apply(rec=rec,
dm=dm,
u0=u0,
v0=v0,
du=du,
dv=dv,
dt=dt,
**kwargs)
return dm, summary
def run(space_order=4,
kernel='OT2',
nbpml=40,
autotune=False,
filename='',
chunk=1000000,
algo=None,
shuffle="SHUFFLE",
**kwargs):
solver = overthrust_setup(filename=filename,
nbpml=nbpml,
space_order=space_order,
kernel=kernel,
**kwargs)
m = solver.model.m
dt = solver.dt
u = TimeFunction(name='u',
grid=solver.model.grid,
time_order=2,
space_order=solver.space_order)
v = TimeFunction(name='v',
grid=solver.model.grid,
time_order=2,
space_order=solver.space_order)
rec = Receiver(name='rec',
grid=solver.model.grid,
time_range=solver.receiver.time_range,
coordinates=solver.receiver.coordinates.data)
grad = Function(name='grad', grid=solver.model.grid)
cp = CompressionCheckpoint([u])
n_checkpoints = 60
wrap_fw = CheckpointOperator(solver.op_fwd(save=False),
src=solver.source,
u=u,
m=m,
dt=solver.dt,
rec=rec)
wrap_rev = CheckpointOperator(solver.op_grad(save=False),
u=u,
v=v,
m=m,
rec=rec,
dt=dt,
grad=grad)
# Run forward
wrp = Revolver(cp,
wrap_fw,
wrap_rev,
n_checkpoints,
rec.data.shape[0] - 2,
compression='blosc',
compression_params={
CHUNK_SIZE: chunk,
CNAME: algo,
SHUFFLE: shuffle
})
info("Applying Forward")
solver.forward(time=100)
#raw_fw(dt=dt)
info("Again")
wrp.apply_forward()
print(np.linalg.norm(u.data))
print(np.linalg.norm(rec.data))
info("Applying Gradient")
summary = wrp.apply_reverse()
print("Gradient is: %d" % np.linalg.norm(grad.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
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 forward_modeling_single_shot(record, table, par_files):
"Serial modeling function"
worker = get_worker() # The worker on which this task is running
strng = '{} : {} =>'.format(worker.address, str(record).zfill(5))
filename = 'logfile_{}.txt'.format(str(record).zfill(5))
g = open(filename, 'w')
g.write("This will show up in the worker logs")
# Read velocity model
f = segyio.open(par_files[-1], iline=segyio.tracefield.TraceField.FieldRecord,
xline=segyio.tracefield.TraceField.CDP)
xl, il, t = f.xlines, f.ilines, f.samples
dz = t[1] - t[0]
dx = f.header[1][segyio.TraceField.SourceX]-f.header[0][segyio.TraceField.SourceX]
if len(il) != 1:
dims = (len(xl), len(il), len(f.samples))
else:
dims = (len(xl), len(f.samples))
vp = f.trace.raw[:].reshape(dims)
vp *= 1./1000 # convert to km/sec
epsilon = np.empty(dims)
delta = np.empty(dims)
theta = np.empty(dims)
params = [epsilon, delta, theta]
# Read Thomsem parameters
for segyfile, par in zip(par_files, params):
f = segyio.open(segyfile, iline=segyio.tracefield.TraceField.FieldRecord,
xline=segyio.tracefield.TraceField.CDP)
par[:] = f.trace.raw[:].reshape(dims)
theta *= (np.pi/180.) # use radians
g.write('{} Parameter model dims: {}\n'.format(strng, vp.shape))
origin = (0., 0.)
shape = vp.shape
spacing = (dz, dz)
# Get a single shot as a numpy array
filename = table[record]['filename']
position = table[record]['Trace_Position']
traces_in_shot = table[record]['Num_Traces']
src_coord = np.array(table[record]['Source']).reshape((1, len(dims)))
rec_coord = np.array(table[record]['Receivers'])
start = time.time()
f = segyio.open(filename, ignore_geometry=True)
num_samples = len(f.samples)
samp_int = f.bin[segyio.BinField.Interval]/1000
retrieved_shot = np.zeros((traces_in_shot, num_samples))
shot_traces = f.trace[position:position+traces_in_shot]
for i, trace in enumerate(shot_traces):
retrieved_shot[i] = trace
g.write('{} Shot loaded in: {} seconds\n'.format(strng, time.time()-start))
# Only keep receivers within the model'
xmin = origin[0]
idx_xrec = np.where(rec_coord[:, 0] < xmin)[0]
is_empty = idx_xrec.size == 0
if not is_empty:
g.write('{} in {}\n'.format(strng, rec_coord.shape))
idx_tr = np.where(rec_coord[:, 0] >= xmin)[0]
rec_coord = np.delete(rec_coord, idx_xrec, axis=0)
# For 3D shot records, scan also y-receivers
if len(origin) == 3:
ymin = origin[1]
idx_yrec = np.where(rec_coord[:, 1] < ymin)[0]
is_empty = idx_yrec.size == 0
if not is_empty:
rec_coord = np.delete(rec_coord, idx_yrec, axis=0)
if rec_coord.size == 0:
g.write('all receivers outside of model\n')
return np.zeros(vp.shape)
space_order = 8
g.write('{} before: {} {} {}\n'.format(strng, params[0].shape, rec_coord.shape, src_coord.shape))
model = limit_model_to_receiver_area(rec_coord, src_coord, origin, spacing,
shape, vp, params, space_order=space_order, nbl=80)
g.write('{} shape_vp: {}\n'.format(strng, model.vp.shape))
model.smooth(('epsilon', 'delta', 'theta'))
# Geometry for current shot
geometry = AcquisitionGeometry(model, rec_coord, src_coord, 0, (num_samples-1)*samp_int, f0=0.018, src_type='Ricker')
g.write("{} Number of samples modelled data & dt: {} & {}\n".format(strng, geometry.nt, model.critical_dt))
g.write("{} Samples & dt: {} & {}\n".format(strng, num_samples, samp_int))
# Set up solver.
solver_tti = AnisotropicWaveSolver(model, geometry, space_order=space_order)
# Create image symbol and instantiate the previously defined imaging operator
image = Function(name='image', grid=model.grid)
itemsize = np.dtype(np.float32).itemsize
full_fld_mem = model.vp.size*itemsize*geometry.nt*2.
checkpointing = True
if checkpointing:
op_imaging = ImagingOperator(geometry, image, space_order, save=False)
n_checkpoints = 150
ckp_fld_mem = model.vp.size*itemsize*n_checkpoints*2.
g.write('Mem full fld: {} == {} use ckp instead\n'.format(full_fld_mem, humanbytes(full_fld_mem)))
g.write('Number of checkpoints/timesteps: {}/{}\n'.format(n_checkpoints, geometry.nt))
g.write('Memory saving: {}\n'.format(humanbytes(full_fld_mem-ckp_fld_mem)))
u = TimeFunction(name='u', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
v = TimeFunction(name='v', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
vv = TimeFunction(name='vv', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
uu = TimeFunction(name='uu', grid=model.grid, staggered=None,
time_order=2, space_order=space_order)
cp = DevitoCheckpoint([u, v])
op_fwd = solver_tti.op_fwd(save=False)
op_fwd.cfunction
op_imaging.cfunction
wrap_fw = CheckpointOperator(op_fwd, src=geometry.src,
u=u, v=v, vp=model.vp, epsilon=model.epsilon,
delta=model.delta, theta=model.theta, dt=model.critical_dt)
time_range = TimeAxis(start=0, stop=(num_samples-1)*samp_int, step=samp_int)
dobs = Receiver(name='dobs', grid=model.grid, time_range=time_range, coordinates=geometry.rec_positions)
if not is_empty:
dobs.data[:] = retrieved_shot[idx_tr, :].T
else:
dobs.data[:] = retrieved_shot[:].T
dobs_resam = dobs.resample(num=geometry.nt)
g.write('Shape of residual: {}\n'.format(dobs_resam.data.shape))
wrap_rev = CheckpointOperator(op_imaging, u=u, v=v, vv=vv, uu=uu, vp=model.vp,
epsilon=model.epsilon, delta=model.delta, theta=model.theta,
dt=model.critical_dt, residual=dobs_resam.data)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints, dobs_resam.shape[0]-2)
g.write('Revolver storage: {}\n'.format(humanbytes(cp.size*n_checkpoints*itemsize)))
wrp.apply_forward()
g.write('{} run finished\n'.format(strng))
summary = wrp.apply_reverse()
form = 'image_{}.bin'.format(str(record).zfill(5))
h = open(form, 'wb')
g.write('{}\n'.format(str(image.data.shape)))
np.transpose(image.data).astype('float32').tofile(h)
else:
# For illustrative purposes, assuming that there is enough memory
g.write('enough memory to save full fld: {} == {}\n'.format(full_fld_mem, humanbytes(full_fld_mem)))
op_imaging = ImagingOperator(geometry, image, space_order)
vv = TimeFunction(name='vv', grid=model.grid, staggered=None, time_order=2, space_order=space_order)
uu = TimeFunction(name='uu', grid=model.grid, staggered=None, time_order=2, space_order=space_order)
time_range = TimeAxis(start=0, stop=(num_samples-1)*samp_int, step=samp_int)
dobs = Receiver(name='dobs', grid=model.grid, time_range=time_range, coordinates=geometry.rec_positions)
if not is_empty:
dobs.data[:] = retrieved_shot[idx_tr, :].T
else:
dobs.data[:] = retrieved_shot[:].T
dobs_resam = dobs.resample(num=geometry.nt)
u, v = solver_tti.forward(vp=model.vp, epsilon=model.epsilon, delta=model.delta,
theta=model.theta, dt=model.critical_dt, save=True)[1:-1]
op_imaging(u=u, v=v, vv=vv, uu=uu, epsilon=model.epsilon, delta=model.delta,
theta=model0.theta, vp=model.vp, dt=model.critical_dt, residual=dobs_resam)
full_image = extend_image(origin, vp, model, image)
return full_image
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 test_index_alignment(const):
""" A much simpler test meant to ensure that the forward and reverse indices are
correctly aligned (i.e. u * v , where u is the forward field and v the reverse field
corresponds to the correct timesteps in u and v). The forward operator does u = u + 1
which means that the field a will be equal to nt (0 -> 1 -> 2 -> 3), the number of
timesteps this operator is run for. The field at the last time step of the forward is
used to initialise the field v for the reverse pass. The reverse operator does
v = v - 1, which means that if the reverse operator is run for the same number of
timesteps as the forward operator, v should be 0 at the last time step
(3 -> 2 -> 1 -> 0). There is also a grad = grad + u * v accumulator in the reverse
operator. If the alignment is correct, u and v should have the same value at every
time step:
0 -> 1 -> 2 -> 3 u
0 <- 1 <- 2 <- 3 v
and hence grad = 0*0 + 1*1 + 2*2 + 3*3 = sum(n^2) where n -> [0, nt]
If the test fails, the resulting number can tell you how the fields are misaligned
"""
n = 4
grid = Grid(shape=(2, 2))
order_of_eqn = 1
modulo_factor = order_of_eqn + 1
nt = n - order_of_eqn
u = TimeFunction(name='u', grid=grid, save=n)
# Increment one in the forward pass 0 -> 1 -> 2 -> 3
fwd_op = Operator(Eq(u.forward, u + 1.*const))
# Invocation 1
fwd_op(time=nt-1, constant=1)
last_time_step_v = nt % modulo_factor
# Last time step should be equal to the number of timesteps we ran
assert(np.allclose(u.data[nt, :, :], nt))
v = TimeFunction(name='v', grid=grid, save=None)
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Decrement one in the reverse pass 3 -> 2 -> 1 -> 0
adj_eqn = Eq(v, v.forward - 1.*const)
adj_op = Operator(adj_eqn)
# Invocation 2
adj_op(time=nt-1, constant=1)
# Last time step should be back to 0
assert(np.allclose(v.data[0, :, :], 0))
# Reset v to run the backward again
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
prod = Function(name="prod", grid=grid)
# Multiply u and v and add them
# = 3*3 + 2*2 + 1*1 + 0*0
prod_eqn = Eq(prod, prod + u * v)
comb_op = Operator([adj_eqn, prod_eqn])
# Invocation 3
comb_op(time=nt-1, constant=1)
final_value = sum([n**2 for n in range(nt)])
# Final value should be sum of squares of first nt natural numbers
assert(np.allclose(prod.data, final_value))
# Now reset to repeat all the above tests with checkpointing
prod.data[:] = 0
v.data[last_time_step_v, :, :] = u.data[nt, :, :]
# Checkpointed version doesn't require to save u
u_nosave = TimeFunction(name='u_n', grid=grid)
# change equations to use new symbols
fwd_eqn_2 = Eq(u_nosave.forward, u_nosave + 1.*const)
fwd_op_2 = Operator(fwd_eqn_2)
cp = DevitoCheckpoint([u_nosave])
wrap_fw = CheckpointOperator(fwd_op_2, constant=1)
prod_eqn_2 = Eq(prod, prod + u_nosave * v)
comb_op_2 = Operator([adj_eqn, prod_eqn_2])
wrap_rev = CheckpointOperator(comb_op_2, constant=1)
wrp = Revolver(cp, wrap_fw, wrap_rev, None, nt)
# Invocation 4
wrp.apply_forward()
assert(np.allclose(u_nosave.data[last_time_step_v, :, :], nt))
# Invocation 5
wrp.apply_reverse()
assert(np.allclose(v.data[0, :, :], 0))
assert(np.allclose(prod.data, final_value))
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 process_shot_checkpointed(i,
solver,
shots_container,
auth,
exclude_boundaries=True,
checkpoint_params=None):
model = solver.model
so = solver.space_order
geometry = solver.geometry
time_order = 2
dt = solver.model.critical_dt # 1.75
nt = solver.geometry.time_axis.num - time_order
if checkpoint_params is not None:
n_checkpoints = checkpoint_params.pop('n_checkpoints', 1000)
compression_params = checkpoint_params
else:
n_checkpoints = 1000
compression_params = None
# Create symbols to hold the gradient and residual
grad = Function(name="grad", grid=model.grid)
smooth_d = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
residual = Receiver(name='rec',
grid=model.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
u = TimeFunction(name='u',
grid=model.grid,
time_order=time_order,
space_order=so)
v = TimeFunction(name='v',
grid=model.grid,
time_order=time_order,
space_order=so)
fwd_op = solver.op_fwd(save=False)
rev_op = solver.op_grad(save=False)
cp = DevitoCheckpoint([u])
true_d, source_location, old_dt = load_shot(i,
auth=auth,
container=shots_container)
true_d = reinterpolate(true_d, solver.geometry.nt, old_dt)
# Update source location
solver.geometry.src_positions[0, :] = source_location[:]
wrap_fw = CheckpointOperator(fwd_op,
src=solver.geometry.src,
u=u,
rec=smooth_d,
dt=dt,
vp=solver.model.vp)
wrap_rev = CheckpointOperator(rev_op,
u=u,
v=v,
rec=residual,
grad=grad,
dt=dt,
vp=solver.model.vp)
wrp = Revolver(cp,
wrap_fw,
wrap_rev,
n_checkpoints,
nt,
compression_params=compression_params)
wrp.apply_forward()
# Compute gradient from data residual and update objective function
residual.data[:] = smooth_d.data[:] - true_d[:]
objective = .5 * np.linalg.norm(residual.data.ravel())**2
wrp.apply_reverse()
# Prepare for serialization before returning
if exclude_boundaries:
grad = trim_boundary(grad, solver.model.nbl)
else:
grad = grad.data
grad = np.array(grad, dtype=solver.model.dtype)
return objective, grad
def jacobian_adjoint(self,
rec,
u,
v=None,
grad=None,
precon=None,
vp=None,
checkpointing=False,
**kwargs):
"""
Gradient modelling function for computing the adjoint of the
Linearized Born modelling function, ie. the action of the
Jacobian adjoint on an input data.
Parameters
----------
rec : SparseTimeFunction
Receiver data.
u : TimeFunction
Full wavefield `u` (created with save=True).
v : TimeFunction, optional
Stores the computed wavefield.
grad : Function, optional
Stores the gradient field.
vp : Function or float, optional
The time-constant velocity.
Returns
-------
Gradient field and performance summary.
"""
dt = kwargs.pop('dt', self.dt)
# Gradient symbol
grad = grad or Function(name='grad', grid=self.model.grid)
# Create the forward wavefield
v = v or TimeFunction(name='v',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
# Pick vp from model unless explicitly provided
vp = vp or self.model.vp
if checkpointing:
u = TimeFunction(name='u',
grid=self.model.grid,
time_order=2,
space_order=self.space_order)
cp = DevitoCheckpoint([u])
n_checkpoints = None
wrap_fw = CheckpointOperator(self.op_fwd(save=False),
src=self.geometry.src,
u=u,
vp=vp,
dt=dt)
wrap_rev = CheckpointOperator(self.op_grad(save=False),
u=u,
v=v,
vp=vp,
rec=rec,
dt=dt,
grad=grad)
# Run forward
wrp = Revolver(cp, wrap_fw, wrap_rev, n_checkpoints,
rec.data.shape[0] - 2)
wrp.apply_forward()
summary = wrp.apply_reverse()
else:
if precon is not None:
precon = precon #don't know if this is needded but I'll just follow style
summary = self.op_grad_precon().apply(rec=rec,
grad=grad,
precon=precon,
v=v,
u=u,
vp=vp,
dt=dt,
**kwargs)
else:
summary = self.op_grad().apply(rec=rec,
grad=grad,
v=v,
u=u,
vp=vp,
dt=dt,
**kwargs)
return grad, summary
