Generating banded matrices using numpy

Question

I'm using the following piece of code to create a banded matrix from a generator g:

def banded(g, N):
    """Creates a `g` generated banded matrix with 'N' rows"""
    n=len(g)
    T = np.zeros((N,N+n-1))
    for x in range(N):
        T[x][x:x+n]=g
    return T

The usage is simple as:

banded([1,2,3], 3)

and it returns

[1, 2, 3, 0, 0]
[0, 1, 2, 3, 0]
[0, 0, 1, 2, 3]

It will used mostly to solve a finite difference model with a given stencil for example (-1, 1)

Is there a better way to generate that stencil? I could not find any good NumPy Function for that.

By better I mean, faster, using less memory, removing the loop from python and sending to Numpy stack. Any (or all) of those are improvements.

Answer 1

Here's one with np.lib.stride_tricks.as_strided to give us a 2D view into a zeros padded 1D version of the input and as such pretty memory efficient and hence performant too. This trick had been explored numerous times - 1,2.

Thus, the implementation would be -

def sliding_windows(a, W):
    a = np.asarray(a)
    p = np.zeros(W-1,dtype=a.dtype)
    b = np.concatenate((p,a,p))
    s = b.strides[0]
    strided = np.lib.stride_tricks.as_strided
    return strided(b[W-1:], shape=(W,len(a)+W-1), strides=(-s,s))

Sample runs -

In [99]: a = [1,2,3]

In [100]: sliding_windows(a, W=3)
Out[100]: 
array([[1, 2, 3, 0, 0],
       [0, 1, 2, 3, 0],
       [0, 0, 1, 2, 3]])

In [101]: a = [1,2,3,4,5]

In [102]: sliding_windows(a, W=3)
Out[102]: 
array([[1, 2, 3, 4, 5, 0, 0],
       [0, 1, 2, 3, 4, 5, 0],
       [0, 0, 1, 2, 3, 4, 5]])

With the same philosophy, but less messier version, we can also leverage np.lib.stride_tricks.as_strided based scikit-image's view_as_windows to get sliding windows. More info on use of as_strided based view_as_windows.

from skimage.util.shape import view_as_windows

def sliding_windows_vw(a, W):
    a = np.asarray(a)
    p = np.zeros(W-1,dtype=a.dtype)
    b = np.concatenate((p,a,p))
    return view_as_windows(b,len(a)+W-1)[::-1]