ChrisRackauckas / universal_differential_equations

Repository for the Universal Differential Equations for Scientific Machine Learning paper, describing a computational basis for high performance SciML
https://arxiv.org/abs/2001.04385
MIT License
219 stars 59 forks source link

type SparseIdentificationResult has no field basis #24

Open smollada opened 4 years ago

smollada commented 4 years ago

When I try to run the code, I get the following error: https://github.com/ChrisRackauckas/universal_differential_equations/blob/dd836890e7e09923a0ae9ed8b6b49d2ee64b2e6a/SEIR_exposure/seir_exposure.jl#L223

type SparseIdentificationResult has no field basis

Thanks!

ChrisRackauckas commented 4 years ago

Are you instantiating the manifest?

smollada commented 4 years ago

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

Everything until line 222 https://github.com/ChrisRackauckas/universal_differential_equations/blob/dd836890e7e09923a0ae9ed8b6b49d2ee64b2e6a/SEIR_exposure/seir_exposure.jl#L222 runs smoothly and I get the output figures.

ChrisRackauckas commented 4 years ago

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

https://github.com/ChrisRackauckas/universal_differential_equations/blob/dd836890e7e09923a0ae9ed8b6b49d2ee64b2e6a/SEIR_exposure/seir_exposure.jl#L1-L2

smollada commented 4 years ago

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 commented 4 years ago

Are you using Julia v1.4?

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

smollada commented 4 years ago

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 commented 4 years ago

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 commented 4 years ago

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 https://github.com/ChrisRackauckas/universal_differential_equations/blob/dd836890e7e09923a0ae9ed8b6b49d2ee64b2e6a/SEIR_exposure/seir_exposure.jl#L161-L166

Thanks!

ChrisRackauckas commented 4 years ago

Yeah exactly, you can write a loop on that.

smollada commented 4 years ago

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) https://github.com/ChrisRackauckas/universal_differential_equations/blob/dd836890e7e09923a0ae9ed8b6b49d2ee64b2e6a/SEIR_exposure/seir_exposure.jl#L160 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 commented 4 years ago

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

smollada commented 4 years ago

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!

universalode_extrapolation universalodesindy_extrapolation

ChrisRackauckas commented 4 years ago

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 commented 4 years ago

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 commented 4 years ago

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