Search code examples
pythonmatrixscipyblockdiag

Create a block diagonal matrix in python using loops


Some basic help on how to construct large block diagonal matrices would be welcome. Preferably solving this with the block_diag function in python.

I would like to construct two matrices: one of the form diag(A,A...A) and one of the form diag(k,A,A...A,k), where A is a 2x2 matrix and k is a real number. The matrices should have the same size and I would like to give an input N for how many A's I want in there.

So far this code works (although I don't know what it does exactly; if I just use the block_diag it does not seem to work):

me = sparse.block_diag(A for _ in range(int(N)))

Doing the second matrix this way errors (in different ways depending on how I put my brackets, like 'expression should be paranthesized' or 'int object is not iterable'):

mo = sparse.block_diag(k,A for _ in range(int(N)),k)

Could I get some help how to do this in an understandable way?


Solution

  • block_diag expects a sequence of matrices as its first input (see the documentation). In your first line, (A for _ in range(N)) creates a generator expression (note that the parentheses in your code are from the function call, and the generator expression is automatically created inside it. On its own, you'll need parentheses around the expression).

    In the second line, you basically pass three arguments to block_diag: k, (A for _ in range(N)) and k again as third argument. You need to combine them. There are a few ways to do that:

    • [k] + list(A for _ in range(N)) + [k]
    • [k, *[A for _ in range(N)], k]

    and probably some variations, where the latter uses list unpacking to create multiple arguments for the list.

    So I ended up with

    mo = scipy.sparse.block_diag([k, *[A for _ in range(N)], k])