Solve Ax=b with sparse matrix A

The following code demonstrate the usage of cupy to solve Ax=b with sparse matrix A

'''
This program gives an example of using cupy to solve
large scale linear equation Ax=b with sparse matrix A.
'''

import fenics as fe
import numpy as np
import matplotlib.pyplot as plt
import cupy
import cupyx
import scipy.sparse as sps
import scipy.sparse.linalg as spsl
import time

## ------------------------------------------------------------------------
fe.parameters['linear_algebra_backend'] = 'Eigen'
def tran2SparseMatrix(A):
    row, col, val = fe.as_backend_type(A).data()
    return sps.csr_matrix((val, col, row))
## ------------------------------------------------------------------------
num_grid = 100
mesh = fe.UnitSquareMesh(num_grid, num_grid)
V = fe.FunctionSpace(mesh, 'P', 2)

u_D = fe.Expression('1 + x[0]*x[0] + 2*x[1]*x[1]', degree=2)

def boundary(x, on_boundary):
    return on_boundary

bc = fe.DirichletBC(V, u_D, boundary)

du, dv = fe.TrialFunction(V), fe.TestFunction(V)
f = fe.Expression('exp(-(pow(x[0]-0.5, 2)+pow(x[1]-0.5, 2)))', degree=3)
a = fe.dot(fe.grad(du), fe.grad(dv))*fe.dx
L = f*dv*fe.dx

u = fe.Function(V)
A, b = fe.assemble(a), fe.assemble(L)
bc.apply(A, b)
A_2 = tran2SparseMatrix(A)
b_2 = b[:]

start = time.time()
u.vector()[:] = spsl.lsqr(A_2, b_2)[:0]
end = time.time()
print(end - start)

u2 = fe.Function(V)
start = time.time()
A_2g = cupyx.scipy.sparse.csr_matrix(A_2)
b_2g = cupy.array(b_2)
u2.vector()[:] = cupy.asnumpy(cupyx.scipy.sparse.linalg.lsqr(A_2g, b_2g)[:1][0])
end = time.time()
print(end - start)

eng = fe.assemble(fe.inner(u-u2, u-u2)*fe.dx)
print(eng)