如何使用numpy更快地制作循环

时间:2018-06-15 10:29:41

标签: python python-3.x numpy

如果输入了一个numpy 2d数组,例如

                           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

应该获得这样的输出

                            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

除了边界之外的每个数字都成为其相邻数字的平均值。

我当前的代码有效,但我需要在没有for循环的情况下使用它,并使用numpy对其进行矢量化。

我目前的代码:

import numpy as np
def evolve_heat_slow(u):      
    u2 = np.copy(u)
    x=u2.shape[0]
    y=u2.shape[1]
    for i in range(1,x-1):
        for s in range(1,y-1):
            u2[i,s]=(u[i-1,s]+u[i+1,s]+u[i,s+1]+u[i,s-1])/4
    return u2

3 个答案:

答案 0 :(得分:6)

这几乎是2D卷积的定义。 scipy让你满意。我复制a以保留边界; valid模式下的卷积将形成一个较小的数组(没有边框),然后我将其粘贴到准备好的“框架”中。

import numpy as np
from scipy.signal import convolve2d

a = np.array([[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]])
b = np.array([[0, 0.25, 0], [0.25, 0, 0.25], [0, 0.25, 0]])
r = np.copy(a)
r[1:-1, 1:-1] = convolve2d(a, b, mode='valid')
r
# => array([[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]])

答案 1 :(得分:2)

虽然Amadan的scipy回答对于这种情况最有意义,但是这是另一种做法"手动":

import numpy as np

# Create your array
data = np.ones((5,5)) * 100
data[1:-1,1:-1] = 0

def evolve_heat_slow(m, should_copy=True):
    if should_copy: m = m.copy()

    components = (
        m[:-2,  1:-1],  # N
        m[2:,   1:-1],  # S
        m[1:-1, 2:],    # E
        m[1:-1, :-2],   # W
    )

    m[1:-1, 1:-1] = np.mean(np.stack(components), axis=0)
    return m

for _ in range(2):
    data = evolve_heat_slow(data)
    print(data)

在这里,我们首先采用中央3x3"窗口"定义组件。并在每个方向上将它移动1。然后我们堆叠移动的窗口,取平均值,并用这些值替换中心窗口。

经过1次迭代:

[[ 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.]]

经过2次迭代:

[[ 100.   100.   100.   100.   100. ]
 [ 100.    62.5   50.    62.5  100. ]
 [ 100.    50.    25.    50.   100. ]
 [ 100.    62.5   50.    62.5  100. ]
 [ 100.   100.   100.   100.   100. ]]

答案 2 :(得分:-1)

并不比其他方法更好,但您也可以在每个方向使用np.roll来做同样的事情:

def evolve_heat_slow(u):
    u2 = u.copy()
    u2[1:-1, 1:-1] = ((np.roll(u2,1,0) + np.roll(u2,-1,0) 
                        + np.roll(u2,1,1) + np.roll(u2,-1,1))/4)[1:-1, 1:-1]
    return u2

现在使用u2 = evolve_heat_slow(u)

u2 =
array([[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]])