Skip to main content

Wave propagation modules for PyTorch.

Project description

Deepwave

Build StatusCodacy Badge

Deepwave provides wave propagation modules for PyTorch (currently only for the constant density acoustic / scalar wave equation). It is primarily aimed at seismic imaging/inversion. One interesting advantage of this is that it allows you to chain operations together. You could, for example, use Deepwave to perform FWI using a custom cost function (as long as PyTorch is able to automatically differentiate it), or add some other operations before and after wave propagation and PyTorch will backpropagate through all of them.

Wave propagation and FWI can be implemented using deep neural networks. Deep learning tools such as TensorFlow and PyTorch are currently too slow and memory inefficient to do this with realistic seismic datasets, however. Deepwave extends PyTorch with higher performance wave propagation modules (written in C and CUDA) so that you can benefit from the conveniences of PyTorch without sacrificing performance.

Deepwave runs on CPUs and NVIDIA GPUs. It should automatically detect compatible GPUs during installation and compile the GPU code if any are found. To run on the GPU you must transfer the model, source amplitude, and source and receiver location Tensors to the GPU in the standard PyTorch way (using .cuda() or .to(device)).

Installation

Deepwave needs NumPy, so installing Anaconda first is highly recommended. It also needs PyTorch, so it is also best to install that (using something like conda install pytorch-cpu -c pytorch) before installing Deepwave. You can then install the latest release of Deepwave using

pip install deepwave

or the latest development version using

pip install git+https://github.com/ar4/deepwave.git

Usage

Deepwave can do two things: forward modeling and backpropagation.

Forward modeling

We first need to create a wave propagation module. If you are familiar with training neural networks, this is the same as defining your network (or a portion of it) and initializing the weights.

prop = deepwave.scalar.Propagator({'vp': model}, dx)

The two required arguments are the wave speed model (model), and the cell size used in the model (dx). The model should be provided as a PyTorch Float Tensor of shape [nz, (ny, (nx))]. 1D, 2D, and 3D propagators are available, with the model shape used to choose between them. If you want 1D wave propagation, then the model shape would be [nz], for example. The cell size should also be provided as a single float (if it is the same in every dimension) or a Float Tensor containing the cell size in each spatial dimension; [dz, dy, dx] in 3D, [dz] in 1D.

Now we can call our propagator.

receiver_amplitudes = prop(source_amplitudes, x_s, x_r, dt)

The required arguments are the source amplitudes, the source coordinates (x_s), the receiver coordinates (x_r), and the time interval between source samples (dt). Receiver amplitudes are returned. The source amplitudes should be a Float Tensor of shape [nt, num_shots, num_sources_per_shot], where nt is the number of time samples. Note that you can have more than one source per shot, although in active source seismic experiments there is normally only one. The source coordinates array is a Float Tensor of shape [num_shots, num_sources_per_shot, num_dimensions], where num_dimensions corresponds to the number of dimensions in the model. Similarly, the receiver coordinates array is a Float Tensor of shape [num_shots, num_receivers_per_shot, num_dimensions]. The coordinates are specified in the same units as dx and are with respect to the origin of the model. The time step size dt should be a regular float. The returned Float Tensor of shape [nt, num_shots, num_receivers_per_shot] contains receiver amplitudes.

You can run forward propagation multiple times with different inputs, for example to forward model each shot in a dataset. You may have noticed that the propagator is also designed to forward model multiple separate shots simultaneously (the num_shots dimension in the inputs). Doing so may give better performance than propagating each shot separately, but you might also run out of memory.

Here is a full example:

import math
import torch
import deepwave

dx = 5.0 # 5m in each dimension
dt = 0.004 # 4ms
nz = 20
ny = 40
nt = int(1 / dt) # 1s
peak_freq = 4
peak_source_time = 1/peak_freq

# constant 1500m/s model
model = torch.ones(nz, ny) * 1500

# one source and receiver at the same location
x_s = torch.Tensor([[[0, 20 * dx]]])
x_r = x_s.clone()

source_amplitudes = deepwave.wavelets.ricker(peak_freq, nt, dt,
                                             peak_source_time).reshape(-1, 1, 1)

prop = deepwave.scalar.Propagator({'vp': model}, dx)
receiver_amplitudes = prop(source_amplitudes, x_s, x_r, dt)

Backpropagation/Inversion (FWI)

FWI attempts to match the predicted and true receiver amplitudes by modifying the model. This is similar to modifying the weights in a neural network during training. To achieve this with Deepwave, we first specify that we will need to calculate gradients with respect to our model.

model.requires_grad = True

We also need to specify which optimization algorithm and loss/cost/objective function we wish to use.

optimizer = torch.optim.Adam([model], lr=10)
criterion = torch.nn.MSELoss()

I have chosen to use the Adam optimizer with a learning rate of 10, and mean squared error (MSE) as the cost function. The MSE cost function is provided for us by PyTorch, but I could also have made-up my own cost function and used that instead (as long as PyTorch can automatically differentiate it).

With our propagator initialized as in the forward modeling section above, we can then iteratively optimize the model.

for iter in range(num_iter):
    optimizer.zero_grad()
    receiver_amplitudes_pred = prop(source_amplitudes, x_s, x_r, dt)
    loss = criterion(receiver_amplitudes_pred, receiver_amplitudes_true)
    loss.backward()
    optimizer.step()

In each iteration, we zero the gradients, calculate the predicted receiver amplitudes from the current model, compare these with the true data using our cost function, backpropagate to calculate the gradient of the loss with respect to the model, and then update the model.

The flexibility of PyTorch means that it is easy to modify this to suit your own needs.

For a full example of using Deepwave for FWI, see this notebook.

It is also possible to perform source inversion with Deepwave. Simply set requires_grad to True on the source amplitude Tensor.

Contributing

Your help to improve Deepwave would be greatly appreciated. If you encounter any difficulty using Deepwave, please report it as a Github Issue so that it can be fixed. If you have feature suggestions or other ideas to make Deepwave better, please also report those as Github Issues. If you want to help with the coding, that would be especially wonderful. The Github Issues contain a list of things that need work. If you see one that you would like to attempt, please ask for it to be assigned to you.

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Files for deepwave, version 0.0.2
Filename, size File type Python version Upload date Hashes
Filename, size deepwave-0.0.2.tar.gz (22.9 kB) File type Source Python version None Upload date Hashes View

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring DigiCert DigiCert EV certificate Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Microsoft Microsoft PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page