from numpy import *
def solver(I, V, f, c, L, Nx, C, T, user_action=None):
"""Solve u_tt=c^2*u_xx + f on (0,L)x(0,T]."""
x = linspace(0, L, Nx+1) # Mesh points in space
dx = x[1] - x[0]
dt = C*dx/c
Nt = int(round(T/dt))
t = linspace(0, Nt*dt, Nt+1) # Mesh points in time
C2 = C**2 # Help variable in the scheme
if f is None or f == 0 :
f = lambda x, t: 0
if V is None or V == 0:
V = lambda x: 0
u = zeros(Nx+1) # Solution array at new time level
u_1 = zeros(Nx+1) # Solution at 1 time level back
u_2 = zeros(Nx+1) # Solution at 2 time levels back
import time; t0 = time.clock() # for measuring CPU time
# Load initial condition into u_1
for i in range(0,Nx+1):
u_1[i] = I(x[i])
if user_action is not None:
user_action(u_1, x, t, 0)
# Special formula for first time step
n = 0
for i in range(1, Nx):
u[i] = u_1[i] + dt*V(x[i]) + \
0.5*C2*(u_1[i-1] - 2*u_1[i] + u_1[i+1]) + \
0.5*dt**2*f(x[i], t[n])
u[0] = 0; u[Nx] = 0
if user_action is not None:
user_action(u, x, t, 1)
# Switch variables before next step
u_2[:], u_1[:] = u_1, u
for n in range(1, Nt):
# Update all inner points at time t[n+1]
for i in range(1, Nx):
u[i] = - u_2[i] + 2*u_1[i] + \
C2*(u_1[i-1] - 2*u_1[i] + u_1[i+1]) + \
dt**2*f(x[i], t[n])
# Insert boundary conditions
u[0] = 0; u[Nx] = 0
if user_action is not None:
if user_action(u, x, t, n+1):
break
# Switch variables before next step
u_2[:], u_1[:] = u_1, u
cpu_time = t0 - time.clock()
return u, x, t, cpu_time