Writing compiled functions as fast as Python's Numba

Question

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

Answer 1

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. )