NNODE for DAE Problems

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,

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)

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 @@
 
 ## 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 @@
  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 @@
  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 @@
  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 @@
  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 @@
  dense_errors = false)
  sol
 end #solve
+
+function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem,
+ alg::NNDAE,
+ 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