Numpy:快速设置数组元素

时间:2016-02-02 11:43:05

标签: python numpy

我刚刚遇到了一部分代码的showstopper,我不确定我做错了什么......

我只是拥有一个大型数据立方体,并希望将z轴上的最大值更改为其他数字:

import numpy as np
from time import time

x, y, z = 100, 100, 10
a = np.arange(x*y*z).reshape((z, y, x))

t = time()
a[np.argmax(a, axis=0)] = 1
print(time() - t)

这需要大约0.02秒,这对于如此小的阵列来说有点慢,但还可以。我的问题是,我需要使用与(32, 4096, 4096)一样大的数组来执行此操作,并且我没有耐心让上述代码完成...它效率太低,但实际上应该是非常快!我是否在设置数组元素时出错?

2 个答案:

答案 0 :(得分:4)

您基本上使用包含数字的numpy数组索引numpy数组。我认为这就是为什么它如此缓慢的原因(我不确定它是否真的能按照你的意愿去做)。

如果你创建一个布尔numpy数组并将其用作切片,那么它的数量级会更快。

例如:

pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
pos_max_indices = np.arange(a.shape[0]).reshape(10,1,1) == pos_max
a[pos_max_indices] = 1

比原版快20倍,并且也是如此。

答案 1 :(得分:2)

我不认为数字的索引会减慢它的速度。通常使用布尔向量索引单个维度比使用相应的np.where索引要慢。

此处还有其他事情发生。看看这些形状:

In [14]: a.shape
Out[14]: (10, 100, 100)
In [15]: np.argmax(a,axis=0).shape
Out[15]: (100, 100)
In [16]: a[np.argmax(a,axis=0)].shape
Out[16]: (100, 100, 100, 100)

索引a远远大于原始版本1000x。

@ MSeifert的解决方案更快,但我无法感觉它比需要的更复杂。

In [35]: %%timeit 
   ....: a=np.arange(x*y*z).reshape((z,y,x))
   ....: pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
   ....: pos_max_indices = np.arange(a.shape[0]).reshape(10,1,1) == pos_max
   ....: a[pos_max_indices]=1
   ....: 
1000 loops, best of 3: 1.28 ms per loop

我还在努力改进。

示例数组不是一个好的 - 它太大而无法显示,以及最后一个z平面上的所有最大值:

In [46]: x,y,z=4,2,3
In [47]: a=np.arange(x*y*z).reshape((z,y,x))
In [48]: a
Out[48]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[16, 17, 18, 19],
        [20, 21, 22, 23]]])

In [49]: a[np.argmax(a,axis=0)]=1
In [50]: a
Out[50]: 
array([[[ 0,  1,  2,  3],
        [ 4,  5,  6,  7]],

       [[ 8,  9, 10, 11],
        [12, 13, 14, 15]],

       [[ 1,  1,  1,  1],
        [ 1,  1,  1,  1]]])

我可以使用:

访问相同的argmax
In [51]: a[-1,...]
Out[51]: 
array([[1, 1, 1, 1],
       [1, 1, 1, 1]])

让我们尝试一个随机数组,其中argmax可以在任何平面上:

In [57]: a=np.random.randint(2,10,(z,y,x))
In [58]: a
Out[58]: 
array([[[9, 7, 6, 5],
        [6, 3, 5, 2]],

       [[5, 6, 2, 3],
        [7, 9, 6, 9]],

       [[7, 7, 8, 9],
        [2, 4, 9, 7]]])
In [59]: a[np.argmax(a,axis=0)]=0
In [60]: a
Out[60]: 
array([[[0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0]],

       [[0, 0, 0, 0],
        [0, 0, 0, 0]]])

哎呀 - 我把一切都变成了0.这就是你想要的吗?

让我们试试pos_max方法:

In [61]: a=np.random.randint(0,10,(z,y,x))
In [62]: a
Out[62]: 
array([[[9, 3, 9, 0],
        [6, 6, 2, 4]],

       [[9, 9, 4, 9],
        [5, 9, 7, 9]],

       [[1, 8, 1, 7],
        [1, 0, 2, 3]]])
In [63]: pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
In [64]: pos_max
Out[64]: 
array([[[0, 1, 0, 1],
        [0, 1, 1, 1]]], dtype=int32)
In [66]: pos_max_indices = np.arange(a.shape[0]).reshape(z,1,1) == pos_max
In [67]: pos_max_indices
Out[67]: 
array([[[ True, False,  True, False],
        [ True, False, False, False]],

       [[False,  True, False,  True],
        [False,  True,  True,  True]],

       [[False, False, False, False],
        [False, False, False, False]]], dtype=bool)
In [68]: a[pos_max_indices]=0
In [69]: a
Out[69]: 
array([[[0, 3, 0, 0],
        [0, 6, 2, 4]],

       [[9, 0, 4, 0],
        [5, 0, 0, 0]],

       [[1, 8, 1, 7],
        [1, 0, 2, 3]]])

这看起来更合理。第二架飞机上仍然有一个9,但那是因为第一架飞机上还有9

这仍然需要清理,但这是一个非布尔掩码解决方案:

In [98]: a=np.random.randint(0,10,(z,y,x))
In [99]: a1=a.reshape(z,-1)   # it's easier to work with a 2d view
In [100]: ind=np.argmax(a1,axis=0)
In [101]: ind
Out[101]: array([2, 2, 1, 0, 2, 0, 1, 2], dtype=int32)
In [102]: a1[ind,np.arange(a1.shape[1])]  # the largest values
Out[102]: array([9, 8, 7, 4, 9, 7, 9, 6])
In [104]: a1
Out[104]: 
array([[3, 1, 5, 4, 2, 7, 4, 5],
       [4, 4, 7, 1, 3, 7, 9, 4],
       [9, 8, 3, 3, 9, 1, 2, 6]])

In [105]: a1[ind,np.arange(a1.shape[1])]=0
In [106]: a
Out[106]: 
array([[[3, 1, 5, 0],
        [2, 0, 4, 5]],

       [[4, 4, 0, 1],
        [3, 7, 0, 4]],

       [[0, 0, 3, 3],
        [0, 1, 2, 0]]])

使用a1 2d视图更容易; x,y维度的确切形状对此问题并不重要。我们正在改变单个值,而不是列或平面。我还是想在没有'a1。

的情况下让它工作

以下是两个替换最大值的函数(在第一个平面中)。我使用copy因为它使重复的时间测试更容易。

def setmax0(a, value=-1):
    # @MSeifert's
    a = a.copy()
    z = a.shape[0]
    # a=np.arange(x*y*z).reshape((z,y,x))
    pos_max = np.expand_dims(np.argmax(a, axis=0), axis=0)
    pos_max_indices = np.arange(z).reshape(z,1,1) == pos_max
    a[pos_max_indices]=value
    return a

def setmax1(a, value=-2):
    a = a.copy()
    z = a.shape[0]
    a1 = a.reshape(z, -1)
    ind = np.argmax(a1, axis=0)
    a1[ind, np.arange(a1.shape[1])] = value
    return a

它们在如下测试中产生相同的结果:

ab = np.random.randint(0,100,(20,1000,1000))
test = np.allclose(setmax1(ab,-1),setmax0(ab,-1))

计时(使用ipython timeit)基本相同。

他们会按不同的顺序分配值,因此setmax0(ab,-np.arange(...))会有所不同。