1D-wave-simulation

from math import sin, pi
from numpy import zeros, linspace
from scitools.numpyutils import iseq
import time

import matplotlib.pyplot as plt
import numpy as np
fig, ax = plt.subplots()



def I(x):
    return sin(2*x*pi/L)

def f(x,t):
    return 0

def solver0(I, f, c, L, n, dt, tstop):

    ##########
    # PLOT
    ##########

    x = []
    y = []
    for i in range(n+1):
        x.append(i)
        y.append(3)
    points, = ax.plot(x, y, marker='o', linestyle='None', alpha = 0.3)
    ax.set_xlim(0, 20)
    ax.set_ylim(-20, 20)


    # f is a function of x and t, I is a function of x
    x = linspace(0, L, n+1) # grid points in x dir
    dx = L/float(n)
    if dt <= 0: dt = dx/float(c) # max time step
    C2 = (c*dt/dx)**2 # help variable in the scheme
    dt2 = dt*dt

    up = zeros(n+1) # NumPy solution array
    u = up.copy() # solution at t-dt
    um = up.copy() # solution at t-2*dt

    t = 0.0
    for i in iseq(0,n):
        u[i] = I(x[i])*4        # I = initial force
    for i in iseq(1,n-1):
        um[i] = u[i] + 0.5*C2*(u[i-1] - 2*u[i] + u[i+1]) + \
                dt2*f(x[i], t)

    um[0] = 0; um[n] = 0

    while t <= tstop:
          t_old = t; t += dt
          # update all inner points:
          for i in iseq(start=1, stop=n-1):
              up[i] = - um[i] + 2*u[i] + \
                      C2*(u[i-1] - 2*u[i] + u[i+1]) + \
                      dt2*f(x[i], t_old)

          #print "len(up)", len(up)

          # insert boundary conditions:
          up[0] = 0; up[n] = 0; #up[70] = 0
          # update data structures for next step
          um = u.copy(); u = up.copy()

          # PLOT
          #print len(x), len(up[0:101])
          new_y = up[0:n+1]
          points.set_data(x, new_y)
          plt.draw()
          time.sleep(.1)

    return u

if __name__ == '__main__':

# When choosing the parameters you should also check that the units are correct

   c = 5100     # wave velocity
   L = 20       # length of the array
   n = 100      # amount of points
   dt = 0.00004
   tstop = 3
   a = solver0(I, f, c, L, n, dt, tstop)