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

Question

我不知道这个标题是否有意义。通常单位矩阵是一个二维矩阵，如

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 单位矩阵”，因为前两个维度上的所有对角线元素都是 1。例如，对于形状 (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：沿第三维堆叠二维单位矩阵

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 分配给对角线元素的位置（第 1 维和第 2 维）：

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 作为最后一个轴的长度，我们可以这样做 -

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

运行示例让事情更清晰 -

In [154]: i = np.identity(3); i # Create an identity matrix
Out[154]: 
array([[ 1.,  0.,  0.],
       [ 0.,  1.,  0.],
       [ 0.,  0.,  1.]])

# Extend it to 3D. This helps us broadcast to reqd. shape later on
In [152]: i[...,None]
Out[152]: 
array([[[ 1.],
        [ 0.],
        [ 0.]],

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

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

# Broadcast to (n,n,r) shape for the 3D identity matrix    
In [153]: np.broadcast_to(i[...,None], (3,3,3))
Out[153]: 
array([[[ 1.,  1.,  1.],
        [ 0.,  0.,  0.],
        [ 0.,  0.,  0.]],

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

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

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

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

In [140]: n,r = 100,300

In [141]: %timeit np.broadcast_to(np.identity(n)[...,None], (n,n,r))
100000 loops, best of 3: 9.29 µs per loop

Answer 3

类似于@Psidom，使用高级np.einsum

%%timeit
shape = (100,100,300)
identity_3d = np.zeros(shape)
np.einsum('iij->ij', identity_3d)[:] = 1

1000 loops, best of 3: 251 µs per loop

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

1000 loops, best of 3: 320 µs per loop

%timeit np.broadcast_to(np.identity(100)[...,None], (100,100,300)).copy()

100 loops, best of 3: 12.1 ms per loop

我假设您希望写入单位矩阵的copy()，因为anyarray.dot(identity) 更容易由np.broadcast_to(anyarray[..., None], a.shape + (300,)) 计算。如果您真的只想要一个裸identity 矩阵，那么@Divakar 的解决方案比其他任何解决方案都要快，

%timeit np.broadcast_to(np.identity(100)[...,None], (100,100,300))

The slowest run took 5.16 times longer than the fastest. This could mean that an intermediate result is being cached.
10000 loops, best of 3: 20.8 µs per loop

但broadcast_to(anyarray[..., None], . . . ) 可能比这更快。