Извлечь и установить толщину диагонали матрицы numpy

Question

У меня вопрос простой numpy. Как я могу извлечь и, следовательно, установить диагональ, имеющую «толщину», равную 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 [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)