MATLAB / vs Python np.linalg.solve, heat diffusion

Question

I am solving a simple 1D steady heat question using spectral method. I am a long time MATLAB and Mathematica user and trying to learn Python. I compare only linear solving times and Python is way slower than MATLAB, which doesn't feel right. Can you please tell me why Python is an order of magnitude slower for only even linear solve?

MATLAB code

N=5000;
% HERE I AM JUST CREATING THE MATRIX FOR d/dx
Dhat=zeros(N+1,N+1);

for j=1:N;
for i=mod(j-1,2):2:j;
Dhat(i+1,j+1)=2*j;
end
end
Dhat(1,:)=Dhat(1,:)/2;
cbar=[2; ones(N-1,1); 2];
p=0:1:N;
pn=cos((pi/N)*kron(p',p));
I=ones(N+1,1);
II=(-1).^(0:N);
G=kron(I,II);
x=-cos(pi*(0:N)/N)';
T=G.*pn;
Tinv=(2/N)*(G'./kron(cbar,cbar')).*pn;

D=T*Dhat*Tinv;     % derivative matrix
% We have d/dx matrix

D2=D*D;            % d2/dx2 to solve d^2T/dx^2 = F

T0= 0;             % Temp at left boundary
Tn= 10;            % Temp at right boundary
F= zeros(N+1,1);   % source term
F(1)=0;            % Boundary condition
F(end)=10;         % Boundary condition

D2(1,:)=zeros;
D2(end,:)=zeros;
D2(1,1)=1;          % At the left boundary T = 0
D2(end,end)=1;      % At the right boundary T = 10
tic
T=D2/F';
toc

Python code

def chebgl(N):
    import numpy as np
    cbar = np.array([1] * (N - 1))
    cbar = np.insert(cbar, 0, 2)
    cbar = np.insert(cbar, N, 2)

    p = np.array(range(N + 1))
    pkp = np.kron(p, p)
    pkp = pkp.reshape(N + 1, N + 1)
    pn = np.cos((np.pi / N) * pkp)

    I = np.ones(N + 1)
    Ia = -1 * np.ones(N + 1)
    II = np.power(Ia, range(N + 1))
    G = np.kron(I, II)
    G = G.reshape(N + 1, N + 1)

    T = np.multiply(G, pn)
    GTrans = np.transpose(G)

    cbarK = np.kron(cbar, cbar)
    cbarK = cbarK.reshape(N + 1, N + 1)

    Tinv = (2 / N) * np.multiply((np.divide(GTrans, cbarK)), pn)

    dhat = np.zeros((N + 1, N + 1))

    for j in range(1, N + 1, 1):

        for i in range((j - 1) % 2, j, 2):
            dhat[i, j] = 2 * j

    dhat[0] = dhat[0] / 2
    # This is the operator for d/dx

    D = np.matmul(np.matmul(T, dhat), Tinv)
    # predefined x-locations
    x = -np.cos(np.pi * (p / N))

    return D, x

Main code

import numpy as np
import chebGL
import matplotlib.pyplot as plt
from time import process_time

# Because I want to see all double precision digits
np.set_printoptions(precision=15)
# Number of points (d^2/dx^2) T = F(x) is going to be solved  
N = 5000

D, x = chebGL.chebgl(N)

T0 = 0  # Left boundary condition
Tend = 10  # Right boundary condition
F = 0 * np.ones(N + 1)  # Source Term
F[0] = T0  # inputting the boundary condition
F[N] = Tend  # inputting the boundary condition

D2T = np.matmul(D, D)  # Creating d^2/dx^2 operator
D2T[0] = abs(0 * D2T[0])
D2T[N] = abs(0 * D2T[0])

D2T[0, 0] = 1  # inputting the left boundary condition
D2T[N, N] = 1  # inputting the right boundary condition

t1_start = process_time()

T = np.linalg.solve(D2T, F)
t1_stop = process_time()

print("Elapsed time during the whole program in seconds:", t1_stop -
      t1_start)

Answer 1

Can you please tell me why Python is an order of magnitude slower for only even linear solve?

Answering this generically is perhaps a little off-topic for review, but: first, it's important that we check the documentation for solve(), which explains

The solutions are computed using LAPACK routine _gesv.

For this solution step only, if MATLAB is indeed much faster, look at the fundamentals:

• Which flavour of BLAS does your instance of Numpy use? Have you used ATLAS? If not, why not?
• When you run your program, is there obvious occupancy of more than one CPU core? Does this change between MATLAB and Numpy?

Note that the slow parts of your program are not only between your t1_start and t1_stop. Any time that you do a matrix multiplication in chebGL, that's slow too. Since the dominant CPU time is spent in the lower-level matrix libraries and not Numpy itself, you can't blame Python itself for being slow - not even Numpy, necessarily; I would look to BLAS and make sure that the best implementation is selected and compiled correctly.

Though it doesn't make much of a difference in runtime, there's a long list of things you should work on in your usage of Numpy:

Don't call np.insert, and definitely don't use built-in list repetition. Create cbar at the correct size using np.ones and then overwrite the first and last entries with 2.

Use np.arange instead of range.

I don't think that your use of the Kronecker product is appropriate. You can just multiply a vector by its transpose; no reshaping necessary.

Don't use np.power to get a vector of alternating 1, -1. Just use slice assignment on a vector and then broadcast to a matrix.

Prefer .T instead of np.transpose.

Prefer *, / and @ operators instead of their corresponding functions np.multiply, np.divide and np.matmul.

You should not be recalculating x. Just negate the second row from pn which is equivalent.

Don't 0 * np.ones; just call np.zeros.

abs(0 * is strange and unnecessary; just write 0.

Add regression tests.

This:

Elapsed time during the whole program in seconds:

is a lie, and is only the elapsed time during np.linalg.solve.

Suggested

I believe this to be equivalent but you're going to want to scrutinise.

import numpy as np
from time import process_time


def chebgl(N: int) -> tuple[
    np.ndarray,  # D: (n+1)*(n+1)
    np.ndarray,  # x: (n+1)
]:
    cbar = np.ones((N + 1, 1))
    cbar[(0, -1), :] = 2
    cbarK = cbar * cbar.T

    p = np.arange(N + 1)[np.newaxis, :]
    pkp = p * p.T
    pn = np.cos(np.pi / N * pkp)

    g = np.ones_like(p)
    g[:, 1::2] = -1
    G = np.broadcast_to(g, pkp.shape)
    GTrans = G.T

    T = G * pn
    Tinv = 2 / N * GTrans / cbarK * pn

    # This is the operator for d/dx
    dhat = np.zeros_like(G)
    dhat[ ::2, 1::2] = np.arange(2,   2*N, 4)
    dhat[1::2,  ::2] = np.arange(0, 2+2*N, 4)
    dhat = np.triu(dhat)
    dhat[0, :] //= 2

    D = T @ dhat @ Tinv  # quite slow
    # predefined x-locations
    x = -pn[1, :]

    return D, x


def main() -> None:
    # Because I want to see all double precision digits
    np.set_printoptions(precision=15)
    # Number of points (d^2/dx^2) T = F(x) is going to be solved
    N = 5000

    D, x = chebgl(N)
    assert D.shape == (5001, 5001)
    assert x.shape == (5001,)

    T0 = 0     # Left boundary condition
    Tend = 10  # Right boundary condition
    F = np.zeros(N + 1)  # Source Term
    F[0] = T0    # inputting the boundary condition
    F[N] = Tend  # inputting the boundary condition

    D2T = D @ D  # Creating d^2/dx^2 operator
    D2T[(0, N), :] = 0
    D2T[0, 0] = 1  # inputting the left boundary condition
    D2T[N, N] = 1  # inputting the right boundary condition

    t1_start = process_time()
    T = np.linalg.solve(D2T, F)
    t1_stop = process_time()
    print("Elapsed time during the solve() in seconds:", t1_stop - t1_start)

    assert T.shape == (5001,)
    assert T.argmin() == 0
    assert T.argmax() == 5000
    assert np.isclose(0, T.min(),  rtol=0, atol=1e-7)
    assert np.isclose(5, T.mean(), rtol=0, atol=1e-7)
    assert np.isclose(5, T[2500],  rtol=0, atol=1e-7)
    assert np.isclose(10, T.max(), rtol=0, atol=1e-7)
    assert np.isclose(1.1434317494703112, T[1098], rtol=0, atol=1e-7)
    non_monotonic = np.count_nonzero(np.diff(T) < 0)
    assert non_monotonic == 0


if __name__ == '__main__':
    main()