Block tridiagonal matrix python

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)

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

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. E.g. 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 (e.g. 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]]

@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.]])

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