Crank-Nicolson. Advection - Diffusion

Question

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[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

    C = adv_diff(D, w, t)

Answer 1

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(..)
• Use docstrings instead of comments to document your functions.

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 
    C = np.asarray(list(adv_diff(D, w, t)))

Answer 2

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):
    """Solver for AD PDE equation""" 
    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
    C = np.array(list(adv_diff(D, w, t)))

EDIT:I also updated cosmetics of the code according to the comments of @MathiasEttinger (no increasing performance here).