solving sparse Ax=b in scipy

1.3k Views Asked by At

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:

  1. What's the best way to solve Ax=b in scipy? I use the spsolve in the sp.sparse.linalg library, which uses the LU-decomposition. I tried using the standard sp.linalg.solve for dense matrix, but it's considerably slower. I also tried using all the other iterative solvers in the sp.sparse.linalg library, but they are also slower (for 1000x1000, they all take a couple seconds, compared to less than half a second for spsolve, and my A 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?

  1. I have installed umfpack. I tried setting use_umfpack True and False to test it, but surprisingly use_umfpack=True takes nearly twice longer than use_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.

0

There are 0 best solutions below