Block tridiagonal matrix python

Question

I would like to create a block tridiagonal matrix starting from three numpy.ndarray. Is there any (direct) way to do that in python?

Thank you in advance!

Cheers

Answer 1

With "regular" numpy arrays, using numpy.diag :

def tridiag(a, b, c, k1=-1, k2=0, k3=1):
    return np.diag(a, k1) + np.diag(b, k2) + np.diag(c, k3)

a = [1, 1]; b = [2, 2, 2]; c = [3, 3]
A = tridiag(a, b, c)

Answer 2

Use the function scipy.sparse.diags .

Example:

from scipy.sparse import diags
import numpy as np

n = 10
k = [np.ones(n-1),-2*np.ones(n),np.ones(n-1)]
offset = [-1,0,1]
A = diags(k,offset).toarray()

This returns:

array([[-2.,  1.,  0.,  0.,  0.],
       [ 1., -2.,  1.,  0.,  0.],
       [ 0.,  1., -2.,  1.,  0.],
       [ 0.,  0.,  1., -2.,  1.],
       [ 0.,  0.,  0.,  1., -2.]])

Answer 3

You can also do this with "regular" numpy arrays through fancy indexing:

import numpy as np
data = np.zeros((10,10))
data[np.arange(5), np.arange(5)+2] = [5, 6, 7, 8, 9]
data[np.arange(3)+4, np.arange(3)] = [1, 2, 3]
print data

(You could replace those calls to np.arange with np.r_ if you wanted to be more concise. Eg instead of data[np.arange(3)+4, np.arange(3)] , use data[np.r_[:3]+4, np.r_[:3]] )

This yields:

[[0 0 5 0 0 0 0 0 0 0]
 [0 0 0 6 0 0 0 0 0 0]
 [0 0 0 0 7 0 0 0 0 0]
 [0 0 0 0 0 8 0 0 0 0]
 [1 0 0 0 0 0 9 0 0 0]
 [0 2 0 0 0 0 0 0 0 0]
 [0 0 3 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]
 [0 0 0 0 0 0 0 0 0 0]]

However, if you're going to be using sparse matrices anyway, have a look at scipy.sparse.spdiags . (Note that you'll need to prepend fake data onto your row values if you're placing data into a diagonal position with a positive value (eg the 3's in position 4 in the example))

As a quick example:

import numpy as np
import scipy as sp
import scipy.sparse

diag_rows = np.array([[1, 1, 1, 1, 1, 1, 1],
                      [2, 2, 2, 2, 2, 2, 2],
                      [0, 0, 0, 0, 3, 3, 3]])
positions = [-3, 0, 4]
print sp.sparse.spdiags(diag_rows, positions, 10, 10).todense()

This yields:

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

Answer 4

@TheCorwoodRep 's answer can actually be done in a single line. No need for a seperate function.

np.eye(3,3,k=-1) + np.eye(3,3)*2 + np.eye(3,3,k=1)*3

This produces:

array([[ 2.,  3.,  0.],
       [ 1.,  2.,  3.],
       [ 0.,  1.,  2.]])

Answer 5

For building a block-wise tridiagonal matrix from the three individual blocks (and repeat the blocks for N times), one solution can be:

import numpy as np
from scipy.linalg import block_diag

def tridiag(c, u, d, N): 
    # c, u, d are center, upper and lower blocks, repeat N times
    cc = block_diag(*([c]*N)) 
    shift = c.shape[1]
    uu = block_diag(*([u]*N)) 
    uu = np.hstack((np.zeros((uu.shape[0], shift)), uu[:,:-shift]))
    dd = block_diag(*([d]*N)) 
    dd = np.hstack((dd[:,shift:],np.zeros((uu.shape[0], shift))))
    return cc+uu+dd

For example:

c = np.matrix([[1,1],[1,1]])
u = np.matrix([[2,2],[2,2]])
d = -1*u
N =4
H = tridiag(c,u,d,N)
print(H)

gives the answer

[[ 1.  1.  2.  2.  0.  0.  0.  0.]
 [ 1.  1.  2.  2.  0.  0.  0.  0.]
 [-2. -2.  1.  1.  2.  2.  0.  0.]
 [-2. -2.  1.  1.  2.  2.  0.  0.]
 [ 0.  0. -2. -2.  1.  1.  2.  2.]
 [ 0.  0. -2. -2.  1.  1.  2.  2.]
 [ 0.  0.  0.  0. -2. -2.  1.  1.]
 [ 0.  0.  0.  0. -2. -2.  1.  1.]]

Answer 6

My answer builds of @TheCorwoodRep's answer. I am just posting it because I made a few changes to make it more modular so that it would work for different orders of matrices and also changing the values of k1 , k2 , k3 ie which decide where the diagonal appears, will take care of the overflow automatically. While calling the function you can specify what values should appear on the diagonals.

import numpy as np
def tridiag(T,x,y,z,k1=-1, k2=0, k3=1):
    a = [x]*(T-abs(k1)); b = [y]*(T-abs(k2)); c = [z]*(T-abs(k3))
    return np.diag(a, k1) + np.diag(b, k2) + np.diag(c, k3)

D=tridiag(10,-1,2,-1)

Answer 7

For better or worse, all the other answers seem to answer about tridiagonal matrices and not block tridiagonal matrices.

I don't think there is native support for tridiagonal matrices, so I wrote my own code. I had zeros on the main diagonal and my matrix was symmetric.

Here is my code.

n1 = 784
n2 = 256
n3 = 128
n4 = 10
M1 = np.ones((n1,n2))
M2 = np.ones((n2,n3))
M3 = np.ones((n3, n4))

def blockTri(Ms):
    #Takes in a list of matrices (not square) and returns a tridiagonal block matrix with zeros on the diagonal
    count = 0
    idx = []
    for M in Ms:
        #print(M.shape)
        count += M.shape[0]
        idx.append(count)
    count += Ms[-1].shape[-1]
    mat = np.zeros((count,count))
    count = 0
    for i, M in enumerate(Ms):
        mat[count:count+M.shape[0],idx[i]:idx[i]+M.shape[1]] = M
        count = count + M.shape[0]
    mat = mat + mat.T    
    return mat

M = blockTri([M1, M2, M3])

Hopefully this can help future people looking for block tridiagonal matrices.

Answer 8

Since tridiagonal matrix is a sparse matrix using a sparse package could be a nice option, see http://pysparse.sourceforge.net/spmatrix.html#matlab-implementation , there are some examples and comparisons with MATLAB even...

Answer 9

I just needed something similar. I share my workaround too.

from  functools import reduce
import numpy as np
size = 5
vals = [1,3,4]
res = reduce(lambda a, b: a+np.eye(size,k=b[0])*b[1], enumerate(vals), np.zeros((size, size)))
print(res)
# [[1. 3. 4. 0. 0.]
#  [0. 1. 3. 4. 0.]
#  [0. 0. 1. 3. 4.]
#  [0. 0. 0. 1. 3.]
#  [0. 0. 0. 0. 1.]]

size = 7
vals = [1,3,4,7,6]
offset = -2
res = reduce(lambda a, b: a+np.eye(size,k=b[0]+offset)*b[1], enumerate(vals), np.zeros((size, size)))
print(res)
# [[4. 7. 6. 0. 0. 0. 0.]
#  [3. 4. 7. 6. 0. 0. 0.]
#  [1. 3. 4. 7. 6. 0. 0.]
#  [0. 1. 3. 4. 7. 6. 0.]
#  [0. 0. 1. 3. 4. 7. 6.]
#  [0. 0. 0. 1. 3. 4. 7.]
#  [0. 0. 0. 0. 1. 3. 4.]]