10
$\begingroup$

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
$\endgroup$
10
  • 2
    $\begingroup$ I think it may be AbsoluteTiming that you want, not Timing. $\endgroup$
    – C. E.
    Jul 4, 2020 at 20:57
  • 1
    $\begingroup$ Right, or RepatedTiming. And with adding SetOptions[Compile, CompilationTarget->"C" ,RuntimeOptions->"Speed"] I get AbsoluteTiming down to about 3 ms, on Windows. $\endgroup$ Jul 4, 2020 at 21:13
  • 1
    $\begingroup$ Maybe someone could try the newer FunctionCompile. Though probably it will be slower than Compile. $\endgroup$ Jul 4, 2020 at 21:25
  • 1
    $\begingroup$ Why do you need this? If Python runs fine, just use Python. $\endgroup$ Jul 4, 2020 at 21:27
  • 2
    $\begingroup$ 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 Native`Range 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? $\endgroup$
    – Carl Lange
    Jul 4, 2020 at 22:38

1 Answer 1

18
$\begingroup$

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

$\endgroup$
6
  • 1
    $\begingroup$ 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 $\endgroup$
    – C. E.
    Jul 5, 2020 at 8:59
  • 2
    $\begingroup$ @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. ) $\endgroup$
    – xzczd
    Jul 5, 2020 at 9:25
  • $\begingroup$ @xzczd great! Can we apply this method for PDEs (FEM) in MMA? $\endgroup$
    – ABCDEMMM
    Jul 8, 2020 at 15:02
  • $\begingroup$ @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. $\endgroup$
    – xzczd
    Jul 8, 2020 at 15:10
  • $\begingroup$ @xzczd it means that if we use build-in NDSolve, every simulation will run directly in c-code? $\endgroup$
    – ABCDEMMM
    Jul 8, 2020 at 22:11

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.