[Visualpython-users] Speeding up a wave propagation simulation

[Andy D. <dou...@la...>]

I've been considering using Vpython in a course to illustrate the
propagation of wave packets with dispersion.  The problem is it's just too
slow.  I'm hoping someone on this list with experience with the Numeric
module might be able to help me speed it up.  I think I'm spending way
too much time explicitly looping over arrays.

Sorry for the length of what follows, but it might be helpful to be
specific:

Here's my original naive attempt:

    #!/usr/bin/python
    # Show propagation of a wave pulse with dispersion.
    from visual.graph import *

    def w(k):
	"Dispersion relation w(k), gives v ~ 1/sqrt(k)."
	return sqrt(k)

    def yt(x,t):
	"Total wave obtained by superimposing sinusoids with k from 1 to 4."
	ytotal = 0.0
	for k in arange(1.0, 4.0, 0.05):
	    ytotal = ytotal + sin(k * x - w(k) * t)
	return ytotal

    gdisplay(xtitle='x', ytitle='ytotal', title='Wave',
	    xmin = -10.0, ymin = -60.0, xmax =100.0, ymax = 60.0)
    pulse = gcurve()

    t = 0.0
    # Make initial plot
    for x in arange(-10.0, 100.0, 0.2):
	pulse.plot(pos=(x, yt(x,t)))
    npts = len(pulse.gcurve.pos)
    
    while 1:
	t = t + 0.5
	for n in range(npts):
	    x = pulse.gcurve.pos[n][0]
	    pulse.gcurve.pos[n][1] = yt(x, t)

This would be rather simple for my students to follow (I hope!), but it
is painfully slow.  I was able to speed it up considerably by doing
some obvious things:
    1.  Precomputing the dispersion relationship w(k) once and storing it
    2.  storing each set of points in a temporary array and updating the 
    display only once per time step (instead of every x value)
    3.  Calling math.sin() directly instead of the overloaded Numeric sin()
but it's still too slow.  (Obviously, I can reduce the spatial and
temporal resolution, but that's beside the point here.)

To be specific, here's what I ended up with.

    #!/usr/bin/python
    # Show propagation of a wave pulse with dispersion.
    from visual.graph import *

    def w(k):
	"Dispersion relation w(k), gives v ~ 1/sqrt(k)."
	return sqrt(k)

    # Range of wavenumbers 'k' to include.
    ka = arange(1.0, 4.0, 0.05)
    wk = w(ka)
    nkmax = len(ka)
    def yt(x,t):
	"Total wave obtained by superimposing sinusoids with k from 1 to 4."
	ytotal = 0.0
	for nk in range(nkmax):
	    ytotal = ytotal + math.sin(ka[nk] * x - wk[nk] * t)
	return ytotal

    gdisplay(xtitle='x', ytitle='ytotal', title='Wave',
	    xmin = -10.0, ymin = -60.0, xmax =100.0, ymax = 60.0)
    pulse = gcurve()

    t = 0.0
    for x in arange(-10.0, 100.0, 0.2):
	pulse.plot(pos=(x, yt(x,t)))

    # Make a temporary copy -- see below
    temp = array(pulse.gcurve.pos)
    npts = len(pulse.gcurve.pos)

    while 1:
	t = t + 0.5
	# Compute the new values in a temp array that isn't displayed, then
	# copy to the pulse.gcurve.pos array to update the display all at once.
	for n in arange(npts):
	    temp[n][1] = yt(temp[n][0], t)
	pulse.gcurve.pos = temp

Does anyone have any suggestions for how to significantly speed this up?
All the looping over arrays seems to me to be the sort of thing that I
ought to be able to get the Numeric extension to do for me efficiently,
but I confess I'm new to it and haven't seen how to do it yet.

Thanks for any advice,

    Andrew Dougherty		dou...@la...
    Dept. of Physics
    Lafayette College, Easton PA 18042