如何使用numpy在2D数组中构建对角数组？

Question

使用np.diag，我可以构造一个二维数组，在对角线上输入一个一维数组。但是如果我有n维数组作为输入，怎么做呢？

这个可以工作。

foo = np.random.randint(2, size=(36))
print foo
print np.diag(foo)
[1 1 1 1 1 1 0 1 0 0 1 1 0 1 1 1 1 1 0 0 0 0 0 0 0 1 0 1 0 1 0 0 0 1 1 0]
[[1 0 0 ..., 0 0 0]
 [0 1 0 ..., 0 0 0]
 [0 0 1 ..., 0 0 0]
 ..., 
 [0 0 0 ..., 1 0 0]
 [0 0 0 ..., 0 1 0]
 [0 0 0 ..., 0 0 0]]

这不会。

foo = np.random.randint(2, size=(2,36))
print foo
[[1 1 0 1 1 0 1 1 1 1 1 0 0 0 0 0 1 0 0 1 0 0 0 0 1 0 0 0 1 0 0 0 0 1 0 0]
 [0 1 1 1 0 0 1 1 1 1 0 1 0 1 0 1 0 0 0 0 1 0 1 1 1 1 0 0 1 0 1 1 1 1 0 1]]
do_something(foo)

应返回

array([[[ 1.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  1.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        ..., 
        [ 0.,  0.,  0., ...,  1.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.]],

       [[ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  1.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  1., ...,  0.,  0.,  0.],
        ..., 
        [ 0.,  0.,  0., ...,  1.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  0.],
        [ 0.,  0.,  0., ...,  0.,  0.,  1.]]])

编辑

这篇文章基于Alan和ajcr在这篇帖子中的回答，以及Saulo Castro和jaime的贡献，其中ajcr提到了链接。通常情况下，一切都取决于您的输入。我的输入通常具有以下形式：

M = np.random.randint(2, size=(1000, 36))

具有以下功能：

def Alan(M):
    M = np.asarray(M)
    depth, size = M.shape
    x = np.zeros((depth,size,size))
    for i in range(depth):
        x[i].flat[slice(0,None,1+size)] = M[i]
    return x

def ajcr(M):
    return np.eye(M.shape[1]) * M[:,np.newaxis,:]

def Saulo(M):
    b = np.zeros((M.shape[0], M.shape[1], M.shape[1]))
    diag = np.arange(M.shape[1])
    b[:, diag, diag] = M
    return b

def jaime(M):
    b = np.zeros((M.shape[0], M.shape[1]*M.shape[1]))
    b[:, ::M.shape[1]+1] = M
    return b.reshape(M.shape[0], M.shape[1], M.shape[1])

给我以下的结果。

%timeit Alan(M)
100 loops, best of 3: 2.22 ms per loop    
%timeit ajcr(M)
100 loops, best of 3: 5.1 ms per loop    
%timeit Saulo(M)
100 loops, best of 3: 4.33 ms per loop    
%timeit jaime(M)
100 loops, best of 3: 2.07 ms per loop

- Mattijn

Answer 1

在纯NumPy中完成此操作的简单方法是执行以下数组乘法：

np.eye(foo.shape[1]) * foo[:, np.newaxis]

其中 foo 是二维对角线数组。

这将NxN单位矩阵与foo的每一行相乘，得到所需的三维矩阵。

由于此方法的语法非常简单，您可以轻松扩展至更高维度。例如：

>>> foo = np.array([[0, 1], [1, 1]])
>>> d = np.eye(foo.shape[1]) * foo[:, np.newaxis] # 2D to 3D
>>> d
array([[[ 0.,  0.],
        [ 0.,  1.]],

       [[ 1.,  0.],
        [ 0.,  1.]]])

>>> np.eye(d.shape[1]) * d[:, :, np.newaxis] # 3D to 4D
array([[[[ 0.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]],

       [[[ 1.,  0.],
         [ 0.,  0.]],

        [[ 0.,  0.],
         [ 0.,  1.]]]])

这个问题可能与编程有关；它还展示了从2D数组中推导所需对角矩阵的更快（但略微冗长）的方法。

- Alex Riley

---

这种方法会创建额外的单位矩阵，并进行许多不必要的乘法运算。如果 foo 很大，这可能会变得很昂贵。 - Alan

当然 - 这种方法的主要优势在于简单的语法使得它相对容易地扩展到更高维度。对于较小的数组，这种方法仍然非常快速，但是对于较大的数组来说，利用索引的方法肯定会更高效（尤其是如果你使用的是NumPy 1.9）。 - Alex Riley

Answer 2

def makediag3d(a):
    a = np.asarray(a)
    depth, size = a.shape
    x = np.zeros((depth,size,size))
    for i in range(depth):
        x[i].flat[slice(0,None,1+size)] = a[i]
    return x

- Alan

---

我试图进行一些高级索引，并使用numpy.diag_indices_from，没想到我需要在这里使用for循环。 - Mattijn