使用numpy高效地计算标准基向量

Question

给定一个索引和一个大小，有没有更高效的方法来生成标准基向量：

import numpy as np
np.array([1.0 if i == index else 0.0 for i in range(size)])

- Neil G

Answer 1

In [2]: import numpy as np

In [9]: size = 5

In [10]: index = 2

In [11]: np.eye(1,size,index)
Out[11]: array([[ 0.,  0.,  1.,  0.,  0.]])

哎呀，不幸的是，使用np.eye进行此操作相当慢：

In [12]: %timeit np.eye(1,size,index)
100000 loops, best of 3: 7.68 us per loop

In [13]: %timeit a = np.zeros(size); a[index] = 1.0
1000000 loops, best of 3: 1.53 us per loop

将np.zeros(size); a[index] = 1.0封装成一个函数只会带来一点点改善，但仍然比np.eye要快得多：

In [24]: def f(size, index):
   ....:     arr = np.zeros(size)
   ....:     arr[index] = 1.0
   ....:     return arr
   ....: 

In [27]: %timeit f(size, index)
1000000 loops, best of 3: 1.79 us per loop

- unutbu

---

哇，我看了一眼，因为文档上说：“返回一个二维数组，对角线上为1，其他地方为0”，所以就跳过了它。谢谢。 - Neil G

是的，我认为这样的事情是可能的，但我有点惊讶numpy使用eye来实现这一点（因为名称暗示它是单位矩阵，而这绝对不是一个单位矩阵）。比我的答案更好的回答加1。 - mgilson

@NeilG：嗯，也许这不是一个好的解决方案。timeit结果显示它比mgilson和JoranBeasley的解决方案慢5倍。 - unutbu

@unutbu: 你介意比较一下，另一个解决方案是包装在一个接受大小和索引的函数中吗？（我刚刚比较了一下，在我的机器上只慢了2倍）。 - Neil G

@unutbu：非常感谢您对这个答案（以及您其他好的答案）的认真解答。 - Neil G

Answer 2

x = np.zeros(size)
x[index] = 1.0

至少我认为是这样的...

>>> t = timeit.Timer('np.array([1.0 if i == index else 0.0 for i in range(size)]
)','import numpy as np;size=10000;index=5123')
>>> t.timeit(10)
0.039461429317952934  #original method
>>> t = timeit.Timer('x=np.zeros(size);x[index]=1.0','import numpy as np;size=10000;index=5123')
>>> t.timeit(10)
9.4077963240124518e-05 #zeros method
>>> t = timeit.Timer('x=np.eye(1.0,size,index)','import numpy as np;size=10000;index=5123')
>>> t.timeit(10)
0.0001398340635319073 #eye method

看起来np.zeros是最快的...

- Joran Beasley

---

有趣的是，我们两个都没有在“=”周围保持一致的空格（我已经更新了我的）。 - mgilson

嘿，已经修复了空格问题 :P - Joran Beasley

Answer 3

我不确定这样是否更快，但对于我来说肯定更清晰易懂。

a = np.zeros(size)
a[index] = 1.0

- mgilson

Answer 4

通常，你不仅需要一个基向量，而是需要所有的基向量。如果是这种情况，请考虑使用np.eye：

basis = np.eye(3)
for vector in basis:
  ...

虽然不完全相同，但有密切关联：通过一些技巧，这甚至可以用来获取一组基矩阵：

>>> d, e = 2, 3    # want 2x3 matrices
>>> basis = np.eye(d*e,d*e).reshape((d*e,d,e))
>>> print(basis)
[[[ 1.  0.  0.]
  [ 0.  0.  0.]]

 [[ 0.  1.  0.]
  [ 0.  0.  0.]]

 [[ 0.  0.  1.]
  [ 0.  0.  0.]]

 [[ 0.  0.  0.]
  [ 1.  0.  0.]]

 [[ 0.  0.  0.]
  [ 0.  1.  0.]]

 [[ 0.  0.  0.]
  [ 0.  0.  1.]]]

等等。

- coldfix

Answer 5

虽然不是最快的方法，但scipy.signal.unit_impulse这种方法将上述概念推广到任意形状的numpy数组中。

- Neil G

Answer 6

另一种实现方法如下：

>>> def f(size, index):
...     return (np.arange(size) == index).astype(float)
... 

这种方法的执行速度略慢：

>>> timeit.timeit('f(size, index)', 'from __main__ import f, size, index', number=1000000)
2.2554846050043125

- Oussa