From 61c6c3c73178f2b13cd043c3897a60e36eb485eb Mon Sep 17 00:00:00 2001 From: Samedh Desai Date: Tue, 27 Jun 2023 00:45:18 -0400 Subject: [PATCH 1/5] Initial changes to ode_solve, added NNDAE to get this solver working for DAE problems --- src/ode_solve.jl | 287 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 275 insertions(+), 12 deletions(-) diff --git a/src/ode_solve.jl b/src/ode_solve.jl index 84e6daf29..c6bcd22c4 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -1,4 +1,5 @@ abstract type NeuralPDEAlgorithm <: DiffEqBase.AbstractODEAlgorithm end +abstract type NeuralPDEAlgorithmDAE <: DiffEqBase.AbstractDAEAlgorithm end """ ```julia @@ -29,7 +30,8 @@ of the physics-informed neural network which is used as a solver for a standard ## Example -```julia +```juliap = [1.5, 1.0, 3.0, 1.0] +u0 = [1.0, 1.0] ts=[t for t in 1:100] (u_, t_) = (analytical_func(ts), ts) function additional_loss(phi, θ) @@ -96,6 +98,25 @@ function NNODE(chain, opt, init_params = nothing; NNODE(chain, opt, init_params, autodiff, batch, strategy, additional_loss, kwargs) end +struct NNDAE{C, O, P, B, K, AL <: Union{Nothing, Function}, + S <: Union{Nothing, AbstractTrainingStrategy} + } <: + NeuralPDEAlgorithmDAE + chain::C + opt::O + init_params::P + autodiff::Bool + batch::B + strategy::S + additional_loss::AL + kwargs::K +end +function NNDAE(chain, opt, init_params = nothing; + strategy = nothing, + autodiff = false, batch = nothing, additional_loss = nothing, kwargs...) + NNDAE(chain, opt, init_params, autodiff, batch, strategy, additional_loss, kwargs) +end + """ ```julia ODEPhi(chain::Lux.AbstractExplicitLayer, t, u0, st) @@ -120,23 +141,21 @@ mutable struct ODEPhi{C, T, U, S} end end -function generate_phi_θ(chain::Lux.AbstractExplicitLayer, t, u0, init_params::Nothing) - θ, st = Lux.setup(Random.default_rng(), chain) - ODEPhi(chain, t, u0, st), ComponentArrays.ComponentArray(θ) -end - function generate_phi_θ(chain::Lux.AbstractExplicitLayer, t, u0, init_params) θ, st = Lux.setup(Random.default_rng(), chain) - ODEPhi(chain, t, u0, st), ComponentArrays.ComponentArray(init_params) -end - -function generate_phi_θ(chain::Flux.Chain, t, u0, init_params::Nothing) - θ, re = Flux.destructure(chain) - ODEPhi(re, t, u0), θ + if init_params === nothing + init_params = ComponentArrays.ComponentArray(θ) + else + init_params = ComponentArrays.ComponentArray(init_params) + end + ODEPhi(chain, t, u0, st), init_params end function generate_phi_θ(chain::Flux.Chain, t, u0, init_params) θ, re = Flux.destructure(chain) + if init_params === nothing + init_params = θ + end ODEPhi(re, t, u0), init_params end @@ -252,6 +271,35 @@ function inner_loss(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, sum(abs2, dxdtguess .- fs) / length(t) end +""" +L2 inner loss for DAEProblems +""" + +function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, + p) where {C, T, U <: Number} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(phi(t, θ), p, t)) +end + +function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, + p) where {C, T, U <: Number} + out = phi(t, θ) + dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) + sum(abs2, f(dxdtguess[i], p, t[i]) for i in 1:size(out, 2)) / length(t) +end + +function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, + p) where {C, T, U} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(phi(t, θ), p, t)) +end + +function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, + p) where {C, T, U} + out = Array(phi(t, θ)) + arrt = Array(t) + dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) + sum(abs2, f(dxdtguess[:, i], p, arrt[i]) for i in 1:size(out, 2)) / length(t) +end + """ Representation of the loss function, parametric on the training strategy `strategy` """ @@ -334,6 +382,89 @@ function generate_loss(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, ts error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") end +""" +Representation of the loss function, parametric on the training strategy `strategy` for DAE problems +""" +function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, + batch) + integrand(t::Number, θ) = abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) + integrand(ts, θ) = [abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) for t in ts] + @assert batch == 0 # not implemented + + function loss(θ, _) + intprob = IntegralProblem(integrand, tspan[1], tspan[2], θ) + sol = solve(intprob, QuadGKJL(); abstol = strategy.abstol, reltol = strategy.reltol) + sol.u + end + + return loss +end + +function generate_loss_DAE(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, batch) + ts = tspan[1]:(strategy.dx):tspan[2] + + # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken + function loss(θ, _) + if batch + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + else + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + end + end + return loss +end + +function generate_loss_DAE(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, + batch) + # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken + function loss(θ, _) + ts = adapt(parameterless_type(θ), + [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)]) + if batch + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + else + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + end + end + return loss +end + +function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p, + batch) + minT = tspan[1] + maxT = tspan[2] + + weights = strategy.weights ./ sum(strategy.weights) + + N = length(weights) + samples = strategy.samples + + difference = (maxT - minT) / N + + data = Float64[] + for (index, item) in enumerate(weights) + temp_data = rand(1, trunc(Int, samples * item)) .* difference .+ minT .+ + ((index - 1) * difference) + data = append!(data, temp_data) + end + + ts = data + + function loss(θ, _) + if batch + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + else + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + end + end + return loss +end + +function generate_loss_DAE(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan) + error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") +end + + struct NNODEInterpolation{T <: ODEPhi, T2} phi::T θ::T2 @@ -483,3 +614,135 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, dense_errors = false) sol end #solve + +function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, + alg::NNODE, + args...; + dt = nothing, + timeseries_errors = true, + save_everystep = true, + adaptive = false, + abstol = 1.0f-6, + reltol = 1.0f-3, + verbose = false, + saveat = nothing, + maxiters = nothing) + + u0 = prob.u0 + tspan = prob.tspan + f = prob.f + p = prob.p + t0 = tspan[1] + + #hidden layer + chain = alg.chain + opt = alg.opt + autodiff = alg.autodiff + + #train points generation + init_params = alg.init_params + + if chain isa Lux.AbstractExplicitLayer || chain isa Flux.Chain + phi, init_params = generate_phi_θ(chain, t0, u0, init_params) + else + error("Only Lux.AbstractExplicitLayer and Flux.Chain neural networks are supported") + end + + # if isinplace(prob) + # throw(error("The NNODE solver only supports out-of-place ODE definitions, i.e. du=f(u,p,t).")) + # end + + try + phi(t0, init_params) + catch err + if isa(err, DimensionMismatch) + throw(DimensionMismatch("Dimensions of the initial u0 and chain should match")) + else + throw(err) + end + end + + strategy = if alg.strategy === nothing + if dt !== nothing + GridTraining(dt) + else + QuadratureTraining(; quadrature_alg = QuadGKJL(), + reltol = convert(eltype(u0), reltol), + abstol = convert(eltype(u0), abstol), maxiters = maxiters, + batch = 0) + end + else + alg.strategy + end + + batch = if alg.batch === nothing + if strategy isa QuadratureTraining + strategy.batch + else + true + end + else + alg.batch + end + + inner_f = generate_loss_DAE(strategy, phi, f, autodiff, tspan, p, batch) + additional_loss = alg.additional_loss + + # Creates OptimizationFunction Object from total_loss + function total_loss(θ, _) + L2_loss = inner_f(θ, phi) + if !(additional_loss isa Nothing) + return additional_loss(phi, θ) + L2_loss + end + L2_loss + end + + # Choice of Optimization Algo for Training Strategies + opt_algo = if strategy isa QuadratureTraining + Optimization.AutoForwardDiff() + else + Optimization.AutoZygote() + end + + # Creates OptimizationFunction Object from total_loss + optf = OptimizationFunction(total_loss, opt_algo) + + iteration = 0 + callback = function (p, l) + iteration += 1 + verbose && println("Current loss is: $l, Iteration: $iteration") + l < abstol + end + + optprob = OptimizationProblem(optf, init_params) + res = solve(optprob, opt; callback, maxiters, alg.kwargs...) + + #solutions at timepoints + if saveat isa Number + ts = tspan[1]:saveat:tspan[2] + elseif saveat isa AbstractArray + ts = saveat + elseif dt !== nothing + ts = tspan[1]:dt:tspan[2] + elseif save_everystep + ts = range(tspan[1], tspan[2], length = 100) + else + ts = [tspan[1], tspan[2]] + end + + if u0 isa Number + u = [first(phi(t, res.u)) for t in ts] + else + u = [phi(t, res.u) for t in ts] + end + + sol = DiffEqBase.build_solution(prob, alg, ts, u; + k = res, dense = true, + interp = NNODEInterpolation(phi, res.u), + calculate_error = false, + retcode = ReturnCode.Success) + DiffEqBase.has_analytic(prob.f) && + DiffEqBase.calculate_solution_errors!(sol; timeseries_errors = true, + dense_errors = false) + sol +end #solve \ No newline at end of file From c56922c7f291dc07aa5b9b48cc29a5115510fbcb Mon Sep 17 00:00:00 2001 From: Samedh Desai Date: Tue, 27 Jun 2023 01:23:45 -0400 Subject: [PATCH 2/5] Added test file to commit --- src/NeuralPDE.jl | 2 +- src/dae_problem_test.jl | 25 +++++++++++++++++++++++++ src/ode_solve.jl | 2 +- 3 files changed, 27 insertions(+), 2 deletions(-) create mode 100644 src/dae_problem_test.jl diff --git a/src/NeuralPDE.jl b/src/NeuralPDE.jl index f61544274..9ad913d9b 100644 --- a/src/NeuralPDE.jl +++ b/src/NeuralPDE.jl @@ -48,7 +48,7 @@ include("transform_inf_integral.jl") include("discretize.jl") include("neural_adapter.jl") -export NNODE, TerminalPDEProblem, NNPDEHan, NNPDENS, NNRODE, +export NNODE, NNDAE, TerminalPDEProblem, NNPDEHan, NNPDENS, NNRODE, KolmogorovPDEProblem, NNKolmogorov, NNStopping, ParamKolmogorovPDEProblem, KolmogorovParamDomain, NNParamKolmogorov, PhysicsInformedNN, discretize, diff --git a/src/dae_problem_test.jl b/src/dae_problem_test.jl new file mode 100644 index 000000000..f9a6eb84b --- /dev/null +++ b/src/dae_problem_test.jl @@ -0,0 +1,25 @@ +using DAEProblemLibrary, Sundials, Optimisers, OptimizationOptimisers, DifferentialEquations +using NeuralPDE, Lux, Test, Statistics, Plots + +prob = DAEProblemLibrary.prob_dae_resrob +true_sol = solve(prob, IDA(), saveat = 0.01) +# sol = solve(prob, IDA()) + +u0 = [1.0, 1.0, 1.0] +func = Lux.σ +N = 12 +chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func), + Lux.Dense(N, N, func), Lux.Dense(N, length(u0))) + +opt = Optimisers.Adam(0.01) +dx = 0.05 +alg = NeuralPDE.NNDAE(chain, opt, autodiff = false, strategy = NeuralPDE.GridTraining(dx)) +sol = solve(prob, alg, verbose=true, maxiters = 100000, saveat = 0.01) + +println(abs(mean(true_sol .- sol))) + +using Plots + +plot(sol) +plot!(true_sol) +# ylims!(0,8) \ No newline at end of file diff --git a/src/ode_solve.jl b/src/ode_solve.jl index c6bcd22c4..ab40c87c2 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -616,7 +616,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, end #solve function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, - alg::NNODE, + alg::NNDAE, args...; dt = nothing, timeseries_errors = true, From 595359f4b3cb4e4d96a80aecddbebc39a79b7d57 Mon Sep 17 00:00:00 2001 From: Samedh Desai Date: Tue, 18 Jul 2023 14:36:40 -0400 Subject: [PATCH 3/5] Added changes to make DAE function definitions include u, but it still throws an error --- src/dae_problem_test.jl | 43 ++++++++++++++++++----- src/ode_solve.jl | 48 ++++++++++++++------------ src/weighted_interval_training_test.jl | 31 +++++++++++++++++ 3 files changed, 90 insertions(+), 32 deletions(-) create mode 100644 src/weighted_interval_training_test.jl diff --git a/src/dae_problem_test.jl b/src/dae_problem_test.jl index f9a6eb84b..ed77520fd 100644 --- a/src/dae_problem_test.jl +++ b/src/dae_problem_test.jl @@ -1,9 +1,34 @@ using DAEProblemLibrary, Sundials, Optimisers, OptimizationOptimisers, DifferentialEquations using NeuralPDE, Lux, Test, Statistics, Plots -prob = DAEProblemLibrary.prob_dae_resrob -true_sol = solve(prob, IDA(), saveat = 0.01) -# sol = solve(prob, IDA()) +f = function (r, yp, y, p, tres) + r[1] = -0.04 * y[1] + 1.0e4 * y[2] * y[3] + r[2] = -r[1] - 3.0e7 * y[2] * y[2] - yp[2] + r[1] -= yp[1] + r[3] = y[1] + y[2] + y[3] - 1.0 +end +u0 = [1.0, 0, 0] +du0 = [-0.04, 0.04, 0.0] + +""" +The Robertson biochemical reactions in DAE form + +```math +\frac{dy₁}{dt} = -k₁y₁+k₃y₂y₃ +``` +```math +\frac{dy₂}{dt} = k₁y₁-k₂y₂^2-k₃y₂y₃ +``` +```math +1 = y₁ + y₂ + y₃ +``` +where ``k₁=0.04``, ``k₂=3\times10^7``, ``k₃=10^4``. For details, see: +Hairer Norsett Wanner Solving Ordinary Differential Equations I - Nonstiff Problems Page 129 +Usually solved on ``[0,1e11]`` +""" + +prob_oop = DAEProblem(f, du0, u0, (0.0, 100000.0)) +true_sol = solve(prob_oop, IDA(), saveat = 0.01) u0 = [1.0, 1.0, 1.0] func = Lux.σ @@ -14,12 +39,12 @@ chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, opt = Optimisers.Adam(0.01) dx = 0.05 alg = NeuralPDE.NNDAE(chain, opt, autodiff = false, strategy = NeuralPDE.GridTraining(dx)) -sol = solve(prob, alg, verbose=true, maxiters = 100000, saveat = 0.01) +sol = solve(prob_oop, alg, verbose=true, maxiters = 100000, saveat = 0.01) -println(abs(mean(true_sol .- sol))) +# println(abs(mean(true_sol .- sol))) -using Plots +# using Plots -plot(sol) -plot!(true_sol) -# ylims!(0,8) \ No newline at end of file +# plot(sol) +# plot!(true_sol) +# # ylims!(0,8) \ No newline at end of file diff --git a/src/ode_solve.jl b/src/ode_solve.jl index ab40c87c2..8fc621cda 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -275,29 +275,31 @@ end L2 inner loss for DAEProblems """ +function inner_loss_DAE end + function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, - p) where {C, T, U <: Number} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(phi(t, θ), p, t)) + p, u) where {C, T, U <: Number} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(phi(t, θ), u, p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, - p) where {C, T, U <: Number} + p, u) where {C, T, U <: Number} out = phi(t, θ) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[i], p, t[i]) for i in 1:size(out, 2)) / length(t) + sum(abs2, f(dxdtguess[i], u, p, t[i]) for i in 1:size(out, 2)) / length(t) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, - p) where {C, T, U} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(phi(t, θ), p, t)) + p, u) where {C, T, U} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(phi(t, θ), u, p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, - p) where {C, T, U} + p, u) where {C, T, U} out = Array(phi(t, θ)) arrt = Array(t) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[:, i], p, arrt[i]) for i in 1:size(out, 2)) / length(t) + sum(abs2, f(dxdtguess[:, i], u, p, arrt[i]) for i in 1:size(out, 2)) / length(t) end """ @@ -385,10 +387,10 @@ end """ Representation of the loss function, parametric on the training strategy `strategy` for DAE problems """ -function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, +function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, u, batch) - integrand(t::Number, θ) = abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) - integrand(ts, θ) = [abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) for t in ts] + integrand(t::Number, θ) = abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p, u)) + integrand(ts, θ) = [abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p, u)) for t in ts] @assert batch == 0 # not implemented function loss(θ, _) @@ -400,36 +402,36 @@ function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, return loss end -function generate_loss_DAE(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, batch) +function generate_loss_DAE(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, u, batch) ts = tspan[1]:(strategy.dx):tspan[2] # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken function loss(θ, _) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, +function generate_loss_DAE(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, u, batch) # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken function loss(θ, _) ts = adapt(parameterless_type(θ), [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)]) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p, +function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p, u, batch) minT = tspan[1] maxT = tspan[2] @@ -452,15 +454,15 @@ function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff: function loss(θ, _) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan) +function generate_loss_DAE(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan, p, u, batch) error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") end @@ -683,9 +685,9 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, end else alg.batch - end + endx - inner_f = generate_loss_DAE(strategy, phi, f, autodiff, tspan, p, batch) + inner_f = generate_loss_DAE(strategy, phi, f, autodiff, tspan, p, u0, batch) additional_loss = alg.additional_loss # Creates OptimizationFunction Object from total_loss diff --git a/src/weighted_interval_training_test.jl b/src/weighted_interval_training_test.jl new file mode 100644 index 000000000..5d81d8960 --- /dev/null +++ b/src/weighted_interval_training_test.jl @@ -0,0 +1,31 @@ + +using NeuralPDE, OrdinaryDiffEq, OptimizationPolyalgorithms, Lux, OptimizationOptimJL, Test, Statistics, Plots, Optimisers + +function f(u, p, t) + [p[1] * u[1] - p[2] * u[1] * u[2], -p[3] * u[2] + p[4] * u[1] * u[2]] +end + +p = [1.5, 1.0, 3.0, 1.0] +u0 = [1.0, 1.0] +prob_oop = ODEProblem{false}(f, u0, (0.0, 3.0), p) +true_sol = solve(prob_oop, Tsit5(), saveat = 0.01) +func = Lux.σ +N = 12 +chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func), + Lux.Dense(N, N, func), Lux.Dense(N, length(u0))) + +opt = Optimisers.Adam(0.01) +weights = [0.7, 0.2, 0.1] +samples = 200 +alg = NeuralPDE.NNODE(chain, opt, autodiff = false, strategy = NeuralPDE.WeightedIntervalTraining(weights, samples)) +sol = solve(prob_oop, alg, verbose=true, maxiters = 100000, saveat = 0.01) + +println(abs(mean(true_sol .- sol))) +# println(abs(mean(sol) - mean(true_sol))) +# @test abs(mean(sol) - mean(true_sol)) < 0.2 + +using Plots + +plot(sol) +plot!(true_sol) +ylims!(0,8) \ No newline at end of file From 41c90840a7817b0eedbf63d214b5079a9f03476d Mon Sep 17 00:00:00 2001 From: Samedh Desai Date: Wed, 26 Jul 2023 19:31:10 -0400 Subject: [PATCH 4/5] Fixed some of the issues with problem definitions and loss function generation, still have an error --- src/ode_solve.jl | 46 ++++++++++++++++++------------------ test/dae_problem_test.jl | 50 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 73 insertions(+), 23 deletions(-) create mode 100644 test/dae_problem_test.jl diff --git a/src/ode_solve.jl b/src/ode_solve.jl index 8fc621cda..a70d453d5 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -278,28 +278,28 @@ L2 inner loss for DAEProblems function inner_loss_DAE end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, - p, u) where {C, T, U <: Number} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(phi(t, θ), u, p, t)) + p) where {C, T, U <: Number} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(ode_dfdx(phi, t, θ, autodiff), phi, p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, - p, u) where {C, T, U <: Number} + p) where {C, T, U <: Number} out = phi(t, θ) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[i], u, p, t[i]) for i in 1:size(out, 2)) / length(t) + sum(abs2, f(dxdtguess[i], phi, p, t[i]) for i in 1:size(out, 2)) / length(t) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, - p, u) where {C, T, U} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(phi(t, θ), u, p, t)) + p) where {C, T, U} + sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(ode_dfdx(phi, t, θ, autodiff), phi, p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, - p, u) where {C, T, U} + p) where {C, T, U} out = Array(phi(t, θ)) arrt = Array(t) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[:, i], u, p, arrt[i]) for i in 1:size(out, 2)) / length(t) + sum(abs2, f(dxdtguess[:, i], phi, p, arrt[i]) for i in 1:size(out, 2)) / length(t) end """ @@ -387,10 +387,10 @@ end """ Representation of the loss function, parametric on the training strategy `strategy` for DAE problems """ -function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, u, +function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, batch) - integrand(t::Number, θ) = abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p, u)) - integrand(ts, θ) = [abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p, u)) for t in ts] + integrand(t::Number, θ) = abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) + integrand(ts, θ) = [abs2(inner_loss_DAE(phi, f, autodiff, t, θ, p)) for t in ts] @assert batch == 0 # not implemented function loss(θ, _) @@ -402,36 +402,36 @@ function generate_loss_DAE(strategy::QuadratureTraining, phi, f, autodiff::Bool, return loss end -function generate_loss_DAE(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, u, batch) +function generate_loss_DAE(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, batch) ts = tspan[1]:(strategy.dx):tspan[2] # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken function loss(θ, _) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, u, +function generate_loss_DAE(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, batch) # sum(abs2,inner_loss(t,θ) for t in ts) but Zygote generators are broken function loss(θ, _) ts = adapt(parameterless_type(θ), [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)]) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p, u, +function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p, batch) minT = tspan[1] maxT = tspan[2] @@ -454,15 +454,15 @@ function generate_loss_DAE(strategy::WeightedIntervalTraining, phi, f, autodiff: function loss(θ, _) if batch - sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p, u)) + sum(abs2, inner_loss_DAE(phi, f, autodiff, ts, θ, p)) else - sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p, u) for t in ts]) + sum(abs2, [inner_loss_DAE(phi, f, autodiff, t, θ, p) for t in ts]) end end return loss end -function generate_loss_DAE(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan, p, u, batch) +function generate_loss_DAE(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan, p, batch) error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.") end @@ -685,9 +685,9 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, end else alg.batch - endx + end - inner_f = generate_loss_DAE(strategy, phi, f, autodiff, tspan, p, u0, batch) + inner_f = generate_loss_DAE(strategy, phi, f, autodiff, tspan, p, batch) additional_loss = alg.additional_loss # Creates OptimizationFunction Object from total_loss diff --git a/test/dae_problem_test.jl b/test/dae_problem_test.jl new file mode 100644 index 000000000..1ca7eb096 --- /dev/null +++ b/test/dae_problem_test.jl @@ -0,0 +1,50 @@ +using DAEProblemLibrary, Sundials, Optimisers, OptimizationOptimisers, DifferentialEquations +using NeuralPDE, Lux, Test, Statistics, Plots + +f = function (yp, y, p, tres) + [-0.04 * y[1] + 1.0e4 * y[2] * y[3] - yp[1], + -(-0.04 * y[1] + 1.0e4 * y[2] * y[3]) - 3.0e7 * y[2] * y[2] - yp[2], + y[1] + y[2] + y[3] - 1.0] +end +u0 = [1.0, 0, 0] +du0 = [-0.04, 0.04, 0.0] + +println("f defined") +""" +The Robertson biochemical reactions in DAE form + +```math +\frac{dy₁}{dt} = -k₁y₁+k₃y₂y₃ +``` +```math +\frac{dy₂}{dt} = k₁y₁-k₂y₂^2-k₃y₂y₃ +``` +```math +1 = y₁ + y₂ + y₃ +``` +where ``k₁=0.04``, ``k₂=3\times10^7``, ``k₃=10^4``. For details, see: +Hairer Norsett Wanner Solving Ordinary Differential Equations I - Nonstiff Problems Page 129 +Usually solved on ``[0,1e11]`` +""" + +prob_oop = DAEProblem{false}(f, du0, u0, (0.0, 100000.0)) +true_sol = solve(prob_oop, IDA(), saveat = 0.01) + +u0 = [1.0, 1.0, 1.0] +func = Lux.σ +N = 12 +chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func), + Lux.Dense(N, N, func), Lux.Dense(N, length(u0))) + +opt = Optimisers.Adam(0.01) +dx = 0.05 +alg = NeuralPDE.NNDAE(chain, opt, autodiff = false, strategy = NeuralPDE.GridTraining(dx)) +sol = solve(prob_oop, alg, verbose=true, maxiters = 100000, saveat = 0.01) + +# println(abs(mean(true_sol .- sol))) + +# using Plots + +# plot(sol) +# plot!(true_sol) +# # ylims!(0,8) \ No newline at end of file From fe9274f8c55d711d389cfacc8ec933c61fffded0 Mon Sep 17 00:00:00 2001 From: Samedh Desai Date: Mon, 14 Aug 2023 01:56:22 -0400 Subject: [PATCH 5/5] New edits to DAE loss function generation, solving runs forever --- src/dae_problem_test.jl | 50 ------------------- src/my_test.jl | 27 ++++++++++ src/ode_solve.jl | 13 ++--- test/dae_problem_test.jl | 14 +++--- .../weighted_interval_training_test.jl | 6 +-- 5 files changed, 42 insertions(+), 68 deletions(-) delete mode 100644 src/dae_problem_test.jl create mode 100644 src/my_test.jl rename {src => test}/weighted_interval_training_test.jl (74%) diff --git a/src/dae_problem_test.jl b/src/dae_problem_test.jl deleted file mode 100644 index ed77520fd..000000000 --- a/src/dae_problem_test.jl +++ /dev/null @@ -1,50 +0,0 @@ -using DAEProblemLibrary, Sundials, Optimisers, OptimizationOptimisers, DifferentialEquations -using NeuralPDE, Lux, Test, Statistics, Plots - -f = function (r, yp, y, p, tres) - r[1] = -0.04 * y[1] + 1.0e4 * y[2] * y[3] - r[2] = -r[1] - 3.0e7 * y[2] * y[2] - yp[2] - r[1] -= yp[1] - r[3] = y[1] + y[2] + y[3] - 1.0 -end -u0 = [1.0, 0, 0] -du0 = [-0.04, 0.04, 0.0] - -""" -The Robertson biochemical reactions in DAE form - -```math -\frac{dy₁}{dt} = -k₁y₁+k₃y₂y₃ -``` -```math -\frac{dy₂}{dt} = k₁y₁-k₂y₂^2-k₃y₂y₃ -``` -```math -1 = y₁ + y₂ + y₃ -``` -where ``k₁=0.04``, ``k₂=3\times10^7``, ``k₃=10^4``. For details, see: -Hairer Norsett Wanner Solving Ordinary Differential Equations I - Nonstiff Problems Page 129 -Usually solved on ``[0,1e11]`` -""" - -prob_oop = DAEProblem(f, du0, u0, (0.0, 100000.0)) -true_sol = solve(prob_oop, IDA(), saveat = 0.01) - -u0 = [1.0, 1.0, 1.0] -func = Lux.σ -N = 12 -chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func), - Lux.Dense(N, N, func), Lux.Dense(N, length(u0))) - -opt = Optimisers.Adam(0.01) -dx = 0.05 -alg = NeuralPDE.NNDAE(chain, opt, autodiff = false, strategy = NeuralPDE.GridTraining(dx)) -sol = solve(prob_oop, alg, verbose=true, maxiters = 100000, saveat = 0.01) - -# println(abs(mean(true_sol .- sol))) - -# using Plots - -# plot(sol) -# plot!(true_sol) -# # ylims!(0,8) \ No newline at end of file diff --git a/src/my_test.jl b/src/my_test.jl new file mode 100644 index 000000000..29bfdd25d --- /dev/null +++ b/src/my_test.jl @@ -0,0 +1,27 @@ + +using OrdinaryDiffEq, OptimizationPolyalgorithms, Lux, OptimizationOptimJL, Test, Statistics, Plots, Optimisers + +function fu(u, p, t) + [p[1] * u[1] - p[2] * u[1] * u[2], -p[3] * u[2] + p[4] * u[1] * u[2]] +end + +p = [1.5, 1.0, 3.0, 1.0] +u0 = [1.0, 1.0] +prob_oop = ODEProblem{false}(fu, u0, (0.0, 3.0), p) +true_sol = solve(prob_oop, Tsit5(), saveat = 0.01) +func = Lux.σ +N = 12 +chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, length(u0))) + +opt = Optimisers.Adam(0.01) +dx=0.05 +alg = NeuralPDE.NNODE(chain, opt, autodiff = false, strategy = NeuralPDE.GridTraining(dx)) +sol = solve(prob_oop, alg, verbose=true, maxiters = 3, saveat = 0.01) + +@test abs(mean(sol) - mean(true_sol)) < 0.2 + +# using Plots + +# plot(sol) +# plot!(true_sol) +# ylims!(0,8) diff --git a/src/ode_solve.jl b/src/ode_solve.jl index a70d453d5..39aaeff6d 100644 --- a/src/ode_solve.jl +++ b/src/ode_solve.jl @@ -279,19 +279,20 @@ function inner_loss_DAE end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, p) where {C, T, U <: Number} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) - f(ode_dfdx(phi, t, θ, autodiff), phi, p, t)) + sum(abs2,f(ode_dfdx(phi, t, θ, autodiff), phi(t, θ), p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, p) where {C, T, U <: Number} out = phi(t, θ) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[i], phi, p, t[i]) for i in 1:size(out, 2)) / length(t) + fs = reduce(hcat, [f(dxdtguess[:, i], out, p, arrt[i]) for i in 1:size(out, 2)]) + sum(abs2, fs) / length(t) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::Number, θ, p) where {C, T, U} - sum(abs2, ode_dfdx(phi, t, θ, autodiff) .- f(ode_dfdx(phi, t, θ, autodiff), phi, p, t)) + sum(abs2,f(ode_dfdx(phi, t, θ, autodiff), phi(t, θ), p, t)) end function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector, θ, @@ -299,7 +300,8 @@ function inner_loss_DAE(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVect out = Array(phi(t, θ)) arrt = Array(t) dxdtguess = Array(ode_dfdx(phi, t, θ, autodiff)) - sum(abs2, f(dxdtguess[:, i], phi, p, arrt[i]) for i in 1:size(out, 2)) / length(t) + fs = reduce(hcat, [f(dxdtguess[:, i], out, p, arrt[i]) for i in 1:size(out, 2)]) + sum(abs2, fs) / length(t) end """ @@ -573,7 +575,6 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractODEProblem, else Optimization.AutoZygote() end - # Creates OptimizationFunction Object from total_loss optf = OptimizationFunction(total_loss, opt_algo) @@ -698,7 +699,6 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, end L2_loss end - # Choice of Optimization Algo for Training Strategies opt_algo = if strategy isa QuadratureTraining Optimization.AutoForwardDiff() @@ -717,6 +717,7 @@ function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem, end optprob = OptimizationProblem(optf, init_params) + println("attempting to solve") res = solve(optprob, opt; callback, maxiters, alg.kwargs...) #solutions at timepoints diff --git a/test/dae_problem_test.jl b/test/dae_problem_test.jl index 1ca7eb096..ef9c64e28 100644 --- a/test/dae_problem_test.jl +++ b/test/dae_problem_test.jl @@ -1,15 +1,15 @@ -using DAEProblemLibrary, Sundials, Optimisers, OptimizationOptimisers, DifferentialEquations -using NeuralPDE, Lux, Test, Statistics, Plots +using Optimisers, OptimizationOptimisers, Sundials +using Lux, Test, Statistics, Plots -f = function (yp, y, p, tres) +function fu(yp, y, p, tres) [-0.04 * y[1] + 1.0e4 * y[2] * y[3] - yp[1], -(-0.04 * y[1] + 1.0e4 * y[2] * y[3]) - 3.0e7 * y[2] * y[2] - yp[2], y[1] + y[2] + y[3] - 1.0] end u0 = [1.0, 0, 0] du0 = [-0.04, 0.04, 0.0] +p = [1.5, 1.0, 3.0, 1.0] -println("f defined") """ The Robertson biochemical reactions in DAE form @@ -27,14 +27,12 @@ Hairer Norsett Wanner Solving Ordinary Differential Equations I - Nonstiff Probl Usually solved on ``[0,1e11]`` """ -prob_oop = DAEProblem{false}(f, du0, u0, (0.0, 100000.0)) +prob_oop = DAEProblem{false}(fu, du0, u0, (0.0, 100000.0), p) true_sol = solve(prob_oop, IDA(), saveat = 0.01) -u0 = [1.0, 1.0, 1.0] func = Lux.σ N = 12 -chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func), - Lux.Dense(N, N, func), Lux.Dense(N, length(u0))) +chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, length(u0))) opt = Optimisers.Adam(0.01) dx = 0.05 diff --git a/src/weighted_interval_training_test.jl b/test/weighted_interval_training_test.jl similarity index 74% rename from src/weighted_interval_training_test.jl rename to test/weighted_interval_training_test.jl index 5d81d8960..921f88592 100644 --- a/src/weighted_interval_training_test.jl +++ b/test/weighted_interval_training_test.jl @@ -1,5 +1,5 @@ -using NeuralPDE, OrdinaryDiffEq, OptimizationPolyalgorithms, Lux, OptimizationOptimJL, Test, Statistics, Plots, Optimisers +using OrdinaryDiffEq, OptimizationPolyalgorithms, Lux, OptimizationOptimJL, Test, Statistics, Plots, Optimisers function f(u, p, t) [p[1] * u[1] - p[2] * u[1] * u[2], -p[3] * u[2] + p[4] * u[1] * u[2]] @@ -20,9 +20,7 @@ samples = 200 alg = NeuralPDE.NNODE(chain, opt, autodiff = false, strategy = NeuralPDE.WeightedIntervalTraining(weights, samples)) sol = solve(prob_oop, alg, verbose=true, maxiters = 100000, saveat = 0.01) -println(abs(mean(true_sol .- sol))) -# println(abs(mean(sol) - mean(true_sol))) -# @test abs(mean(sol) - mean(true_sol)) < 0.2 +@test abs(mean(true_sol .- sol)) < 0.2 using Plots