[Tutor] Solve wave equation

[Jose Amoreira]

On Thursday, February 23, 2012 12:57:39 PM David Craig wrote:
> Hi,
> I am trying to write some code that will solve the 2D wave equation by
> the finite difference method, but it just returns an array full of zeros
> and NaN's. Not sure where I am going wrong, the code is attached so if
> could someone point me in the right direction I'd appreciate this.
> Thanks
> D

Let me add my 2 cents to Steven's suggestions.

The main cicle of your program can be reorganized, pulling out all constant 
calculations. Where you write
>for t in range(2,nsteps-1):
>
>    
>    for z in range(1,nz-1):
>
>        for x in range(2,nx-1):
>
>            p[xs,zs,t] = s[t]
>
>            k = (c*dt/h)**2
>
>            p[x,z,t] = 2*p[x,z,t-1] - p[x,z,t-2] + [...]

I'd write

k = (c*dt/h)**2
for t in range(2,nsteps-1):
	p[xs,zs,t] = s[t]
	for z in range(1,nz-1):
		for x in range(2,nx-1):
			p[x,z,t] = 2*p[x,z,t-1] - p[x,z,t-2] + [...]

Like that you don't have to compute the value of k (wich is constant) for each 
cell in your mesh and for every time slice. I didn't quite understand the way 
you integrate the source in the calculation, but if it works the way you do 
it, it should also work putting it out of the x and z loops; like that, you 
just have to compute it once for each time slice.

Also, whenever possible, try to implement the iterations (for loops) over 
array elements as whole array operations, wich are way faster then python for 
loops. For instance, the laplacian of the wave function,

>  k*(p[x+1,z,t-1]-4*p[x,z,t-1]+p[x-1,z,t-1]+p[x,z+1,t-1]+p[x,z-1,t-1])

can be computed at once (without for loops) with something like (haven't tryed 
it, take care, read the docs)

>k*(roll(p[:,:,t-1],-1,axis=1) - 4*p[:,:,t-1] + roll(p[:,:,t-1],1,axis=1) +
	roll(p[:,:,t-1],-1,axis=0) + roll(p[:,:,t-1],1,axis=0))

(mind the linebreak). This expression returns an array with the dimensions of 
your pressure array p. It may have to be tweaked a bit because of boundary 
behaviour of the roll function.
roll() is a numpy function (must be imported from numpy) that shifts the array 
elements. See 
http://docs.scipy.org/doc/numpy/reference/generated/numpy.roll.html
Some examples:

In [1]: from numpy import *

In [2]: m=array([[11,12,13],[21,22,23],[31,32,33]])

In [3]: print m
[[11 12 13]
 [21 22 23]
 [31 32 33]]

In [4]: roll(m,-1,axis=1)
Out[4]: 
array([[12, 13, 11],
       [22, 23, 21],
       [32, 33, 31]])

In [5]: roll(m,1,axis=0)
Out[5]: 
array([[31, 32, 33],
       [11, 12, 13],
       [21, 22, 23]])

Finally, about the plot, I find the matplotlib contour function too slow (it's 
not trivial to compute the contour lines) for animations. I prefer a method 
that directly maps values into colors. Someone in this list suggested pygame's 
surfarray methods, and that's what I've been using.
I hope this helps. Cheers,
Ze
-------------- next part --------------
An HTML attachment was scrubbed...
URL: <http://mail.python.org/pipermail/tutor/attachments/20120224/0436c430/attachment.html>