创建包含单位矩阵的多维数组

Question

假设我有一个 Numpy 数组

n

索引，其中第一个

n-2

代表一些计数索引，最后 2 个索引代表一个方

MxM

矩阵。我想初始化这个结构，以便它将包含单位矩阵的副本。

示例（此处 N=3，M=2）：

A = numpy.zeros((3,2,2))
for k in range(3):
     A[k,:,:] = numpy.eye(2)

另一个例子（这里 N=4，M=5）：

B = numpy.zeros((3,4,5,5))
for k1 in range(3):
     for k2 in range(4):
         B[k1,k2,:,:] = numpy.eye(5)

是否有更紧凑的方法来代替（嵌套）循环来执行此操作？

Answer 1

你可以

repeat

:

A = np.repeat(np.eye(2)[None], 3, axis=0)

对于更复杂的情况，结合

reshape

：

extra = (3, 4)
M = 5
B = np.repeat(np.eye(M)[None], np.prod(extra), axis=0).reshape(extra+(M, M))