Question

我试图在python中进行一些矩阵计算，当我尝试使用堆栈数组而不是简单的for循环来加快代码速度时遇到一个问题。我需要创建一个2D数组，在对角线上带有值（给定为1D数组），但无法找到一种聪明的方法来处理堆叠数组。

在旧的（循环）版本中，我使用了np.diag（）方法，如果我将值作为1D数组作为输入，该方法将精确返回我需要的值（在这种情况下为2D数组）。但是，当我切换到堆叠数组时，我的输入不再是一维数组，因此np.diag（）方法将返回我的二维输入对角线的副本。

具有1D输入的旧版本：

import numpy as np
vals = np.array([1,2,3])
mat = np.diag(vals)
print(mat.shape)

Out: (3, 3)

具有2D输入的新版本：

vals_stack = np.repeat(np.expand_dims(vals, axis=0), 5, axis=0) 
# btw: is there a better way to repeat/stack my array? 
mat_stack = np.diag(vals_stack)
print(mat_stack.shape)

Out: (3,)

因此您可以看到np.diag（）返回一个1D数组（如文档所期望的那样），但实际上我需要一堆2D数组。因此，mat_stack的形状必须为（7,3,3）而不是（3，）。 numpy中有任何功能吗？还是我必须像这样遍历其他维度：

def mydiag(stack):
    diag = np.zeros([stack.shape[0], stack.shape[1], stack.shape[1]])
    for i in np.arange(stack.shape[0]):
        diag[i,:,:] = np.diag([stack[i,:].ravel()])
    return diag

Answer 1

在numpy中，您应该使用apply_along_axis。文档末尾甚至有一个针对您的特定案例（here）的示例。答案是：

np.apply_along_axis(np.diag, -1, vals_stack)

更Python化的方式是这样的：

[np.diag(row) for row in vals_stack]

Answer 2

这是您的初衷吗？

In [499]: x = np.arange(12).reshape(4,3)                                                                     
In [500]: X = np.zeros((4,3,3),int)                                                                          
In [501]: X[np.arange(4)[:,None],np.arange(3), np.arange(3)] = x                                             
In [502]: X                                                                                                  
Out[502]: 
array([[[ 0,  0,  0],
        [ 0,  1,  0],
        [ 0,  0,  2]],

       [[ 3,  0,  0],
        [ 0,  4,  0],
        [ 0,  0,  5]],

       [[ 6,  0,  0],
        [ 0,  7,  0],
        [ 0,  0,  8]],

       [[ 9,  0,  0],
        [ 0, 10,  0],
        [ 0,  0, 11]]])

X[0,np.arange(3), np.arange(3)]在第一个平面上索引对角线。 np.arange(4)[:,None]是一个（4,1）数组，它与（3，）一起广播以索引（4,3）块，与x的大小匹配。

使用numpy根据对角线值创建数组堆栈

2 个答案: