TL; DR,我的单线:
mydiag=np.vectorize(np.diag, signature='(n)->(n,n)')
[[1, 2, 3],
[9, 8, 7]]
您想要的输出矩阵是
[[[1, 0, 0],
[0, 2, 0],
[0, 0, 3]],
[[9, 0, 0],
[0, 8, 0],
[0, 0, 7]]]
如果不是,您可以忽略此答案[编辑:与此同时,您确切地解释了这一点。因此,您可能会继续阅读]。
有很多方法可以做到这一点。 我的衬里会是
mydiag=np.vectorize(np.diag, signature='(n)->(n,n)')
构建一个可以执行您想要的新函数(它将输入解释为1d array的列表,请致电每个函数的np.diag,以获取2D阵列,然后将每个2D阵列放入numpy阵列中,因此获得3D阵列),然后你只打电话
mydiag(M)
矢量化的一种优势是它使用numpy广播。换句话说,循环是在C中执行的,而不是在Python中执行。换句话说,它更快。好吧,它应该是(在小矩阵上,实际上比迈克尔的方法慢 - 在评论中;在大型矩阵上,它具有完全相同的速度。这是令人沮丧的,因为Einsum doc本身指定它牺牲了广播)。
Plus,它是一个单线,除了在论坛上吹牛之外,别无其他兴趣。但是,我们在这里。
在这里是索引的一种方法:
out = np.zeros(data.shape+(data.shape[-1],), dtype=data.dtype)
x,y = np.indices(data.shape).reshape(2, -1)
out[x,y,y] = data.ravel()
输出:
array([[[1, 0, 0],
[0, 2, 0],
[0, 0, 3]],
[[4, 0, 0],
[0, 5, 0],
[0, 0, 6]]])
index1 = np.arange(2)[:, None] # 2 is the number of arrays
index2 = np.arange(3)[None, :] # 3 is the square size of each matrix
result = np.zeros((2, 3, 3))
result[index1, index2, index2] = data
result = np.zeros((2, 3*3)) result[:,::4]=data[:,:] result.shape=(2,3,3)