如何在numpy中使用2d数组构造对角线数组？

Question

使用np.diag我能够构建一个二维数组,其中输入一维数组在对角线上返回.但如果我有nD数组作为输入,如何做同样的事情？

这有效

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和ajcr 所指的 jaime .像往常一样,这一切都取决于你的输入.我的输入通常有以下形状:

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

Answer 1

一个简单的方法是在纯NumPy中执行以下数组乘法:

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

foo对角线的2D数组在哪里.

这将NxN标识数组与每行相乘foo以产生所需的3D矩阵.

由于此方法中的语法非常简单,因此您可以轻松地将其扩展到更高的维度.例如:

>>> 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阵列导出所需对角矩阵的更快(但稍微更详细)的方法.