Compiling FoldList implementation for RK4

Question

Original

I'm looking to write an integrator for a function of two variables. Here is my implementation for the RK4 update rule using FoldList.

RK4Update[fx_, x_, prms_, numsteps_] :=
Block[{},

1.0/(6.0 numsteps)
Dot[{1.0, 2.0, 2.0, 1.0},
 Rest[
  FoldList[
   fx @@ {x + #1 #2, prms} &,
   x,
   Divide[{0.0, 0.5, 0.5, 1.0}, 1.0 numsteps]
   ]
  ]
 ]

];

This function evaluates correctly when given g (below) and numerical list arguments for x and prms. I'm attempting to compile to gain speed:

Compile[
{{sts, _Real, 1}, {pms, _Real, 1}, {tsteps, _Real}},
Evaluate[RK4Update[g, sts, pms, tsteps]],
CompilationOptions -> {"InlineExternalDefinitions" -> True}]

Here g is an already compiled function of two vector arguments. The compiler is unhappy, returning errors of the form:

Part::partd: Part specification pms[[2]] is longer than depth of object. 
Part specification pms[[1]] is longer than depth of object
Part::partw: "Part 3 of 0. +sts does not exist"

and further errors from the attempted evaluation of g. I'm not sure if there is any way to proceed here. Is there a simple fix? Or is an overhaul needed?

If it's of any help, a sample function g is polynomial in elements of its arguments

(#2[[2]] #1[[3]] (1- #1[[1]])+ #1[[2]]^2(1 - #1[[1]])
+ #1[[1]](2 - #1[[1]]))/#2[[1]]

Curiously, CompilePrint returns what appears to be a fully compiled code. However, this code (perhaps unsurprisingly) does not produce the correct result.

Update

The function RK4Update can be condensed to

RK4Update[fx_, x_, prms_,numsteps_] :=
  1.0/(6.0 numsteps)
  Dot[{1.0, 2.0, 2.0, 1.0},
   FoldList[
    fx[x + #1 #2, prms] &,
    fx[x, prms],
    Divide[{0.5, 0.5, 1.0}, 1.0 numsteps]
   ]
  ]

The following compilation command outputs several errors, but returns a successfully compiled function. A Quiet prefix suppresses display of these errors:

 cf= Quiet@Compile[
 {{st, _Real, 1}, {pm, _Real, 1}, {tstep, _Real}},
 Evaluate[RK4Update[h, st, pm, tstep]],
 CompilationOptions -> {"InlineCompiledFunctions" -> True}]

where, importantly, h[x,y] = {g1[x,y],g2[x,y],g3[x,y]} with x,y each a numeric list and g_i a set of compiled functions. In the original question, g in the FoldList call returns a single number, whereas x is a triplet. A function in FoldList must output a quantity with the same dimension as its argument.

For some reason, cf outputs a 3x3 matrix whose diagonal is the desired result. My current working (yet unsatisfactorily ad hoc) solution is given by calling Diagonal[cf[Range[3],Range[2]]

Answer 1

The problem is the use of the Evaluate function inside the Compile function. The Evaluate function requires that RK4Update, and therefore g, be passed the correct form of the x and pms parameters. The pure function definition of g provided dictates that these should be a lists of length at least three and at least two, respectively.

However, within the scope of the Compile function RK4Update, and therefore g, are passed undefined variables sts and pms, which have length 0. When RK4Update is evaluated in this context, you see the Part specification errors in the original problem.

This simple modification to the code should work (no Evaluation):

h = Compile[{{sts, _Real, 1}, {pms, _Real, 1}, {tsteps, _Real}}, 
        RK4Update[g, sts, pms, tsteps], 
        CompilationOptions -> {"InlineExternalDefinitions" -> True}]