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

时间:2014-11-30 13:59:34

标签: python arrays numpy

使用np.diag我能够在对角线上返回输入1-D数组的2-D数组。但是如果我有n-D数组作为输入怎么做呢?

这有效

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 refers来进行测试。像往常一样,这一切都取决于你的输入。我的输入通常具有以下形状:

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

2 个答案:

答案 0 :(得分:7)

一种简单的方法是在纯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.]]]])

This question可能具有相关性;它还显示了从2D数组中获得所需对角矩阵的更快(但稍微更详细)的方法。

答案 1 :(得分: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