[FIXED] Efficient way of creating a band (or banded) matrix

Issue

I'm trying to create a certain style of band(ed) matrix (see wikipedia). The following code works, but for large M (~300 or so) it becomes quite slow because of the for loop. Is there a way to vectorize it/make better use of numpy and/or scipy? I am having trouble figuring out the mathematical operation that this corresponds to, and hence I have not succeeded thus far.

The code I have is as follows

def banded_matrix(M):
    phis = np.linspace(0, 2*np.pi, M)
    i = 0
    ham = np.zeros((int(2*M), int(2*M)))
    for phi in phis:
        ham_phi = np.array([[1, 1],
                            [1, -1]])*(1+np.cos(phi))
        array_phi = np.zeros(M)
        array_phi[i] = 1
        mat_phi = np.diag(array_phi)
        ham += np.kron(mat_phi, ham_phi)
        i += 1
    return ham

With %timeit banded_matrix(M=300) it takes about 4 seconds on my computer.

Since the code is a bit opaque, what I want to do is construct a large 2M by 2M matrix. In a sense it has M entries on it's 'width 2' diagonal, where the entries are 2x2 matrices ham_phi that depend on phi. The matrix will afterwards be diagonalized, so perhaps one could even make use of its structure/the fact that it is rather sparse to speed that up, but of that I am not sure.

If anyone has an idea where to go with this, I'd be happy to follow up on that!

Solution

Your matrix is diagonal by blocks, so you can use scipy.linalg.block_diag:

import numpy as np
from scipy.linalg import block_diag

def banded_matrix_scipy(M):
    ham = np.array([[1, 1], [1, -1]])
    phis = np.linspace(0, 2 * np.pi, M)
    ham_phis = ham * (1 + np.cos(phis))[:, None, None]
    return block_diag(*ham_phis)

Let's check that it works and is faster:

b1 = banded_matrix(300)
b2 = banded_matrix_scipy(300)

np.all(b1 == b2)  # True

>>> %timeit banded_matrix(300)
>>> %timeit banded_matrix_scipy(300)
1.51 s ± 57 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
1.24 ms ± 4.57 µs per loop (mean ± std. dev. of 7 runs, 1,000 loops each)

Answered By - paime