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
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[0] = 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[1] = 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[0] = 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


• after profiling the code, I found out that the slowest part is "np.linalg.solve". Thus, I realised that my matrices were not sparse and simple optimisation of AL AR matrices resolved the problem: 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. Jan 15, 2017 at 18:19
• Removing numba decorators also improved the performance. Jan 15, 2017 at 18:45
• Could you provide sample data? It is hard to speed up code without being able to run it. Jan 15, 2017 at 22:01
• You can write your ideas for improvement as a self-answer. Alternatively, if no answer has been posted yet, you can just replace the code in the question with your faster code. Jan 15, 2017 at 22:45
• Within the time loop, the solver is the only thing that takes a significant amount of time. I'd be temped to clean up the creation of AL and AR, but that's only happening once, so isn't a time consumer. numba can't touch functions imported from scipy or numpy functions that are already compiled. So it's no use here. Jan 16, 2017 at 8:03

I know nothing about your business logic but you're trying to transpose matlab in Python and it doesn't look like anything: there are some useless parameters, you are building arrays only to take their size. I'll try to provide advices to clean up the code and maybe others will provide runtime improvements.

• Remove unused variables like F_bottom or tend which is only an alias to years.
• Replaces calls like arr = np.linspace(begin, end, size); var = arr.size with var = size if you don't plan on using the array at all.
• Use slice assignment to eliminate visual clutter.
• Remove update_bc_dirichlet as you already perform the same kind of operation directly in the loop: stay consistent and see previous point.
• Remove linalg_solver as it only indirects a single function call. If you want to avoid lookups in the loop, create an alias as local variable.
• Some of the constants should either be extracted out of the function as real constants or defined as parameter with default values.
• 0 * np.ones(..) is np.zeros(..)

I would also turn adv_diff into a generator, but only because I don't know numpy enough to come up with a more idiomatic code:

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, L=30, phi=1, dx=0.1, dt=0.001):
"""What does this function do?"""
N = int(L / dx) + 1
C1 = np.zeros((N, 1))
C1[0] = 1

linalg_solver = np.linalg.solve
AL, AR = AL_AR_dirichlet(D, w, phi, dt, dx, N)

yield C1
for _ in range(int(year/dt)):
C1 = linalg_solver(AL, AR.dot(C1))
yield C1
C1[:2, 0] = 1, 1

@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, :2] = 1, 0
AL[-1, -2:] = -s, 1 + s
AR[0, :2] = 1, 0
AR[-1, -2:] = s, 1 - s

return AL, AR

if __name__ == '__main__':
D = 0.5
w = 3
t = 10

• Thanks for advices. However, when I run your code I got Segmentation fault: 11 Jan 16, 2017 at 1:29
• @IgorMarkelov Not very instructive. Anything else? As I can't test the code for a few days. Jan 16, 2017 at 7:52
• Ok. I found the problem. There was a typo. After a fix, a receive a different error at C1[:2] = 1, 1 "could not broadcast input array from shape (2) into shape (2,1)". I have the question regarding the yield. Would it be better to allocate the size of the resulting array before the actual computation (in terms of memory managment)? Jan 16, 2017 at 17:20
• @IgorMarkelov The error is easily fixed using C1[:2, 0] = 1, 1 instead. For the memory management, I don't really know, you'd have to copy each output anyway, either right after computing it or at the last np.asarray call. Both would use numpy's operations that are already compiled and optimized/vectorized. Jan 16, 2017 at 17:31
• Some minor improvement: replace divisions by multiplications (/2 --> *0.5. Etc). dt and dx (and L and N) seem to be known(=constant), maybe just set them as global variables and get rid of some function arguments? Jun 16, 2017 at 23:30

After a bit of investigation of the issue, I came up with the following. I changed the solver type which gave 5x increase in performance with enabled umfpack (you may do so pip install scikit-umfpack. However, it is still 4x times slower than MATLAB. It seems that MATLAB uses an advantage of tridiagonal matrices and scipy/numpy do not. I used tridiagonal solver, but it was much slower than regular sparse solver. Now I am looking into scipy.linalg.solve_banded.

If you have any further notes regarding the possible increase in performance - please let me know.

import numpy as np
import scipy as sp
from scipy.sparse import spdiags

def adv_diff(D, w, years, L=30, phi=1, dx=0.1, dt=0.001, bc_top=1):
N = int(L / dx) + 1

[AL, AR] = AL_AR_dirichlet(D, w, phi, dt, dx, N)

res = np.zeros((N))
res[0] = bc_top

yield res

for _ in range(int(years / dt)):
res = sp.sparse.linalg.spsolve(AL, AR.dot(res), use_umfpack=True)
res[0] = bc_top
yield res

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, format='csc')
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, format='csc')
AL[0, :2] = 1, 0
AL[-1, -2:] = -s, 1 + s
AR[0, :2] = 1, 0
AR[-1, -2:] = s, 1 - s
return AL, AR

if __name__ == '__main__':
D = 0.5
w = 3
t = 10