Brusselator Stiff Partial Differential Equation Benchmark: Julia DifferentialEquations.jl vs Python SciPy. Julia 935x faster

a_pde_bruss_scaling_benchmark_python_julia.md

Brusselator Stiff Partial Differential Equation Benchmark: Julia DifferentialEquations.jl vs Python SciPy

Tested is DifferentialEquations.jl vs Python's SciPy ODE solvers. Notes:

• Stiff ODE solvers are used since they are required to solve this problem effectively.
• The Python code is vectorized with for maximum performance
• All of the performance features are tried: automatic sparsity detection, preconditioners, etc.

Results

• Julia best (Preconditioned GMRES): 72.949 ms, or 935x faster than fastest SciPy
• Julia automatic algorithm choice: 2.08 s, or 33x faster than fastest SciPy
• Python BDF: 68.2s
• Python LSODA: 444.8 s
• Python Radau: 344 s

bruss.jl

using OrdinaryDiffEq, LinearAlgebra, SparseArrays, BenchmarkTools

const N = 32
const xyd_brusselator = range(0,stop=1,length=N)
brusselator_f(x, y, t) = (((x-0.3)^2 + (y-0.6)^2) <= 0.1^2) * (t >= 1.1) * 5.
limit(a, N) = a == N+1 ? 1 : a == 0 ? N : a
function brusselator_2d_loop(du, u, p, t)
  A, B, alpha, dx = p
  alpha = alpha/dx^2
  @inbounds for I in CartesianIndices((N, N))
    i, j = Tuple(I)
    x, y = xyd_brusselator[I[1]], xyd_brusselator[I[2]]
    ip1, im1, jp1, jm1 = limit(i+1, N), limit(i-1, N), limit(j+1, N), limit(j-1, N)
    du[i,j,1] = alpha*(u[im1,j,1] + u[ip1,j,1] + u[i,jp1,1] + u[i,jm1,1] - 4u[i,j,1]) +
                B + u[i,j,1]^2*u[i,j,2] - (A + 1)*u[i,j,1] + brusselator_f(x, y, t)
    du[i,j,2] = alpha*(u[im1,j,2] + u[ip1,j,2] + u[i,jp1,2] + u[i,jm1,2] - 4u[i,j,2]) +
                A*u[i,j,1] - u[i,j,1]^2*u[i,j,2]
    end
end
p = (3.4, 1., 10., step(xyd_brusselator))

function init_brusselator_2d(xyd)
  N = length(xyd)
  u = zeros(N, N, 2)
  for I in CartesianIndices((N, N))
    x = xyd[I[1]]
    y = xyd[I[2]]
    u[I,1] = 22*(y*(1-y))^(3/2)
    u[I,2] = 27*(x*(1-x))^(3/2)
  end
  u
end
u0 = init_brusselator_2d(xyd_brusselator)
prob_ode_brusselator_2d = ODEProblem(brusselator_2d_loop,u0,(0.,11.5),p)

du = similar(u0)
brusselator_2d_loop(du, u0, p, 0.0)
du[34] # 802.9807693762164
du[1058] # 985.3120721709204
du[2000] # -403.5817880634729
du[end] # 1431.1460373522068
du[521] # -323.1677459142322

du2 = similar(u0)
brusselator_2d_loop(du2, u0, p, 1.3)
du2[34] # 802.9807693762164
du2[1058] # 985.3120721709204
du2[2000] # -403.5817880634729
du2[end] # 1431.1460373522068
du2[521] # -318.1677459142322

using Symbolics
du0 = copy(u0)
jac_sparsity = Symbolics.jacobian_sparsity((du,u)->brusselator_2d_loop(du,u,p,0.0),du0,u0)

f = ODEFunction(brusselator_2d_loop;jac_prototype=float.(jac_sparsity))
prob_ode_brusselator_2d_sparse = ODEProblem(f,u0,(0.,11.5),p,tstops=[1.1])

using IncompleteLU
function incompletelu(W,du,u,p,t,newW,Plprev,Prprev,solverdata)
  if newW === nothing || newW
    Pl = ilu(convert(AbstractMatrix,W), τ = 150.0)
  else
    Pl = Plprev
  end
  Pl,nothing
end

# Required due to a bug in Krylov.jl: https://github.com/JuliaSmoothOptimizers/Krylov.jl/pull/477
Base.eltype(::IncompleteLU.ILUFactorization{Tv,Ti}) where {Tv,Ti} = Tv

sol = solve(prob_ode_brusselator_2d_sparse,KenCarp47(linsolve=KrylovJL_GMRES(),precs=incompletelu,concrete_jac=true),abstol=1e-12,reltol=1e-12)
sol.t[104] # 1.1
sol[104][1] # 0.5305455584661771
sol[end][1] # 3.2723204220222195

@btime solve(prob_ode_brusselator_2d_sparse,KenCarp47(linsolve=KrylovJL_GMRES(),precs=incompletelu,concrete_jac=true),save_everystep=false,abstol=1e-8,reltol=1e-8);
# 1.422 s (304441 allocations: 140.77 MiB)

using DifferentialEquations
@btime solve(prob_ode_brusselator_2d,alg_hints=[:stiff],save_everystep=false,abstol=1e-8,reltol=1e-8);
# 2.076 s (551306 allocations: 23.87 MiB)

using Sundials, ModelingToolkit
prob_ode_brusselator_2d_mtk = ODEProblem(modelingtoolkitize(prob_ode_brusselator_2d_sparse),[],(0.0,11.5),jac=true,sparse=true);

using LinearAlgebra
u0 = prob_ode_brusselator_2d_mtk.u0
p  = prob_ode_brusselator_2d_mtk.p
const jaccache = prob_ode_brusselator_2d_mtk.f.jac(u0,p,0.0)
const WW = I - 1.0*jaccache

prectmp = ilu(WW, τ = 50.0)
const preccache = Ref(prectmp)

function psetupilu(p, t, u, du, jok, jcurPtr, gamma)
  if jok
    prob_ode_brusselator_2d_mtk.f.jac(jaccache,u,p,t)
    jcurPtr[] = true

    # W = I - gamma*J
    @. WW = -gamma*jaccache
    idxs = diagind(WW)
    @. @view(WW[idxs]) = @view(WW[idxs]) + 1

    # Build preconditioner on W
    preccache[] = ilu(WW, τ = 5.0)
  end
end

function precilu(z,r,p,t,y,fy,gamma,delta,lr)
  ldiv!(z,preccache[],r)
end

@btime solve(prob_ode_brusselator_2d_sparse,CVODE_BDF(linear_solver=:GMRES,prec=precilu,psetup=psetupilu,prec_side=1),save_everystep=false);
# 72.949 ms (15412 allocations: 55.75 MiB)

bruss.py

import numpy as np
from scipy import sparse as sp
from scipy import integrate as solver
import math
import scipy

def ff(x,y,t):
    if (((x-0.3)**2+(y-0.6)**2<=0.01) and (t>=1.1)):
        return 5
    else:
        return 0

def r0(x,y):
    return [22*(y*(1-y))**1.5,27*(x*(1-x))**1.5]

alpha, A, B, C = (10,3.4,1,4.4)
nx=30
ny=30
L=1

A_x = sp.diags([[1],[1]*(nx+1),[-2]*(nx+2),[1]*(nx+1),[1]],[-(nx+1),-1,0,1,nx+1])
A_y = sp.diags([[1],[1]*(nx+1),[-2]*(nx+2),[1]*(nx+1),[1]],[-(nx+1),-1,0,1,nx+1])
A_x = A_x / (L/(nx+1))**2
A_y = A_y / (L/(ny+1))**2
gradsq = alpha*sp.kronsum(A_x,A_y)

u0 = []
v0 = []
for j in range(ny+2):
    for i in range(nx+2):
        r0ij = r0(i*L/(nx+1),j*L/(ny+1))
        u0.append(r0ij[0])
        v0.append(r0ij[1])

V0 = u0 + v0

def g(t,V):
    u_non_linearities = np.array([[B+V[(nx+2)*(ny+2)+i]*V[i]**2-C*V[i]+ff(i*L/(nx+1) % 1,(i//(nx+2))*L/(ny+1),t)] for i in range((nx+2)*(ny+2))])
    v_non_linearities = np.array([[A*V[i]-V[(nx+2)*(ny+2)+i]*V[i]**2] for i in range((nx+2)*(ny+2))])
    u_linearities = gradsq*np.transpose(np.array([V[:(nx+2)*(ny+2)]]))
    v_linearities = gradsq*np.transpose(np.array([V[(nx+2)*(ny+2):]]))
    out = np.concatenate((u_non_linearities+u_linearities,v_non_linearities+v_linearities))
    return np.transpose(out)[0]

## Verify the solution

sample = g(0.0,V0)
sample[33] # 802.9807693762165
sample[1057] # 985.3120721709204
sample[1999] # -403.58178806346996
sample[-1] # 1431.1460373522068
sample[535] # -323.16774591424223

sample2 = g(1.3,V0)
sample2[33] # 80.71952692049332
sample2[1057] # 98.90054733889691
sample2[1999] # -40.03152087416721
sample2[-1] # 1431.1460373522068
sample2[535] # -318.1677459142322

solution1 = solver.solve_ivp(g,(0,1.1),V0,method="BDF",rtol=1e-12,atol=1e-12)
V02 = solution1["y"][:,-1]
V02[0] # 0.5305455578390353
solution2 = solver.solve_ivp(g,(1.1,11.5),V02,method="BDF",rtol=1e-12,atol=1e-12)
sol = solution2["y"][:,-1]
sol[0] # 3.272258137972002

## Time it

import timeit

def timebdf():
    solution1 = solver.solve_ivp(g,(0,1.1),V0,method="BDF",t_eval=[0,1.1],rtol=1e-8,atol=1e-8)
    V02 = solution1["y"][:,-1]
    V02[0] # 0.5305455578390353
    solution2 = solver.solve_ivp(g,(1.1,11.5),V02,method="BDF",t_eval=[1.1,11.5],rtol=1e-8,atol=1e-8)
    sol = solution2["y"][:,-1]
    sol[0] # 3.272258137972002

timeit.Timer(timebdf).timeit(number=2)/2 # 68.23427899999842

def timelsoda():
    solution1 = solver.solve_ivp(g,(0,1.1),V0,method="LSODA",t_eval=[0,1.1],rtol=1e-8,atol=1e-8)
    V02 = solution1["y"][:,-1]
    V02[0] # 0.5305455578390353
    solution2 = solver.solve_ivp(g,(1.1,11.5),V02,method="LSODA",t_eval=[1.1,11.5],rtol=1e-8,atol=1e-8)
    sol = solution2["y"][:,-1]
    sol[0] # 3.272258137972002

timeit.Timer(timelsoda).timeit(number=2)/2 # 444.8325969999987

def timeradau():
    solution1 = solver.solve_ivp(g,(0,1.1),V0,method="Radau",t_eval=[0,1.1],rtol=1e-8,atol=1e-8)
    V02 = solution1["y"][:,-1]
    V02[0] # 0.5305455578390353
    solution2 = solver.solve_ivp(g,(1.1,11.5),V02,method="Radau",t_eval=[1.1,11.5],rtol=1e-8,atol=1e-8)
    sol = solution2["y"][:,-1]
    sol[0] # 3.272258137972002

timeit.Timer(timeradau).timeit(number=2)/2 # 343.9781433500011

A = gradsq != 0
sparsity = scipy.sparse.csr_matrix(scipy.linalg.block_diag(A.todense(),A.todense())).astype(float)

def timebdfsparse():
    solution1 = solver.solve_ivp(g,(0,1.1),V0,method="BDF",jac_sparsity=sparsity,t_eval=[0,1.1],rtol=1e-8,atol=1e-8)
    V02 = solution1["y"][:,-1]
    V02[0]
    solution2 = solver.solve_ivp(g,(1.1,11.5),V02,method="BDF",jac_sparsity=sparsity,t_eval=[1.1,11.5],rtol=1e-8,atol=1e-8)
    sol = solution2["y"][:,-1]
    sol[0]

timeit.Timer(timebdf).timeit(number=2)/2 # Errors
# NotImplementedError: subtracting a nonzero scalar from a sparse matrix is not supported