import numpy as np
from matplotlib import animation
from fatiando import utils
from fatiando.seismic import wavefd
from fatiando.vis import mpl

# Set the parameters of the finite difference grid
area = [0, 20000, 0, 5000]
velocity = utils.fromimage('scipy2014.png', ranges=[3500, 4500])
density = utils.fromimage('scipy2014.png', ranges=[2200, 2500])
shape = velocity.shape
mu = wavefd.lame_mu(velocity, density)

# Make a wave source from a mexican hat wavelet
sources = [wavefd.MexHatSource(1000, 0, area, shape, 100000000, 10, delay=0.1),
           wavefd.MexHatSource(19000, 0, area, shape, 100000000, 10, delay=0.1)]

# Get the iterator for the simulation
dt = wavefd.maxdt(area, shape, velocity.max())
duration = 7
maxit = int(duration/dt)
snapshot = int(0.1/dt) # Plot a snapshot of the simulation every 0.5 seconds
simulation = wavefd.elastic_sh(mu, density, area, dt, maxit, sources, None,
                               snapshot, padding=100, taper=0.007)

# This part makes an animation using matplotlibs animation API
fig = mpl.figure(figsize=(6, 1.5))
ax = mpl.subplot(1, 1, 1)
wavefield_plt = mpl.imshow(np.zeros(shape), vmin=-10**(-5), vmax=10**(-5), cmap=mpl.cm.seismic)
mpl.imshow(np.ma.masked_array(velocity, velocity < velocity.max() - 10), cmap=mpl.cm.gray_r, alpha=0.9)
ax.axison = False
mpl.subplots_adjust(top=1, bottom=0, right=1, left=0)
# Update the plot everytime the simulation yields
def animate(i):
    # Grab the iteration number, displacment panel and seismograms
    t, u, seismograms = simulation.next()
    wavefield_plt.set_array(u) # Revert the z axis so that 0 is top
    mpl.savefig('snapshots/frame%05d.png' % (i + 1), dpi=50)
    return wavefield_plt
anim = animation.FuncAnimation(fig, animate, frames=maxit/snapshot, interval=1)
mpl.show()

! convert -coalesce -layers Optimize -delay 10 -size 300x75 -loop 0 snapshots/*.png scipy2014.gif

from IPython.display import Image
Image(url='files/scipy2014.gif')