Delay Differential Equations

Other differential equation problem types from DifferentialEquations.jl are supported. For example, we can build a layer with a delay differential equation like:

using DifferentialEquations, Optimization, OptimizationPolyalgorithms


# Define the same LV equation, but including a delay parameter
function delay_lotka_volterra!(du, u, h, p, t)
  x, y = u
  α, β, δ, γ = p
  du[1] = dx = (α   - β*y) * h(p, t-0.1)[1]
  du[2] = dy = (δ*x - γ)   * y
end

# Initial parameters
p = [2.2, 1.0, 2.0, 0.4]

# Define a vector containing delays for each variable (although only the first
# one is used)
h(p, t) = ones(eltype(p), 2)

# Initial conditions
u0 = [1.0, 1.0]

# Define the problem as a delay differential equation
prob_dde = DDEProblem(delay_lotka_volterra!, u0, h, (0.0, 10.0),
                      constant_lags = [0.1])

function predict_dde(p)
  return Array(solve(prob_dde, MethodOfSteps(Tsit5()),
                              u0=u0, p=p, saveat = 0.1,
                              sensealg = ReverseDiffAdjoint()))
end

loss_dde(p) = sum(abs2, x-1 for x in predict_dde(p))

#using Plots
callback = function (p,l...)
  display(loss_dde(p))
  #display(plot(solve(remake(prob_dde,p=p),MethodOfSteps(Tsit5()),saveat=0.1),ylim=(0,6)))
  return false
end

callback(p,loss_dde(p))

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss_dde(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)
result_dde = Optimization.solve(optprob, PolyOpt(), p, callback=callback)

Notice that we chose sensealg = ReverseDiffAdjoint() to utilize the ReverseDiff.jl reverse-mode to handle the delay differential equation.

We define a callback to display the solution at the current parameters for each step of the training:

#using Plots
callback = function (p,l...)
  display(loss_dde(p))
  #display(plot(solve(remake(prob_dde,p=p),MethodOfSteps(Tsit5()),saveat=0.1),ylim=(0,6)))
  return false
end

callback(p,loss_dde(p))

We use Optimization.solve to optimize the parameters for our loss function:

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss_dde(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)
result_dde = Optimization.solve(optprob, PolyOpt(), p, callback=callback)