# 04 - Full waveform inversion with Devito and Dask¶

## Introduction¶

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.

## scipy.optimize.minimize¶

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.

• 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.

## Setting up (synthetic) data¶

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.

In [1]:
#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.
'nbpml': 40}              # nbpml 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'], nbpml=param['nbpml'])

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'], nbpml=param['nbpml'],
grid=grid)

def wrap_model(x, astype=None):
'''Wrap a flat array as a subsurface model.
'''
model = get_initial_model()
if astype:
model.m.data[:] = x.astype(astype).reshape(model.m.data.shape)
else:
model.m.data[:] = x.reshape(model.m.data.shape)
return model

""" Returns the current model. This is used by the
worker to get the current model.
"""

return pkl['model']

def dump_model(filename, model):
''' Dump model to disk.
'''
pickle.dump({'model':model}, open(filename, "wb"))

''' Load shot data from disk, resampling to the model time step.
'''

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.

# Set up receiver data and geometry.
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(m=true_model.m)

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)

In [2]:
#NBVAL_IGNORE_OUTPUT

cluster = LocalCluster(n_workers=5, death_timeout=600)
client = Client(cluster)

# Generate shot data.
generate_shotdata(param)

Operator Forward run in 0.04 s
Operator Forward run in 0.04 s
Operator Forward run in 0.04 s
Operator Forward run in 0.05 s
Operator Forward run in 0.07 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).

In [3]:
# Define a type to store the functional and gradient.
class fg_pair:
def __init__(self, f, g):
self.f = f
self.g = g

f = self.f + other.f
g = self.g + other.g

return fg_pair(f, g)

if other == 0:
return self
else:


## Create operators for gradient based inversion¶

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.

In [4]:
from devito import Function

# Create FWI gradient kernel for a single shot
from devito import clear_cache

# Need to clear the workers cache.
clear_cache()

# 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.

dt = model0.critical_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
time_range=geometry.time_axis,
coordinates=geometry.rec_positions)

residual.data[:] = d.data[:] - rec.data[:]
f = .5*np.linalg.norm(residual.data.flatten())**2

# backpropagates the data misfit through the model.

# Copying here to avoid a (probably overzealous) destructor deleting

# return the objective functional and gradient.
return fg_pair(f, g)


Define the global functional-gradient operator. This does the following:

• Maps the worklist (shots) to the workers so that the invidual contributions to (f, g) are computed.
• Sum individual contributions to (f, g) and returns the result.
In [5]:
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.

# 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)


## FWI with L-BFGS-B¶

Equipped with a function to calculate the functional and gradient, we are finally ready to define the optimization function.

In [6]:
from scipy import optimize

# Define bounding box constraints on the solution.
def apply_box_constraint(m):
# Maximum possible 'realistic' velocity is 3.5 km/sec
# Minimum possible 'realistic' velocity is 2 km/sec
return np.clip(m, 1/3.5**2, 1/2**2)

# 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().m.data.flatten()
relative_error.append(np.linalg.norm((x-true_x)/true_x))

def fwi(model, param, ftol=0.1, maxiter=5):
model.m.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.

In [7]:
#NBVAL_IGNORE_OUTPUT

model0 = get_initial_model()

# Baby steps
result = fwi(model0, param)

# Print out results of optimizer.
print(result)

      fun: 90.41895049522077
hess_inv: <32761x32761 LbfgsInvHessProduct with dtype=float64>
jac: array([5.31055928e-12, 2.42461554e-11, 6.16747486e-11, ...,
1.02646572e-10, 4.08787067e-11, 9.01362492e-12])
message: b'STOP: TOTAL NO. of ITERATIONS EXCEEDS LIMIT'
nfev: 8
nit: 6
status: 1
success: False
x: array([0.16, 0.16, 0.16, ..., 0.16, 0.16, 0.16])

In [8]:
#NBVAL_SKIP

# Show what the update does to the model
from examples.seismic import plot_image, plot_velocity

model0.m.data[:] = result.x.astype(np.float32).reshape(model0.m.data.shape)
model0.vp = np.sqrt(1. / model0.m.data[40:-40, 40:-40])
plot_velocity(model0)

In [9]:
#NBVAL_SKIP

# Plot percentage error
plot_image(100*np.abs(model0.vp-get_true_model().vp.data)/get_true_model().vp.data, vmax=15, cmap="hot")

In [10]:
#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.