python, calculation on large 3d numpy arrays

Shmulik Edelman 0
Viewed 32

I have a numPy 3d array (for now, it is 200x200x200), and later I plan to use larger arrays (~500x500x500).

For each cell, I want to calculate the average values of its neighbors. I implemented it using loops and NumPy, but it takes a lot of time to do it. I assume it is because I wrote it not efficiently:

def fun1(mat, R, C, D):
    for r in range(1, R - 2):
        for c in range(1, C - 2):
            for d in range(1, D - 2):
                new_val = (mat[r - 1, c, d] + mat[r + 1, c, d] + \
                            mat[r, c - 1, d] + mat[r, c + 1, d] + \
                            mat[r, c, d - 1] + mat[r, c, d + 1]) / 6
                mat[r, c, d] = new_val
    return mat

# R=C=D=200
# mat is a NumPy array with the size of RxCxD 

for iter in range(1000):
    mat = fun1(mat, R, C, D)

Is there any way to write it more efficiently?

The first step of my project was to solve this problem over 2D arrays. Then used the OpenCV function cv.filter2D, which was fast (compared to NumPy). But now, with the 3D matrix size, I am lost.

Thank you all for the help

large-data performance opencv numpy
1 Answers

Loops are probably one of the least efficient ways to write code in python - avoid them if you care about efficiency. Numpy implements many vectorized operations that can speed up your code by orders of magnitude.

There are 6 direct neighbors per cell. We are going to exclude the outermost cells, as they do not have 6 neighbors.

import numpy as np

mat = np.random.rand(200,200,200) # create 3d array of random values

avg = np.random.rand(200,200,200) # array to paste the average values in

avg[1:-1,1:-1,1:-1] = (mat[0:-2,1:-1,1:-1] + mat[1:-1,0:-2,1:-1] + mat[1:-1,1:-1,0:-2,] + mat[2:,1:-1,1:-1] + mat[1:-1,2:,1:-1] + mat[1:-1,1:-1,2:]) / 6

This code will be much quicker than yours, because it vectorizes the operation (it does multiple operations at once!) instead of calculating the values one by one with the loop overhead.

userV5dj10
Related