Neural Hybrid DE

Specific Domains Machine Learning
1

Hey, I am following this article Neural Hybrid Differential Equations | juliabloggers.com concerning Neural Hybrid DE.
Currently, I am stuck at this moment with my code.


using DiffEqFlux, DifferentialEquations, Flux, Optimization
using Plots, DiffEqSensitivity, Zygote, Random, OptimizationFlux
using SciMLSensitivity, OrdinaryDiffEq

u0= Float32[1.;0.]

tspan = (0.0f0, 50.0f0)

dtsave = 0.5f0

t = tspan[1]:dtsave:tspan[2]

function oscilator!(du, u, p, t)
    du[1] = u[2]
    du[2] = -u[1] - 1//10*u[2]
end


prob_data = ODEProblem(oscilator!, u0, tspan)


pl = plot(solve(prob_data, Tsit5(), saveat = t), label =["x(t)" "v(t)"])

kicktimes = tspan[1]:1:tspan[2]

function kick!(integrator)
    integrator.u[end] += one(eltype(integrator.u))
end

cb_= PresetTimeCallback(kicktimes, kick!, save_positions=(true, true))

sol_data = solve(prob_data, Tsit5(), callback=cb_, saveat=t)
t_data = sol_data.t
ode_data = Array(sol_data)

pl1 = plot(t_data, ode_data[1, :], label = "data x(t)")
plot!(pl1, t_data, ode_data[2, :], label = "data v(t)")


pl2 = plot(t_data[1:20], ode_data[1,1:20], label="data x(t)")
plot!(pl2,t_data[1:20], ode_data[2,1:20], label="data v(t)")
pl = plot(pl2, pl1, layout=(1,2), xlabel="t")




nn1 = Chain(Dense(2, 64, tanh), Dense(64, 1))


p_nn1 = Flux.params(nn1)

function f1!(du,u,p,t)
    du[1] = u[2]
    du[2] = - u[1] - 1//10*u[2]
    du[3] = nn1(u[1:2], p)[1]
    return nothing
  end
  
  affect!(integrator) = integrator.u[2] += integrator.u[3]
  cb = PresetTimeCallback(kicktimes,affect!,save_positions=(false,false))
  z0 = Float32[u0;one(u0[1])]
  prob1 = ODEProblem(f1!,z0,tspan,p_nn1)
  
  function loss(p; prob=prob1, sensealg = ReverseDiffAdjoint())
    pred = Array(solve(prob,Tsit5(),callback=cb,
                 saveat=dtsave,sensealg=sensealg))[1:2,:]
    sum(abs2, ode_data .- pred)
  end
  
  function loss(p, prob=prob1, sensealg = ReverseDiffAdjoint())
    _prob = remake(prob,p=p)
    pred = Array(solve(_prob,Tsit5(),callback=cb,
                 saveat=dtsave,sensealg=sensealg))[1:2,:]
    sum(abs2,ode_data .- pred)
  end
  
  loss(p_nn1)

The error I am getting is: MethodError: no method matching (::Chain{Tuple{Dense{typeof(tanh),
Matrix{Float32}, Vector{Float32}}, Dense{typeof(identity), Matrix{Float32}, Vector{Float32}}}})(::Vector{Float32}, ::Params{Zygote.Buffer{Any, Vector{Any}}}).
Any help would be appreciated :slight_smile:

2

You may want to instead look at the tutorials in the library which are tested and maintained:

https://docs.sciml.ai/SciMLSensitivity/dev/examples/hybrid_jump/hybrid_diffeq/

https://docs.sciml.ai/SciMLSensitivity/dev/examples/hybrid_jump/bouncing_ball/

@frankschae maybe link to them from the blog post?

1 Like
3

FWIW, I think you want to do:

p_nn1, re = Flux.destructure(nn1)
function f1!(du, u, p, t)
    du[1] = u[2]
    du[2] = -u[1] - 1 // 10 * u[2]
    du[3] = re(p)(u[1:2])[1]
    return nothing
end

@frankschae maybe link to them from the blog post?

Good idea. I’ll update the post.

1 Like