def callback(final_solution_basename, vec):
callback.call_count += 1
fwi_iteration = callback.call_count
filename = "%s_%d.h5" % (final_solution_basename, fwi_iteration)
with profiler.get_timer('io', 'write_progress'):
to_hdf5(vec2mat(vec, model.shape), filename)
print(profiler.summary())
def callback(progress_dir, intermediates_dir, model, exclude_boundaries, vec):
global plot_model_to_file
callback.call_count += 1
if not hasattr(callback, "obj_fn_history"):
callback.obj_fn_history = []
callback.obj_fn_history.append(fwi_gradient.obj_fn_cache[vec.tobytes()])
fwi_iteration = callback.call_count
filename = os.path.join(intermediates_dir, "solution%d.h5" % fwi_iteration)
if exclude_boundaries:
to_hdf5(vec2mat(vec, model.shape), filename)
else:
to_hdf5(vec2mat(vec, model.vp.shape), filename)
progress_filename = os.path.join(progress_dir, "fwi-iter%d.pdf" % (fwi_iteration))
plot_model_to_file(solver.model, progress_filename)
def test_vec2mat():
shape = (2, 2, 2)
vec = np.arange(8)
mat = vec.reshape(shape)
assert (np.array_equal(mat, vec2mat(vec, shape)))
back = mat2vec(mat)
assert (np.array_equal(back, vec))
def fwi_gradient(vp_in, model, geometry, nshots, client, solver_params):
start_time = time.time()
vp_in = vec2mat(vp_in, model.shape)
f_vp_in = client.scatter(vp_in) # Dask enforces this for large arrays
assert (model.shape == vp_in.shape)
futures = []
for i in range(nshots):
futures.append(
client.submit(
fwi_gradient_shot,
f_vp_in,
i,
solver_params,
resources={
'tasks': 1
})) # Ensure one task per worker (to run two, tasks=0.5)
shape = model.shape
def reduction(*args):
grad = np.zeros(shape) # Closured from above
objective = 0.
for a in args:
o, g = a
objective += o
grad += g
return objective, grad
reduce_future = client.submit(reduction, *futures)
wait(reduce_future)
objective, grad = reduce_future.result()
elapsed_time = time.time() - start_time
print("Objective function evaluation completed in %f seconds" %
elapsed_time)
# Scipy LBFGS misbehaves if type is not float64
grad = mat2vec(np.array(grad)).astype(np.float64)
return objective, grad
def fwi_gradient_local(vp_in, nshots, solver, shots_container):
model = solver.model
vp_in = np.array(vec2mat(vp_in, solver.model.vp.shape),
dtype=solver.model.dtype)
assert (model.vp.shape == vp_in.shape)
solver.model.update("vp", vp_in)
objective = 0.
grad = np.zeros(model.vp.shape)
for i in range(nshots):
o, g = process_shot(i,
solver,
shots_container,
exclude_boundaries=False)
objective += o
grad += g
return objective, -mat2vec(grad).astype(np.float64)
def run(initial_model_filename, results_dir, tn, nshots, shots_container, so, nbl, kernel, scale_gradient, max_iter,
checkpointing, n_checkpoints, compression, tolerance, reference_solution, dtype):
if dtype == 'float32':
dtype = np.float32
elif dtype == 'float64':
dtype = np.float64
else:
raise ValueError("Invalid dtype")
water_depth = 20 # Number of points at the top of the domain that correspond to water
exclude_boundaries = True # Exclude the boundary regions from the optimisation problem
mute_water = True # Mute the gradient in the water region
initial_model_filename, datakey = initial_model_filename
model, geometry, bounds = initial_setup(initial_model_filename, tn, dtype, so, nbl,
datakey=datakey, exclude_boundaries=exclude_boundaries, water_depth=water_depth)
client = setup_dask()
if not os.path.exists(results_dir):
os.mkdir(results_dir)
intermediates_dir = os.path.join(results_dir, "intermediates")
if not os.path.exists(intermediates_dir):
os.mkdir(intermediates_dir)
progress_dir = os.path.join(results_dir, "progress")
if not os.path.exists(progress_dir):
os.mkdir(progress_dir)
auth = default_auth()
solver_params = {'h5_file': Blob("models", initial_model_filename, auth=auth), 'tn': tn,
'space_order': so, 'dtype': dtype, 'datakey': datakey, 'nbl': nbl,
'opt': ('noop', {'openmp': True, 'par-dynamic-work': 1000})}
if kernel in ['OT2', 'OT4']:
solver_params['kernel'] = kernel
solver = overthrust_solver_iso(**solver_params)
elif kernel == "rho":
solver_params['water_depth'] = water_depth
solver_params['calculate_density'] = False
solver = overthrust_solver_density(**solver_params)
solver._dt = 1.75
solver.geometry.resample(1.75)
f_args = [nshots, client, solver, shots_container, auth, scale_gradient, mute_water, exclude_boundaries, water_depth]
if checkpointing:
f_args += [checkpointing, {'n_checkpoints': n_checkpoints, 'scheme': compression,
'tolerance': tolerance}]
if exclude_boundaries:
v0 = mat2vec(trim_boundary(model.vp, model.nbl)).astype(np.float64)
else:
v0 = mat2vec(model.vp).astype(np.float64)
def callback(progress_dir, intermediates_dir, model, exclude_boundaries, vec):
global plot_model_to_file
callback.call_count += 1
if not hasattr(callback, "obj_fn_history"):
callback.obj_fn_history = []
callback.obj_fn_history.append(fwi_gradient.obj_fn_cache[vec.tobytes()])
fwi_iteration = callback.call_count
filename = os.path.join(intermediates_dir, "solution%d.h5" % fwi_iteration)
if exclude_boundaries:
to_hdf5(vec2mat(vec, model.shape), filename)
else:
to_hdf5(vec2mat(vec, model.vp.shape), filename)
progress_filename = os.path.join(progress_dir, "fwi-iter%d.pdf" % (fwi_iteration))
plot_model_to_file(solver.model, progress_filename)
callback.call_count = 0
partial_callback = partial(callback, progress_dir, intermediates_dir, model, exclude_boundaries)
fwi_gradient.call_count = 0
fwd_op = solver.op_fwd(save=False)
rev_op = solver.op_grad(save=False)
fwd_op.ccode
rev_op.ccode
solution_object = minimize(fwi_gradient,
v0,
args=tuple(f_args),
jac=True, method='L-BFGS-B',
callback=partial_callback, bounds=bounds,
options={'disp': True, 'maxiter': max_iter})
if exclude_boundaries:
final_model = vec2mat(solution_object.x, model.shape)
else:
final_model = vec2mat(solution_object.x, model.vp.shape)
solver.model.update("vp", final_model)
# Save plot of final model
final_plot_filename = os.path.join(results_dir, "final_model.pdf")
plot_model_to_file(solver.model, final_plot_filename)
# Save objective function values to CSV
obj_fn_history = callback.obj_fn_history
obj_fn_vals_filename = os.path.join(results_dir, "objective_function_values.csv")
with open(obj_fn_vals_filename, "w") as vals_file:
vals_writer = csv.writer(vals_file)
for r in obj_fn_history:
vals_writer.writerow([r])
# Plot objective function values
plt.title("FWI convergence")
plt.xlabel("Iteration number")
plt.ylabel("Objective function value")
obj_fun_plt_filename = os.path.join(results_dir, "convergence.tex")
obj_fun_plt_pdf_filename = os.path.join(results_dir, "convergence.pdf")
plt.clf()
if reference_solution is None:
plt.plot(obj_fn_history)
else:
# Load reference solution convergence history
with open(reference_solution, 'r') as reference_file:
vals_reader = csv.reader(reference_file)
reference_solution_values = []
for r in vals_reader:
reference_solution_values.append(float(r[0]))
# Pad with 0s to ensure same size
# Plot with legends
plt.plot(obj_fn_history, label="Lossy FWI")
plt.plot(reference_solution_values, label="Reference FWI")
# Display legend
plt.legend()
plt.savefig(obj_fun_plt_pdf_filename)
tikzplotlib.save(obj_fun_plt_filename)
true_model = overthrust_model_iso("overthrust_3D_true_model_2D.h5", datakey="m", dtype=dtype, space_order=so, nbl=nbl)
true_model_vp = trim_boundary(true_model.vp, true_model.nbl)
error_norm = np.linalg.norm(true_model_vp - final_model)
print("L2 norm of final solution vs true solution: %f" % error_norm)
data = {'error_norm': error_norm,
'checkpointing': checkpointing,
'compression': compression,
'tolerance': tolerance,
'ncp': n_checkpoints}
write_results(data, "fwi_experiment.csv")
# Final solution
final_solution_filename = os.path.join(results_dir, "final_solution.h5")
to_hdf5(final_model, final_solution_filename)
def fwi_gradient(vp_in, nshots, client, solver, shots_container, auth, scale_gradient=None, mute_water=True,
exclude_boundaries=True, water_depth=20, checkpointing=False, checkpoint_params=None):
start_time = time.time()
reset_cluster(client)
if not hasattr(fwi_gradient, "obj_fn_cache"):
fwi_gradient.obj_fn_cache = {}
if exclude_boundaries:
vp = np.array(vec2mat(vp_in, solver.model.shape), dtype=solver.model.dtype)
else:
vp = np.array(vec2mat(vp_in, solver.model.vp.shape), dtype=solver.model.dtype)
solver.model.update("vp", vp)
# Dask enforces this for large objects
f_solver = client.scatter(solver, broadcast=True)
futures = []
for i in range(nshots):
if checkpointing:
futures.append(client.submit(process_shot_checkpointed, i, f_solver, shots_container, auth, exclude_boundaries,
checkpoint_params, resources={'tasks': 1}))
else:
futures.append(client.submit(process_shot, i, f_solver, shots_container, auth, exclude_boundaries,
resources={'tasks': 1})) # Ensure one task per worker (to run two, tasks=0.5)
if exclude_boundaries:
gradient_shape = solver.model.shape
else:
gradient_shape = solver.model.vp.shape
def reduction(*args):
grad = np.zeros(gradient_shape) # Closured from above
objective = 0.
for a in args:
o, g = a
objective += o
grad += g
return objective, grad
reduce_future = client.submit(reduction, *futures)
wait(reduce_future)
objective, grad = reduce_future.result()
if mute_water:
if exclude_boundaries:
muted_depth = water_depth
else:
muted_depth = water_depth + solver.model.nbl
grad[:, 0:muted_depth] = 0
# Scipy LBFGS misbehaves if type is not float64
grad = mat2vec(grad).astype(np.float64)
if scale_gradient is not None:
if scale_gradient == "W":
if not hasattr(fwi_gradient, "gradient_scaling_factor"):
fwi_gradient.gradient_scaling_factor = np.max(np.abs(grad))
grad /= fwi_gradient.gradient_scaling_factor
elif scale_gradient == "L":
grad /= np.max(np.abs(grad))
else:
raise ValueError("Invalid value %s for gradient scaling. Allowed: None, L, W" % scale_gradient)
fwi_gradient.obj_fn_cache[vp_in.tobytes()] = objective
elapsed_time = time.time() - start_time
eprint("Objective function evaluation completed in %f seconds. F=%f" % (elapsed_time, objective))
return objective, -grad
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 run(initial_model_filename, final_solution_basename, tn, nshots,
shots_container, so, nbl, kernel, checkpointing, n_checkpoints,
compression, tolerance):
dtype = np.float32
model, geometry, bounds = initial_setup(initial_model_filename, tn, dtype,
so, nbl)
solver_params = {
'filename': initial_model_filename,
'tn': tn,
'space_order': so,
'dtype': dtype,
'datakey': 'm0',
'nbl': nbl,
'origin': model.origin,
'spacing': model.spacing,
'shots_container': shots_container
}
client = setup_dask()
f_args = [model, geometry, nshots, client, solver_params]
# if checkpointing:
# f_g = fwi_gradient_checkpointed
# compression_params = {'scheme': compression, 'tolerance': 10**(-tolerance)}
# f_args.append(n_checkpoints)
# f_args.append(compression_params)
# else:
f_g = fwi_gradient
clipped_vp = mat2vec(clip_boundary_and_numpy(model.vp.data, model.nbl))
def callback(final_solution_basename, vec):
callback.call_count += 1
fwi_iteration = callback.call_count
filename = "%s_%d.h5" % (final_solution_basename, fwi_iteration)
with profiler.get_timer('io', 'write_progress'):
to_hdf5(vec2mat(vec, model.shape), filename)
print(profiler.summary())
callback.call_count = 0
partial_callback = partial(callback, final_solution_basename)
solution_object = minimize(f_g,
clipped_vp,
args=tuple(f_args),
jac=True,
method='L-BFGS-B',
callback=partial_callback,
bounds=bounds,
options={
'disp': True,
'maxiter': 60
})
final_model = vec2mat(solution_object.x, model.shape)
true_model = overthrust_model_iso("overthrust_3D_true_model_2D.h5",
datakey="m",
dtype=dtype,
space_order=so,
nbl=nbl)
true_model_vp = clip_boundary_and_numpy(true_model.vp.data, true_model.nbl)
error_norm = np.linalg.norm(true_model_vp - final_model)
print(error_norm)
data = {
'error_norm': error_norm,
'checkpointing': checkpointing,
'compression': compression,
'tolerance': tolerance,
'ncp': n_checkpoints
}
write_results(data, "fwi_experiment.csv")
to_hdf5(final_model, '%s_final.h5' % final_solution_basename)
def test_gradientFWI(self):
r"""
This test ensures that the FWI gradient computed with devito
satisfies the Taylor expansion property:
.. math::
\Phi(m0 + h dm) = \Phi(m0) + \O(h) \\
\Phi(m0 + h dm) = \Phi(m0) + h \nabla \Phi(m0) + \O(h^2) \\
\Phi(m0) = .5* || F(m0 + h dm) - D ||_2^2
where
.. math::
\nabla \Phi(m0) = <J^T \delta d, dm> \\
\delta d = F(m0+ h dm) - D \\
with F the Forward modelling operator.
"""
initial_model_filename = "overthrust_3D_initial_model_2D.h5"
true_model_filename = "overthrust_3D_true_model_2D.h5"
tn = 4000
dtype = np.float32
so = 6
nbl = 40
shots_container = "shots-iso"
shot_id = 10
##########
model0 = overthrust_model_iso(initial_model_filename,
datakey="m0",
dtype=dtype,
space_order=so,
nbl=nbl)
model_t = overthrust_model_iso(true_model_filename,
datakey="m",
dtype=dtype,
space_order=so,
nbl=nbl)
_, geometry, _ = initial_setup(initial_model_filename,
tn,
dtype,
so,
nbl,
datakey="m0")
# rec, source_location, old_dt = load_shot(shot_id,
# container=shots_container)
source_location = geometry.src_positions
solver_params = {
'h5_file': initial_model_filename,
'tn': tn,
'space_order': so,
'dtype': dtype,
'datakey': 'm0',
'nbl': nbl,
'origin': model0.origin,
'spacing': model0.spacing,
'shots_container': shots_container,
'src_coordinates': source_location
}
solver = overthrust_solver_iso(**solver_params)
true_solver_params = solver_params.copy()
true_solver_params['h5_file'] = true_model_filename
true_solver_params['datakey'] = "m"
solver_true = overthrust_solver_iso(**true_solver_params)
rec, _, _ = solver_true.forward()
v0 = mat2vec(clip_boundary_and_numpy(model0.vp.data, model0.nbl))
v_t = mat2vec(clip_boundary_and_numpy(model_t.vp.data, model_t.nbl))
dm = np.float64(v_t**(-2) - v0**(-2))
print("dm", np.linalg.norm(dm), dm.shape)
F0, gradient = fwi_gradient_shot(vec2mat(v0, model0.shape), shot_id,
solver_params, source_location)
G = np.dot(gradient.reshape(-1), dm.reshape(-1))
# FWI Gradient test
H = [0.5, 0.25, .125, 0.0625, 0.0312, 0.015625, 0.0078125]
error1 = np.zeros(7)
error2 = np.zeros(7)
for i in range(0, 7):
# Add the perturbation to the model
vloc = np.sqrt(v0**2 * v_t**2 /
((1 - H[i]) * v_t**2 + H[i] * v0**2))
m = Model(vp=vloc,
nbl=nbl,
space_order=so,
dtype=dtype,
shape=model0.shape,
origin=model0.origin,
spacing=model0.spacing,
bcs="damp")
# Data for the new model
d = solver.forward(vp=m.vp)[0]
# First order error Phi(m0+dm) - Phi(m0)
F_i = .5 * linalg.norm((d.data - rec.data).reshape(-1))**2
error1[i] = np.absolute(F_i - F0)
# Second order term r Phi(m0+dm) - Phi(m0) - <J(m0)^T \delta d, dm>
error2[i] = np.absolute(F_i - F0 - H[i] * G)
print(i, F0, F_i, H[i] * G)
# Test slope of the tests
p1 = np.polyfit(np.log10(H), np.log10(error1), 1)
p2 = np.polyfit(np.log10(H), np.log10(error2), 1)
info('1st order error, Phi(m0+dm)-Phi(m0): %s' % (p1))
info(
r'2nd order error, Phi(m0+dm)-Phi(m0) - <J(m0)^T \delta d, dm>: %s'
% (p2))
print("Error 1:")
print(error1)
print("***")
print("Error 2:")
print(error2)
assert np.isclose(p1[0], 1.0, rtol=0.1)
assert np.isclose(p2[0], 2.0, rtol=0.1)