创建块三对角矩阵

时间:2018-10-10 10:03:03

标签: python

假设我有一个方形的numpy数组B和一个相同大小的标识数组I。我要在以下意义上创建一个给定整数n的块三对角矩阵:

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

以此类推。

其中0只是大小为c的零数组。如何才能做到这一点?

1 个答案:

答案 0 :(得分:0)

您可以尝试使用我构建的此功能,也许可以使用一些我不知道的numpy或scipy内置函数来实现。

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

尝试一些例子:

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

enter image description here