# Writing compiled functions as fast as Python's Numba

I want to write some code to simulate a damped oscillator that is just as fast as code written using Numba's @njit decorator. I've written the mathematica code and mine is 20-40x slower than the python code written by YouTuber Jack of Some.

Here is the code from Jack of Some's video on speeding up Python code with Numba; I've changed it a bit to run in just one jupyter cell:

import numpy as np
from numba import jit, njit, types, vectorize

@njit
def friction_fn(v, vt):
if v > vt:
return - v * 3
else:
return - vt * 3 * np.sign(v)

@njit
def simulate_spring_mass_funky_damper(x0, T=10, dt=0.0001, vt=1.0):
times = np.arange(0, T, dt)
positions = np.zeros_like(times)

v = 0
a = 0
x = x0
positions[0] = x0/x0

for ii in range(len(times)):
if ii == 0:
continue
t = times[ii]
a = friction_fn(v, vt) - 100*x
v = v + a*dt
x = x + v*dt
positions[ii] = x/x0
return times, positions

_ = simulate_spring_mass_funky_damper(0.1)

%time _ = simulate_spring_mass_funky_damper(0.1)


The output is

CPU times: user 1.38 ms, sys: 337 µs, total: 1.72 ms
Wall time: 1.72 ms


vs my Mathematica code

ClearAll[friction, simulateSpring, jacksSpring];

friction = Compile[{{v, _Real}, {vt, _Real}},
If[v > vt,
-v*3.0,
-vt*3.0*Sign[v]]
];

simulateSpring =
Compile[{{x0, _Real}, {t, _Real}, {dt, _Real}, {vt, _Real}},

Module[{τ, times, positions, v = 0.0, a = 0.0, x = x0},

τ = t;
times = Range[0.0, t, dt];
positions = ConstantArray[0.0, Length@times];
positions[[1]] = x0/x0;

Do[
τ = times[[i]];
a = friction[v, vt] - 100*x;
v = v + a*dt;
x = x + v*dt;
positions[[i]] = x/x0;
,
{i, 2, Length@times}];
{times, positions}
]
];
jacksSpring[x_] := simulateSpring[x, 10.0, 0.0001, 1.0];

Print["CPU time: ", Timing[jacksSpring[0.1]][[1]]*1000, " ms"]


from which we have

CPU time: 27.703 ms

• I think it may be AbsoluteTiming that you want, not Timing. – C. E. Jul 4 at 20:57
• Right, or RepatedTiming. And with adding SetOptions[Compile, CompilationTarget->"C" ,RuntimeOptions->"Speed"] I get AbsoluteTiming down to about 3 ms, on Windows. – Rolf Mertig Jul 4 at 21:13
• Maybe someone could try the newer FunctionCompile. Though probably it will be slower than Compile. – Rolf Mertig Jul 4 at 21:25
• Why do you need this? If Python runs fine, just use Python. – Rolf Mertig Jul 4 at 21:27
• I checked out FunctionCompile, and the results were so close as to not be worth it on my machine (~6.1ms with Compile and the options above, ~6.2ms with FunctionCompile). I had to wrap Range in a KernelFunction as the built-in NativeRange does not support Real arguments, but since that's only called once I would be surprised if that were the bottleneck. Is there some way to replace the Do with some matrix maths? – Carl Lange Jul 4 at 22:38

Based on the experience obtained here:

friction = Compile[{{v, _Real}, {vt, _Real}}, If[v > vt, -v*3.0, -vt*3.0*Sign[v]]];

simulateSpring =
Compile[{{x0, _Real}, {t, _Real}, {dt, _Real}, {vt, _Real}},
Module[{τ, times, positions, v = 0.0, a = 0.0, x = x0}, τ = t;
times = Range[0.0, t, dt];
positions = Table[0.0, {Length@times}];
positions[[1]] = x0/x0;
Do[τ = times[[i]];
a = friction[v, vt] - 100*x;
v = v + a*dt;
x = x + v*dt;
positions[[i]] = x/x0;, {i, 2, Length@times}];
{times, positions}], RuntimeOptions -> Speed, CompilationTarget -> C,
CompilationOptions -> {InlineExternalDefinitions -> True,
InlineCompiledFunctions -> True}];

lib = simulateSpring[[-1]]

Print["CPU time: ", RepeatedTiming[lib[0.1, 10.0, 0.0001, 1.0]][[1]]*1000, " ms"]

CPU time: 1.6 ms


Compiled with TDM GCC 5.1.0-2 64bit, "SystemCompileOptions"->"-Ofast".

RepeatedTiming of the original jacksSpring[0.1] is 20 ms here. I tried timing the python code but failed. (Seems that Numba is not configured correctly on my laptop, the timing is about 0.35 s. )

• I'm surprised it even works to write RuntimeOptions -> Speed and CompilationTarget -> C, usually they're written "Speed" and "C". Nevertheless, I tried this and it gives me 1.26 both with and without quotation marks, whereas Numba gives me 1.23 ms. So it's very close. +1 – C. E. Jul 5 at 8:59
• @C.E. String options and option values of Compile don't need to be string, so do those of NDSolve, AFAIK. (But properties of FittedModel and methods of InterpolatingFunction need to be string. ) – xzczd Jul 5 at 9:25
• @xzczd great! Can we apply this method for PDEs (FEM) in MMA? – ABCDEMMM Jul 8 at 15:02
• @ABCDEMMM If you mean FEM built in NDSolve, then high level functions like NDSolve already makes use of Compile internally AFAIK; if you mean self implementation of FEM, just search in this site, you'll see most (if not all) of the answers implementing FEM in a low level way makes use of Compile`, here is an example. – xzczd Jul 8 at 15:10
• @xzczd it means that if we use build-in NDSolve, every simulation will run directly in c-code? – ABCDEMMM Jul 8 at 22:11