confused
confused

Reputation: 49

Python numpy vectorization for heat dispersion

I'm supposed to write a code to represent heat dispersion using the finite difference formula given below.

๐‘ข(๐‘ก)๐‘–๐‘—=(๐‘ข(๐‘กโˆ’1)[๐‘–+1,๐‘—] + ๐‘ข(๐‘กโˆ’1) [๐‘–โˆ’1,๐‘—] +๐‘ข(๐‘กโˆ’1)[๐‘–,๐‘—+1] + ๐‘ข(๐‘กโˆ’1)[๐‘–,๐‘—โˆ’1])/4 The formula is supposed to produce the result only for a time step of 1. So, if an array like this was given:

100   100   100   100   100  
100    0     0     0    100  
100    0     0     0    100  
100    0     0     0    100  
100   100   100   100   100

The resulting array at time step 1 would be:

100   100   100   100   100  
100    50    25    50   100   
100    25     0    25   100  
100    50    25    50   100  
100   100   100   100   100

I know the representation using for loops would be as follows, where the array would have a minimum of 2 rows and 2 columns as a precondition:

h = np.copy(u)
for i in range(1,h.shape[0]-1):
    for j in range (1, h.shape[1]-1):
        num = u[i+1][j] + u[i-1][j] + u[i][j+1] + u[i][j-1]
        h[i][j] = num/4

But I cannot figure out how to vectorize the code to represent heat dispersion. I am supposed to use numpy arrays and vectorization and am not allowed to use for loops of any kind, and I think I am supposed to rely on slicing, but I cannot figure out how to write it out and have started out with.

r, c = h.shape
if(c==2 or r==2):
    return h

I'm sure that if the rows=2 or columns =2 then the array is returned as is, but correct me if Im wrong. Any help would be greatly appreciated. Thank you!

Upvotes: 2

Views: 1035

Answers (1)

tstanisl
tstanisl

Reputation: 14147

Try:

h[1:-1,1:-1] = (h[2:,1:-1] + h[:-2,1:-1] + h[1:-1,2:] + h[1:-1,:-2]) / 4

This solution uses slicing where:

  • 1:-1 stays for indices 1,2, ..., LAST - 1
  • 2: stays for 2, 3, ..., LAST
  • :-2 stays for 0, 1, ..., LAST - 2

During each iteration only the inner elements (indices 1..LAST-1) are updated

Upvotes: 4

Related Questions