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

Question

我有一个简单的问题。如何提取并随后设置“厚度”等于 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 中移动此函数，但我会将其保留为最后的选择。

Answer 1

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的符号(并使用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)

一旦有了这个 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)