Create Block Tridiagonal Matrix

Lets say i have a square numpy array B and an identity array I of the same size. I want to create a block tridiagonal matrix , given an integer n, in the following sense: n=1 retur

Solution 1:

You can try this function I built, probably there is a simpler way to do it using some numpy or scipy built-in functions I am not aware off.

import numpy as np
defcreate_block_tridiagonal(n, block_size):
    if block_size < 2:
        raise ValueError('block_size should be at leat 2')
    if n < 1:
        raise ValueError('n cannot be less than 1')


    B = np.ones((block_size, block_size))*2# matrix of all 2
    I = np.identity(block_size)

    if n == 1:
        return B

    else:
        if n%block_size !=0:
            raise ValueError('Cant broadcast square array sizes of\
                             {} into the diagonal of of a matrix\
                             of size {}'.format(block_size, n))

        a = np.zeros((n*block_size, n*block_size))

        for i inrange(a.shape[0]//block_size):
            m = int(i*block_size)
            a[m:m+block_size, m:m+block_size] = B


            if i == 0:
                k = m+block_size
                a[m:m+block_size, k:k+block_size] = I
            elif i == a.shape[0]//block_size - 1:
                k=m-block_size
                a[m:m+block_size, k:k+block_size] = I

            else:
                a[m:m+block_size, m-block_size:m] = I
                a[m:m+block_size, m+block_size:m+2*block_size] = I
        return a

Try some examples:

#n=1
create_block_tridiagonal(1, 2)
array([[2., 2.],
       [2., 2.]])
#n=2
create_block_tridiagonal(2, 2)

array([[2., 2., 1., 0.],
       [2., 2., 0., 1.],
       [1., 0., 2., 2.],
       [0., 1., 2., 2.]])
#n=4
create_block_tridiagonal(4, 2)
array([[2., 2., 1., 0., 0., 0., 0., 0.],
       [2., 2., 0., 1., 0., 0., 0., 0.],
       [1., 0., 2., 2., 1., 0., 0., 0.],
       [0., 1., 2., 2., 0., 1., 0., 0.],
       [0., 0., 1., 0., 2., 2., 1., 0.],
       [0., 0., 0., 1., 2., 2., 0., 1.],
       [0., 0., 0., 0., 1., 0., 2., 2.],
       [0., 0., 0., 0., 0., 1., 2., 2.]])
#n=3, block size=3
create_block_tridiagonal(3, 3)
array([[2., 2., 2., 1., 0., 0., 0., 0., 0.],
       [2., 2., 2., 0., 1., 0., 0., 0., 0.],
       [2., 2., 2., 0., 0., 1., 0., 0., 0.],
       [1., 0., 0., 2., 2., 2., 1., 0., 0.],
       [0., 1., 0., 2., 2., 2., 0., 1., 0.],
       [0., 0., 1., 2., 2., 2., 0., 0., 1.],
       [0., 0., 0., 1., 0., 0., 2., 2., 2.],
       [0., 0., 0., 0., 1., 0., 2., 2., 2.],
       [0., 0., 0., 0., 0., 1., 2., 2., 2.]])
# n=6, block size=3
print(create_block_tridiagonal(6, 3))