Improve perfromance of DifferentialEquations, when odefun is evaluated repeatedly and many timepoints are required

General Usage Performance
1

Is it possible to improve the performance of an ODE system solving, if we have to evaluate the equation many times with different parameters (not in a Monte Carlo manner, but sequentially), and if we want to evaluate part of the solution at many time points? For instance, if we consider the following Lotka-Volterra system (just as an example), we’ll use only the first component of the solution (sol[1,:]), but we want to know its values at 10000 time points between [0,10]. Then we would do something with that and reassign the parameters for the system for further iterations.

In the example, we do 1000 iterations. If I run main() multiple times after the initial compilation, @time tells

1.549323 seconds (10.09 M allocations: 1.368 GiB, 11.56% gc time)

The number of allocations look big to me. Can we improve the performance, for instance by pre-allocating the solution variable beforehand, or specifying that we do not need to know the values of the second component? And indeed, we do not need to store the solution after we calculate new parameters for the system.

using LinearAlgebra;
using Plots;
using Random;
using DifferentialEquations;
using BenchmarkTools;

@inbounds function odefun!(du, u, p, t)
    du[1] = u[1]*(p[1]-p[2]*u[2])
    du[2] = -u[2]*(p[3]-p[4]*u[1])
    return nothing
end

function main()

    u0 = [0.7, 0.3]
    p = [0.15,0.2,0.5,0.9]
    tspan = (0.0, 10)
    prob = ODEProblem{true}(odefun!, u0, tspan, p)
    sol = solve(
        prob,
        DP5(),
        saveat = 0.001,
        reltol = 1e-3,
        abstol = 1e-6,
        alg_hints = [:nonstiff],
    )
    @time for i::Int32 = 1:1000
        p = rand(4,) # Just to simulate new parameters.
        prob = remake(prob; p = p)
        sol = solve(
            prob,
            DP5(),
            saveat = 0.001,
            reltol = 1e-3,
            abstol = 1e-6,
            alg_hints = [:nonstiff],
        )

        # Do something with sol[1,:], and set new value for variable p:
        # p = fancyfunc(sol[1,:])
    end
    plot(sol.t,sol[1,:])
end

main()
2

What processing are you doing that requires 10000 time points? How are you updating the parameters? If this is some kind of optimization/fitting process, for example, there are probably more efficient ways to go about it.

1 Like
3

No, this is not related to optimization. For example, I might need many time points in MCMC parameter estimation, if I have already a long noisy timeseries of some measurements related to an ODE - in other words I evaluate likelihood with the parameters. The parameter update is done by the MCMC algorithm to propose new parameter values according to some rule. But in case of a simple MCMC proposal step, it’s very cheap to generate new parameters out of the current one compared to the likelihood evaluation.

4

Generally when using DiffEq solvers I let the solver determine the explicit save times based on the desired tolerances. To get the solution at other times I then use the interpolation form of sol, see https://docs.sciml.ai/dev/basics/solution/#Interpolations-1. Have you tried seeing if this helps with allocations?

Sounds like I mixed up saveat and another saving option…

5

Also, is there is a specific reason you are using DP5? I would suggest using one of the recommended solvers from https://docs.sciml.ai/dev/solvers/ode_solve/#Recommended-Methods-1 depending on the stiffness of your actual problem and desired tolerances.

6

(I would say that MCMC parameter estimation is closely related to optimization; indeed, I seem to recall seeing papers that explicitly formulated it as an optimization problem.)

You could take a page from machine learning — typically, if you have a large training set, it is expensive to evaluate your loss function (or, in your case, the likelihood) with the full data set for every set of parameters. Instead, you choose a random sample of the training set for every step of the learning process (which leads people to use stochastic optimization methods).

Couldn’t you do the same thing? Estimate the likelihood using a different random subset of your training point for each parameter update, so that you don’t need to evaluate your ODE solution at 10000 points every time.

7

saveat takes as large steps as possible and then interpolates, so post interpolation won’t help here. One thing that could help is to add save_idxs=1, to tell it only save the timeseries of the first variable.

But secondly, I agree with @stevengj here that you probably want to do is take random subsets from your training dataset, i.e. “minibatching”, and evaluate the cost on those subsets. Normally when the data is too dense it isn’t helpful.

But thirdly, if your dataset is that dense, then you probably want to fit it in a completely different way. You can instead interpolate your data and calculate the approximate derivatives u'(t) from cubic spline approximations. You can then write your cost/likelihoods as ||u' - f(u(t),p,t)||. This doesn’t require running differential equation solves at all and is a lot cheaper. You can take a look at how https://docs.sciml.ai/latest/analysis/parameter_estimation/#Two-Stage-method-(Non-Parametric-Collocation)-1 is implemented for one way to do the interpolation, or just use something like https://github.com/kbarbary/Dierckx.jl

1 Like
8

Thanks for the tips! Using save_idxs=1 helped a lot to reduce allocations, and increased the performance by almost 40%. Now @times gives

  1.000711 seconds (84.00 k allocations: 485.321 MiB, 1.23% gc time)

I plan to use MCMC to sample from the parameters’ joint posterior to do full-level UQ, so for now I stick with ODE approach at least to obtain reference results with minimum possible of bias as possible. In the batch-likelihood approach, do I have to just evaluate the likelihood indeed only at the batch points and otherwise do everything the same, or do I have to do take the batch-likelihood into account somehow?
How about the spline method? It seems that the method fits a curve to the smoothed data and uses an extra penalty function with respect to the ODE function to do that. For me, it sounds that the likelihood of the spline method would be greatly different, or at least it would be difficult to express the posterior in the form P(obs | param)P(param), where likelihood P(obs | param) is the same as in the ODE approach.

Nevertheless, solving this example system 1000 times at 10000 points in 1 second is a good result already :grinning:

9

One other thing that could be helpful is to use StaticArrays and an out-of-place ODE definition, since then the solver cache wouldn’t need to be created.