Search code examples
pythonpython-3.xscipysparse-matrix

Create arbitrary dimension sparse matrix in Python 3


So I'm using Python 3 to create a matrix of the form

L=[B 0 0
   0 B 0
   0 0 B]

where

B=[4 -1  0
  -1  4 -1
   0 -1  4]

But instead of repeating B three times, I want to do it N times (depending on input). My shot so far is the following

import numpy as np
import scipy.sparse as sp

one=np.ones(N)
four=4*np.ones(N)
data = np.array([-one, four, -one])
diags = np.array([-1,0, 1])
B=sp.spdiags(data, diags, N, N).toarray() # Create matrix B

L=np.kron(np.eye(N), B)

However, it takes a lot of time when N is big (which is necessary since this is to solve a differential equation). Is there a more effective way to do this?


Solution

  • No timings here (no performance guarantees from me), but for me the most natural approach (don't have much experience with kronecker) would be scipy's block_diag although i always wonder if i'm using it correctly (in this case: list-comprehension):

    Code

    import numpy as np
    import scipy.sparse as sp
    
    N = 2
    B = np.array([[4,-1,0],[-1,4,-1],[0,-1,4]])
    L = sp.block_diag([B for i in range(N)])
    
    print(L.todense())
    

    Out

    [[ 4 -1  0  0  0  0]
     [-1  4 -1  0  0  0]
     [ 0 -1  4  0  0  0]
     [ 0  0  0  4 -1  0]
     [ 0  0  0 -1  4 -1]
     [ 0  0  0  0 -1  4]]