This is a prototype of classes to run FD simulations of the wave equation for Fatiando a Terra.
Below each class implementation there are a few example use cases of what it can do. This was also an oportunity for me to play around with h5py and IPython's rich display.
Skip below to the example usage for some figures and animations.
from __future__ import division
from tempfile import NamedTemporaryFile
import time
import sys
import os
import cPickle as pickle
from abc import ABCMeta, abstractmethod
import six
import numpy as np
import scipy.sparse
from IPython.display import Image, HTML, display, display_png
from IPython.html import widgets
from IPython.core.pylabtools import print_figure
from matplotlib import animation
from fatiando import utils
from fatiando.seismic import wavefd
from fatiando.gravmag import prism
from fatiando.mesher import PrismMesh
from fatiando.vis import mpl, myv
import h5py
The class Source
represents a source function with the given amplitude, central frequency, and time delay. It is used to build the source functions, for example Ricker
.
class Source(six.with_metaclass(ABCMeta)):
def __init__(self, amp, cf, delay):
self.amp = amp
self.cf = cf
self.delay = delay
def __mul__(self, scalar):
return self.__class__(self.amp*scalar, self.cf, self.delay)
def __rmul__(self, scalar):
return self.__mul__(scalar)
def _repr_png_(self):
t = self.delay + np.linspace(-2/self.cf, 2/self.cf, 200)
fig = mpl.figure(figsize=(6, 4), facecolor='white')
fig.set_figheight(0.5*fig.get_figwidth())
mpl.title('{} wavelet'.format(self.__class__.__name__))
mpl.plot(t, self(t), '-k')
mpl.xlabel('Time')
mpl.ylabel('Amplitude')
mpl.xlim(t.min(), t.max())
mpl.grid(True)
mpl.tight_layout()
data = print_figure(fig, dpi=70)
mpl.close(fig)
return data
def __call__(self, time):
return self.value(time)
@abstractmethod
def value(self, time):
"""
Return the source functions value at a given time
"""
pass
class Ricker(Source):
def __init__(self, amp, cf, delay=0):
super(Ricker, self).__init__(amp, cf, delay)
def value(self, time):
t = (time - self.delay)
return self.amp*(1 - 2*(np.pi*self.cf*t)**2)*np.exp(-(np.pi*self.cf*t)**2)
They have a rich display feature for IPython notebooks that plots the source function. So just putting as the last element of a cell will plot the source from [-2T, 2T], where T is the period.
w = Ricker(amp=10, cf=20)
w
Applying a delay shifts the wavelet in time:
w = Ricker(amp=-1, cf=50, delay=1)
w
Sources can be multiplied by a scalar to alter their amplitude. So this:
w = Ricker(amp=5, cf=10)
w
... is the same as this:
w_half = Ricker(amp=2, cf=10)
w = 2.5*w_half
w
To get the value of the source function at a given time, simply call it as function and pass the desired time:
w(0)
5.0
You can also pass a numpy array to get the value at many times:
w(np.linspace(0, 0.1, 10))
array([ 5.00000000e+00, 3.34772766e+00, 7.74673176e-02, -1.99270307e+00, -2.06326111e+00, -1.21046286e+00, -4.83647425e-01, -1.39659766e-01, -2.99569989e-02, -4.84625793e-03])
The source objects should be picklable:
pickle.loads(pickle.dumps(w))
This is the generic base class that runs a 2D simulation. It implements the run
method that time steps the simulation for the given amount of iterations. The actual timestepping will be implemented by the sub class. The displacement panels and all metadata associated with the simulation are cached in an HDF5 file using h5py. You can give the cache file name or a temporary file will be created for you automatically.
class WaveFD2D(six.with_metaclass(ABCMeta)):
def __init__(self, cachefile, spacing, dt, shape, verbose=True):
self.spacing = spacing
self.shape = shape
self.set_verbose(verbose)
self.sources = []
self.it = -1
self.size = 0
if cachefile is None:
cachefile = self.create_tmp_cache()
self.cachefile = cachefile
self.dt = dt
def create_tmp_cache(self):
tmpfile = NamedTemporaryFile(
suffix='.h5',
prefix='{}-'.format(self.__class__.__name__),
dir=os.path.curdir,
delete=False)
fname = tmpfile.name
tmpfile.close()
return fname
def set_verbose(self, verbose):
self.verbose = verbose
if verbose:
self.stream = sys.stderr
else:
self.stream = None
def get_cache(self, mode='r'):
return h5py.File(self.cachefile, mode)
def __getitem__(self, args):
with self.get_cache() as f:
data = f['panels'][args]
return data
@abstractmethod
def plot_snapshot(self, frame, **kwargs):
pass
def snapshot(self, frame, embed=False, raw=False, ax=None, **kwargs):
if ax is None:
fig = mpl.figure(facecolor='white')
ax = mpl.subplot(111)
if frame < 0:
title = self.size + frame
else:
title = frame
mpl.sca(ax)
fig = ax.get_figure()
mpl.title('Time frame {:d}'.format(title))
self.plot_snapshot(frame, **kwargs)
nz, nx = self.shape
dx, dz = nx*self.spacing, nz*self.spacing
ax.set_xlim(0, dx)
ax.set_ylim(0, dz)
ax.set_xlabel('x')
ax.set_ylabel('z')
ax.invert_yaxis()
# Check the aspect ratio of the plot and adjust figure size to match
aspect = min(self.shape)/max(self.shape)
try:
aspect /= ax.get_aspect()
except TypeError:
pass
if nx > nz:
width = 10
height = width*aspect*0.8
else:
height = 8
width = height*aspect*1.5
fig.set_size_inches(width, height)
mpl.tight_layout()
if raw or embed:
png = print_figure(fig, dpi=70)
mpl.close(fig)
if raw:
return png
elif embed:
return Image(png)
else:
return fig
def _repr_png_(self):
return self.snapshot(-1, raw=True)
def explore(self, **kwargs):
plotargs = kwargs
def plot(Frame):
image = Image(self.snapshot(Frame, raw=True, **plotargs))
display(image)
return image
slider = widgets.IntSliderWidget(min=0, max=self.it, step=1, value=self.it,
description="Frame")
widget = widgets.interact(plot, Frame=slider)
return widget
@abstractmethod
def init_cache(self, iterations):
pass
@abstractmethod
def expand_cache(self, iterations):
pass
@abstractmethod
def init_panels(self):
pass
@abstractmethod
def timestep(self, panels, tm1, t, tp1):
pass
@abstractmethod
def cache_panels(self, u, tp1, iteration, simul_size):
pass
def run(self, iterations):
nz, nx = self.shape
dz, dx = self.spacing, self.spacing
# Initialize the cache on the first run
if self.size == 0:
self.init_cache(iterations)
else:
self.expand_cache(iterations)
u = self.init_panels()
if self.verbose:
# The size of the progress status bar
places = 50
self.stream.write(''.join(['|', '-'*places, '|', ' 0%']))
self.stream.flush()
nprinted = 0
start_time = time.clock()
for iteration in xrange(iterations):
t, tm1 = iteration%2, (iteration + 1)%2
tp1 = tm1
self.it += 1
self.timestep(u, tm1, t, tp1, self.it)
self.size += 1
self.cache_panels(u, tp1, self.it, self.size)
# Update the status bar
if self.verbose:
percent = int(round(100*(iteration + 1)/iterations))
n = int(round(0.01*percent*places))
if n > nprinted:
self.stream.write(''.join(['\r|', '#'*n, '-'*(places - n), '|', '%3d%s' % (percent, '%')]))
self.stream.flush()
nprinted = n
# Make sure the progress bar ends in 100 percent
if self.verbose:
self.stream.write(''.join(
['\r|', '#'*places, '|', '100%',
' Ran {:d} iterations in {:g} seconds.'.format(iterations, time.clock() - start_time)]))
self.stream.flush()
Now I can use the WaveFD2D
class to make an elastic horizontal S wave solver. It has to implement the timestep
method, among others. It uses the Cython functions in fatiando.seismic.wavefd
to do the timestepping and applying the boundary conditions.
class ElasticSH(WaveFD2D):
def __init__(self, velocity, density, spacing, cachefile=None, dt=None, padding=50, taper=0.007, verbose=True):
super(ElasticSH, self).__init__(cachefile, spacing, dt, velocity.shape, verbose)
self.density = density
self.velocity = velocity
self.mu = wavefd.lame_mu(velocity, density)
self.padding = padding
self.taper = taper
if self.dt is None:
self.dt = self.maxdt()
def maxdt(self):
nz, nx = self.shape
ds = self.spacing
return 0.6*wavefd.maxdt([0, nx*ds, 0, nz*ds], self.shape, self.velocity.max())
def plot_snapshot(self, frame, **kwargs):
with h5py.File(self.cachefile) as f:
data = f['panels'][frame]
scale = np.abs(data).max()
nz, nx = self.shape
dx, dz = nx*self.spacing, nz*self.spacing
extent = [0, dx, dz, 0]
if 'cmap' not in kwargs:
kwargs['cmap'] = mpl.cm.seismic
mpl.imshow(data, extent=extent, vmin=-scale, vmax=scale, **kwargs)
mpl.colorbar(pad=0, aspect=30).set_label('Displacement')
def add_point_source(self, position, wavelet):
self.sources.append([position, wavelet])
@staticmethod
def from_cache(fname, verbose=True):
with h5py.File(fname, 'r') as f:
vel = f['velocity']
dens = f['density']
panels = f['panels']
spacing = panels.attrs['spacing']
dt = panels.attrs['dt']
padding = panels.attrs['padding']
taper = panels.attrs['taper']
sim = ElasticSH(vel[:], dens[:], spacing, dt=dt, padding=padding,
taper=taper, cachefile=fname)
sim.size = panels.attrs['size']
sim.it = panels.attrs['iterations']
sim.sources = pickle.loads(f['sources'].value.tostring())
sim.set_verbose(verbose)
return sim
def init_cache(self, panels, chunks=None, compression='lzf', shuffle=True):
nz, nx = self.shape
if chunks is None:
chunks = (1, nz//10, nx//10)
with self.get_cache(mode='w') as f:
nz, nx = self.shape
dset = f.create_dataset('panels', (panels, nz, nx),
maxshape=(None, nz, nx),
chunks=chunks,
compression=compression,
shuffle=shuffle,
dtype=np.float)
dset.attrs['shape'] = self.shape
dset.attrs['size'] = self.size
dset.attrs['iterations'] = self.it
dset.attrs['spacing'] = self.spacing
dset.attrs['dt'] = self.dt
dset.attrs['padding'] = self.padding
dset.attrs['taper'] = self.taper
dset = f.create_dataset('velocity', data=self.velocity)
dset.attrs['spacing'] = self.spacing
dset.attrs['shape'] = self.shape
dset = f.create_dataset('density', data=self.density)
dset.attrs['spacing'] = self.spacing
dset.attrs['shape'] = self.shape
dset = f.create_dataset('sources', data=np.void(pickle.dumps(self.sources)))
dset.attrs['len'] = len(self.sources)
def expand_cache(self, iterations):
with self.get_cache(mode='a') as f:
cache = f['panels']
cache.resize(self.size + iterations, axis=0)
def cache_panels(self, u, tp1, iteration, simul_size):
# Save the panel to disk
with self.get_cache(mode='a') as f:
cache = f['panels']
cache[simul_size - 1] = u[tp1]
cache.attrs['size'] = simul_size
cache.attrs['iterations'] = iteration
def init_panels(self):
if self.size == 0:
nz, nx = self.shape
u = np.zeros((2, nz, nx), dtype=np.float)
else:
# Get the last two panels from the cache
with self.get_cache() as f:
cache = f['panels']
u = cache[self.size - 2 : self.size][::-1]
return u
def timestep(self, u, tm1, t, tp1, iteration):
nz, nx = self.shape
ds = self.spacing
wavefd._step_elastic_sh(u[tp1], u[t], u[tm1], 3, nx - 3, 3, nz - 3,
self.dt, ds, ds, self.mu, self.density)
wavefd._apply_damping(u[t], nx, nz, self.padding, self.taper)
wavefd._nonreflexive_sh_boundary_conditions(
u[tp1], u[t], nx, nz, self.dt, ds, ds, self.mu, self.density)
wavefd._apply_damping(u[tp1], nx, nz, self.padding, self.taper)
for pos, src in self.sources:
i, j = pos
u[tp1, i, j] += src(iteration*self.dt)
def animate(self, every=1, cutoff=None, ax=None, cmap=mpl.cm.seismic, embed=False, fps=10, dpi=70, writer='avconv', **kwargs):
if ax is None:
mpl.figure(facecolor='white')
ax = mpl.subplot(111)
ax.set_xlabel('x')
ax.set_ylabel('z')
fig = ax.get_figure()
nz, nx = self.shape
# Check the aspect ratio of the plot and adjust figure size to match
aspect = min(self.shape)/max(self.shape)
try:
aspect /= ax.get_aspect()
except TypeError:
pass
if nx > nz:
width = 10
height = width*aspect*0.8
else:
height = 10
width = height*aspect*1.5
fig.set_size_inches(width, height)
# Separate the arguments for imshow
imshow_args = dict(cmap=cmap)
if cutoff is not None:
imshow_args['vmin'] = -cutoff
imshow_args['vmax'] = cutoff
wavefield = ax.imshow(np.zeros(self.shape), **imshow_args)
fig.colorbar(wavefield, pad=0, aspect=30).set_label('Displacement')
ax.set_title('iteration: 0')
frames = self.size//every
def plot(i):
ax.set_title('iteration: {:d}'.format(i*every))
u = self[i*every]
wavefield.set_array(u)
return wavefield
anim = animation.FuncAnimation(fig, plot, frames=frames, **kwargs)
if embed:
return display_animation(anim, fps=fps, dpi=dpi, writer=writer)
else:
mpl.show()
return anim
This code (adapted from the yt project docs) saves a matplotlib animation and embeds it into the notebook. Magic. Pure magic.
def anim_to_html(anim, fps=6, dpi=30, writer='avconv'):
VIDEO_TAG = """
<video controls>
<source src="data:video/webm;base64,{0}" type="video/webm">
Your browser does not support the video tag.
</video>"""
if not hasattr(anim, '_encoded_video'):
with NamedTemporaryFile(suffix='.webm') as f:
anim.save(f.name, fps=fps, dpi=dpi, writer=writer, extra_args=['-vcodec', 'libvpx'])
video = open(f.name, "rb").read()
anim._encoded_video = video.encode("base64")
return VIDEO_TAG.format(anim._encoded_video)
def display_animation(anim, fps=6, dpi=30, writer='avconv'):
mpl.close(anim._fig)
return HTML(anim_to_html(anim, fps, dpi, writer))
First, I'll need to make the S wave velocity and density grids for the simulation.
shape = (300, 800)
velocity = 3000*np.ones(shape)
density = 2200*np.ones(shape)
Then, I can create my simulation class and add a Ricker wavelet point source to the middle of the grid.
sim = ElasticSH(velocity, density, spacing=10)
sim.add_point_source((shape[0]//2, shape[1]//2), Ricker(5, 20, 1/20))
The HDF5 cache file is created in the current directory. It will have the class name and a random-looking string at the end. For me, the cache file is:
sim.cachefile
'/home/leo/src/prototypes/ElasticSH-wAJcgJ.h5'
Running the simulation is as simple as calling run
with the desired number of iterations. It will even print an ASCII progress bar. You can turn this off by setting verbose=False
when making the simulation instance.
sim.run(200)
|##################################################|100% Ran 200 iterations in 6.56431 seconds.
The cache file will be a bit large because it stores all of the frames for the simulation. But not too large because of the awesome compression available in HDF5.
!du -h $sim.cachefile
170M /home/leo/src/prototypes/ElasticSH-wAJcgJ.h5
The WaveFD
class defines some rich display for the simulations. By default, the display for the simulation object is the last time frame computed.
sim
You can browse each time frame of the simulation using the explore
method. This creates an IPython widget that allows you to interactively explore each time frame of the simulation using a slider. This will only work if you're running the notebook localy. In nbviewer, only the last frame should appear.
sim.explore()
<function __main__.plot>
You can inspect individual frames by calling snapshot
. Passing embed=True
will display the image in the notebook. Otherwise, it will call matplotlib.pyplot.show()
to make the plot window popup (if you're not using %matplotlib inline
).
sim.snapshot(frame=100, embed=True)
You can get the displacement data for each frame of the simulation by indexing the simulation object.
sim[100]
array([[ 0., 0., 0., ..., 0., 0., 0.], [ 0., 0., 0., ..., 0., 0., 0.], [ 0., 0., 0., ..., 0., 0., 0.], ..., [ 0., 0., 0., ..., 0., 0., 0.], [ 0., 0., 0., ..., 0., 0., 0.], [ 0., 0., 0., ..., 0., 0., 0.]])
The first dimension is time, so the above is a 2D array corresponding to a time step. The second dimension is z and the third is x.
You can slice the simulation object as if it where a 3D numpy array:
part = sim[100, 120:180, 370:430]
part
array([[ 2.60317298e-07, 8.74915863e-07, -6.45496603e-07, ..., -8.18761962e-06, -6.45496603e-07, 8.74915863e-07], [ 8.74915863e-07, -7.26305098e-07, -8.73484007e-06, ..., -1.31911982e-05, -8.73484007e-06, -7.26305098e-07], [ -6.45496603e-07, -8.73484007e-06, -1.31072982e-05, ..., 3.14237808e-05, -1.31072982e-05, -8.73484007e-06], ..., [ -8.18761962e-06, -1.31911982e-05, 3.14237808e-05, ..., 1.59627536e-04, 3.14237808e-05, -1.31911982e-05], [ -6.45496603e-07, -8.73484007e-06, -1.31072982e-05, ..., 3.14237808e-05, -1.31072982e-05, -8.73484007e-06], [ 8.74915863e-07, -7.26305098e-07, -8.73484007e-06, ..., -1.31911982e-05, -8.73484007e-06, -7.26305098e-07]])
fig = mpl.figure(facecolor='white')
mpl.imshow(part, cmap=mpl.cm.seismic)
mpl.colorbar(pad=0)
# Some magic to embed the image in the notebook
mpl.close(fig)
Image(print_figure(fig, dpi=70))
To animate the simulation, use the animate
method. The code below will animate using matplotlib's animation package. It will show every 10 frames of the simulation. (A plot window should popup with the animation)
anim = sim.animate(every=10, cutoff=1, interval=100)
sim.animate(every=5, cutoff=1, blit=True, fps=20, dpi=50, embed=True)
To resume the simulation, call run
again.
sim.run(150)
|##################################################|100% Ran 150 iterations in 6.75071 seconds.
sim
You can reload the simulation from the cache file at a later time. So as long as you have the cache file, your simulation is safe.
fname = sim.cachefile
del sim
reloaded = ElasticSH.from_cache(fname)
reloaded
It will contain the whole simulation history, just like it did before.
reloaded.snapshot(frame=100, embed=True)
And the best thing is, you can even resume the simulation as if nothing had happened! (warning: this will write to the cache file so the old sim
object might not work anymore).
reloaded.run(200)
|##################################################|100% Ran 200 iterations in 7.85226 seconds.
reloaded
Make sure that simulation classes can be pickled (important for running simulations in parallel).
pickle.loads(pickle.dumps(reloaded))
--------------------------------------------------------------------------- TypeError Traceback (most recent call last) <ipython-input-26-2a8a38a06fc4> in <module>() ----> 1 pickle.loads(pickle.dumps(reloaded)) /home/leo/bin/anaconda/lib/python2.7/copy_reg.pyc in _reduce_ex(self, proto) 68 else: 69 if base is self.__class__: ---> 70 raise TypeError, "can't pickle %s objects" % base.__name__ 71 state = base(self) 72 args = (self.__class__, base, state) TypeError: can't pickle StringIO objects
Turns out you can't pickle the simulation because of the reference to sys.stderr
that it keeps for the progress bar. To pickle the simulation, first turn off verbosity.
reloaded.set_verbose(False)
pickle.loads(pickle.dumps(reloaded))
The cache file stores the source functions as well. So even if you interupt the simulation before a source goes off, you can be sure it will be triggered once you reload the simulation from cache.
sim = ElasticSH(velocity, density, spacing=10)
sim.add_point_source((shape[0]//2, shape[1]//2), Ricker(5, 20, 1/20))
# Add a delayed source
sim.add_point_source((shape[0]//3, shape[1]//3), Ricker(5, 20, 0.3))
sim.run(200)
|##################################################|100% Ran 200 iterations in 6.91455 seconds.
Check that the second source has not been fired:
sim
fname = sim.cachefile
del sim
reloaded = ElasticSH.from_cache(fname)
reloaded.run(200)
|##################################################|100% Ran 200 iterations in 8.33649 seconds.
reloaded
This class simulates P and SV elastic waves and behaves the same as ElasticSH
. The code is larger and slower because it has to calculate the x and z displacements (ElasticSH
only calculates y displacements).
class ElasticPSV(WaveFD2D):
def __init__(self, pvel, svel, density, spacing, cachefile=None, dt=None, padding=50, taper=0.007, verbose=True):
super(ElasticPSV, self).__init__(cachefile, spacing, dt, pvel.shape, verbose)
self.pvel = pvel
self.svel = svel
self.density = density
self.mu = wavefd.lame_mu(svel, density)
self.lamb = wavefd.lame_lamb(pvel, svel, density)
self.padding = padding
self.taper = taper
self.make_free_surface_matrices()
if self.dt is None:
self.dt = self.maxdt()
def maxdt(self):
nz, nx = self.shape
ds = self.spacing
return 0.6*wavefd.maxdt([0, nx*ds, 0, nz*ds], self.shape, self.pvel.max())
def __getitem__(self, args):
with self.get_cache() as f:
ux = f['xpanels'][args]
uz = f['zpanels'][args]
return [ux, uz]
def add_blast_source(self, position, wavelet):
nz, nx = self.shape
i, j = position
amp = 1/(2**0.5)
locations = [
[i - 1, j , 0, -1],
[i + 1, j , 0, 1],
[i , j - 1, -1, 0],
[i , j + 1, 1, 0],
[i - 1, j - 1, -amp, -amp],
[i + 1, j - 1, -amp, amp],
[i - 1, j + 1, amp, -amp],
[i + 1, j + 1, amp, amp],
]
for k, l, xamp, zamp in locations:
if k >= 0 and k < nz and l >= 0 and l < nx:
xwav = xamp*wavelet
zwav = zamp*wavelet
self.sources.append([[k, l], xwav, zwav])
def add_point_source(self, position, dip, wavelet):
d2r = np.pi/180
xamp = np.cos(d2r*dip)
zamp = np.sin(d2r*dip)
self.sources.append([position, xamp*wavelet, zamp*wavelet])
def plot_snapshot(self, frame, **kwargs):
with h5py.File(self.cachefile) as f:
ux = f['xpanels'][frame]
uz = f['zpanels'][frame]
plottype = kwargs.get('plottype', ['wavefield'])
nz, nx = self.shape
dx, dz = nx*self.spacing, nz*self.spacing
if 'wavefield' in plottype:
extent = [0, dx, dz, 0]
cmap = kwargs.get('cmap', mpl.cm.seismic)
p = np.empty(self.shape, dtype=np.float)
s = np.empty(self.shape, dtype=np.float)
wavefd._xz2ps(ux, uz, p, s, nx, nz, self.spacing, self.spacing)
data = p + s
scale = np.abs(data).max()
vmin = kwargs.get('vmin', -scale)
vmax = kwargs.get('vmax', scale)
mpl.imshow(data, cmap=cmap, extent=extent, vmin=vmin, vmax=vmax)
mpl.colorbar(pad=0, aspect=30).set_label('Divergence + Curl')
if 'particles' in plottype:
every_particle = kwargs.get('every_particle', 5)
markersize = kwargs.get('markersize', 1)
scale = kwargs.get('scale', 1)
xs = np.linspace(0, dx, nx)[::every_particle]
zs = np.linspace(0, dz, nz)[::every_particle]
x, z = np.meshgrid(xs, zs)
x += scale*ux[::every_particle, ::every_particle]
z += scale*uz[::every_particle, ::every_particle]
mpl.plot(x, z, '.k', markersize=markersize)
if 'vectors' in plottype:
every_particle = kwargs.get('every_particle', 5)
scale = kwargs.get('scale', 1)
linewidth = kwargs.get('linewidth', 0.1)
xs = np.linspace(0, dx, nx)[::every_particle]
zs = np.linspace(0, dz, nz)[::every_particle]
x, z = np.meshgrid(xs, zs)
mpl.quiver(x, z,
ux[::every_particle, ::every_particle],
uz[::every_particle, ::every_particle],
scale=1/scale, linewidth=linewidth,
pivot='tail', angles='xy', scale_units='xy')
def init_cache(self, panels, chunks=None, compression='lzf', shuffle=True):
nz, nx = self.shape
if chunks is None:
chunks = (1, nz//10, nx//10)
with self.get_cache(mode='w') as f:
nz, nx = self.shape
dset = f.create_dataset('xpanels', (panels, nz, nx),
maxshape=(None, nz, nx),
chunks=chunks,
compression=compression,
shuffle=shuffle,
dtype=np.float)
dset.attrs['shape'] = self.shape
dset.attrs['size'] = self.size
dset.attrs['iterations'] = self.it
dset.attrs['spacing'] = self.spacing
dset.attrs['dt'] = self.dt
dset.attrs['padding'] = self.padding
dset.attrs['taper'] = self.taper
dset = f.create_dataset('zpanels', (panels, nz, nx),
maxshape=(None, nz, nx),
chunks=chunks,
compression=compression,
shuffle=shuffle,
dtype=np.float)
dset.attrs['shape'] = self.shape
dset.attrs['size'] = self.size
dset.attrs['iterations'] = self.it
dset.attrs['spacing'] = self.spacing
dset.attrs['dt'] = self.dt
dset.attrs['padding'] = self.padding
dset.attrs['taper'] = self.taper
dset = f.create_dataset('pvel', data=self.pvel,
chunks=chunks[1:],
compression=compression,
shuffle=shuffle)
dset.attrs['spacing'] = self.spacing
dset.attrs['shape'] = self.shape
dset = f.create_dataset('svel', data=self.svel,
chunks=chunks[1:],
compression=compression,
shuffle=shuffle)
dset.attrs['spacing'] = self.spacing
dset.attrs['shape'] = self.shape
dset = f.create_dataset('density', data=self.density,
chunks=chunks[1:],
compression=compression,
shuffle=shuffle)
dset.attrs['spacing'] = self.spacing
dset.attrs['shape'] = self.shape
dset = f.create_dataset('sources', data=np.void(pickle.dumps(self.sources)))
dset.attrs['len'] = len(self.sources)
@staticmethod
def from_cache(fname):
with h5py.File(fname, 'r') as f:
pvel = f['pvel']
svel = f['svel']
dens = f['density']
panels = f['xpanels']
spacing = panels.attrs['spacing']
dt = panels.attrs['dt']
padding = panels.attrs['padding']
taper = panels.attrs['taper']
sim = ElasticPSV(pvel[:], svel[:], dens[:], spacing, dt=dt, padding=padding,
taper=taper, cachefile=fname)
sim.size = panels.attrs['size']
sim.it = panels.attrs['iterations']
sim.sources = pickle.loads(f['sources'].value.tostring())
return sim
def cache_panels(self, u, tp1, iteration, simul_size):
# Save the panel to disk
with self.get_cache(mode='a') as f:
ux, uz = u
xpanels = f['xpanels']
zpanels = f['zpanels']
xpanels[simul_size - 1] = ux[tp1]
zpanels[simul_size - 1] = uz[tp1]
xpanels.attrs['size'] = simul_size
xpanels.attrs['iterations'] = iteration
zpanels.attrs['size'] = simul_size
zpanels.attrs['iterations'] = iteration
def expand_cache(self, iterations):
with self.get_cache(mode='a') as f:
f['xpanels'].resize(self.size + iterations, axis=0)
f['zpanels'].resize(self.size + iterations, axis=0)
def init_panels(self):
if self.size == 0:
nz, nx = self.shape
ux = np.zeros((2, nz, nx), dtype=np.float)
uz = np.zeros((2, nz, nx), dtype=np.float)
else:
# Get the last two panels from the cache
with self.get_cache() as f:
# Reverse the array becase the later time comes first in ux, uz
# Need to copy them and reorder because the timestep function takes
# the whole ux array and complains that it isn't C contiguous
# Could change the timestep to pass ux[tp1], etc, like in ElasticSH
ux = np.copy(f['xpanels'][self.size - 2 : self.size][::-1], order='C')
uz = np.copy(f['zpanels'][self.size - 2 : self.size][::-1], order='C')
return [ux, uz]
def make_free_surface_matrices(self):
# Pre-compute the matrices required for the free-surface BC
nz, nx = self.shape
dzdx = 1
identity = scipy.sparse.identity(nx)
B = scipy.sparse.eye(nx, nx, k=1) - scipy.sparse.eye(nx, nx, k=-1)
gamma = scipy.sparse.spdiags(self.lamb[0]/(self.lamb[0] + 2*self.mu[0]), [0], nx, nx)
Mx1 = identity - 0.0625*(dzdx**2)*B*gamma*B
Mx2 = identity + 0.0625*(dzdx**2)*B*gamma*B
Mx3 = 0.5*dzdx*B
Mz1 = identity - 0.0625*(dzdx**2)*gamma*B*B
Mz2 = identity + 0.0625*(dzdx**2)*gamma*B*B
Mz3 = 0.5*dzdx*gamma*B
self.Mx = [Mx1, Mx2, Mx3]
self.Mz = [Mz1, Mz2, Mz3]
def timestep(self, u, tm1, t, tp1, iteration):
nz, nx = self.shape
ds = self.spacing
ux, uz = u
wavefd._step_elastic_psv(ux, uz, tp1, t, tm1, 1, nx - 1, 1, nz - 1,
self.dt, ds, ds,
self.mu, self.lamb, self.density)
wavefd._apply_damping(ux[t], nx, nz, self.padding, self.taper)
wavefd._apply_damping(uz[t], nx, nz, self.padding, self.taper)
# Free-surface boundary conditions
Mx1, Mx2, Mx3 = self.Mx
Mz1, Mz2, Mz3 = self.Mz
ux[tp1,0,:] = scipy.sparse.linalg.spsolve(Mx1, Mx2*ux[tp1,1,:] + Mx3*uz[tp1,1,:])
uz[tp1,0,:] = scipy.sparse.linalg.spsolve(Mz1, Mz2*uz[tp1,1,:] + Mz3*ux[tp1,1,:])
wavefd._nonreflexive_psv_boundary_conditions(
ux, uz, tp1, t, tm1, nx, nz, self.dt, ds, ds,
self.mu, self.lamb, self.density)
wavefd._apply_damping(ux[tp1], nx, nz, self.padding, self.taper)
wavefd._apply_damping(uz[tp1], nx, nz, self.padding, self.taper)
for pos, xsrc, zsrc in self.sources:
i, j = pos
ux[tp1, i, j] += xsrc(iteration*self.dt)
uz[tp1, i, j] += zsrc(iteration*self.dt)
def animate(self, every=1, plottype=['wavefield'], cutoff=None, cmap=mpl.cm.seismic,
scale=1, every_particle=5,
ax=None, interval=100, embed=False, blit=False,
fps=10, dpi=70, writer='avconv', **kwargs):
nz, nx = self.shape
dx, dz = nx*self.spacing, nz*self.spacing
if ax is None:
mpl.figure(facecolor='white')
ax = mpl.subplot(111)
ax.set_xlabel('x')
ax.set_ylabel('z')
ax.set_xlim(0, dx)
ax.set_ylim(0, dz)
ax.invert_yaxis()
fig = ax.get_figure()
wavefield = None
particles = None
vectors = None
if 'wavefield' in plottype:
extent = [0, dx, dz, 0]
p = np.empty(self.shape, dtype=np.float)
s = np.empty(self.shape, dtype=np.float)
imshow_args = dict(cmap=cmap, extent=extent)
if cutoff is not None:
imshow_args['vmin'] = -cutoff
imshow_args['vmax'] = cutoff
wavefield = ax.imshow(np.zeros(self.shape), **imshow_args)
fig.colorbar(wavefield, pad=0, aspect=30).set_label('Divergence + Curl')
if 'particles' in plottype or 'vectors' in plottype:
xs = np.linspace(0, dx, nx)[::every_particle]
zs = np.linspace(0, dz, nz)[::every_particle]
x, z = np.meshgrid(xs, zs)
if 'particles' in plottype:
markersize = kwargs.get('markersize', 1)
stype = kwargs.get('style', '.k')
particles, = mpl.plot(x.ravel(), z.ravel(), style, markersize=markersize)
if 'vectors' in plottype:
linewidth = kwargs.get('linewidth', 0.1)
vectors = mpl.quiver(x, z, np.zeros_like(x), np.zeros_like(z),
scale=1/scale, linewidth=linewidth,
pivot='tail', angles='xy',
scale_units='xy')
# Check the aspect ratio of the plot and adjust figure size to match
aspect = min(self.shape)/max(self.shape)
try:
aspect /= ax.get_aspect()
except TypeError:
pass
if nx > nz:
width = 10
height = width*aspect*0.8
else:
height = 8
width = height*aspect*1.5
fig.set_size_inches(width, height)
def plot(i):
ax.set_title('iteration: {:d}'.format(i*every))
ux, uz = self[i*every]
if wavefield is not None:
wavefd._xz2ps(ux, uz, p, s, nx, nz, self.spacing, self.spacing)
wavefield.set_array(p + s)
if particles is not None or vectors is not None:
ux = ux[::every_particle, ::every_particle]
uz = uz[::every_particle, ::every_particle]
if particles is not None:
particles.set_data(x.ravel() + scale*ux.ravel(),
z.ravel() + scale*uz.ravel())
if vectors is not None:
vectors.set_UVC(ux, uz)
return wavefield, particles, vectors
frames = self.size//every
anim = animation.FuncAnimation(fig, plot, frames=frames, blit=blit, interval=interval)
if embed:
return display_animation(anim, fps=fps, dpi=dpi, writer=writer)
else:
mpl.show()
return anim
For this, we'll need both P and S wave velocities.
shape = (500, 300)
pvel = 4000*np.ones(shape)
svel = 3000*np.ones(shape)
density = 2200*np.ones(shape)
The source function for this kind of wave needs to have an x and z component. When adding a point source, you only have to specify the dip of the source (with respect to the horizontal) and the source wavelet. The class will take care of projecting the source onto the x and z directions.
sim = ElasticPSV(pvel, svel, density, spacing=10)
sim.add_point_source((shape[0]//2, shape[1]//2), dip=45, wavelet=Ricker(5, 20, 1/20))
sim.cachefile
'/home/leo/src/prototypes/ElasticPSV-BeqfqS.h5'
Running the simulation is the same as it was with ElasticSH
.
sim.run(300)
|##################################################|100% Ran 300 iterations in 13.2192 seconds.
The cachefile will be larger because it has to store 2 panels per iteration (ux and uz, the x and z displacements).
!du -h $sim.cachefile
352M /home/leo/src/prototypes/ElasticPSV-BeqfqS.h5
The rich display features all work for PSV as well. However, PSV simulations would have 2 panels to show. It would be better to condense them into a single representation. A very convinient way to show P and S waves is to plot the divergence plus the curl of the displacement vector field. The curl will show the S waves and the divergence the P waves.
sim
You can also use the interactive widget by calling explore
.
sim.explore()
And you can plot each time frame by calling snapshot
.
sim.snapshot(250, embed=True)
However, plotting methods in ElasticPSV
can take some extra arguments. For example, you can use plottype=['vectors']
to plot the displacement vector field instead of the divergence+curl (plottype=['wavefield']
).
sim.snapshot(250, embed=True, plottype=['vectors'], scale=200)
The vector field is a good way of showing the difference in the vibration direction of P and S waves. Notice that the outer ring has displacement along the propagation direction, while in the inner ring it is perpendicular.
You can also combine different plot types.
sim.snapshot(250, embed=True, plottype=['vectors', 'wavefield'], scale=200)
The same arguments will work for explore
and animate
as well.
sim.explore(plottype=['vectors', 'wavefield'], scale=200)
sim.animate(every=10, plottype=['vectors', 'wavefield'], scale=100, cutoff=0.1, blit=True, embed=True, fps=10, dpi=50)
Calling run
again resumes the simulation.
sim.run(100)
|##################################################|100% Ran 100 iterations in 5.51067 seconds.
sim
Reloading from the cache also works and preserves your sources.
sim = ElasticPSV(pvel, svel, density, spacing=10)
sim.add_point_source((shape[0]//2, shape[1]//2), dip=45, wavelet=Ricker(5, 20, 1/20))
# Add a delayed source deeper down
sim.add_point_source((2*shape[0]//3, shape[1]//3), dip=-45, wavelet=Ricker(5, 20, 0.25))
sim.run(200)
|##################################################|100% Ran 200 iterations in 8.26454 seconds.
Second source hasn't fired yet:
sim
Reloading from cache and resuming will eventualy trigger the second source.
fname = sim.cachefile
del sim
reloaded = ElasticPSV.from_cache(fname)
reloaded.run(300)
|##################################################|100% Ran 300 iterations in 17.0418 seconds.
reloaded.snapshot(-1, embed=True, plottype=['wavefield', 'vectors'], scale=200)
ElasticPSV
also has the a blast source. This source explodes in all directions, creating only P waves.
sim = ElasticPSV(pvel, svel, density, spacing=10)
sim.add_blast_source((shape[0]//2, shape[1]//2), wavelet=Ricker(3, 20, 1/20))
sim.run(300)
|##################################################|100% Ran 300 iterations in 13.4258 seconds.
sim.snapshot(-1, embed=True, plottype=['wavefield', 'vectors'], every_particle=3, scale=100)
sim.animate(every=10, cutoff=0.05, plottype=['wavefield', 'vectors'], scale=100,
blit=True, embed=True, fps=10, dpi=50)
Check that ElasticPSV can be pickled.
sim.set_verbose(False)
pickle.loads(pickle.dumps(sim))