Question on ODE system with inhomogeneous coefficients, parameters, code design


I have a question relating to code design. I’m using some ODE solvers to calculate a system of ODEs with time-varying coefficients which I am interested in. They take on the form F, G, and dG in the MWE and depend on some parameters which I’ve defined in a separate struct. The fact that the functions are listed twice seems to be an obvious case of code smell, but there’s a couple issues

  • ForwardDiff.derivative only takes unary arguments, but e.g. G depends on parameters b and c (the actual function also depends on something like u[1] + u[2])
  • The actual function to differentiate depends on t in a much more complicated manner than G below, with several layers of functions in t. G is needed in the solver, but once the solver is done, I don’t have access to G without copying the function definition.

In the real code, there are about a dozen parameters and a dozen functions like the three I’ve given, so I’d like to only write them out once. Is this a case of using the integrator interface? Any assistance is greatly appreciated.

using OrdinaryDiffEq, Plots
import ForwardDiff

struct Parameters{T}

function f(du, u, p, t)
    F(t) = sqrt(p.a * t)
    G(t) = p.b - exp(-p.c * t)
    dG(t) = ForwardDiff.derivative(G, t)

    du[1] = -F(t) * u[2] + dG(t)
    du[2] = +G(t) * u[1]

p = Parameters(0.5, 0.7, 0.2)
u0 = [1.0, 0.0]; tspan = (0.0, 1.0)
prob = ODEProblem(f, u0, tspan, p)
sol = solve(prob, Tsit5())

F(t) = sqrt(p.a * t)
G(t) = p.b - exp(-p.c * t)
dG(t) = ForwardDiff.derivative(G, t)

# plot F, G, dG, etc...
plt = plot(F, tspan...)
plot!(G, tspan...)

An easy alternative: modify your functions to take your parameter struct as an argument:

F(p, t) = sqrt(p.a * t)
G(p, t) = p.b - exp(-p.c * t)
dG(p, t) = ForwardDiff.derivative(t -> G(p, t), t)

function f(du, u, p, t)
    du[1] = -F(p, t) * u[2] + dG(p, t)
    du[2] = G(p, t) * u[1]

dG(p, t) = ForwardDiff.derivative(t -> G(p, t), t)

Thanks, I will keep this type of line in mind.

However I must ask what to do in the case that G or dG depends on u as well. For example if I multiplied G by u[1] + u[2].

Just pass another argument:

G(u, p, t) = (u[1] + u[2])*(p.b - exp(-p.c * t))
dG(u, p, t) = ForwardDiff.derivative(t -> G(u, p, t), t)

A separate note: for small state vectors, you can get a decent speed-up by using a StaticVector instead of a Vector.

function g(u, p, t)
    du1 = -F(p, t) * u[2] + dG(p, t)
    du2 = G(p, t) * u[1]
    SA[du1, du2]

prob_f = ODEProblem(f, u0, tspan, (a = 0.5, b = 0.7, c = 0.2))
prob_g = ODEProblem(g, SA[1.0, 0.0], tspan, (a = 0.5, b = 0.7, c = 0.2))
julia> @btime solve($prob_f, Tsit5());
  8.700 μs (94 allocations: 12.17 KiB)

julia> @btime solve($prob_g, Tsit5());
  3.900 μs (26 allocations: 5.56 KiB)

Much appreciated. I’ll take a look at the static vectors; I suspect performance may become important in the future.

While this thread is fresh, can I ask how one would implement e.g. if F was given by an integral, something like F(t) = \int_{0}^{t}u_{1}^{2}(t')\,\mathrm{d}t', which requires information from the initial time up to the present time during each stop of the solve? It’s a straightforward problem if I coded the time integration manually and preallocated a bunch of vectors, but it’d be useful to know how to do it with ODEProblem.

The easiest thing is to include it in your state vector like [u1, u2, F] - you’re integrating the ODE anyways, so why not just take care of it there?