Develop Reference

Monte Carlo integration in case of coordinate transformation with Python

I am trying to evaluate the integral using the Monte Carlo method where the integrand underwent a transformation from cylindrical to cartesian coordinates. The integrand itself is quite simple and could be calculated using scipy.integrate.quad, but I need it to be in cartesian coordinates for a specific purpose later on.
So here is the integrand: rho*k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2) d rho
Here the kn(i,rho) is modified Bessel function of 2nd kind.
Solving it with quad gives the following result:
from scipy.special import kn
from scipy.integrate import quad
import random
k_r = 6.2e-2
k_n = k_r/1.05
C_factor = 2*np.pi*1e5
lmax,lmin = 80,50
def integration_polar():
def K_int(rho):
return rho*k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2)
rho = np.linspace(lmin,lmax,200)
I,_ = quad(K_int,lmin,lmax)
Gamma = I*C_factor
print("expected=",Gamma)
Output: Expected = 7.641648442007296
Now the same integral using Monte Carlo method (hit-or-miss method looked up from here) gives almost the same result:
def integration_polar_MC():
random.seed(1)
n = 100000
def K_int(rho):
return rho*k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2)
def sampler():
x = random.uniform(lmin,lmax)
y = random.uniform(0,c_lim)
return x,y
c_lim = 2*K_int(50) #Upper limit of integrand
sum_I = 0
for i in range(n):
x,y = sampler()
func_Int = K_int(x)
if y>func_Int:
I = 0
elif y<=func_Int:
I = 1
sum_I += I
Gamma = C_factor*(lmax-lmin)*c_lim*sum_I/n
print("MC_integral_polar:",Gamma)
Output: MC_integral_polar = 7.637391399699502
Since Monte Carlo worked with this example, I thought the cartesian case would go smoothly as well but I couldn't get the right answer.
For the cartesian case, similarly as in previous case I've employed the hit-or-miss method, with rho = np.sqrt(x**2+y**2) and integrand becoming k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2) dx dy where domain over x and y:
-80 <= x <= 80
-80 <= y <= 80
50 <= np.sqrt(x**2+y**2) <= 80
Here is my attempt:
def integration_cartesian_MCtry():
random.seed(1)
lmin,lmax = -100,100
n = 100000
def K_int(x,y):
rho = np.sqrt(x**2+y**2)
if rho>=50 and rho<=80:
return k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2)
else:
return 0
def sampler():
x = random.uniform(lmin,lmax)
y = random.uniform(lmin,lmax)
z = random.uniform(0,c_lim)
return x,y,z
c_lim = K_int(50,0)
sum_I = 0
for i in range(n):
x,y,z = sampler()
func_Int = K_int(x,y)
if z>func_Int:
I = 0
elif z<=func_Int:
I = 1
sum_I += I
Gamma = C_factor*(lmax-lmin)**2*c_lim*sum_I/n
print("MC_integral_cartesian:",Gamma)
Output: MC_integral_cartesian = 48.83166430996952
As you can see Monte Carlo in cartesian overestimates the integral. I am not sure why it is happening but think that it may be related to the incorrect limits or domain over which I should integrate the function.
Any help appreciated as I am stuck without any progress for a few days.

Problem, as I said, is with jacobian. In case of polar, you have integration over
f(ρ)*ρ*dρ*dφ
You integrate over dφ analytically (your f(ρ) doesn't depend on φ), and get 2π
In case of cartesian there are no analytical integration, so it is over dx*dy, no factor
of 2π. Code to illustrate it, Python 3.9.1, Windows 10 x64, and it produced pretty much the same answer
import numpy as np
from scipy.special import kn
k_r = 6.2e-2
k_n = k_r/1.05
C_factor = 2*np.pi*1e5
lmin = 50
lmax = 80
def integration_polar_MC(rng, n):
def K_int(rho):
if rho>=50 and rho<=80:
return rho*k_r**2*(k_r**2*kn(0, k_r*rho)**2 + k_n**2*kn(1, k_r*rho)**2)
return 0.0
def sampler():
x = rng.uniform(lmin, lmax)
y = rng.uniform(0.0, c_lim)
return x,y
c_lim = 2*K_int(50) # Upper limit of integrand
sum_I = 0
for i in range(n):
x,y = sampler()
func_Int = K_int(x)
I = 1
if y>func_Int:
I = 0
sum_I += I
Gamma = C_factor*(lmax-lmin)*c_lim*sum_I/n
return Gamma
def integration_cartesian_MC(rng, n):
def K_int(x,y):
rho = np.hypot(x, y)
if rho>=50 and rho<=80:
return k_r**2*(k_r**2*kn(0,k_r*rho)**2+k_n**2*kn(1,k_r*rho)**2)
return 0.0
def sampler():
x = rng.uniform(lmin,lmax)
y = rng.uniform(lmin,lmax)
z = rng.uniform(0,c_lim)
return x,y,z
lmin,lmax = -100,100
c_lim = K_int(50, 0)
sum_I = 0
for i in range(n):
x,y,z = sampler()
func_Int = K_int(x,y)
I = 1
if z>func_Int:
I = 0
sum_I += I
Gamma = C_factor*(lmax-lmin)**2*c_lim*sum_I/n
return Gamma/(2.0*np.pi) # to compensate for 2π in the constant
rng = np.random.default_rng()
q = integration_polar_MC(rng, 100000)
print("MC_integral_polar:", q)
q = integration_cartesian_MC(rng, 100000)
print("MC_integral_cart:", q)

Multigrid Poisson Solver

I am trying to make my own CFD solver and one of the most computationally expensive parts is solving for the pressure term. One way to solve Poisson differential equations faster is by using a multigrid method. The basic recursive algorithm for this is:
function phi = V_Cycle(phi,f,h)
% Recursive V-Cycle Multigrid for solving the Poisson equation (\nabla^2 phi = f) on a uniform grid of spacing h
% Pre-Smoothing
phi = smoothing(phi,f,h);
% Compute Residual Errors
r = residual(phi,f,h);
% Restriction
rhs = restriction(r);
eps = zeros(size(rhs));
% stop recursion at smallest grid size, otherwise continue recursion
if smallest_grid_size_is_achieved
eps = smoothing(eps,rhs,2*h);
else
eps = V_Cycle(eps,rhs,2*h);
end
% Prolongation and Correction
phi = phi + prolongation(eps);
% Post-Smoothing
phi = smoothing(phi,f,h);
end
I've attempted to implement this algorithm myself (also at the end of this question) however it is very slow and doesn't give good results so evidently it is doing something wrong. I've been trying to find why for too long and I think it's just worthwhile seeing if anyone can help me.
If I use a grid size of 2^5 by 2^5 points, then it can solve it and give reasonable results. However, as soon as I go above this it takes exponentially longer to solve and basically get stuck at some level of inaccuracy, no matter how many V-Loops are performed. at 2^7 by 2^7 points, the code takes way too long to be useful.
I think my main issue is that my implementation of a jacobian iteration is using linear algebra to calculate the update at each step. This should, in general, be fast however, the update matrix A is an n*m sized matrix, and calculating the dot product of a 2^7 * 2^7 sized matrix is expensive. As most of the cells are just zeros, should I calculate the result using a different method?
if anyone has any experience in multigrid methods, I would appreciate any advice!
Thanks
my code:
# -*- coding: utf-8 -*-
"""
Created on Tue Dec 29 16:24:16 2020
#author: mclea
"""
import numpy as np
import matplotlib.pyplot as plt
from scipy.signal import convolve2d
from mpl_toolkits.mplot3d import Axes3D
from scipy.interpolate import griddata
from matplotlib import cm
def restrict(A):
"""
Creates a new grid of points which is half the size of the original
grid in each dimension.
"""
n = A.shape[0]
m = A.shape[1]
new_n = int((n-2)/2+2)
new_m = int((m-2)/2+2)
new_array = np.zeros((new_n, new_m))
for i in range(1, new_n-1):
for j in range(1, new_m-1):
ii = int((i-1)*2)+1
jj = int((j-1)*2)+1
# print(i, j, ii, jj)
new_array[i,j] = np.average(A[ii:ii+2, jj:jj+2])
new_array = set_BC(new_array)
return new_array
def interpolate_array(A):
"""
Creates a grid of points which is double the size of the original
grid in each dimension. Uses linear interpolation between grid points.
"""
n = A.shape[0]
m = A.shape[1]
new_n = int((n-2)*2 + 2)
new_m = int((m-2)*2 + 2)
new_array = np.zeros((new_n, new_m))
i = (np.indices(A.shape)[0]/(A.shape[0]-1)).flatten()
j = (np.indices(A.shape)[1]/(A.shape[1]-1)).flatten()
A = A.flatten()
new_i = np.linspace(0, 1, new_n)
new_j = np.linspace(0, 1, new_m)
new_ii, new_jj = np.meshgrid(new_i, new_j)
new_array = griddata((i, j), A, (new_jj, new_ii), method="linear")
return new_array
def adjacency_matrix(rows, cols):
"""
Creates the adjacency matrix for an n by m shaped grid
"""
n = rows*cols
M = np.zeros((n,n))
for r in range(rows):
for c in range(cols):
i = r*cols + c
# Two inner diagonals
if c > 0: M[i-1,i] = M[i,i-1] = 1
# Two outer diagonals
if r > 0: M[i-cols,i] = M[i,i-cols] = 1
return M
def create_differences_matrix(rows, cols):
"""
Creates the central differences matrix A for an n by m shaped grid
"""
n = rows*cols
M = np.zeros((n,n))
for r in range(rows):
for c in range(cols):
i = r*cols + c
# Two inner diagonals
if c > 0: M[i-1,i] = M[i,i-1] = -1
# Two outer diagonals
if r > 0: M[i-cols,i] = M[i,i-cols] = -1
np.fill_diagonal(M, 4)
return M
def set_BC(A):
"""
Sets the boundary conditions of the field
"""
A[:, 0] = A[:, 1]
A[:, -1] = A[:, -2]
A[0, :] = A[1, :]
A[-1, :] = A[-2, :]
return A
def create_A(n,m):
"""
Creates all the components required for the jacobian update function
for an n by m shaped grid
"""
LaddU = adjacency_matrix(n,m)
A = create_differences_matrix(n,m)
invD = np.zeros((n*m, n*m))
np.fill_diagonal(invD, 1/4)
return A, LaddU, invD
def calc_RJ(rows, cols):
"""
Calculates the jacobian update matrix Rj for an n by m shaped grid
"""
n = int(rows*cols)
M = np.zeros((n,n))
for r in range(rows):
for c in range(cols):
i = r*cols + c
# Two inner diagonals
if c > 0: M[i-1,i] = M[i,i-1] = 0.25
# Two outer diagonals
if r > 0: M[i-cols,i] = M[i,i-cols] = 0.25
return M
def jacobi_update(v, f, nsteps=1, max_err=1e-3):
"""
Uses a jacobian update matrix to solve nabla(v) = f
"""
f_inner = f[1:-1, 1:-1].flatten()
n = v.shape[0]
m = v.shape[1]
A, LaddU, invD = create_A(n-2, m-2)
Rj = calc_RJ(n-2,m-2)
update=True
step = 0
while update:
v_old = v.copy()
step += 1
vt = v_old[1:-1, 1:-1].flatten()
vt = np.dot(Rj, vt) + np.dot(invD, f_inner)
v[1:-1, 1:-1] = vt.reshape((n-2),(m-2))
err = v - v_old
if step == nsteps or np.abs(err).max()<max_err:
update=False
return v, (step, np.abs(err).max())
def MGV(f, v):
"""
Solves for nabla(v) = f using a multigrid method
"""
# global A, r
n = v.shape[0]
m = v.shape[1]
# If on the smallest grid size, compute the exact solution
if n <= 6 or m <=6:
v, info = jacobi_update(v, f, nsteps=1000)
return v
else:
# smoothing
v, info = jacobi_update(v, f, nsteps=10, max_err=1e-1)
A = create_A(n, m)[0]
# calculate residual
r = np.dot(A, v.flatten()) - f.flatten()
r = r.reshape(n,m)
# downsample resitdual error
r = restrict(r)
zero_array = np.zeros(r.shape)
# interploate the correction computed on a corser grid
d = interpolate_array(MGV(r, zero_array))
# Add prolongated corser grid solution onto the finer grid
v = v - d
v, info = jacobi_update(v, f, nsteps=10, max_err=1e-6)
return v
sigma = 0
# Setting up the grid
k = 6
n = 2**k+2
m = 2**(k)+2
hx = 1/n
hy = 1/m
L = 1
H = 1
x = np.linspace(0, L, n)
y = np.linspace(0, H, m)
XX, YY = np.meshgrid(x, y)
# Setting up the initial conditions
f = np.ones((n,m))
v = np.zeros((n,m))
# How many V cyles to perform
err = 1
n_cycles = 10
loop = True
cycle = 0
# Perform V cycles until converged or reached the maximum
# number of cycles
while loop:
cycle += 1
v_new = MGV(f, v)
if np.abs(v - v_new).max() < err:
loop = False
if cycle == n_cycles:
loop = False
v = v_new
print("Number of cycles " + str(cycle))
plt.contourf(v)

I realize that I'm not answering your question directly, but I do note that you have quite a few loops that will contribute some overhead cost. When optimizing code, I have found the following thread useful - particularly the line profiler thread. This way you can focus in on "high time cost" lines and then start to ask more specific questions regarding opportunities to optimize.
How do I get time of a Python program's execution?

ValueError in > comparison

# -*- coding: utf-8 -*-
"""
Created on Sat Feb 29 14:56:52 2020
#author: 12064
"""
import numpy as np
k = 0
### size of function
n = 5
######## creating test function
A = np.random.randn(n,n)
B = np.transpose(A)
C = np.matmul(B,A) # A,B,C working
gradf = np.gradient(C, axis=0)
def myFcn(x):
return np.dot(C, x)
def gradientmyFcn(x):
return np.dot(gradf, x)
######## C = Q Rick defines
######## define the point
xT = np.random.randn((n))
x = xT.reshape((n,1))
print('myFcn(x)')
###### define gradient
#####initial alpha, c
###### define descent direction, and set of descent directions
p = np.zeros((n,1))
pT = np.zeros((1,n))
pT = gradientmyFcn(x)
pT.reshape(1,n)
p = np.transpose(pT)
p.reshape(n,1)
#P = {p[0]:p[k]}
#p[0] = gradf initial descent direction
def ArmijoLineSearch(alpha):
tao = .1
alpha = 1
c = 1.e-4
backtrack_factor = .5
gradfx = gradientmyFcn(x)
gradfx = gradfx
gradTx = np.transpose(gradfx)
xk1 = (x + (alpha*p))
check_step = np.dot(gradfx, p)
LessThan = myFcn(x) + (c*alpha)*check_step
if myFcn(xk1) > LessThan:
alpha = backtrack_factor * alpha
else:
if (np.transpose(mygradientFcn(xk1)) # p) < tao * np.dot(np.transpose(gradfx), p):
print("weak curvature condition does not hold")
else:
return alpha
ArmijoLineSearch(alpha)
The error occurs on this line: if myFcn(xk1) > LessThan:
ValueError: The truth value of an array with more than one element is ambiguous. Use a.any() or a.all()
I'm struggling to understand this error since it seems to be taking the numpy arrays of random numbers (at least initially random) and trying to operate with booleans. What is the problem and how should I fix it?

Fastest way to perform sparse matrix multiplications in Python

BACKGROUND: I am trying to build a real-time drum simulation model, for which I need really fast matrix-vector products. My matrices are of the size ~5000-10000 rows/cols, out of which only 6 entries per row are non-zero, hence I am inclined to use sparse matrices. I am using scipy.sparse module. The iterations are as below.
Vjk_plus_sparse = Vjk_minus_sparse.transpose()
Vj = Vjk_plus_sparse.dot(constant)
np.put(Vj, Nr, 0.0)
Uj[t] = Uj[t-1] + np.transpose(Vj)/fs
Vj_mat = adj_mat_sparse.multiply(Vj)
Vjk_minus_sparse = Vj_mat-Vjk_plus_sparse.multiply(end_gain)
Here, Vjk_plus_sparse, Vjk_minus_sparse and Vj_mat are sparse CSR matrices, Vj is a numpy array, and Uj is a numpy matrix where each row represents Uj(t). end_gain is an array which is a static numpy array for dampening of vibrations.
THE ISSUE: A single iteration takes about 3 ms for size = 4250. With the most significant
steps being the last 2 lines. They together take about 2.5 ms. I would ideally need it to run in 0.1 ms, which would be more than a 10x speedup. This is the maximum extent of vectorization possible for the problem, and I cannot parallelize as I am marching in time, at least physically it won't be accurate.
ATTEMPTS: I tried fiddling with the sparse data structures, and found best performance with all of them being CSR (Compressed Sparse Row), with the values as quoted above. I also tried to replace the multiply() method with a matrix multiplication, by repeating Vj, but that worsened the time, as the resultant operation would be a sparse*dense operation.
How can I speed this up within python itself? I am open to trying c++ as well, though migrating now would be a major pain. Also, since scipy is essentially based in c, would it even give that much of a speedup?
Added a complete runnable example
import matplotlib.pyplot as plt
import matplotlib as mpl
import matplotlib.patches
import math
from mpl_toolkits import mplot3d
import numpy as np
import scipy.sparse as sp
import scipy.fftpack as spf
import matplotlib.animation as animation
import time
sqrt_3 = 1.73205080757
class Pt:
def __init__(self,x_0,y_0):
self.x_0 = x_0
self.y_0 = y_0
self.id = -1
self.neighbours = []
self.distance = (x_0**2 + y_0**2)**0.5
class Circle:
def __init__(self,radius,center):
self.radius = radius
self.center = center
self.nodes = []
def construct_mesh(self, unit):
queue = [self.center]
self.center.distance = 0
curr_id = 0
delta = [(1.,0.), (1./2, (3**0.5)/2),(-1./2, (3**0.5)/2),(-1.,0.), (-1./2,-(3**0.5)/2), (1./2,- (3**0.5)/2)]
node_dict = {}
node_dict[(self.center.x_0,self.center.y_0)] = curr_id
self.nodes.append(self.center)
curr_id+=1
while len(queue)!=0:
curr_pt = queue[0]
queue.pop(0)
# self.nodes.append(curr_pt)
# curr_id+=1
for i in delta:
temp_pt = Pt(curr_pt.x_0 + 2*unit*i[0], curr_pt.y_0 + 2*unit*i[1])
temp_pt.id = curr_id
temp_pt.distance = (temp_pt.x_0 ** 2 + temp_pt.y_0 ** 2)**0.5
# curr_id+=1
if (round(temp_pt.x_0,5), round(temp_pt.y_0,5)) not in node_dict and temp_pt.distance <= self.radius:
# print(temp_pt.x_0, temp_pt.y_0)
self.nodes.append(temp_pt)
node_dict[(round(temp_pt.x_0,5), round(temp_pt.y_0,5))] = curr_id
curr_id+=1
queue.append(temp_pt)
curr_pt.neighbours.append(temp_pt.id)
elif temp_pt.distance <= self.radius:
curr_pt.neighbours.append(node_dict[round(temp_pt.x_0,5), round(temp_pt.y_0,5)])
# print(node_dict)
def plot_neighbours(self, pt):
x = []
y = []
x.append(pt.x_0)
y.append(pt.y_0)
for i in (pt.neighbours):
x.append(self.nodes[i].x_0)
y.append(self.nodes[i].y_0)
plt.scatter(x,y)
plt.axis('scaled')
def boundary_node_ids(self):
boundary_nodes = []
for j in range(len(self.nodes)):
if(len(self.nodes[j].neighbours) < 6):
boundary_nodes.append(j)
return boundary_nodes
def add_rim(self, boundary_node_ids, unit):
c = self.center
rim_ids = []
N = len(self.nodes)
for i in range(len(boundary_node_ids)):
d = self.nodes[boundary_node_ids[i]].distance
xp = self.nodes[boundary_node_ids[i]].x_0
yp = self.nodes[boundary_node_ids[i]].y_0
xnew = xp + xp*unit/d
ynew = yp + yp*unit/d
new_point = Pt(xnew, ynew)
new_point.id = N + i
rim_ids.append(N+i)
self.nodes.append(new_point)
self.nodes[boundary_node_ids[i]].neighbours.append(new_point.id)
self.nodes[N+i].neighbours.append(boundary_node_ids[i])
return rim_ids
def find_nearest_point(mesh, pt):
distances_from_center = np.zeros(len(mesh.nodes))
for i in xrange(len(mesh.nodes)):
distances_from_center[i] = mesh.nodes[i].distance
target_distance = pt.distance
closest_point_id = np.argmin(np.abs(distances_from_center-target_distance))
return closest_point_id
def init_impulse(mesh, impulse, Vj, poi, roi):
data = []
for i in range(len(Vj)):
r = ((mesh.nodes[i].x_0 - mesh.nodes[poi].x_0)**2 + (mesh.nodes[i].y_0 - mesh.nodes[poi].y_0)**2)**0.5
Vj[i] = max(0, impulse*(1. - (r/roi)))
if i in Nr:
Vj[i] = 0.
for k in mesh.nodes[i].neighbours:
data.append(np.asscalar(Vj[i])/2.)
return Vj, data
r = 0.1016 #Radius of drum head
# rho = 2500 #Density of drum head
thickness = 0.001 #Thickness of membrane
# tension = 1500 #Tension in membrane in N
param = 0.9
c = (param/thickness)**(0.5) #Speed of wave in string
duration = 0.25
fs = 4000
delta = c/fs
center = Pt(0,0)
point_of_impact = Pt(r/2., 0)
center.id = 0
mesh = Circle(r,center)
mesh.construct_mesh(delta)
N = len(mesh.nodes)
Nb = []
for j in range(N):
if len(mesh.nodes[j].neighbours) < 6:
Nb.append(j)
Nr = mesh.add_rim(Nb, delta)
N = len(mesh.nodes)
print(N)
row_ind = []
col_ind = []
for j in range(N):
for k in mesh.nodes[j].neighbours:
row_ind.append(j)
col_ind.append(k)
data = np.ones(len(col_ind))
adj_mat_sparse = sp.csr_matrix((data, (row_ind, col_ind)), shape = (N,N))
Vjk_plus = sp.csr_matrix([N, N])
Vj = np.zeros([N,1])
Uj = np.zeros([int(duration*fs), N])
Vj_mat = sp.csc_matrix([N,N])
closest_point_id = find_nearest_point(mesh, point_of_impact)
Vj, Vjk_data = init_impulse(mesh, -10.0, Vj, closest_point_id, r/10.)
Vjk_minus_sparse = sp.csr_matrix((Vjk_data, (row_ind, col_ind)), shape = (N,N))
constant = (1./3)*np.ones([N,1])
Vjk_plus = Vjk_minus_sparse.transpose()
np.put(Vj, Nr, 0.0)
Uj[1] = Uj[0] + np.transpose(Vj)/fs
Vj_mat = adj_mat_sparse.multiply(Vj)
Vjk_minus_sparse = Vj_mat - Vjk_plus
end_gain = np.ones([N,1])
end_gain[Nr] = 1.0
for t in range(2,int(duration*fs)):
Vjk_plus = Vjk_minus_sparse.transpose()
Vj = Vjk_plus.dot(constant)
np.put(Vj, Nr, 0.0)
Uj[t] = Uj[t-1] + np.transpose(Vj)/fs
Vj_mat = adj_mat_sparse.multiply(Vj)
Vjk_minus_sparse = Vj_mat-Vjk_plus.multiply(end_gain)

Python: Writing a program to simulate 1D wave motion along a string

I am working on a program that simulates wave motion along a 1-dimensional string to eventually simulate different wave packets. I found a program in the book "Python Scripting for Computational Science" that claims to describe wave motion, though I'm not certain how to implement it (the book was on Google Books and won't show me the text before/after the code).
For example, I understand that "f" is a function of x and t and that "I" is a function of x but what functions are actually needed to produce a wave?
I=
f=
c=
L=
n=
dt=
tstop=
x = linespace(0,L,n+1) #grid points in x dir
dx = L/float(n)
if dt <= 0: dt = dx/float(c) #max step time
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] +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)
#insert boundary conditions
up[0] = 0; up[n] = 0
#updata data structures for next step
um = u.copy(); u = up.copy()

The code below should work:
from math import sin, pi
from numpy import zeros, linspace
from scitools.numpyutils import iseq
def I(x):
return sin(2*x*pi/L)
def f(x,t):
return 0
def solver0(I, f, c, L, n, dt, tstop):
# 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])
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)
# insert boundary conditions:
up[0] = 0; up[n] = 0
# update data structures for next step
um = u.copy(); u = up.copy()
return u
if __name__ == '__main__':
# When choosing the parameters you should also check that the units are correct
c = 5100
L = 1
n = 10
dt = 0.1
tstop = 1
a = solver0(I, f, c, L, n, dt, tstop)
It returns an array with the values of the wave at time tstop and at all the points in our solution grid.
Before you apply it to a practical situation, you should read up both about the wave equation and the finite element method to understand what the code does. It can be used to find the numerical solutions of the wave equation:
Utt + beta*Ut = c^2*Uxx + f(x,t)
which is one of most important differential equations in physics. The solution of this PDE or wave, is given by a function which is a function of space and time u(x,t).
To visualize the concept of wave, consider two dimensions, space and time. If you fix the time, e.g. t1, you will get a function of x:
U(x) = U(x,t=t1)
However, at a particular point of space, x1, the wave is a function of time:
U(t) = U(x=x1, t)
This should help you to understand how the wave propagates. In order to find a solution, you need to impose some initial and boundary conditions to restrict all the possibles waves to the one you are interested in. For this particular case:
I = I(xi) is the initial force that we will apply to get the
perturbation/wave going.
The term f = f(x,t) accounts for any external force that generates
waves.
c is wave velocity; it is a constant (assuming the medium is
homogeneous).
L is the length of the domain where we want to solve the PDE; also
a constant.
n is the number of grid points in space.
dt is a time step.
tstop is the stop time.