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

Question

我正在尝试在 python 中进行一些矩阵计算，当我尝试使用堆叠数组而不是简单的 for 循环来加速我的代码时遇到了一个问题。我需要在对角线上创建一个带有值（以一维数组形式给出）的二维数组，但无法找到一种使用堆叠数组的聪明方法。

在旧的（循环）版本中，我使用了 np.diag() 方法，如果我将值作为一维数组作为输入。但是，当我切换到堆叠数组时，我的输入不再是一维数组，因此 np.diag() 方法 returns 是我的二维输入对角线的副本。

具有一维输入的旧版本：

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() returns 一个一维数组（正如文档中预期的那样），但我实际上需要一堆二维数组。所以 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)

更像 pythonic 的方式是这样的：

[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 从对角线值创建数组堆栈

Create stack of arrays from diagonal values using numpy

python

numpy

numpy-ndarray