提取并设置numpy数组的粗对角线

时间:2013-04-21 14:06:16

标签: optimization numpy scipy diagonal

我有一个简单的问题。如何提取并因此将具有等于width的“厚度”的对角线设置为具有恒定值?我知道fill_diagonal函数用给定的值填充主对角线。类似于我想要填充主对角线及其周围的对角线。请参阅banded diagonal matrix

例如:

In [293]: a = np.random.randint(1, 100, (5,5)) % 2 == 0

In [294]: a
Out[294]: 
array([[ True,  True, False, False, False],
       [ True,  True, False,  True, False],
       [ True,  True, False, False,  True],
       [False, False, False,  True, False],
       [False, False, False, False,  True]], dtype=bool)

In [295]: fill_banded(a, val=True, width=3) # width must be odd number (?)

In [296]: a
Out[296]: 
array([[ True,  True, False, False, False],
       [ True,  True,  True,  True, False],
       [ True,  True,  True,  True,  True],
       [False, False,  True,  True,  True],
       [False, False, False,  True,  True]], dtype=bool)

到目前为止,我可以通过以下方式实现fill_banded(有效):

def fill_banded(a, val, width=1):
    # TODO: Add some error checking
    for i in range(width // 2):
        a[range(0,a.shape[0]-(i+1)),range(i+1,a.shape[1])] = val
        a[range(i+1,a.shape[0]),range(0,a.shape[1]-(i+1))] = val
    np.fill_diagonal(a, val)

但我确信在numpy / scipy中有更好的方法。我可以在Cython中移动这个函数,但我会保留它作为最后一个选择。

1 个答案:

答案 0 :(得分:4)

In [53]: a = np.arange(25).reshape(5,5)

In [54]: a
Out[54]: 
array([[ 0,  1,  2,  3,  4],
       [ 5,  6,  7,  8,  9],
       [10, 11, 12, 13, 14],
       [15, 16, 17, 18, 19],
       [20, 21, 22, 23, 24]])

In [55]: mask = np.abs(np.add.outer(np.arange(5), -np.arange(5))) < 3

In [56]: mask
Out[56]: 
array([[ True,  True,  True, False, False],
       [ True,  True,  True,  True, False],
       [ True,  True,  True,  True,  True],
       [False,  True,  True,  True,  True],
       [False, False,  True,  True,  True]], dtype=bool)

In [57]: a[mask] = 100

In [58]: a
Out[58]: 
array([[100, 100, 100,   3,   4],
       [100, 100, 100, 100,   9],
       [100, 100, 100, 100, 100],
       [ 15, 100, 100, 100, 100],
       [ 20,  21, 100, 100, 100]])

说明: np.add.outer可用于制作附加表:

In [59]: np.add.outer(np.arange(5), np.arange(5))
Out[59]: 
array([[0, 1, 2, 3, 4],
       [1, 2, 3, 4, 5],
       [2, 3, 4, 5, 6],
       [3, 4, 5, 6, 7],
       [4, 5, 6, 7, 8]])

通过更改其中一个arange s(并使用np.abs)的符号,您可以测量距对角线的距离:

In [61]: np.abs(np.add.outer(np.arange(5), -np.arange(5)))
Out[61]: 
array([[0, 1, 2, 3, 4],
       [1, 0, 1, 2, 3],
       [2, 1, 0, 1, 2],
       [3, 2, 1, 0, 1],
       [4, 3, 2, 1, 0]])

所以你可以通过写一个简单的不等式来“选择”与对角线相距一定距离的所有元素:

In [62]: np.abs(np.add.outer(np.arange(5), -np.arange(5))) < 3
Out[62]: 
array([[ True,  True,  True, False, False],
       [ True,  True,  True,  True, False],
       [ True,  True,  True,  True,  True],
       [False,  True,  True,  True,  True],
       [False, False,  True,  True,  True]], dtype=bool)

拥有此布尔掩码后,您可以使用

将新值分配给a
a[mask] = val

因此,fill_banded看起来像这样:

import numpy as np

def fill_banded(a, val, width=1):
    mask = np.abs(np.add.outer(np.arange(a.shape[0]), -np.arange(a.shape[1]))) < width
    a[mask] = val

a = np.arange(30).reshape(6,5)
fill_banded(a, val=True, width=3)
print(a)