Compile a program with Fold function

Question

I have the following small program:

(*Parameters*)
x0 = 4.;
g = 10.;
k = 0.5;
(*Define function*)    
fun[nTot_?IntegerQ, dt_?NumberQ] :=
 (fac = 1. - dt*k/g;
  list = Reap[Fold[Sow[fac*#] &, x0, Range[nTot]]][[2, 1]];
  list2 = ({(# - 1.)*dt, list[[#]]} & /@ Range[Length[list]]);)
(*Parameters*)
nTot = 10^7; dt = 10.^-3;
(*Computefunction*)
fun[nTot, dt]

which takes on my machine about 13 seconds. I want to speed it up and thought about Mathematica's Compile capabilities. Is there any way to Compile the function fun to C? For example here it is done for a Do loop, but I cannot apply this to my code. Any suggestions?

[Edit]

Thank you again for your answer. To illustrate what I mean see the following change of my initial program:

f[nTot_?IntegerQ, nMem_?IntegerQ, 
  dt_?NumberQ] :=
 (fun[t_] = -163.6 E^(-2. t) + 40. E^(-1. t);
  fac1 = 1. - dt*(41.8/0.1 + 1.);
  fac2 = dt^2/0.1;
  listExp = 
   Developer`ToPackedArray[
    fun[#*dt] & /@ Range[0, nMem - 1] // Reverse];
  list = Prepend[
    Reap[Fold[Append[Rest@#, Sow[fac1*Last[#] - fac2*listExp.#]] &, 
       ConstantArray[1., nMem], Range[nTot]]][[2, 1]], 1.];
  list2 = ({(# - 1.)*dt, list[[#]]} & /@ Range[Length[list]]);)
(*Define parameters*)
nTot = 10^4; nMem = 10^5; dt = 10.^-4;
(*Compute function*)
f[nTot, nMem, dt]

which takes 2.3 secs on my machine. And now the (equivalent) compiled version which takes 8.4 secs on my machine:

fun[t_] = -163.6 E^(-2.* t) + 40. E^(-1.* t);
cfun = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {nMem, \
_Integer}}, 
   Block[{}, x = ConstantArray[x0, nMem]; 
    fac1 = 1. - dt*(41.8/0.1 + 1.);
    fac2 = dt^2/0.1; 
    listExp = Reverse[fun[#*dt] & /@ Range[0, nMem - 1]];
    Table[{i*dt, 
      Last[x = Append[Rest@x, fac1*Last[x] - fac2*(listExp.x)]]}, {i, 
      1, nTot}]], CompilationTarget -> "C"];
clist2 = Prepend[cfun[1., 10^4, 10.^-4, 10^5], {0., 1.}];

Why is the compiled version now so slow? Is there a way to speed it up? Or did I do some simple mistakes? Help is highly appreciated again!

[Edit2]:

Ok, a small modification in the compiled version speeds it up to around 3 secs on my machine. But still it is slower compared to the solution with Fold:

fac1 = 1. - dt*(41.8/0.1 + 1.);
fac2 = dt^2/0.1;
fun[t_] = -163.6 E^(-2.* t) + 40. E^(-1.* t);
cfun = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {nMem, \
_Integer}}, Block[{fac2 = fac2, fac1 = fac1, x, listExp},
    x = Table[x0, {i, 1, 10^5}];
    listExp = Reverse[fun[#*dt] & /@ Range[0, nMem - 1]];
    Table[{i*dt, 
      Last[x = Append[Rest@x, fac1*Last[x] - fac2*listExp.x]]}, {i, 1,
       nTot}]], {{listExp, _Real, 1}}, 
   CompilationOptions -> {"InlineExternalDefinitions" -> True, 
     "InlineCompiledFunctions" -> True, 
     "ExpressionOptimization" -> True}, 
   RuntimeAttributes -> {Listable}, Parallelization -> True];
clist2 = Prepend[cfun[1, 10^4, 10^-4, 10^5], {0, 1}];

[Edit3]

Getting rid of Append and using Part instead is a little bit helpful:

dt = 10.^-4;
nMem = 10^5;
nTot = 10^5;
fac1 = 1. - dt*(41.8/0.1 + 1.);
fac2 = dt^2/0.1;
x0 = 1.;
x = Array[x0 &, nTot + nMem];
fun[t_] = -163.6 E^(-2.* t) + 40. E^(-1.* t);
listExp = 
  Developer`ToPackedArray[Reverse[fun[#*dt] & /@ Range[0, nMem - 1]]];
list = Table[{i*dt, 
    x[[i + nMem]] = 
     fac1*x[[i + nMem - 1]] - 
      fac2*listExp.x[[i ;; nMem + i - 1]]}, {i, 1, nTot}];

and a compiled version:

fac1 = 1. - dt*(41.8/0.1 + 1.);
fac2 = dt^2/0.1;
fun[t_] = -163.6 E^(-2.* t) + 40. E^(-1.* t);
cfun = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {nMem, \
_Integer}}, Block[{fac2 = fac2, fac1 = fac1, x, listExp},
    x = Array[x0 &, nTot + nMem];
    listExp = Reverse[fun[#*dt] & /@ Range[0, nMem - 1]];
    Table[{i*dt, 
      x[[i + nMem]] = 
       fac1*x[[i + nMem - 1]] - 
        fac2*listExp.x[[i ;; nMem + i - 1]]}, {i, 1, 
      nTot}]], {{listExp, _Real, 1}}, 
   CompilationOptions -> {"InlineExternalDefinitions" -> True, 
     "InlineCompiledFunctions" -> True, 
     "ExpressionOptimization" -> True}, 
   RuntimeAttributes -> {Listable}, Parallelization -> True];
clist2 = cfun[1, 10^5, 10^-4, 10^5];

[Edit4]

Managed to get rid of the MainEvaluate code by using pure function definition:

Clear["Global`*"]
fac1 = 1. - dt*(41.8/0.1 + 1.);
fac2 = dt^2/0.1;
fun = -163.6 E^(-2.* #) + 40. E^(-1.* #) &;
cfun = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {nMem, \
_Integer}}, Block[{fac2 = fac2, fac1 = fac1, x, listExp},
    x = Array[x0 &, nTot + nMem];
    listExp = Reverse[fun[#*dt] & /@ Range[0, nMem - 1]];
    Table[{i*dt, 
      x[[i + nMem]] = 
       fac1*x[[i + nMem - 1]] - 
        fac2*listExp.x[[i ;; nMem + i - 1]]}, {i, 1, 
      nTot}]], {{listExp, _Real, 1}}, 
   CompilationOptions -> {"InlineExternalDefinitions" -> True, 
     "InlineCompiledFunctions" -> True, 
     "ExpressionOptimization" -> True}];
clist3 = cfun[1., 10^4, 10.^-4.5, 4*10^5];

Any further recommendations how to speed it up?

Answer 1

This new task with dotting listExp against a moving window in x would have deserved a new question post.

Your problem as a very specific structure that I am going to exploit: The list listExp is a sum of two lists listExp1 and listExp2 that are both of the form Table[a base^i, {i,1,nMem - 1}]. These list are dotted against a window moving along x resulting in numbers sum1 and sum2. The main obserservation for speeding up your code is that after a moving step, sum1 can be computed as sum1 = (sum1 + a1 xout)/base1 + b1 xin where xin is the element of x entering the window and xout is the one leaving the window. Hence, we don't have to compute the expensive dot product any more! Along with some other tweaks (replacing Part in read operations by Compile`GetElement, compiling into binaries via "C", and switching off some secrity checks with RuntimeOptions -> "Speed"), the code looks like this:

cfunDotLess = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {nMem, _Integer}}, 
   Block[{fac2, fac1, x, listExp1, listExp2, base1, base2, sum1, sum2,
      xnew, a1, a2, b1, b2, xout, xin, j}, 
    x = Table[x0, {i, 1, nMem}];
    fac1 = 1. - dt (41.8/0.1 + 1.);
    fac2 = dt^2/0.1;
    base1 = E^(-2. dt);
    base2 = E^(-1. dt);
    a1 = -163.6 fac2;
    a2 = 40. fac2;
    sum1 = 0.;
    sum2 = 0.;
    b1 = (-a1);
    b2 = (-a2);
    Do[
     sum1 += b1 x0;
     sum2 += b2 x0;
     b1 *= base1;
     b2 *= base2;, {i, 0, nMem - 1}];
    xnew = x0;
    j = 0;
    Table[
     If[j >= nMem, j = 1, j++];
     xout = Compile`GetElement[x, j];
     xin = xnew;
     sum1 = (sum1 + a1 xout)/base1 + b1 xin;
     sum2 = (sum2 + a2 xout)/base2 + b2 xin;
     x[[j]] = xnew = Plus[fac1 xin, sum1, sum2];
     {i dt, xnew}, {i, 1, nTot}]
    ],
   CompilationTarget -> "C",
   RuntimeOptions -> "Speed"
   ];

The following shows that clistDotLess is about 8000 times faster than cfun from Edit 4 at a cost of an acceptable loss of accuracy:

clistDotLess = cfunDotLess[1., 10^5, 10.^-4.5, 4.*10^5]; // AbsoluteTiming // First
clist = cfun[1., 10^5, 10.^-4.5, 4.*10^5]; // AbsoluteTiming // First
Max[Abs[clistDotLess - clist]]

0.003614

37.028

4.63854*10^-9

Edit: I removed some superfluous indexing into x and shortened x to be an array of length nMem by using indices j modulo nMem.

Answer 2

Your use of Reap and Sow along with Fold could be better formulated with FoldList, or even better with NestList. Moreover, the second sweep over Range[Length[list]] can be avoided at all. "Unfolding" your code and compiling it, this could look like this:

cfun = Compile[{{x0, _Real}, {nTot, _Integer}, {dt, _Real}, {g, _Real}, {k, _Real}},
   Block[{fac = 1. - dt*k/g, x = x0},
    Table[x *= fac; {(i - 1.)*dt, x}, {i, 1, nTot}]
    ],
   CompilationTarget -> "C"
   ];

clist2 = cfun[x0, nTot, dt, g, k]; // AbsoluteTiming // First

0.49941

Note that constructs like Fold and Nest are bound to evaluate sequentially while your result can actually obtained in parallel. A simple vectorized version of your code is this:

vfun[x0_?NumberQ, nTot_Integer, dt_?NumberQ, g_?NumberQ, k_?NumberQ] :=
   With[{ran = N@Range[1, nTot]},
   Transpose[{(ran - 1.) dt, x0 ((1. - dt*k/g)^ran)}]
   ];

On my Haswell quad core, this performs like this

vlist2 = vfun[x0, nTot, dt, g, k]; // AbsoluteTiming // First

0.397826

Note that the most expensive operation in vfun is Transpose; skipping will half the cost.

Comparing the results:

Max[Abs[clist2 - vlist2]]

2.79776*10^-14