Numpy partition diagonal matrix

Question

I want to generate a partition diagonal matrix A like this

and the matrix B is given

B = -np.diag(np.ones(n - 2), -1) - np.diag(np.ones(n - 2), 1) + 4 * np.diag(np.ones(n - 1))

For example,

Is there a way to do it without using loop?

---

Sorry to upload the figure of matrix A and B incorrectly in the first time.

Answer 1

You can stack the building blocks into a lookup table and then build A by indexing into it:

>>> from scipy import sparse
>>> 
>>> n = 5
>>> B = sparse.diags([-1, 4, -1], [-1, 0, 1], (n-1, n-1), dtype=int).A
>>> A = sparse.diags([1, 2, 1], [-1, 0, 1], (n-1, n-1), dtype=int).A
# 0 means 0 0 0 ...,
# 1 means -I
# 2 means B
>>>
# next line builds the lookup table (using np.stack)
# does the lookup ...[A]
# and flattens the resulting 4D array after swapping
# the middle axes; the swap reorders the entries from
# Vert, Horz, vert, horz   to   Vert, vert, Horz, horz
>>> A = np.stack([np.zeros_like(B), -np.identity(n-1, int), B])[A].swapaxes(1, 2).reshape((n-1)*(n-1), -1)
>>> A
array([[ 4, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [-1,  4, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0, -1,  4, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  0,  0],
       [ 0,  0, -1,  4,  0,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0,  0],
       [-1,  0,  0,  0,  4, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0,  0],
       [ 0, -1,  0,  0, -1,  4, -1,  0,  0, -1,  0,  0,  0,  0,  0,  0],
       [ 0,  0, -1,  0,  0, -1,  4, -1,  0,  0, -1,  0,  0,  0,  0,  0],
       [ 0,  0,  0, -1,  0,  0, -1,  4,  0,  0,  0, -1,  0,  0,  0,  0],
       [ 0,  0,  0,  0, -1,  0,  0,  0,  4, -1,  0,  0, -1,  0,  0,  0],
       [ 0,  0,  0,  0,  0, -1,  0,  0, -1,  4, -1,  0,  0, -1,  0,  0],
       [ 0,  0,  0,  0,  0,  0, -1,  0,  0, -1,  4, -1,  0,  0, -1,  0],
       [ 0,  0,  0,  0,  0,  0,  0, -1,  0,  0, -1,  4,  0,  0,  0, -1],
       [ 0,  0,  0,  0,  0,  0,  0,  0, -1,  0,  0,  0,  4, -1,  0,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0, -1,  0,  0, -1,  4, -1,  0],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0, -1,  0,  0, -1,  4, -1],
       [ 0,  0,  0,  0,  0,  0,  0,  0,  0,  0,  0, -1,  0,  0, -1,  4]])

Please note that the sparse constructor is only used for its convenience. The sparse matrices are immediately converted to dense (using .A property).