I am writing an advection-diffusion solver in Python. I am quite experienced in MATLAB and, therefore, the code implementation looks very close to possible implementation in MATLAB. I implemented the same code in MATLAB and execution time there is much faster. AS you may note, I am also (as in MATLAB) trying to use JIT in Python, however, it does not give me any improvements. Thus, I was wondering if you could review the code in terms of optimization for computational speed and give me some advices for future.
import numpy as np import scipy as sp from scipy.sparse import spdiags from scipy import special import numba as nb @nb.jit def adv_diff(D, w, years): BC1_top = 1 F_bottom = 0 L = 30 tend = years C1_init = 0 phi = 1 dx = 0.1 dt = 0.001 x = np.linspace(0, L, L / dx + 1) N = x.size C1_init = C1_init * np.ones((N, 1)) C1_init = BC1_top [AL, AR] = AL_AR_dirichlet(D, w, phi, dt, dx, N) C1_old = C1_init time = np.linspace(0, tend, tend / dt + 1) C1_res = np.zeros((N, time.size)) C1_res[:, 0] = C1_init[:, 0] for i in np.arange(1, len(time)): C1_old = update_bc_dirichlet(C1_old, BC1_top) B = AR.dot(C1_old) C1_new = linalg_solver(AL, B) C1_res[:, i] = C1_new[:, 0] C1_old = C1_new C1_old = BC1_top return C1_res @nb.jit def linalg_solver(A, b): # linalg_solver: x = A \ b return np.linalg.solve(A, b) # @nb.jit def update_bc_dirichlet(C, BC_top): # update_bc_dirichlet: function description C = BC_top return C @nb.jit def AL_AR_dirichlet(D, w, phi, dt, dx, N): # AL_AR_dirichlet: creates AL and AR matrices with Dirichlet BC s = phi * D * dt / dx / dx # q = phi * w * dt / dx # e1 = np.ones((N, 1)) # AL = spdiags(np.concatenate((e1 * (-s / 2 - q / 4), e1 * (1 + s), e1 * (-s / 2 + q / 4)), axis=1).T, [-1, 0, 1], N, N).toarray() AR = spdiags(np.concatenate((e1 * (s / 2 + q / 4), e1 * (1 - s), e1 * (s / 2 - q / 4)), axis=1).T, [-1, 0, 1], N, N).toarray() AL[0, 0] = 1 AL[0, 1] = 0 AL[N - 1, N - 1] = 1 + s AL[N - 1, N - 1 - 1] = -s AR[0, 0] = 1 AR[0, 1] = 0 AR[N - 1, N - 1] = 1 - s AR[N - 1, N - 1 - 1] = s return AL, AR if __name__ == '__main__': D = 0.5 w = 3 t = 10 C = adv_diff(D, w, t)
spdiags(np.concatenate((e1 * (s / 2 + q / 4), e1 * (1 - s), e1 * (s / 2 - q / 4)), axis=1).T, [-1, 0, 1], N, N, format = 'csc')and
sp.sparse.linalg.spsolve(A, b)produced very similar to MATLAB results. \$\endgroup\$
solveris the only thing that takes a significant amount of time. I'd be temped to clean up the creation of
AR, but that's only happening once, so isn't a time consumer.
numbacan't touch functions imported from
numpyfunctions that are already compiled. So it's no use here. \$\endgroup\$