Извлечь и установить толстую диагональ массива numpy

У меня простой вопрос. Как я могу извлечь и, следовательно, установить диагональ, имеющую «толщину», равную width, с постоянным значением? Я знаю функцию fill_diagonal, которая заполняет главную диагональ данное значение. Аналогично этому я хочу заполнить основную диагональ и окружающие ее диагонали. См. диагональную матрицу с полосами.

Например:

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, но я бы оставил ее в качестве последней альтернативы.


person mg007    schedule 21.04.2013    source источник


Ответы (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]])

Изменив знак одной из aranges (и используя 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)
person unutbu    schedule 21.04.2013