在Numpy中创建“ 3D身份矩阵”的最佳方法是什么？

Question

我不知道标题是否有意义。通常，单位矩阵是2D矩阵，例如

In [1]: import numpy as np

In [2]: np.identity(2)
Out[2]: 
array([[ 1.,  0.],
       [ 0.,  1.]])

而且没有第三维。

Numpy可以为我提供全零的3D矩阵

In [3]: np.zeros((2,2,3))
Out[3]: 
array([[[ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

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

但是我想要一个“ 3D身份矩阵”，因为前两个维度上的所有对角元素均为1s。例如，对于形状（2,2,3），应为

array([[[ 1.,  1.,  1.],
        [ 0.,  0.,  0.]],

       [[ 0.,  0.,  0.],
        [ 1.,  1.,  1.]]])

有什么优雅的方法可以产生这种效果吗？

Answer 1

从2d身份矩阵开始，可以使用以下两个选项来制作“ 3d身份矩阵”：

import numpy as np    
i = np.identity(2)

选项1：沿三维堆叠2D身份矩阵

np.dstack([i]*3)
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

选项2：重复值，然后重塑

np.repeat(i, 3, axis=1).reshape((2,2,3))
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

选项3：创建零数组，并使用高级索引将1分配给对角线元素（第一维和第二维）：

shape = (2,2,3)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1  


identity_3d
#array([[[ 1.,  1.,  1.],
#        [ 0.,  0.,  0.]],

#       [[ 0.,  0.,  0.],
#        [ 1.,  1.,  1.]]])

时间：

%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.repeat(i, shape[2], axis=1).reshape(shape)

# 10 loops, best of 3: 10.1 ms per loop


%%timeit
shape = (100,100,300)
i = np.identity(shape[0])
np.dstack([i] * shape[2])

# 10 loops, best of 3: 47.2 ms per loop

%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
idx = np.arange(shape[0])
identity_3d[idx, idx, :] = 1

# 100 loops, best of 3: 6.31 ms per loop

Answer 2

一种方法是初始化2D单位矩阵，然后将其广播到3D. 因此，对于沿前两个轴和最后一个轴的n长度，我们可以这样做 -r

\n\n

np.broadcast_to(np.identity(n)[...,None], (n,n,r))\n

\n\n

示例运行以使事情更清楚 -

\n\n

In [154]: i = np.identity(3); i # Create an identity matrix\nOut[154]: \narray([[ 1.,  0.,  0.],\n       [ 0.,  1.,  0.],\n       [ 0.,  0.,  1.]])\n\n# Extend it to 3D. This helps us broadcast to reqd. shape later on\nIn [152]: i[...,None]\nOut[152]: \narray([[[ 1.],\n        [ 0.],\n        [ 0.]],\n\n       [[ 0.],\n        [ 1.],\n        [ 0.]],\n\n       [[ 0.],\n        [ 0.],\n        [ 1.]]])\n\n# Broadcast to (n,n,r) shape for the 3D identity matrix    \nIn [153]: np.broadcast_to(i[...,None], (3,3,3))\nOut[153]: \narray([[[ 1.,  1.,  1.],\n        [ 0.,  0.,  0.],\n        [ 0.,  0.,  0.]],\n\n       [[ 0.,  0.,  0.],\n        [ 1.,  1.,  1.],\n        [ 0.,  0.,  0.]],\n\n       [[ 0.,  0.,  0.],\n        [ 0.,  0.,  0.],\n        [ 1.,  1.,  1.]]])\n

\n\n

这种方法可以提高性能，因为它只是生成单位矩阵的视图。因此，在这种形式下，输出将是一个只读数组。如果您需要一个具有自己的内存空间的可写数组，只需.copy()在此处附加一个即可。

\n\n

断言性能，这里有一个创建 3D形状为单位矩阵的时序测试：(100, 100, 300) -

\n\n

In [140]: n,r = 100,300\n\nIn [141]: %timeit np.broadcast_to(np.identity(n)[...,None], (n,n,r))\n100000 loops, best of 3: 9.29 \xc2\xb5s per loop\n

\n