In [1]:
    
import numpy as np
import math
import matplotlib.pyplot as plt
import numba
%matplotlib inline
from IPython.html import widgets
    
    
In [2]:
    
N = 1024
L = 10.0
v = 1.0
dx = L / N
dt = dx / v
    
In [3]:
    
u1 = np.empty(N, np.float32)
u2 = np.empty_like(u1)
x = np.linspace(0, L, N)
    
In [4]:
    
# Pluck the string at point x0 with a gausian with width sigma
x0, sigma = 3.57, 0.5
u1[:] = np.exp(-(x-x0)**2/(2*sigma**2))
u2[:] = u1
    
In [5]:
    
plt.plot(x, u1)
plt.plot(x, u2)
    
    Out[5]:
    
In [6]:
    
@numba.jit(nopython=True)
def propagate(nsteps, recinterval, u1, u2):
    nrec = nsteps // recinterval + 1
    tstore = np.empty(nrec)
    ustore = np.empty((nrec, N), np.float32)
    istore, t = 0, 0.0
    # Store the initial snapsho
    tstore[istore] = t
    for i in range(N):
        ustore[istore, i] = u1[i]
    a = v * dt / dx
    for istep in range(nsteps):
        # Decide which is the old and new wavefunction array
        if (istep % 2 == 0):
            uold, unew = u1, u2
        else:
            uold, unew = u2, u1
        # Update new wave function
        for i in range(1, N - 1):
            unew[i] = 2 * (1 - a) * uold[i] - unew[i] \
                    + a * (uold[i + 1] + uold[i - 1])
        t = t + dt
        # Store a snapshot every recinterval
        if ((istep + 1) % recinterval == 0):
            istore += 1
            tstore[istore] = t
            for i in range(N):
                ustore[istore, i] = unew[i]
    return tstore, ustore
    
In [7]:
    
NSTEPS, RECINT = 2000, 30
tstore, ustore = propagate(NSTEPS, RECINT, u1, u2)
    
In [8]:
    
def plot_wave(t):
    # First find the step
    istep = np.where(tstore >= t)[0][0]
    istep = min(istep, NSTEPS - 1)
    plt.plot(x, ustore[istep,:])
    plt.ylim(-1.2, 1.2)
    
In [9]:
    
widgets.interact(plot_wave, t=widgets.FloatSliderWidget(min=0, max=tstore[-1], step=tstore[1], value=0));