I need to solve Ax=b
where A
is the matrix that represents finite difference method for PDEs. Typical size of A
for a 2D problem is around (256^2)x(256^2), and it consists of a few diagonals. The following sample code is how I construct A:
N = Nx*Ny # Nx is no. of cols (size in x-direction), Ny is no. rows (size in y-direction)
# finite difference in x-direction
up1 = (0.5)*c
up1[Nx-1::Nx] = 0
down1 = (-0.5)*c
down1[::Nx] = 0
matX = diags([down1[1:], up1[:-1]], [-1,1], format='csc')
# finite difference in y-direction
up1 = (0.5)*c
down1 = (-0.5)*c
matY = diags([down1[Nx:], up1[:N-Nx]], [-Nx,Nx], format='csc')
Adding matX
and matY
together results in four diagonals. The above is for second-order discretization. For fourth-order discretization, then I have eight diagonals. If I have second derivative, then the main diagonal is nonzero as well.
I use the following code to do the actual solving:
# Initialize A_fixed, B_fixed
if const is True: # the potential term V(x) is time-independent
A = A_fixed + sp.sparse.diags(V_func(x))
B = B_fixed + sp.sparse.diags(V_func(x))
A_factored = sp.sparse.linalg.factorized(A)
for idx, t in enumerate(t_steps[1:],1):
# Solve Ax=b=Bu
if const in False: #
A = A_fixed + sp.sparse.diags(V_func(x,t))
B = B_fixed + sp.sparse.diags(V_func(x,t))
psi_n = B.dot(psi_old)
if const is True:
psi_new = A_factored(psi_n)
else:
psi_new = sp.sparse.linalg.spsolve(A,psi_n,use_umfpack=False)
psi_old = psi_new.copy()
I have a couple questions:
- What's the best way to solve
Ax=b
in scipy? I use thespsolve
in thesp.sparse.linalg
library, which uses the LU-decomposition. I tried using the standardsp.linalg.solve
for dense matrix, but it's considerably slower. I also tried using all the other iterative solvers in thesp.sparse.linalg
library, but they are also slower (for 1000x1000, they all take a couple seconds, compared to less than half a second forspsolve
, and myA
is likely to be a lot bigger). Are there any alternative ways to do the computation?
Edit: The problem I'm trying to solve is actually the time-dependent Schrodinger Equation. If the potential term is time-independent, then as suggested I can pre-factorize the matrix A
first to speed up the code, but this doesn't work if the potential is time-varying, as I need to change the diagonal term of both matrices A
and B
at each time step. For this specific problem, are there any ways to speed up the code using method similar to pre-factorization or other ways?
- I have installed
umfpack
. I tried settinguse_umfpack
True and False to test it, but surprisinglyuse_umfpack=True
takes nearly twice longer thanuse_umfpack=False
. I thought having this package should give a speed up. Any idea what that's the case? (PS: I am using Anaconda Spyder to run the code if that makes any difference)
I have use cProfile
to profile my codes, and nearly all the time is spent on that line with spsolve
. So I think the remaining part of the code (matrix /problem initialization) is pretty much optimized, and it's the matrix solving part that needs to be improved.
Thanks.