TheWaveLad
TheWaveLad

Reputation: 1016

Tridiagonal Matrix with Matrix Blocks

I read this, but I wasn't able to create a (N^2 x N^2) - matrix A with (N x N) - matrices I on the lower and upper side-diagonal and T on the diagonal. I tried this

def prep_matrix(N):
    I_N = np.identity(N)
    NV = zeros(N*N - 1)
    # now I want NV[0]=NV[1]=...=NV[N-1]:=I_N

but I have no idea how to fill NV with my matrices. What can I do? I found a lot on how to create tridiagonal matrices with scalars, but not with matrix blocks.

Upvotes: 1

Views: 1140

Answers (2)

JuniorCompressor
JuniorCompressor

Reputation: 20015

You could initialize like this:

n = 3
I, T, A = np.identity(n), np.ones(n), np.zeros([n * n, n * n])
for i in range(n):
    a, b, c = i * n, (i + 1) * n, (i + 2) * n
    A[a:b, a:b] = T
    if i < n - 1: A[b:c, a:b] = A[a:b, b:c] = I

The above example has the following output:

[[ 1.  1.  1.  1.  0.  0.  0.  0.  0.]
 [ 1.  1.  1.  0.  1.  0.  0.  0.  0.]
 [ 1.  1.  1.  0.  0.  1.  0.  0.  0.]
 [ 1.  0.  0.  1.  1.  1.  1.  0.  0.]
 [ 0.  1.  0.  1.  1.  1.  0.  1.  0.]
 [ 0.  0.  1.  1.  1.  1.  0.  0.  1.]
 [ 0.  0.  0.  1.  0.  0.  1.  1.  1.]
 [ 0.  0.  0.  0.  1.  0.  1.  1.  1.]
 [ 0.  0.  0.  0.  0.  1.  1.  1.  1.]] 

Upvotes: 1

Dietrich
Dietrich

Reputation: 5531

For the sake of art, another approach based on bmat():

Z = np.zeros((n, n))
I = np.eye(n)
T = np.ones((n, n))*2

B = np.diag([1]*n) + np.diag([2]*np.ones(n-1), 1)  + np.diag(2*np.ones(n-1), -1)
# B = array([[ 1.,  2.,  0.],
#            [ 2.,  1.,  2.],
#            [ 0.,  2.,  1.]])
# build 2d list `Bm` replacing 0->Z, 1->T, 2->I:
bdict = {0.:Z, 1.:T, 2.:I}
Bm = [[bdict[i] for i in rw] for rw in B]
# Use the power of bmat to construct matrix:
A = np.asarray(np.bmat(Bm))

Upvotes: 1

Related Questions