Hybrid ODE with ContinuousCallback for models of changing sizes?

natalieisenberg
Tue, 04 Jun 2024 21:14:12 +0000
Machine Learning

I am solving a system of ODEs with a ContinuousCallback where once a condition!() is met, certain members of the state vector are removed (i.e., elements disappear during simulation). I would like to solve this system using DiffEqFlux.jl as a Neural Differential Equation, specifically a Universal Differential Equation. I am having trouble getting this working, and I believe the problem is my custom ContinuousCallback due to there being a “cannot resize array with shared data” error. Does anyone know if the functionality I’ve described is supported in DiffEqFlux.jl?

Here’s a minimal example which I have adapted from @ChrisRackauckas’s universal_differential_equations repo (hope that’s okay):

using OrdinaryDiffEq, ModelingToolkit, LinearAlgebra, ComponentArrays, Optimization, OptimizationOptimisers, OptimizationOptimJL, Lux,
Plots, ComponentArrays, DiffEqFlux, JLD2, FileIO, Statistics, Random, Distributions, LinearSolve

# Set a random seed for reproduceable behaviour
using Random
rng = Random.default_rng()
Random.seed!(1234)

#### NOTE
# Since the recent release of DataDrivenDiffEq v0.6.0 where a complete overhaul of the optimizers took
# place, SR3 has been used. Right now, STLSQ performs better and has been changed.

# Create a name for saving ( basically a prefix )
svname = "Scenario_1_"

# === Add a ContinuousCallback to test behaviour
# Callback condition for removing elements when value is at threshold
function condition(u, t, integrator)
    # Trigger the event if the order of magnitude of state variable goes below some arbitrary limit...need to pick something that will for-sure get triggered within the integration
    minimum(u) < 0.5 ? 0 : 1
end

# Callback to modify the state vector if above condition is met
function affect!(integrator)
    original_size = length(integrator.u)
    idxs = findall(x -> x <= 0.5, integrator.u)
    new_size = original_size - length(idxs)

    # Remove the identified elements that have below our threshold
    deleteat!(integrator.u, idxs)
    deleteat_non_user_cache!(integrator, idxs)

    resize!(integrator, new_size)
    resize_non_user_cache!(integrator, new_size)
    nothing
end

## Data generation
function lotka!(du, u, p, t)
    α, β, γ, δ = p
    du[1] = α*u[1] - β*u[2]*u[1]
    du[2] = γ*u[1]*u[2]  - δ*u[2]
end

# Define the experimental parameter
disappearing_callback = ContinuousCallback(condition, affect!)
tspan = (0.0,3.0)
u0 = [0.44249296,4.6280594]
p_ = [1.3, 0.9, 0.8, 1.8]

# This works fine
prob = ODEProblem(lotka!, u0,tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 0.1, callback=disappearing_callback)

# Ideal data
X = Array(solution)
t = solution.t
DX = Array(solution(solution.t, Val{1}))

# Add noise in terms of the mean
x̄ = mean(X, dims = 2)
noise_magnitude = 5e-3
Xₙ = X .+ (noise_magnitude*x̄) .* randn(eltype(X), size(X))

plot(solution, alpha = 0.75, color = :black, label = ["True Data" nothing])
display(scatter!(t, transpose(Xₙ), color = :red, label = ["Noisy Data" nothing]))

## Define the network
# Gaussian RBF as activation
rbf(x) = exp.(-(x.^2))

# Multilayer FeedForward
U = Lux.Chain(
    Lux.Dense(2,5,rbf), Lux.Dense(5,5, rbf), Lux.Dense(5,5, rbf), Lux.Dense(5,2)
)
# Get the initial parameters and state variables of the model
p, st = Lux.setup(rng, U)

# Define the hybrid model
function ude_dynamics!(du,u, p, t, p_true)
    û = U(u, p, st)[1] # Network prediction
    du[1] = p_true[1]*u[1] + û[1]
    du[2] = -p_true[4]*u[2] + û[2]
end

# Closure with the known parameter
nn_dynamics!(du,u,p,t) = ude_dynamics!(du,u,p,t,p_)
# Define the problem
disappearing_callback = ContinuousCallback(condition, affect!)
prob_nn = ODEProblem(nn_dynamics!,Xₙ[:, 1], tspan, p)

## Function to train the network
# Define a predictor
# Add the callbacks
function predict!(θ, X = Xₙ[:,1], T = t)
    _prob = remake(prob_nn, u0 = X, tspan = (T[1], T[end]), p = θ)
    Array(solve(_prob, Rosenbrock23(linsolve = LUFactorization()), saveat = T,
                abstol=1e-6, reltol=1e-6, callback=disappearing_callback))
end

# Simple L2 loss
function loss(θ)
    X̂ = predict!(θ)
    sum(abs2, Xₙ .- X̂)
end

# Container to track the losses
losses = Float64[]

callback = function (p, l)
  push!(losses, l)
  if length(losses)%50==0
      println("Current loss after $(length(losses)) iterations: $(losses[end])")
  end
  return false
end

## Training

# First train with ADAM for better convergence -> move the parameters into a
# favourable starting positing for BFGS
adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x,p)->loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, ComponentVector{Float64}(p))

# This throws the error
res1 = Optimization.solve(optprob, ADAM(0.1), callback=callback, maxiters = 5)
println("Training loss after $(length(losses)) iterations: $(losses[end])")

26 posts - 3 participants