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$
  • 2
    $\begingroup$ I think it may be AbsoluteTiming that you want, not Timing. $\endgroup$ – C. E. Jul 4 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$ – Rolf Mertig Jul 4 at 21:13
  • 1
    $\begingroup$ Maybe someone could try the newer FunctionCompile. Though probably it will be slower than Compile. $\endgroup$ – Rolf Mertig Jul 4 at 21:25
  • 1
    $\begingroup$ Why do you need this? If Python runs fine, just use Python. $\endgroup$ – Rolf Mertig Jul 4 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 at 22:38
17
$\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. )

| improve this answer | |
$\endgroup$
  • 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 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 at 9:25
  • $\begingroup$ @xzczd great! Can we apply this method for PDEs (FEM) in MMA? $\endgroup$ – ABCDEMMM Jul 8 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 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 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.