Weighted sum of adjacent values in numpy array

Question

What is the easiest/fastest way to take a weighted sum of values in a numpy array?

Example: Solving the heat equation with the Euler method

length_l=10
time_l=10
u=zeros((length_l,length_l))# (x,y)
u[:, 0]=1
u[:,-1]=1
print(u)
def dStep(ALPHA=0.1):
    for position,value in ndenumerate(u):
        D2u= (u[position+(1,0)]-2*value+u[position+(-1, 0)])/(1**2) \
            +(u[position+(0,1)]-2*value+u[position+( 0,-1)])/(1**2)
        value+=ALPHA*D2u()
while True:
    dStep()
    print(u)

D2u should be the second central difference in two dimensions. This would work if I could add indexes like (1,4)+(1,3)=(2,7). Unfortunately, python adds them as (1,4)+(1,3)=(1,4,1,3).

Note that computing D2u is equivalent to taking a dot product with this kernel centered around the current position:

 0, 1, 0
 1,-4, 1
 0, 1, 0

Can this be vectorised as a dot product?

Answer 1

I think you want something like:

import numpy as np
from scipy.ndimage import convolve

length_l = 10
time_l = 10
u = np.zeros((length_l, length_l))# (x,y)
u[:,  0] = 1
u[:, -1] = 1

alpha = .1
weights = np.array([[ 0,  1,  0],
                    [ 1, -4,  1],
                    [ 0,  1,  0]])

for i in range(5):
    u += alpha * convolve(u, weights)
    print(u)

You could reduce down a bit by doing:

weights = alpha * weights
weights[1, 1] = weights[1, 1] + 1

for i in range(5):
    u = convolve(u, weights)
    print(u)