type SparseIdentificationResult has no field basis about universal_differential_equations

[ChrisRackauckas]

Are you instantiating the manifest?

[smollada]

How do I do that? Do you mean instantiating the Pkg? Sorry, I am new to Julia!

Everything until line 222

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 222 in dd83689

Ψ_direct = SInDy(X[2:4, :], DX[2:4, :], basis, thresholds, opt = opt, maxiter = 50000) # Fail

runs smoothly and I get the output figures.

[ChrisRackauckas]

Did you run this part to get the specific versions used for the script?

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Lines 1 to 2 in dd83689

	cd(@__DIR__)
	using Pkg; Pkg.activate("."); Pkg.instantiate()

[smollada]

When I run that part, I get the error:
Could not parse project /home/jovyan/Project.toml: CompositeException(Any[Pkg.TOML.ParserError(7, 7, "expected a key but found an empty string")])

For what it's worth, I am running the script in jupyter lab inside a docker container.

[ChrisRackauckas]

Are you using Julia v1.4?

I am not sure about Jupyter Lab and all of that, but it might cause an issue.

[smollada]

Ok, I have got it working locally on my machine with Julia v1.4.

I am trying to train the universal ODE for 120 days and extrapolate the solution up to 240 days, but getting the warning and errors below:
┌ Warning: Interrupted. Larger maxiters is needed.

└ @ DiffEqBase C:\Users\user\.julia\packages\DiffEqBase\ytZvl\src\integrator_interface.jl:329

┌ Warning: dt <= dtmin. Aborting. There is either an error in your model specification or the true solution is unstable.

└ @ DiffEqBase C:\Users\user\.julia\packages\DiffEqBase\ytZvl\src\integrator_interface.jl:343

Any idea on what went wrong and how to proceed?
Thanks!

[ChrisRackauckas]

You have to be a bit careful with very long time integrations. If the random neural network is not stable enough, it'll diverge in the first run and then you can't actually adjust the parameters. That's what's going on here: the random neural network defines a neural network that goes to essentially infinity in that time. The way to handle this is to not start at the full time span, i.e. https://diffeqflux.sciml.ai/dev/examples/local_minima/

[smollada]

That makes sense. Would you for example suggest something like (forgive the syntax please):

tspans = [(0.0, 20.0), (0.0, 40.0), (0.0, 60.0), (0.0, 80.0)]
results_uode = [res1_node, 0.0, 0.0, 0.0]
prob_nn_lst = [prob_nn, 0.0, 0.0, 0.0]
uode_sol_lst = [0.0, 0.0, 0.0]

for i ∈ 1:3: 
    results_uode[i+1] = DiffEqFlux.sciml_train(loss, results_uode[i].minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000)
    loss(results_uode[i + 1].minimizer)
    prob_nn[i+1] = ODEProblem(dudt_,u0, tspans[i], results_uode[i+1].minimizer)
    uode_sol[i] = solve(prob_nn[i+1], Tsit5(), saveat = 1)

to replace

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Lines 161 to 166 in dd83689

	res2_uode = DiffEqFlux.sciml_train(loss, res1_uode.minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000)
	loss(res2_uode.minimizer)
	prob_nn2 = ODEProblem(dudt_,u0, tspan, res2_uode.minimizer)
	uode_sol = solve(prob_nn2, Tsit5(), saveat = 1)

Thanks!

[ChrisRackauckas]

Yeah exactly, you can write a loop on that.

[smollada]

I have implemented that (below), but it does not extrapolate well, even with SInDy. Here is the implementation:

tspans = [(0.0, 20.0), (0.0, 40.0), (0.0, 60.0), (0.0, 80.0), (0.0, 100.0), (0.0, 120.0), (0.0, 140.0), (0.0, 160.0)]
results_uode = Any[]
push!(results_uode, res1_uode)
prob_nn_lst = Any[]
push!(prob_nn_lst, prob_nn)
uode_sol_lst = Any[]

for i = 1:7
    push!(results_uode, DiffEqFlux.sciml_train(loss, results_uode[i].minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000))
    loss(results_uode[i+1].minimizer)
    push!(prob_nn_lst, ODEProblem(dudt_,u0, tspans[i], results_uode[i+1].minimizer))
    push!(uode_sol_lst, solve(prob_nn_lst[i+1], Tsit5(), saveat = 1))
end

In line 160 (before the for loop above)

universal_differential_equations/SEIR_exposure/seir_exposure.jl

Line 160 in dd83689

res1_uode = DiffEqFlux.sciml_train(loss, p, ADAM(0.01), cb=callback, maxiters = 500)

I get the this warning:

* Status: failure (reached maximum number of iterations)

which I guess is related to the convergence and the local minimum that you pointed out before. Is the poor extrapolation related to this warning? And any idea on how to proceed?
Thanks!

[ChrisRackauckas]

I'd need more details. How well did the neural network fit? What kind of basis did it find?

[smollada]

These are the basis:

[ModelingToolkit.Constant(3.100733229880176), 1.5182135321786945 + u₁ * 0.065006601719781, u₂ * 0.16805963923290992]
Operation[u₂ * 0.30686556496954887 + u₁ * 0.000981694582027343]
Operation[u₂ * 0.1025923049193073 + u₁ * u₂ * 0.1140006264568493 + u₁ ^ 2 * u₂ * 0.12667755929642832]

Here is the full code:

cd(@__DIR__)
using Pkg; Pkg.activate("."); Pkg.instantiate()

# Single experiment, move to ensemble further on
# Some good parameter values are stored as comments right now
# because this is really good practice

using OrdinaryDiffEq
using ModelingToolkit
using DataDrivenDiffEq
using LinearAlgebra, DiffEqSensitivity, Optim
using DiffEqFlux, Flux
using Plots
gr()

function corona!(du,u,p,t)
    S,E,I,R,N,D,C = u
    F, β0,α,κ,μ,σ,γ,d,λ = p
    dS = -β0*S*F/N - β(t,β0,D,N,κ,α)*S*I/N -μ*S # susceptible
    dE = β0*S*F/N + β(t,β0,D,N,κ,α)*S*I/N -(σ+μ)*E # exposed
    dI = σ*E - (γ+μ)*I # infected
    dR = γ*I - μ*R # removed (recovered + dead)
    dN = -μ*N # total population
    dD = d*γ*I - λ*D # severe, critical cases, and deaths
    dC = σ*E # +cumulative cases

    du[1] = dS; du[2] = dE; du[3] = dI; du[4] = dR
    du[5] = dN; du[6] = dD; du[7] = dC
end
β(t,β0,D,N,κ,α) = β0*(1-α)*(1-D/N)^κ
S0 = 14e6
u0 = [0.9*S0, 0.0, 0.0, 0.0, S0, 0.0, 0.0]
p_ = [10.0, 0.5944, 0.4239, 1117.3, 0.02, 1/3, 1/5,0.2, 1/11.2]
R0 = p_[2]/p_[7]*p_[6]/(p_[6]+p_[5])
tspan = (0.0, 20.0)
prob = ODEProblem(corona!, u0, tspan, p_)
solution = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 1)

tspan2 = (0.0,160.0)
prob = ODEProblem(corona!, u0, tspan2, p_)
solution_extrapolate = solve(prob, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 1)
# Ideal data
tsdata = Array(solution)
# Add noise to the data
noisy_data = tsdata + Float32(1e-5)*randn(eltype(tsdata), size(tsdata))
plot(abs.(tsdata-noisy_data)')

####################################################### Universal ODE Part 1 ################################################################

ann = FastChain(FastDense(3, 64, tanh),FastDense(64, 64, tanh), FastDense(64, 1))
p = Float64.(initial_params(ann))

function dudt_(u,p,t)
    S,E,I,R,N,D,C = u
    F, β0,α,κ,μ,σ,γ,d,λ = p_
    z = ann([S/N,I,D/N],p) # Exposure does not depend on exposed, removed, or cumulative!
    dS = -β0*S*F/N - z[1] -μ*S # susceptible
    dE = β0*S*F/N + z[1] -(σ+μ)*E # exposed
    dI = σ*E - (γ+μ)*I # infected
    dR = γ*I - μ*R # removed (recovered + dead)
    dN = -μ*N # total population
    dD = d*γ*I - λ*D # severe, critical cases, and deaths
    dC = σ*E # +cumulative cases

    [dS,dE,dI,dR,dN,dD,dC]
end
prob_nn = ODEProblem(dudt_,u0, tspan, p)
s = concrete_solve(prob_nn, Tsit5(), u0, p, saveat = 1)

plot(solution, vars=[2,3,4])
plot(s[2:4,:]')

function predict(θ)
    Array(concrete_solve(prob_nn, Vern7(), u0, θ, saveat = solution.t,
                         abstol=1e-6, reltol=1e-6,
                         sensealg = InterpolatingAdjoint(autojacvec=ReverseDiffVJP())))
end

# No regularisation right now
function loss(θ)
    pred = predict(θ)
    sum(abs2, noisy_data[2:4,:] .- pred[2:4,:]), pred # + 1e-5*sum(sum.(abs, params(ann)))
end

loss(p)

const losses = []
callback(θ,l,pred) = begin
    push!(losses, l)
    if length(losses)%50==0
        println(losses[end])
    end
    false
end

res1_uode = DiffEqFlux.sciml_train(loss, p, ADAM(0.01), cb=callback, maxiters = 500)
#res2_uode = DiffEqFlux.sciml_train(loss, res1_uode.minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000)

tspans = [(0.0, 20.0), (0.0, 40.0), (0.0, 60.0), (0.0, 80.0), (0.0, 100.0), (0.0, 120.0), (0.0, 140.0), (0.0, 160.0)]
results_uode = Any[]
push!(results_uode, res1_uode)
prob_nn_lst = Any[]
push!(prob_nn_lst, prob_nn)
uode_sol_lst = Any[]

for i = 1:7
    push!(results_uode, DiffEqFlux.sciml_train(loss, results_uode[i].minimizer, BFGS(initial_stepnorm=0.01), cb=callback, maxiters = 10000))
    loss(results_uode[i+1].minimizer)
    push!(prob_nn_lst, ODEProblem(dudt_,u0, tspans[i], results_uode[i+1].minimizer))
    push!(uode_sol_lst, solve(prob_nn_lst[i+1], Tsit5(), saveat = 1))
end

#loss(res2_uode.minimizer)

#prob_nn2 = ODEProblem(dudt_,u0, (0.0, 60.0), res2_uode.minimizer)
#uode_sol = solve(prob_nn2, Tsit5(), saveat = 1)

#prob_nn3 = ODEProblem(dudt_,u0, (0.0, 60.0), res3_uode.minimizer)
#uode_sol2 = solve(prob_nn3, Tsit5(), saveat = 1)

plot(solution, vars=[2,3,4])
plot!(uode_sol_lst[7], vars=[2,3,4])

# Plot the losses
plot(losses, yaxis = :log, xaxis = :log, xlabel = "Iterations", ylabel = "Loss")

# Collect the state trajectory and the derivatives
X = noisy_data
# Ideal derivatives
DX = Array(solution(solution.t, Val{1}))

# Extrapolate out
prob_nn2 = ODEProblem(dudt_,u0, tspan2, results_uode[8].minimizer)
_sol_uode = solve(prob_nn2, Vern7(), abstol=1e-12, reltol=1e-12, saveat = 1)
p_uode = scatter(solution_extrapolate, vars=[2,3,4], legend = :topleft, label=["True Exposed" "True Infected" "True Recovered"], title="Universal ODE Extrapolation")
plot!(p_uode,_sol_uode, lw = 5, vars=[2,3,4], label=["Estimated Exposed" "Estimated Infected" "Estimated Recovered"])
plot!(p_uode,[20.99,21.01],[0.0,maximum(hcat(Array(solution_extrapolate[2:4,:]),Array(_sol_uode[2:4,:])))],lw=5,color=:black,label="Training Data End")

savefig("universalode_extrapolation.png")
savefig("universalode_extrapolation.pdf")


################################################################# Universal ODE Part 2: SInDy to Equations ######################################################

# Create a Basis
@variables u[1:3]
# Lots of polynomials
polys = Operation[]
for i ∈ 0:2, j ∈ 0:2, k ∈ 0:2
    push!(polys, u[1]^i * u[2]^j * u[3]^k)
end

# And some other stuff
h = [cos.(u)...; sin.(u)...; unique(polys)...]
basis = Basis(h, u)

X = noisy_data
# Ideal derivatives
DX = Array(solution(solution.t, Val{1}))
S,E,I,R,N,D,C = eachrow(X)
F,β0,α,κ,μ,_,γ,d,λ = p_
L = β.(0:tspan[end],β0,D,N,κ,α).*S.*I./N
L̂ = vec(ann([S./N I D./N]',results_uode[8].minimizer))
X̂ = [S./N I D./N]'

scatter(L,title="Estimated vs Expected Exposure Term",label="True Exposure")
plot!(L̂,label="Estimated Exposure")
savefig("estimated_exposure.png")
savefig("estimated_exposure.pdf")

# Create an optimizer for the SINDY problem
opt = SR3()
# Create the thresholds which should be used in the search process
thresholds = exp10.(-6:0.1:1)

# Test on original data and without further knowledge
Ψ_direct = SInDy(X[2:4, :], DX[2:4, :], basis, thresholds, opt = opt, maxiter = 50000) # Fail
println(Ψ_direct.basis)
# Test on ideal derivative data ( not available )
Ψ_ideal = SInDy(X[2:4, 5:end], L[5:end], basis, thresholds, opt = opt, maxiter = 50000) # Succeed
println(Ψ_ideal.basis)
# Test on uode derivative data
Ψ = SInDy(X̂[:, 2:end], L̂[2:end], basis, thresholds,  opt = opt, maxiter = 10000, normalize = true, denoise = true) # Succeed
println(Ψ.basis)

# Build a ODE for the estimated system
function approx(u,p,t)
    S,E,I,R,N,D,C = u
    F, β0,α,κ,μ,σ,γ,d,λ = p_
    z = Ψ([S/N,I,D/N]) # Exposure does not depend on exposed, removed, or cumulative!
    dS = -β0*S*F/N - z[1] -μ*S # susceptible
    dE = β0*S*F/N + z[1] -(σ+μ)*E # exposed
    dI = σ*E - (γ+μ)*I # infected
    dR = γ*I - μ*R # removed (recovered + dead)
    dN = -μ*N # total population
    dD = d*γ*I - λ*D # severe, critical cases, and deaths
    dC = σ*E # +cumulative cases

    [dS,dE,dI,dR,dN,dD,dC]
end

# Create the approximated problem and solution
a_prob = ODEProblem{false}(approx, u0, tspan2, p_)
a_solution = solve(a_prob, Tsit5())

p_uodesindy = scatter(solution_extrapolate, vars=[2,3,4], legend = :topleft, label=["True Exposed" "True Infected" "True Recovered"])
plot!(p_uodesindy,a_solution, lw = 5, vars=[2,3,4], label=["Estimated Exposed" "Estimated Infected" "Estimated Recovered"])
plot!(p_uodesindy,[20.99,21.01],[0.0,maximum(hcat(Array(solution_extrapolate[2:4,:]),Array(_sol_uode[2:4,:])))],lw=5,color=:black,label="Training Data End")

savefig("universalodesindy_extrapolation.png")
savefig("universalodesindy_extrapolation.pdf")

The loss goes down as expected, but the extrapolation does not look good!

[ChrisRackauckas]

I wouldn't expect anything to be perfect. I'd say that looks pretty good! It's extrapolating from like day 21 to day 80, which is a pretty strong showing!

[smollada]

Ok, good to know. For the SInDy, do you have any suggestions to use data up to more days (e.g. 60 or 90) and extrapolate beyond the peaks? Would a similar strategy to Universal ODE (train incrementally) work for SInDy as well? Sorry, I am new to SInDy too, but it looks promising.

[ChrisRackauckas]

I haven't seen anything, at least if it doesn't have enough data to perfectly learn the generating equation