Search code examples
pythonmathscipylinear-algebrasparse-matrix

How to change elements in sparse matrix in Python's SciPy?


I have built a small code that I want to use for solving eigenvalue problems involving large sparse matrices. It's working fine, all I want to do now is to set some elements in the sparse matrix to zero, i.e. the ones in the very top row (which corresponds to implementing boundary conditions). I can just adjust the column vectors (C0, C1, and C2) below to achieve that. However, I wondered if there is a more direct way. Evidently, NumPy indexing does not work with SciPy's sparse package.

import scipy.sparse as sp
import scipy.sparse.linalg  as la
import numpy as np
import matplotlib.pyplot as plt

#discretize x-axis
N = 11
x = np.linspace(-5,5,N)
print(x)
V = x * x / 2
h = len(x)/(N)
hi2 = 1./(h**2)
#discretize Schroedinger Equation, i.e. build 
#banded matrix from difference equation
C0 = np.ones(N)*30. + V
C1 = np.ones(N) * -16.
C2 = np.ones(N) * 1.
diagonals = np.array([-2,-1,0,1,2])
H = sp.spdiags([C2, C1, C0,C1,C2],[-2,-1,0,1,2], N, N)
H *= hi2 * (- 1./12.) * (- 1. / 2.)
#solve for eigenvalues
EV = la.eigsh(H,return_eigenvectors = False)

#check structure of H
plt.figure()
plt.spy(H)
plt.show()

This is a visualisation of the matrix that is build by the code above. I want so set the elements in the first row zero.enter image description here


Solution

  • As suggested in the comments, I'll post the answer that I found to my own question. There are several matrix classes in in SciPy's sparse package, they are listed here. One can convert sparse matrices from one class to another. So for what I need to do, I choose to convert my sparse matrix to the class csr_matrix, simply by

    H = sp.csr_matrix(H)
    

    Then I can set the elements in the first row to 0 by using the regular NumPy notation:

    H[0,0] = 0
    H[0,1] = 0
    H[0,2] = 0
    

    For completeness, I post the full modified code snippet below.

    #SciPy Sparse linear algebra takes care of sparse matrix computations
    #http://docs.scipy.org/doc/scipy/reference/sparse.linalg.html
    import scipy.sparse as sp
    import scipy.sparse.linalg  as la
    
    import numpy as np
    import matplotlib.pyplot as plt
    
    #discretize x-axis
    N = 1100
    x = np.linspace(-100,100,N)
    V = x * x / 2.
    h = len(x)/(N)
    hi2 = 1./(h**2)
    
    #discretize Schroedinger Equation, i.e. build 
    #banded matrix from difference equation
    C0 = np.ones(N)*30. + V
    C1 = np.ones(N) * -16.
    C2 = np.ones(N) * 1.
    
    H = sp.spdiags([C2, C1, C0, C1, C2],[-2,-1,0,1,2], N, N)
    H *= hi2 * (- 1./12.) * (- 1. / 2.)
    H = sp.csr_matrix(H)
    H[0,0] = 0
    H[0,1] = 0
    H[0,2] = 0
    
    #check structure of H
    plt.figure()
    plt.spy(H)
    plt.show()
    
    EV = la.eigsh(H,return_eigenvectors = False)