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!

Ch

Answer 1

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