Create block tridiagonal matrix

Question

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 return B

n=2 return [[B,I],[I,B]]

n=3 return [[B,I,0],[I,B,I],[0,I,B]]

n=4 return [[B,I,0,0],[I,B,I,0],[0,I,B,I],[0,0,I,B]]

and so on..

where 0 is just a zero array of the same size ofc. How can this be done?

check this out: https://stackoverflow.com/questions/5842903/block-tridiagonal-matrix-python or this: https://stackoverflow.com/questions/44641996/tridiagonal-block-matrix-using-scipy-sparse?rq=1 or this: https://stackoverflow.com/questions/46509744/what-is-the-best-way-to-create-a-block-matrix-form-a-row-vector?rq=1 — Daemon Painter, Oct 10 '18 at 10:13
link number 2 is very close to what i need. The difference is that n should be independet of the size of B. The answer there returns a matrix such that n equals the size of B. Trying to figure out thow to change properly, buty getting weird results — JustANoob, Oct 10 '18 at 10:48
@Sam: Can you explain what do you mean by n is independent of B. The above examples you gave can be solved by the solution posted in link 2. Can you post an example of when n is not equal to B — Rahul Agarwal, Oct 10 '18 at 11:25

Khalil Al Hooti · Answer 1 · 2018-10-10T12:41:54.090

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
def create_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 in range(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))

Create block tridiagonal matrix

1 Answers1

Linked