Python 使用numpy从对角线值创建数组堆栈

我试图用python做一些矩阵计算，当我试图使用堆叠数组而不是简单的for循环来加速代码时遇到了一个问题。我需要创建一个在对角线上有值的二维数组（给定为1D数组），但无法找到一种使用堆叠数组的智能方法

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

具有1D输入的旧版本：

将numpy导入为np
VAL=np.数组（[1,2,3]）
材料=np.诊断（VAL）
印刷品（垫形）

具有2D输入的新版本：

vals\u stack=np.重复（np.展开尺寸（vals，轴=0），5，轴=0）
#顺便问一下：有没有更好的方法来重复/堆叠我的数组？
mat_stack=np.diag（VAL_stack）
打印（材料堆叠形状）

因此，您可以看到np.diag（）返回一个1D数组（正如文档中所预期的），但实际上我需要一个2D数组堆栈。因此，mat_堆栈的形状必须是（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

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

---

在numpy中，应使用

沿_轴应用

。在文档的末尾甚至有一个针对您的特定案例（）的示例。所以答案是：

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

一种更具蟒蛇风格的方式是这样的：

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

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

---

你是这么想的吗

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

的大小相匹配，这正是我想要的，它工作起来像一个符咒！非常感谢。