Optimizing a nested for loop

Question

I'm trying avoid to use for loops to run my calculations. But I don't know how to do it. I have a matrix w with shape (40,100). Each line holds the position to a wave in a t time. For example first line w[0] is the initial condition (also w[1] for reasons that I will show).

To calculate the next line elements I use, for every t and x on shape range:

w[t+1,x] = a * w[t,x] + b * ( w[t,x-1] + w[t,x+1] ) - w[t-1,x]

Where a and b are some constants based on equation solution (it really doesn't matter), a = 2(1-r), b=r, r=(c*(dt/dx))**2. Where c is the wave speed and dt, dx are related to the increment on x and t direction.

Is there any way to avoid a for loop like:

for t in range(1,nt-1):
    for x in range(1,nx-1):
      w[t+1,x] = a * w[t,x] + b * ( w[t,x-1] + w[t,x+1] ) - w[t-1,x]

nt and nx are the shape of w matrix.

Answer 1

I assume you're setting w[:,0] and w[:-1] beforehand (to some constants?) because I don't see it in the loop. If so, you can eliminate for x loop vectorizing this part of code:

for t in range(1,nt-1):
    w[t+1,1:-1] = a*w[t,1:-1] + b*(w[t,:-2] + w[t,2:]) - w[t-1,1:-1]