def test_checkpointed_vs_not_checkpointed(shape, kernel, space_order):
"""
Verifies that the gradients with and without checkpointing are the SpaceDimension
"""
spacing = tuple([15.0 for _ in shape])
tn = 500.
# checkpointing
example = CheckpointingExample(shape, spacing, tn, kernel, space_order)
m0, dm = example.initial_estimate()
gradient, rec = example.gradient(m0)
# No checkpointing
wave = acoustic_setup(shape=shape,
spacing=spacing,
dtype=np.float32,
kernel=kernel,
space_order=space_order,
nbpml=10 + space_order / 2)
m0 = Function(name='m0', grid=wave.model.m.grid, space_order=space_order)
m0.data[:] = smooth10(wave.model.m.data, wave.model.m.shape_domain)
# Compute receiver data for the true velocity
rec, u, _ = wave.forward()
# Compute receiver data and full wavefield for the smooth velocity
rec0, u0, _ = wave.forward(m=m0, save=True)
# Gradient: <J^T \delta d, dm>
residual = Receiver(name='rec',
grid=wave.model.grid,
data=rec0.data - rec.data,
time_range=rec.time_range,
coordinates=rec0.coordinates.data)
grad, _ = wave.gradient(residual, u0, m=m0)
assert np.allclose(grad.data, gradient)
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, 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_acoustic_save_and_nosave(shape=(50, 50), spacing=(15.0, 15.0), tn=500.,
time_order=2, space_order=4, nbpml=10):
solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn,
space_order=space_order, time_order=time_order)
rec, u, summary = solver.forward(save=True)
last_time_step = solver.source.nt-1
field_last_time_step = np.copy(u.data[last_time_step, :, :])
rec_bk = np.copy(rec.data)
rec, u, summary = solver.forward(save=False)
last_time_step = (last_time_step) % (time_order + 1)
assert(np.allclose(u.data[last_time_step, :, :], field_last_time_step))
assert(np.allclose(rec.data, rec_bk))
def test_acoustic_save_and_nosave(shape=(50, 50), spacing=(15.0, 15.0), tn=500.,
time_order=2, space_order=4, nbpml=10):
""" Run the acoustic example with and without save=True. Make sure the result is the
same
"""
solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn,
space_order=space_order, time_order=time_order)
rec, u, summary = solver.forward(save=True)
last_time_step = solver.geometry.nt-1
field_last_time_step = np.copy(u.data[last_time_step, :, :])
rec_bk = np.copy(rec.data)
rec, u, summary = solver.forward(save=False)
last_time_step = (last_time_step) % (time_order + 1)
assert(np.allclose(u.data[last_time_step, :, :], field_last_time_step))
assert(np.allclose(rec.data, rec_bk))
def test_acoustic_save_and_nosave(shape=(50, 50), spacing=(15.0, 15.0), tn=500.,
time_order=2, space_order=4, nbpml=10):
""" Run the acoustic example with and without save=True. Make sure the result is the
same
"""
solver = acoustic_setup(shape=shape, spacing=spacing, nbpml=nbpml, tn=tn,
space_order=space_order, time_order=time_order)
rec, u, summary = solver.forward(save=True)
last_time_step = solver.geometry.nt-1
field_last_time_step = np.copy(u.data[last_time_step, :, :])
rec_bk = np.copy(rec.data)
rec, u, summary = solver.forward(save=False)
last_time_step = (last_time_step) % (time_order + 1)
assert(np.allclose(u.data[last_time_step, :, :], field_last_time_step))
assert(np.allclose(rec.data, rec_bk))
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 setup(self, shape, space_order):
self.solver = acoustic_setup(shape=shape, space_order=space_order,
opt=('advanced', {'openmp': True}))
