Optimizing: `@taylorize`

Here, we describe the use of the macro @taylorize, which parses the functions containing the ODEs to be integrated, allowing to speed up taylorinteg and lyap_taylorinteg.

Warning

The macro @taylorize is still in an experimental phase; be cautious of the resulting integration, which has to be tested carefully.

Some context and the idea

The way in which taylorinteg works by default is by calling repeatedly the function where the ODEs of the problem are defined, in order to compute the recurrence relations that are used to construct the Taylor expansion of the solution. This is done for each order of the series in TaylorIntegration.jetcoeffs!. These computations are not optimized: they waste memory due to allocations of some temporary arrays, and perform some operations whose result has been previously computed.

Here we describe one way to optimize this: The idea is to replace the default method of TaylorIntegration.jetcoeffs! by another (with the same name) which is called by dispatch, that in principle performs better. The new method is constructed specifically for the actual function defining the equations of motion by parsing its expression; the new function performs in principle exactly the same operations, but avoids the extra allocations and the repetition of some operations.

An example

In order to explain how the macro works, we shall use as an example the mathematical pendulum. First, we carry out the integration using the default method, as described before.

using TaylorIntegration

function pendulum!(dx, x, p, t)
    dx[1] = x[2]
    dx[2] = -sin(x[1])
    return dx
end

# Initial time (t0), final time (tf) and initial condition (q0)
t0 = 0.0
tf = 100.0
q0 = [1.3, 0.0]

# The actual integration
t1, x1 = taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500); # warm-up run
e1 = @elapsed taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500);
e1

0.021852419

We note that the initial number of methods defined for TaylorIntegration.jetcoeffs! is 2.

length(methods(TaylorIntegration.jetcoeffs!)) == 2 # initial value

true

Using @taylorize will increase this number by creating a new method.

The macro @taylorize is intended to be used in front of the function that implements the equations of motion. The macro does the following: it first parses the actual function as it is, so the integration can be computed using taylorinteg as above, by explicitly using the keyword argument parse_eqs=false. It then creates and evaluates a new method of TaylorIntegration.jetcoeffs!, which is the specialized method (through Val) on the specific function passed to the macro.

@taylorize function pendulum!(dx, x, p, t)
    dx[1] = x[2]
    dx[2] = -sin(x[1])
    return dx
end

println(methods(pendulum!))

# 1 method for generic function "pendulum!":
[1] pendulum!(dx, x, p, t) in Main.ex-taylorize at none:2

println(methods(TaylorIntegration.jetcoeffs!))

# 3 methods for generic function "jetcoeffs!":
[1] jetcoeffs!(eqsdiff::Function, t::Taylor1{T}, x::Taylor1{U}, params) where {T<:Real, U<:Number} in TaylorIntegration at /home/travis/build/PerezHz/TaylorIntegration.jl/src/explicitode.jl:25
[2] jetcoeffs!(eqsdiff!::Function, t::Taylor1{T}, x::AbstractArray{Taylor1{U},1}, dx::AbstractArray{Taylor1{U},1}, xaux::AbstractArray{Taylor1{U},1}, params) where {T<:Real, U<:Number} in TaylorIntegration at /home/travis/build/PerezHz/TaylorIntegration.jl/src/explicitode.jl:66
[3] jetcoeffs!(::Val{Main.ex-taylorize.pendulum!}, t::Taylor1{_T}, x::AbstractArray{Taylor1{_S},1}, dx::AbstractArray{Taylor1{_S},1}, p) where {_T<:Real, _S<:Number} in Main.ex-taylorize

We see that there is only one method of pendulum!, and there is a new method of TaylorIntegration.jetcoeffs!, whose signature appears in this documentation as Val{Main.ex-taylorize.pendulum!}; it is an specialized version for the function pendulum! (with some extra information about the module where the function was created). This method is selected internally if it exists (default), exploiting dispatch, when calling taylorinteg or lyap_taylorinteg; to integrate using the hard-coded method of TaylorIntegration.jetcoeffs! of the integration above, the keyword argument parse_eqs has to be set to false.

Now we carry out the integration using the specialized method; note that we use the same instruction as above.

t2, x2 = taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500); # warm-up run
e2 = @elapsed taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500);
e2

0.003737381

We note the difference in the performance:

e1/e2

5.846987235178859

We can check that both integrations yield the same results.

t1 == t2 && x1 == x2

true

As stated above, in order to allow to opt-out from using the specialized method created by @taylorize, taylorinteg and lyap_taylorinteg recognize the keyword argument parse_eqs; setting it to false imposes using the standard method.

taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500, parse_eqs=false); # warm-up run

e3 = @elapsed taylorinteg(pendulum!, q0, t0, tf, 25, 1e-20, maxsteps=1500, parse_eqs=false);

e1/e3

0.9488268566171307

We now illustrate the possibility of exploiting the macro when using TaylorIntegration.jl from DifferentialEquations.jl.

using DiffEqBase

prob = ODEProblem(pendulum!, q0, (t0, tf), nothing) # no parameters
solT = solve(prob, TaylorMethod(25), abstol=1e-20, parse_eqs=true); # warm-up run
e4 = @elapsed solve(prob, TaylorMethod(25), abstol=1e-20, parse_eqs=true);

e1/e4

1.1849494245280805

Note that there is a marginal cost of using solve in comparison with taylorinteg.

The speed-up obtained comes from the design of the new (specialized) method of TaylorIntegration.jetcoeffs! as described above: it avoids some allocations and some repeated computations. This is achieved by knowing the specific AST of the function of the ODEs integrated, which is walked through and translated into the actual implementation, where some required auxiliary arrays are created and the low-level functions defined in TaylorSeries.jl are used. For this, we heavily rely on Espresso.jl and some metaprogramming; we thank Andrei Zhabinski for his help and comments.

The new jetcoeffs! method can be inspected by constructing the expression corresponding to the function, and using TaylorIntegration._make_parsed_jetcoeffs:

ex = :(function pendulum!(dx, x, p, t)
    dx[1] = x[2]
    dx[2] = -sin(x[1])
    return dx
end)

new_ex = TaylorIntegration._make_parsed_jetcoeffs(ex)

:(function TaylorIntegration.jetcoeffs!(::Val{pendulum!}, t::Taylor1{_T}, x::AbstractVector{Taylor1{_S}}, dx::AbstractVector{Taylor1{_S}}, p) where {_T <: Real, _S <: Number}
      order = t.order
      dx[1] = Taylor1(identity(constant_term(x[2])), order)
      tmp386 = Taylor1(sin(constant_term(x[1])), order)
      tmp388 = Taylor1(cos(constant_term(x[1])), order)
      dx[2] = Taylor1(-(constant_term(tmp386)), order)
      for __idx = eachindex(x)
          (x[__idx]).coeffs[2] = (dx[__idx]).coeffs[1]
      end
      for ord = 1:order - 1
          ordnext = ord + 1
          TaylorSeries.identity!(dx[1], x[2], ord)
          TaylorSeries.sincos!(tmp386, tmp388, x[1], ord)
          TaylorSeries.subst!(dx[2], tmp386, ord)
          for __idx = eachindex(x)
              (x[__idx]).coeffs[ordnext + 1] = (dx[__idx]).coeffs[ordnext] / ordnext
          end
      end
      return nothing
  end)

This function has a similar structure as the hard-coded method of TaylorIntegration.jetcoeffs!, but uses low-level functions in TaylorSeries (e.g., sincos! above) and explicitly allocates the needed temporary arrays. More complex functions become easily very difficult to read. Note that, if necessary, one can further optimize new_ex manually.

Limitations and some advices

The construction of the internal function obtained by using @taylorize is somewhat complicated and limited. Here we list some limitations and advices.

It is useful to have expressions which involve two arguments at most, which imposes the proper use of parenthesis: For example, res = a+b+c should be written as res = (a+b)+c.
Updating operators such as +=, *=, etc., are not supported. For example, the expression x += y is not recognized by @taylorize. Likewise, expressions such as x = x+y are not supported by @taylorize and should be substituted by equivalent expressions; e.g. z = x+y; x = z.
The macro allows to use array declarations through Array, but other ways (e.g. similar) are not yet implemented.
Avoid using variables prefixed by an underscore, in particular _T and _S; using them may lead to name collisions with some internal variables.
Broadcasting is not recognized by @taylorize.
The macro may be used in combination with the common interface with DifferentialEquations.jl, for functions using the (du, u, p, t) in-place form, as we showed above. Other extensions allowed by DifferentialEquations may not be able to exploit it.
if-else blocks are recognized in its long form, but short-circuit conditional operators (&& and ||) are not.
Expressions which correspond to function calls (so the head field is :call) which are not recognized by the parser are simply copied. The heuristics used, specially for vectors, may not work for all cases.
Use local for internal parameters (simple constant values); this improves performance. Do not use it if the variable is Taylor expanded.

It is recommended to skim test/taylorize.jl, which implements different cases.

Please report any problems you may encounter.

Optimizing: @taylorize

Some context and the idea

An example

Limitations and some advices

Optimizing: `@taylorize`