创建"3D身份矩阵"的最佳方法是什么?在Numpy中? [英] What's the best way to create a "3D identity matrix" in Numpy?

问题描述

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

I don't know if the title makes any sense. Normally an identity matrix is a 2D matrix like

In [1]: import numpy as np

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

，没有第三维.

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

Numpy can give me 3D matrix with all zeros

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)，应为

But I want a "3D identity matrix" in the sense that all diagonal elements along the first 2 dimensions are 1s. For example, for shape (2,2,3) it should be

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

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

是否有任何优雅的方法来生成此代码?

Is there any elegant way to generate this?

推荐答案

从2d身份矩阵开始，可以使用以下两个选项制作"3d身份矩阵" :

Starting from a 2d identity matrix, here are two options you can make the "3d identity matrix":

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

选项1 :沿三维堆叠2d身份矩阵

Option 1: stack the 2d identity matrix along the third dimension

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分配给对角线元素(第一维和第二维):

Option 3: Create an array of zeros and assign 1 to positions of diagonal elements (of the 1st and 2nd dimensions) using advanced indexing:

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

这篇关于创建"3D身份矩阵"的最佳方法是什么?在Numpy中?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！