In this notebook, let's code up some 2D wave propagation.
Inspiration : https://dionhaefner.github.io/2016/11/suck-less-scientific-python-part-2-efficient-number-crunching/
(which in turn comes from here: http://matthewrocklin.com/blog/work/2018/01/30/the-case-for-numba)
Explain that we also want to try all the nice visualization goodness coming from Holoviews.
import numpy as np
from numba import njit
import holoviews as hv
hv.extension('matplotlib')
@njit
def step(u_curr, u_prev, c, dt, dx):
"""Unitary finite difference propagation step."""
u_next = np.zeros_like(u_curr)
for i in range(1, u_curr.shape[0] - 1):
for j in range(1, u_curr.shape[1] - 1):
u_next[i, j] = 2 * u_curr[i, j] - u_prev[i, j] + c**2 * dt**2 / dx**2 * (- 4 * u_curr[i, j] + \
u_curr[i+1, j] + u_curr[i-1, j] + \
u_curr[i, j+1] + u_curr[i, j-1])
return u_next
Let's setup the grid and step sizes.
dx = dy = 0.005
x = np.arange(0, 1 + dx, dx)
y = np.arange(0, 1 + dy, dy)
X, Y = np.meshgrid(x, y)
bounds = (x.min(), y.min(), y.max(), x.max())
source_loc = (0.5, 0.5)
u_curr = np.zeros_like(X)
u_prev = u_curr.copy()
u_curr = np.exp(-((X - source_loc[0])**2 + (Y - source_loc[1])**2) * 5000.)
celerity = 1.
dt = 0.1 * 1/2 * dx / celerity
def make_steps(nstep=10):
global u_curr, u_prev, u_next
for _ in range(nstep):
u_next = step(u_curr, u_prev, celerity, dt, dx)
u_prev = u_curr.copy()
u_curr = u_next.copy()
return u_curr
make_steps(500)
array([[ 0.00000000e+000, 0.00000000e+000, 0.00000000e+000, ..., 0.00000000e+000, 0.00000000e+000, 0.00000000e+000], [ 0.00000000e+000, 5.44361463e-212, 4.95809542e-210, ..., 4.95809542e-210, 5.44361463e-212, 0.00000000e+000], [ 0.00000000e+000, 4.95809542e-210, 4.48537552e-208, ..., 4.48537552e-208, 4.95809542e-210, 0.00000000e+000], ..., [ 0.00000000e+000, 4.95809542e-210, 4.48537552e-208, ..., 4.48537552e-208, 4.95809542e-210, 0.00000000e+000], [ 0.00000000e+000, 5.44361463e-212, 4.95809542e-210, ..., 4.95809542e-210, 5.44361463e-212, 0.00000000e+000], [ 0.00000000e+000, 0.00000000e+000, 0.00000000e+000, ..., 0.00000000e+000, 0.00000000e+000, 0.00000000e+000]])
%%opts Image [colorbar=True] (cmap='seismic')
img = hv.Image(u_curr, bounds=bounds).redim.range(z=(-10, 10))
img
Unexpected plot option 'width' for Image in loaded backend 'matplotlib'. Similar keywords in the currently active 'matplotlib' renderer are: ['cbar_width'] If you believe this keyword is correct, please make sure the backend has been imported or loaded with the hv.extension.Unexpected plot option 'width' for Image in loaded backend 'matplotlib'. Similar keywords in the currently active 'matplotlib' renderer are: ['cbar_width'] If you believe this keyword is correct, please make sure the backend has been imported or loaded with the hv.extension.
Looks nice!
Let's animate this using a HoloMap:
u_curr = np.zeros_like(X)
u_prev = u_curr.copy()
u_curr = np.exp(-((X - source_loc[0])**2 + (Y - source_loc[1])**2) * 500.)
hmap1 = hv.HoloMap({i: hv.Image(make_steps(1000), bounds=bounds).options(colorbar=True, cmap='seismic').redim.range(z=(-10, 10)) for i in list(range(10))}, kdims='index')
hmap1
We can select the best snapshots of the wave below:
(hmap1[4] + hmap1[7] + hmap1[8] + hmap1[9]).cols(2)
By playing with the initial conditions and visualizing method we can get other cool results:
u_curr = np.zeros_like(X)
u_prev = u_curr.copy()
u_curr = np.exp(-((X - source_loc[0])**2 + (Y - source_loc[1])**2) * 5000.)
hmap2 = hv.HoloMap({i: hv.Image(make_steps(2000)).options(colorbar=True, cmap='seismic').redim.range(z=(-2, 2)) for i in list(range(10))}, kdims='index')
hmap2
(hmap2[4] + hmap2[7] + hmap2[8] + hmap2[9]).cols(2)
Another fun variation on this is to put in a source that is a plane wave.
To do this, we have to define a wave vector and also work a little bit on the initial conditions, since we want to analytically define the initial pressure as well as the initial pressure change.
dx = dy = 0.005
x = np.arange(0, 3 + dx, dx)
y = np.arange(0, 1 + dy, dy)
X, Y = np.meshgrid(x, y)
bounds = (x.min(), y.min(), x.max(), y.max())
dt = 0.1 * 1/2 * dx / celerity
@njit
def step_initial(u_curr, du_currdt, c, dt, dx):
"""Initial finite difference propagation step."""
u_next = np.zeros_like(u_curr)
for i in range(1, u_curr.shape[0] - 1):
for j in range(1, u_curr.shape[1] - 1):
u_next[i, j] = u_curr[i, j] + dt * du_currdt[i, j] + 0.5*c**2 * dt**2 / dx**2 * (- 4 * u_curr[i, j] + \
u_curr[i+1, j] + u_curr[i-1, j] + \
u_curr[i, j+1] + u_curr[i, j-1])
return u_next
@njit
def plane_wave(k, X, Y, x0, y0, celerity, dt, dx):
kx, ky = k
absk = np.sqrt(kx**2 + ky**2)
omega = absk * celerity
nx, ny = kx/absk, ky/absk
phase = - kx * (X-x0) - ky * (Y-y0)
envelope = np.exp(- (nx * (X - x0) + ny * (Y - y0))**2 * 10)
u0 = np.sin(phase) * envelope
du0 = omega * np.cos(phase) * envelope
u1 = step_initial(u0, du0, celerity, dt, dx)
return u0, u1
First, let's setup a plane wave travelling towards the bottom.
k = (0., 10.)
u_prev, u_curr = plane_wave(k, X, Y, 0.1, 0.2, celerity, dt, dx)
%%opts Image [colorbar=True fig_size=400] (cmap='seismic')
make_steps(200)
img = hv.Image(u_next, bounds=bounds).redim.range(z=(-1, 1))
img
Let's make an animation of this wave.
k = (0., 10.)
u_prev, u_curr = plane_wave(k, X, Y, 0.1, 0.2, celerity, dt, dx)
hmap3 = hv.HoloMap({i: hv.Image(make_steps(2000), bounds=bounds).options(fig_size=400, colorbar=True, cmap='seismic').redim.range(z=(-1, 1)) for i in list(range(10))}, kdims='index')
hmap3
Let's try a wave travelling to the right:
k = (10., 0.)
u_prev, u_curr = plane_wave(k, X, Y, 0.1, 0.2, celerity, dt, dx)
hmap4 = hv.HoloMap({i: hv.Image(make_steps(2000), bounds=bounds).options(fig_size=400, colorbar=True, cmap='seismic').redim.range(z=(-1, 1)) for i in list(range(10))}, kdims='index')
hmap4
And finally, with a little bit of rotational geometry, we can get oblique wave incidence:
theta = 60
k = 10. * np.array([np.cos(np.deg2rad(theta)), np.sin(np.deg2rad(theta))])
u_prev, u_curr = plane_wave(k, X, Y, 0.1, 0.2, celerity, dt, dx)
hmap5 = hv.HoloMap({i: hv.Image(make_steps(2000), bounds=bounds).options(fig_size=400, colorbar=True, cmap='seismic').redim.range(z=(-1, 1)) for i in list(range(10))}, kdims='index')
hmap5
That's it! I hope you've enjoyed seeing these beautiful wave shapes as much as I did.