In this tutorial we show how Devito and scipy.optimize.minimize are used with Dask to perform full waveform inversion (FWI) on distributed memory parallel computers.
In this tutorial we use scipy.optimize.minimize to solve the FWI gradient based minimization problem rather than the simple grdient decent algorithm in the previous tutorial.
scipy.optimize.minimize(fun, x0, args=(), method=None, jac=None, hess=None, hessp=None, bounds=None, constraints=(), tol=None, callback=None, options=None)
Minimization of scalar function of one or more variables.
In general, the optimization problems are of the form:
minimize f(x) subject to
g_i(x) >= 0, i = 1,...,m h_j(x) = 0, j = 1,...,p where x is a vector of one or more variables. g_i(x) are the inequality constraints. h_j(x) are the equality constrains.
scipy.optimize.minimize provides a wide variety of methods for solving minimization problems depending on the context. Here we are going to focus on using L-BFGS via scipy.optimize.minimize(method=āL-BFGS-Bā)
scipy.optimize.minimize(fun, x0, args=(), method='L-BFGS-B', jac=None, bounds=None, tol=None, callback=None, options={'disp': None, 'maxls': 20, 'iprint': -1, 'gtol': 1e-05, 'eps': 1e-08, 'maxiter': 15000, 'ftol': 2.220446049250313e-09, 'maxcor': 10, 'maxfun': 15000})```
The argument `fun` is a callable function that returns the misfit between the simulated and the observed data. If `jac` is a Boolean and is `True`, `fun` is assumed to return the gradient along with the objective function - as is our case when applying the adjoint-state method.
Dask is a flexible parallel computing library for analytic computing.
Dask is composed of two components:
- Dynamic task scheduling optimized for computation...
- āBig Dataā collections like parallel arrays, dataframes, and lists that extend common interfaces like NumPy, Pandas, or Python iterators to larger-than-memory or distributed environments. These parallel collections run on top of the dynamic task schedulers.
Dask emphasizes the following virtues:
- Familiar: Provides parallelized NumPy array and Pandas DataFrame objects
- Flexible: Provides a task scheduling interface for more custom workloads and integration with other projects.
- Native: Enables distributed computing in Pure Python with access to the PyData stack.
- Fast: Operates with low overhead, low latency, and minimal serialization necessary for fast numerical algorithms
- Scales up: Runs resiliently on clusters with 1000s of cores
- Scales down: Trivial to set up and run on a laptop in a single process
- Responsive: Designed with interactive computing in mind it provides rapid feedback and diagnostics to aid humans
We are going to use it here to parallelise the computation of the functional and gradient as this is the vast bulk of the computational expense of FWI and it is trivially parallel over data shots.
In a real world scenario we work with collected seismic data; for the tutorial we know what the actual solution is and we are using the workers to also generate the synthetic data.
#NBVAL_IGNORE_OUTPUT
# Set up inversion parameters.
param = {'t0': 0.,
'tn': 1000., # Simulation last 1 second (1000 ms)
'f0': 0.010, # Source peak frequency is 10Hz (0.010 kHz)
'nshots': 5, # Number of shots to create gradient from
'm_bounds': (0.08, 0.25), # Set the min and max slowness
'shape': (101, 101), # Number of grid points (nx, nz).
'spacing': (10., 10.), # Grid spacing in m. The domain size is now 1km by 1km.
'origin': (0, 0), # Need origin to define relative source and receiver locations.
'nbl': 40} # nbl thickness.
import numpy as np
import scipy
from scipy import signal, optimize
from devito import Grid
from distributed import Client, LocalCluster, wait
import cloudpickle as pickle
# Import acoustic solver, source and receiver modules.
from examples.seismic import Model, demo_model, AcquisitionGeometry, Receiver
from examples.seismic.acoustic import AcousticWaveSolver
from examples.seismic import AcquisitionGeometry
# Import convenience function for plotting results
from examples.seismic import plot_image
def get_true_model():
''' Define the test phantom; in this case we are using
a simple circle so we can easily see what is going on.
'''
return demo_model('circle-isotropic', vp=3.0, vp_background=2.5,
origin=param['origin'], shape=param['shape'],
spacing=param['spacing'], nbl=param['nbl'])
def get_initial_model():
'''The initial guess for the subsurface model.
'''
# Make sure both model are on the same grid
grid = get_true_model().grid
return demo_model('circle-isotropic', vp=2.5, vp_background=2.5,
origin=param['origin'], shape=param['shape'],
spacing=param['spacing'], nbl=param['nbl'],
grid=grid)
def wrap_model(x, astype=None):
'''Wrap a flat array as a subsurface model.
'''
model = get_initial_model()
if astype:
model.vp = x.astype(astype).reshape(model.vp.data.shape)
else:
model.vp = x.reshape(model.vp.data.shape)
return model
def load_model(filename):
""" Returns the current model. This is used by the
worker to get the current model.
"""
pkl = pickle.load(open(filename, "rb"))
return pkl['model']
def dump_model(filename, model):
''' Dump model to disk.
'''
pickle.dump({'model':model}, open(filename, "wb"))
def load_shot_data(shot_id, dt):
''' Load shot data from disk, resampling to the model time step.
'''
pkl = pickle.load(open("shot_%d.p"%shot_id, "rb"))
return pkl['geometry'].resample(dt), pkl['rec'].resample(dt)
def dump_shot_data(shot_id, rec, geometry):
''' Dump shot data to disk.
'''
pickle.dump({'rec':rec, 'geometry': geometry}, open('shot_%d.p'%shot_id, "wb"))
def generate_shotdata_i(param):
""" Inversion crime alert! Here the worker is creating the
'observed' data using the real model. For a real case
the worker would be reading seismic data from disk.
"""
true_model = get_true_model()
shot_id = param['shot_id']
src_coordinates = np.empty((1, len(param['shape'])))
src_coordinates[0, :] = [30, param['shot_id']*1000./(param['nshots']-1)]
# Number of receiver locations per shot.
nreceivers = 101
# Set up receiver data and geometry.
rec_coordinates = np.empty((nreceivers, len(param['shape'])))
rec_coordinates[:, 1] = np.linspace(0, true_model.domain_size[0], num=nreceivers)
rec_coordinates[:, 0] = 980. # 20m from the right end
# Geometry
geometry = AcquisitionGeometry(true_model, rec_coordinates, src_coordinates,
param['t0'], param['tn'], src_type='Ricker',
f0=param['f0'])
# Set up solver.
solver = AcousticWaveSolver(true_model, geometry, space_order=4)
# Generate synthetic receiver data from true model.
true_d, _, _ = solver.forward(vp=true_model.vp)
dump_shot_data(shot_id, true_d, geometry)
def generate_shotdata(param):
# Define work list
work = [dict(param) for i in range(param['nshots'])]
for i in range(param['nshots']):
work[i]['shot_id'] = i
generate_shotdata_i(work[i])
# Map worklist to cluster
futures = client.map(generate_shotdata_i, work)
# Wait for all futures
wait(futures)
#NBVAL_IGNORE_OUTPUT
# Start Dask cluster
cluster = LocalCluster(n_workers=2, death_timeout=600)
client = Client(cluster)
# Generate shot data.
generate_shotdata(param)
Operator `Forward` run in 0.02 s Operator `Forward` run in 0.06 s Operator `Forward` run in 0.02 s Operator `Forward` run in 0.03 s Operator `Forward` run in 0.04 s
Previously we defined a function to calculate the individual contribution to the functional and gradient for each shot, which was then used in a loop over all shots. However, when using distributed frameworks such as Dask we instead think in terms of creating a worklist which gets mapped onto the worker pool. The sum reduction is also performed in parallel. For now however we assume that the scipy.optimize.minimize itself is running on the master process; this is a reasonable simplification because the computational cost of calculating (f, g) far exceeds the other compute costs.
Because we want to be able to use standard reduction operators such as sum on (f, g) we first define it as a type so that we can define the __add__
(and __rand__
method).
# Define a type to store the functional and gradient.
class fg_pair:
def __init__(self, f, g):
self.f = f
self.g = g
def __add__(self, other):
f = self.f + other.f
g = self.g + other.g
return fg_pair(f, g)
def __radd__(self, other):
if other == 0:
return self
else:
return self.__add__(other)
To perform the inversion we are going to use scipy.optimize.minimize(method=āL-BFGS-Bā).
First we define the functional, f
, and gradient, g
, operator (i.e. the function fun
) for a single shot of data. This is the work that is going to be performed by the worker on a unit of data.
from devito import Function
# Create FWI gradient kernel for a single shot
def fwi_gradient_i(param):
# Load the current model and the shot data for this worker.
# Note, unlike the serial example the model is not passed in
# as an argument. Broadcasting large datasets is considered
# a programming anti-pattern and at the time of writing it
# it only worked relaiably with Dask master. Therefore, the
# the model is communicated via a file.
model0 = load_model(param['model'])
dt = model0.critical_dt
geometry, rec = load_shot_data(param['shot_id'], dt)
geometry.model = model0
# Set up solver.
solver = AcousticWaveSolver(model0, geometry, space_order=4)
# Compute simulated data and full forward wavefield u0
d, u0, _ = solver.forward(save=True)
# Compute the data misfit (residual) and objective function
residual = Receiver(name='rec', grid=model0.grid,
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)
residual.data[:] = d.data[:residual.shape[0], :] - rec.data[:residual.shape[0], :]
f = .5*np.linalg.norm(residual.data.flatten())**2
# Compute gradient using the adjoint-state method. Note, this
# backpropagates the data misfit through the model.
grad = Function(name="grad", grid=model0.grid)
solver.gradient(rec=residual, u=u0, grad=grad)
# Copying here to avoid a (probably overzealous) destructor deleting
# the gradient before Dask has had a chance to communicate it.
g = np.array(grad.data[:])
# return the objective functional and gradient.
return fg_pair(f, g)
Define the global functional-gradient operator. This does the following:
def fwi_gradient(model, param):
# Dump a copy of the current model for the workers
# to pick up when they are ready.
param['model'] = "model_0.p"
dump_model(param['model'], wrap_model(model))
# Define work list
work = [dict(param) for i in range(param['nshots'])]
for i in range(param['nshots']):
work[i]['shot_id'] = i
# Distribute worklist to workers.
fgi = client.map(fwi_gradient_i, work, retries=1)
# Perform reduction.
fg = client.submit(sum, fgi).result()
# L-BFGS in scipy expects a flat array in 64-bit floats.
return fg.f, -fg.g.flatten().astype(np.float64)
Equipped with a function to calculate the functional and gradient, we are finally ready to define the optimization function.
from scipy import optimize
# Define bounding box constraints on the solution.
def apply_box_constraint(vp):
# Maximum possible 'realistic' velocity is 3.5 km/sec
# Minimum possible 'realistic' velocity is 2 km/sec
return np.clip(vp, 2.0, 3.5)
# Many optimization methods in scipy.optimize.minimize accept a callback
# function that can operate on the solution after every iteration. Here
# we use this to apply box constraints and to monitor the true relative
# solution error.
relative_error = []
def fwi_callbacks(x):
# Apply boundary constraint
x.data[:] = apply_box_constraint(x)
# Calculate true relative error
true_x = get_true_model().vp.data.flatten()
relative_error.append(np.linalg.norm((x-true_x)/true_x))
def fwi(model, param, ftol=0.1, maxiter=5):
result = optimize.minimize(fwi_gradient,
model.vp.data.flatten().astype(np.float64),
args=(param, ), method='L-BFGS-B', jac=True,
callback=fwi_callbacks,
options={'ftol':ftol,
'maxiter':maxiter,
'disp':True})
return result
We now apply our FWI function and have a look at the result.
#NBVAL_IGNORE_OUTPUT
model0 = get_initial_model()
# Baby steps
result = fwi(model0, param)
# Print out results of optimizer.
print(result)
fun: 211.63605994835484 hess_inv: <32761x32761 LbfgsInvHessProduct with dtype=float64> jac: array([-6.69854363e-12, -3.05454377e-11, -7.75552331e-11, ..., -1.31146288e-10, -5.26459223e-11, -1.16629224e-11]) message: b'STOP: TOTAL NO. of ITERATIONS REACHED LIMIT' nfev: 6 nit: 5 status: 1 success: False x: array([2.5, 2.5, 2.5, ..., 2.5, 2.5, 2.5])
#NBVAL_SKIP
# Show what the update does to the model
from examples.seismic import plot_image, plot_velocity
model0.vp = result.x.astype(np.float32).reshape(model0.vp.data.shape)
plot_velocity(model0)
#NBVAL_SKIP
# Plot percentage error
plot_image(100*np.abs(model0.vp.data-get_true_model().vp.data)/get_true_model().vp.data, vmax=15, cmap="hot")
#NBVAL_SKIP
import matplotlib.pyplot as plt
# Plot objective function decrease
plt.figure()
plt.loglog(relative_error)
plt.xlabel('Iteration number')
plt.ylabel('True relative error')
plt.title('Convergence')
plt.show()
This notebook is part of the tutorial "Optimised Symbolic Finite Difference Computation with Devito" presented at the IntelĀ® HPC Developer Conference 2017.