scipy.sparse: Set row/column in sparse matrix to the identity without changing sparsity

Question

I'm using the SciPy sparse.csr_matrix format for a finite element code. In applying the essential boundary conditions, I'm setting the desired value in the right hand side vector, and setting the corresponding row and column to the identity (i.e. 1 on the diagonal and zeros elsewhere.)

The code I'm using is similar to:

K = csr_matrix((data,(rows,cols)))
for idx in boundary:
    K[:,idx] = 0.0
    K[idx,:] = 0.0
    K[idx,idx] = 1.0

This code results in inserting values for all of the entries in the idx row and column which is extremely slow.

I'd like to be able to set the diagonal entry to 1 (which shouldn't alter the sparsity), and only zero the entries in the idx row and column which aren't already zero.

Is there a way to do this?

Answer 1

I think it is most efficient if you work directly on the data of the CSR matrix. Have a look at this explanation of the CSR data structure - it is in german, but the example may help you to understand.

Below is a python script with a function that sets a specified row of a CSR matrix to match the identity:

In [13]: run t_csrdiag.py
sparse mat A was before
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 8  9 10 11]
 [12 13 14 15]]

index of the row that is modified: 2

sparse mat A is now
[[ 0  1  2  3]
 [ 4  5  6  7]
 [ 0  0  1  0]
 [12 13 14 15]]

Feel free to loop it or vectorize it.

Here is the code:

import scipy.sparse as sps
import numpy as np


def main():
    # index of the row to become the identity row
    rowind = 2

    A = np.arange(16).reshape((4, 4))
    A = sps.csr_matrix(A)
    print 'sparse mat A was before'
    print A.todense()

    print '\nindex of the row that is modified: {0}'.format(rowind)

    setcsrrow2id(A, rowind)
    print '\nsparse mat A is now'
    print A.todense()


def setcsrrow2id(amat, rowind):
    indptr = amat.indptr
    values = amat.data
    indxs = amat.indices

    # get the range of the data that is changed
    rowpa = indptr[rowind]
    rowpb = indptr[rowind+1]

    # new value and its new rowindex
    values[rowpa] = 1.0
    indxs[rowpa] = rowind

    # number of new zero values
    diffvals = rowpb - rowpa - 1

    # filter the data and indices and adjust the range
    values = np.r_[values[:rowpa+1], values[rowpb:]]
    indxs = np.r_[indxs[:rowpa+1], indxs[rowpb:]]
    indptr = np.r_[indptr[:rowind+1], indptr[rowind+1:]-diffvals]

    # hard set the new sparse data
    amat.indptr = indptr
    amat.data = values
    amat.indices = indxs

if __name__ == '__main__':
    main()

Answer 2

You can use the method nonzero to return an array of all the non-zero entries of a given row or column, for example:

for i in boundary:
    _, J = A[i,:].nonzero()   # ignored the first return value, which is just [i,...,i]
    for j in J:
        if j!=i:
            A[i,j] = 0.0
    A[i,i] = 1.0

There's certainly a clever way of getting rid of the conditional, but you get the idea.

Also, I think you only want to set the row corresponding to an index i in the boundary to be the identity, not the column. If the nodes are re-ordered so that the boundary nodes are at the beginning, your matrix should look like

$A = \left[\begin{matrix}I & 0 \\ B & K\end{matrix}\right]$

where $K$ represents the coupling between internal nodes and $B$ the influence of the boundary conditions on the internal nodes.

Answer 3

This is my very simple function to eliminate all the entries in the bc_id matrix rows, and keeping only the diagonal entry. This how I usually set boundary conditions. I intensively used it. I know it is not optimal, it is not really elegant, but whenever I profiled my code, this function has never been a problem. Probably it would be good idea to include an assertion to check if the matrix is squared.

from scipy import sparse
import numpy as np
import matplotlib.pyplot as plt


def set_diag(A,bc_id):
    ndofs = A.shape[0]
    diago = A.diagonal()

    uno = np.ones((ndofs))
    uno[bc_id] = 0
    uno = sparse.dia_matrix((uno,0), shape = (ndofs,ndofs))
    # up to here I delete the rows
    # I multiply A by an identity matrix
    # where i set to zero the rows I want
    # to delete
    A = uno.dot(A)
    new_diag_entries = np.zeros((ndofs))
    # here I set the diagonal entries
    # equals to the value on the diagonal.
    # Use the second line if you want to
    # to set the diagonal entry to one
    new_diag_entries[bc_id] = diago[bc_id]
    #new_diag_entries[bc_id] = uno[bc_id]
    uno = sparse.dia_matrix((new_diag_entries,0), shape = (ndofs,ndofs))
    A = A+uno # here I set the diagonal entry
    return A

A = sparse.csr_matrix([[1, 2, 3, 10], [4,5, 6,11], [7, 8, 9,12],[13,14,15,16]])

bc_id = np.array([2,3])

A = set_diag(A,bc_id)

print A

Answer 4

You can formulate this at the purely linear algebra level in terms of csr matrix multplications and a matrix addition, which are fast. I have found this to be the fastest and most memory efficient method in scipy on my laptop.

Specifically, let $M$ be the original matrix, and $M'$ be the modified version with rows and columns associated with boundary nodes zeroed out, except for the diagonal entry in which the value of 1 is placed. Then

$$M' = I_\text{interior} ~ M ~ I_\text{interior} + I_\text{boundary},$$

where $I_\text{interior}$ is the diagonal matrix with ones in locations corresponding to interior degrees of freedom and zeros elsewhere, and contrarily $I_\text{boundary}$ is the diagonal matrix with ones in the locations corresponding to the boundary degrees of freedom and zeros elsewhere.

So, the overall strategy is:

1. Form the diagonal matrices $I_\text{interior}$ and $I_\text{boundary}$.

2. Convert $I_\text{interior}$ and $I_\text{boundary}$ to CSR format.

3. Form the modified matrix $M' = I_\text{interior} ~ M ~ I_\text{interior} + I_\text{boundary},$.

Here is the code that I use:

def dirichlet_zero_matrix_modification(matrix, dirichlet_zero_dofs):
    """Enforces dirichlet zero B.C.'s by zeroing-out rows and columns 
    associated with dirichlet dofs, and putting 1 on the diagonal there."""
    N = matrix.shape[1]

    chi_interior = np.ones(N)
    chi_interior[dirichlet_zero_dofs] = 0.0
    I_interior = sps.spdiags(chi_interior, [0], N, N).tocsr()

    chi_boundary = np.zeros(N)
    chi_boundary[dirichlet_zero_dofs] = 1.0
    I_boundary = sps.spdiags(chi_boundary, [0], N, N).tocsr()

    matrix_modified = I_interior * matrix * I_interior + I_boundary
    return matrix_modified

In my experience with scipy, expressing matrix manipulations in terms of higher-level linear algebraic operations (where possible) almost always yields better performance, compared to other methods that actually manipulate matrix entries or slices. In other languages and frameworks this may not be the case.

Answer 5

Although this does not directly answer to your question, let me provide a yet another simple approach to impose the essential boundary conditions that I use in my own FEM library. The following code is adapted to your use case. In the following A is a csr_matrix and b is the right hand side NumPy array.

import scipy.sparse.linalg as spl
import numpy as np
D=boundary
I=np.setdiff1d(np.arange(A.shape[0]),D)
x[I]=spl.spsolve(A[I].T[I].T,b[I]-A[I].T[D].T.dot(x[D]),use_umfpack=use_umfpack)

The advantage is that you do not have to modify the matrix. You simply drop unnecessary rows and columns and modify the right hand side which seems to be rather quick.