Extract and set thick diagonal of numpy array

Question

I have a simple numpy question. How can I extract and consequently set a diagonal having a "thickness" equal to width with a constant value? I know the fill_diagonal function which fills the main diagonal with the given value. Similar to that I want to fill the main diagonal AND its surrounding diagonals. See banded diagonal matrix.

For example:

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)

So far I'm able to implement fill_banded in following way (which works):

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)

But I'm sure there is a better way of doing so in numpy / scipy. I could move this function in Cython, but I would keep that as a last alternative.

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]])

---

Explanation: np.add.outer can be used to make addition tables:

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]])

By changing the sign of one of the aranges (and using np.abs), you can measure the distance from the diagonal:

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]])

So you can "select" all elements that are a certain distance from the diagonal by writing a simple inequality:

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)

Once you have this boolean mask, you can assign new values to a with

a[mask] = val

---

Thus, fill_banded could look something like this:

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)