From 52da385865ba42d1a5b6b967d87647fdf81e76bb Mon Sep 17 00:00:00 2001 From: TT Date: Mon, 9 Oct 2023 16:42:43 +0200 Subject: [PATCH] Revert "adjustments for FMISensitivity.jl" This reverts commit 8684a66d7fcb3bc3b7f80022927eae6fa83bdde8. --- Project.toml | 14 +- examples/src/juliacon_2023.ipynb | 12 +- src/FMIFlux.jl | 64 ++++- src/batch.jl | 3 + src/convert.jl | 4 +- src/flux_overload.jl | 51 +++- src/hotfixes.jl | 19 -- src/layers.jl | 8 +- src/losses.jl | 1 + src/misc.jl | 2 + src/neural.jl | 456 +++++++++++++------------------ test/batching.jl | 1 + test/fmu_params.jl | 41 +-- test/hybrid_ME.jl | 21 +- test/hybrid_ME_dis.jl | 19 +- test/runtests.jl | 71 +++-- test/solution_gradients.jl | 277 ------------------- test/supported_sensitivities.jl | 89 ++---- 18 files changed, 403 insertions(+), 750 deletions(-) delete mode 100644 src/hotfixes.jl delete mode 100644 test/solution_gradients.jl diff --git a/Project.toml b/Project.toml index 6f078bc2..7a2f335b 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "FMIFlux" uuid = "fabad875-0d53-4e47-9446-963b74cae21f" -version = "0.11.0" +version = "0.10.6" [deps] Colors = "5ae59095-9a9b-59fe-a467-6f913c188581" @@ -8,7 +8,6 @@ DiffEqCallbacks = "459566f4-90b8-5000-8ac3-15dfb0a30def" DifferentiableEigen = "73a20539-4e65-4dcb-a56d-dc20f210a01b" DifferentialEquations = "0c46a032-eb83-5123-abaf-570d42b7fbaa" FMIImport = "9fcbc62e-52a0-44e9-a616-1359a0008194" -FMISensitivity = "3e748fe5-cd7f-4615-8419-3159287187d2" Flux = "587475ba-b771-5e3f-ad9e-33799f191a9c" Optim = "429524aa-4258-5aef-a3af-852621145aeb" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" @@ -19,14 +18,13 @@ ThreadPools = "b189fb0b-2eb5-4ed4-bc0c-d34c51242431" [compat] Colors = "0.12.8" -DiffEqCallbacks = "2.33.0" +DiffEqCallbacks = "2.26.0" DifferentiableEigen = "0.2.0" -DifferentialEquations = "7.10.0" -FMIImport = "0.16.0" -FMISensitivity = "0.1.0" -Flux = "0.14" +DifferentialEquations = "7.8.0" +FMIImport = "0.15.8" +Flux = "0.13, 0.14" Optim = "1.7.0" -ProgressMeter = "1.9.0" +ProgressMeter = "1.7.0" Requires = "1.3.0" ThreadPools = "2.1.1" julia = "1.6" diff --git a/examples/src/juliacon_2023.ipynb b/examples/src/juliacon_2023.ipynb index 20d465b0..8848b441 100644 --- a/examples/src/juliacon_2023.ipynb +++ b/examples/src/juliacon_2023.ipynb @@ -6,19 +6,11 @@ "source": [ "# Using NeuralODEs in real life applications\n", "-----\n", - "Tutorial by Tobias Thummerer, Lars Mikelsons | Last edit: 08-04-2023\n", + "Tutorial by Tobias Thummerer | Last edit: 08-04-2023\n", "\n", "This workshop was held at the JuliaCon 2023 | 07-25-2023 | MIT (Boston, USA)\n", "\n", - "Keywords: *#NeuralODE, #NeuralFMU, #PeNODE, #HybridModeling, #SciML*" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Workshop Video\n", - "[![Preview image](https://img.youtube.com/vi/X_u0KlZizD4/0.jpg)](https://www.youtube.com/watch?v=X_u0KlZizD4)" + "Keywords: *#NeuralODE, #NeuralFMU, #PeNODE, #HybridModeling*" ] }, { diff --git a/src/FMIFlux.jl b/src/FMIFlux.jl index 734f56b6..42330e5e 100644 --- a/src/FMIFlux.jl +++ b/src/FMIFlux.jl @@ -4,24 +4,63 @@ # module FMIFlux -import FMISensitivity - -import FMISensitivity.ForwardDiff -import FMISensitivity.Zygote -import FMISensitivity.ReverseDiff -import FMISensitivity.FiniteDiff @debug "Debugging messages enabled for FMIFlux ..." if VERSION < v"1.7.0" - @warn "Training under Julia 1.6 is very slow, please consider using Julia 1.7 or newer." maxlog=1 + @warn "Training under Julia 1.6 is very slow, please consider using Julia 1.7 or newer." end -import FMIImport.FMICore: hasCurrentComponent, getCurrentComponent, unsense -import FMIImport.FMICore.ChainRulesCore: ignore_derivatives +# Overwrite tag printing and limit partials length from ForwardDiff.jl +# import FMIImport.ForwardDiff +# function Base.show(io::IO, d::ForwardDiff.Dual{T,V,N}) where {T,V,N} +# print(io, "Dual(", ForwardDiff.value(d)) +# for i in 1:min(N, 5) +# print(io, ", ", ForwardDiff.partials(d, i)) +# end +# if N > 5 +# print(io, ", [$(N-5) more]...") +# end +# print(io, ")") +# end + +# ToDo: Quick-fixes until patch release SciMLSensitivity v0.7.29 +import FMIImport.SciMLSensitivity: FakeIntegrator, u_modified!, TrackedAffect +import FMIImport.SciMLSensitivity.DiffEqBase: set_u! +function u_modified!(::FakeIntegrator, ::Bool) + return nothing +end +function set_u!(::FakeIntegrator, u) + return nothing +end -using Requires -import Flux +# ToDo: Quick-fixes until patch release SciMLSensitivity v0.7.28 +# function Base.hasproperty(f::TrackedAffect, s::Symbol) +# if hasfield(TrackedAffect, s) +# return true +# else +# _affect = getfield(f, :affect!) +# return hasfield(typeof(_affect), s) +# end +# end +# function Base.getproperty(f::TrackedAffect, s::Symbol) +# if hasfield(TrackedAffect, s) +# return getfield(f, s) +# else +# _affect = getfield(f, :affect!) +# return getfield(_affect, s) +# end +# end +# function Base.setproperty!(f::TrackedAffect, s::Symbol, value) +# if hasfield(TrackedAffect, s) +# return setfield!(f, s, value) +# else +# _affect = getfield(f, :affect!) +# return setfield!(_affect, s, value) +# end +# end + +using Requires, Flux using FMIImport using FMIImport: fmi2ValueReference, FMU, FMU2, FMU2Component @@ -32,9 +71,6 @@ using FMIImport: fmi2SetTime, fmi2CompletedIntegratorStep, fmi2GetEventIndicator using FMIImport: fmi2SampleJacobian, fmi2GetDirectionalDerivative, fmi2GetJacobian, fmi2GetJacobian! using FMIImport: fmi2True, fmi2False -import FMIImport.FMICore: fmi2ValueReferenceFormat - -include("hotfixes.jl") include("convert.jl") include("flux_overload.jl") include("neural.jl") diff --git a/src/batch.jl b/src/batch.jl index 9abeeec0..221d96d7 100644 --- a/src/batch.jl +++ b/src/batch.jl @@ -4,7 +4,10 @@ # import FMIImport: fmi2Real, fmi2FMUstate, fmi2EventInfo, fmi2ComponentState +import FMIImport.ChainRulesCore: ignore_derivatives using DiffEqCallbacks: FunctionCallingCallback +using FMIImport.ForwardDiff +import FMIImport: unsense struct FMULoss{T} loss::T diff --git a/src/convert.jl b/src/convert.jl index d8528fe5..5e088183 100644 --- a/src/convert.jl +++ b/src/convert.jl @@ -3,6 +3,8 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # +import Flux + function is64(model::Flux.Chain) params = Flux.params(model) @@ -18,5 +20,5 @@ function is64(model::Flux.Chain) end function convert64(model::Flux.Chain) - Flux.fmap(Flux.f64, model) + fmap(f64, model) end diff --git a/src/flux_overload.jl b/src/flux_overload.jl index b3e25e5b..11b1c93b 100644 --- a/src/flux_overload.jl +++ b/src/flux_overload.jl @@ -3,5 +3,54 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # +import Flux +import FMIImport.ChainRulesCore +import Flux.Random: AbstractRNG +import Flux.LinearAlgebra: I + # feed through -params = Flux.params \ No newline at end of file +params = Flux.params + +# # exports +# export Adam +# export Parallel + +# # +# Chain = Flux.Chain +# export Chain + +# # Float64 version of the Flux.glorot_uniform +# function glorot_uniform_64(rng::AbstractRNG, dims::Integer...; gain::Real=1) +# scale = Float64(gain) * sqrt(24.0 / sum(Flux.nfan(dims...))) +# (rand(rng, Float64, dims...) .- 0.5) .* scale +# end +# glorot_uniform_64(dims::Integer...; kw...) = glorot_uniform_64(Flux.default_rng_value(), dims...; kw...) +# glorot_uniform_64(rng::AbstractRNG=Flux.default_rng_value(); init_kwargs...) = (dims...; kwargs...) -> glorot_uniform_64(rng, dims...; init_kwargs..., kwargs...) +# ChainRulesCore.@non_differentiable glorot_uniform_64(::Any...) + +# # Float64 version of the Flux.identity_init +# identity_init_64(cols::Integer; gain::Real=1, shift=0) = zeros(Float64, cols) # Assume bias +# identity_init_64(rows::Integer, cols::Integer; gain::Real=1, shift=0) = circshift(Matrix{Float64}(I * gain, rows,cols), shift) +# function identity_init_64(dims::Integer...; gain::Real=1, shift=0) +# nin, nout = dims[end-1], dims[end] +# centers = map(d -> cld(d, 2), dims[1:end-2]) +# weights = zeros(Float64, dims...) +# for i in 1:min(nin,nout) +# weights[centers..., i, i] = gain +# end +# return circshift(weights, shift) +# end +# ChainRulesCore.@non_differentiable identity_init_64(::Any...) + +# """ +# Wrapper for Flux.Dense, that converts all parameters to Float64. +# """ +# function Dense(args...; init=glorot_uniform_64, kwargs...) +# return Flux.Dense(args...; init=init, kwargs...) +# end +# function Dense(W::AbstractMatrix, args...; init=glorot_uniform_64, kwargs...) +# W = Matrix{Float64}(W) +# return Flux.Dense(W, args...; init=init, kwargs...) +# end +# Dense(in::Integer, out::Integer, σ = identity; init=glorot_uniform_64, kwargs...) = Dense(in => out, σ; init=init, kwargs...) +# export Dense \ No newline at end of file diff --git a/src/hotfixes.jl b/src/hotfixes.jl deleted file mode 100644 index a2196309..00000000 --- a/src/hotfixes.jl +++ /dev/null @@ -1,19 +0,0 @@ -# -# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons -# Licensed under the MIT license. See LICENSE file in the project root for details. -# - -# ToDo: Quick-fixes until patch release SciMLSensitivity v0.7.XX -import FMISensitivity.SciMLSensitivity: FakeIntegrator, u_modified! -import FMISensitivity.SciMLSensitivity.DiffEqBase: set_u! -function u_modified!(::FakeIntegrator, ::Bool) - return nothing -end -function set_u!(::FakeIntegrator, u) - return nothing -end - -import FMISensitivity.ReverseDiff: increment_deriv! -function increment_deriv!(t::AbstractArray{<:ReverseDiff.TrackedReal}, x::ReverseDiff.ZeroTangent, args...) - return nothing -end \ No newline at end of file diff --git a/src/layers.jl b/src/layers.jl index 18e2d6b8..0ddb8c1b 100644 --- a/src/layers.jl +++ b/src/layers.jl @@ -19,8 +19,8 @@ struct FMUParameterRegistrator{T} function FMUParameterRegistrator{T}(fmu::FMU2, p_refs::fmi2ValueReferenceFormat, p::AbstractArray{T}) where {T} @assert length(p_refs) == length(p) "`p_refs` and `p` need to be the same length!" p_refs = prepareValueReference(fmu, p_refs) - fmu.default_p_refs = p_refs - fmu.default_p = p + fmu.optim_p_refs = p_refs + fmu.optim_p = p return new(fmu, p_refs, p) end @@ -31,8 +31,8 @@ end export FMUParameterRegistrator function (l::FMUParameterRegistrator)(x) - l.fmu.default_p = l.p - l.fmu.default_p_refs = l.p_refs + l.fmu.optim_p = l.p + l.fmu.optim_p_refs = l.p_refs return x end diff --git a/src/losses.jl b/src/losses.jl index 88f78e1d..b3193614 100644 --- a/src/losses.jl +++ b/src/losses.jl @@ -7,6 +7,7 @@ module Losses using Flux import ..FMIFlux: FMU2BatchElement, NeuralFMU, loss!, run!, ME_NeuralFMU, FMU2Solution +import FMIImport: unsense mse = Flux.Losses.mse mae = Flux.Losses.mae diff --git a/src/misc.jl b/src/misc.jl index e52e2dae..288f4b40 100644 --- a/src/misc.jl +++ b/src/misc.jl @@ -3,6 +3,8 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # +using Flux + """ Compares non-equidistant (or equidistant) datapoints by linear interpolating and comparing at given interpolation points `t_comp`. (Zygote-friendly: Zygote can differentiate through via AD.) diff --git a/src/neural.jl b/src/neural.jl index f519af15..b8ec8882 100644 --- a/src/neural.jl +++ b/src/neural.jl @@ -3,14 +3,18 @@ # Licensed under the MIT license. See LICENSE file in the project root for details. # -import FMIImport.FMICore: assert_integrator_valid, isdual, istracked, undual, unsense, untrack -import FMIImport: finishSolveFMU, handleEvents, prepareSolveFMU +import FMIImport.ChainRulesCore: ignore_derivatives +import FMIImport: assert_integrator_valid, fd_eltypes, fd_set!, finishSolveFMU, + handleEvents, isdual, istracked, prepareSolveFMU, rd_set!, undual, unsense, untrack import Optim import ProgressMeter -import FMISensitivity.SciMLSensitivity.SciMLBase: CallbackSet, ContinuousCallback, ODESolution, ReturnCode, RightRootFind, +import FMIImport.SciMLSensitivity.SciMLBase: CallbackSet, ContinuousCallback, ODESolution, ReturnCode, RightRootFind, VectorContinuousCallback, set_u!, terminate!, u_modified!, build_solution -using FMISensitivity.ReverseDiff: TrackedArray -import FMISensitivity.SciMLSensitivity: InterpolatingAdjoint, ReverseDiffVJP +import FMIImport.SciMLSensitivity.ForwardDiff +import FMIImport.SciMLSensitivity.ReverseDiff +import FMIImport.SciMLSensitivity.FiniteDiff +using FMIImport.SciMLSensitivity.ReverseDiff: TrackedArray +import FMIImport.SciMLSensitivity: InterpolatingAdjoint, ReverseDiffVJP import ThreadPools using DiffEqCallbacks @@ -21,11 +25,12 @@ using FMIImport: FMU2Component, FMU2Event, FMU2Solution, fmi2ComponentState, fmi2ComponentStateInitializationMode, fmi2ComponentStateInstantiated, fmi2ComponentStateTerminated, fmi2StatusOK, fmi2Type, fmi2TypeCoSimulation, fmi2TypeModelExchange, logError, logInfo, logWarning -using FMISensitivity.SciMLSensitivity: +using Flux +using Flux.Zygote +using FMIImport.SciMLSensitivity: ForwardDiffSensitivity, InterpolatingAdjoint, ReverseDiffVJP, ZygoteVJP -DEFAULT_PROGRESS_DESCR="Simulating ME-NeuralFMU ..." -DEFAULT_CHUNK_SIZE = 32 +zero_tgrad(u,p,t) = zero(u) """ The mutable struct representing an abstract (simulation mode unknown) NeuralFMU. @@ -35,16 +40,13 @@ abstract type NeuralFMU end """ Structure definition for a NeuralFMU, that runs in mode `Model Exchange` (ME). """ -mutable struct ME_NeuralFMU{M, R} <: NeuralFMU +mutable struct ME_NeuralFMU{M, P, R} <: NeuralFMU model::M - p::AbstractArray{<:Real} + p::P re::R kwargs - # re_model - # re_p - fmu::FMU tspan @@ -80,19 +82,14 @@ mutable struct ME_NeuralFMU{M, R} <: NeuralFMU execution_start::Real - rd_condition_buffer - - function ME_NeuralFMU{M, R}(model::M, p::AbstractArray{<:Real}, re::R) where {M, R} + function ME_NeuralFMU{M, P, R}(model::M, p::P, re::R) where {M, P, R} inst = new() inst.model = model inst.p = p inst.re = re inst.x0 = nothing - # inst.re_model = nothing - # inst.re_p = nothing - - inst.modifiedState = false + inst.modifiedState = true inst.startState = nothing inst.stopState = nothing @@ -104,7 +101,6 @@ mutable struct ME_NeuralFMU{M, R} <: NeuralFMU inst.customCallbacksAfter = [] inst.execution_start = 0.0 - inst.rd_condition_buffer = nothing return inst end @@ -131,27 +127,29 @@ mutable struct CS_NeuralFMU{F, C} <: NeuralFMU end end -function evaluateModel(nfmu::ME_NeuralFMU, c::FMU2Component, x, p) +function evaluateModel(nfmu::ME_NeuralFMU, c::FMU2Component, x) + @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" + + return nfmu.model(x) +end + +function evaluateModel(nfmu::ME_NeuralFMU, c::FMU2Component, dx, x) + @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" + + dx[:] = nfmu.model(x) + + return nothing +end + +function evaluateReModel(nfmu::ME_NeuralFMU, c::FMU2Component, x, p) @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" - # if isnothing(nfmu.re_model) || p != nfmu.re_p - # nfmu.re_p = p # fast_copy!(nfmu, :re_p, p) - # nfmu.re_model = nfmu.re(p) - # end - # return nfmu.re_model(x) - nfmu.p = p #unsense(p) return nfmu.re(p)(x) end -function evaluateModel(nfmu::ME_NeuralFMU, c::FMU2Component, dx, x, p) +function evaluateReModel(nfmu::ME_NeuralFMU, c::FMU2Component, dx, x, p) @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" - # if isnothing(nfmu.re_model) || p != nfmu.re_p - # nfmu.re_p = p # fast_copy!(nfmu, :re_p, p) - # nfmu.re_model = nfmu.re(p) - # end - # dx[:] = nfmu.re_model(x) - nfmu.p = p #unsense(p) dx[:] = nfmu.re(p)(x) return nothing @@ -221,7 +219,7 @@ function time_choice(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, tStart, t c.solution.evals_timechoice += 1 if c.eventInfo.nextEventTimeDefined == fmi2True - #@debug "time_choice(...): $(c.eventInfo.nextEventTime) at t=$(unsense(integrator.t))" + #@debug "time_choice(...): $(c.eventInfo.nextEventTime) at t=$(ForwardDiff.value(integrator.t))" if c.eventInfo.nextEventTime >= tStart && c.eventInfo.nextEventTime <= tStop #@assert sizeof(integrator.t) == sizeof(c.eventInfo.nextEventTime) "The NeuralFMU/solver are initialized in $(sizeof(integrator.t))-bit-mode, but FMU events are defined in $(sizeof(c.eventInfo.nextEventTime))-bit. Please define your ANN in $(sizeof(c.eventInfo.nextEventTime))-bit mode." @@ -233,7 +231,7 @@ function time_choice(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, tStart, t return nothing end else - #@debug "time_choice(...): nothing at t=$(unsense(integrator.t))" + #@debug "time_choice(...): nothing at t=$(ForwardDiff.value(integrator.t))" return nothing end @@ -241,118 +239,87 @@ function time_choice(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, tStart, t end # Returns the event indicators for an FMU. -# function condition!(nfmu::ME_NeuralFMU, c::FMU2Component, out::SubArray{<:ForwardDiff.Dual{T, V, N}, A, B, C, D}, _x, t, integrator) where {T, V, N, A, B, C, D} # Event when event_f(u,t) == 0 +function condition(nfmu::ME_NeuralFMU, c::FMU2Component, out::SubArray{<:ForwardDiff.Dual{T, V, N}, A, B, C, D}, _x, t, integrator) where {T, V, N, A, B, C, D} # Event when event_f(u,t) == 0 -# @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" -# @debug assert_integrator_valid(integrator) + @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" + @debug assert_integrator_valid(integrator) -# #@assert c.state == fmi2ComponentStateContinuousTimeMode "condition!(...): Must be called in mode continuous time." + #@assert c.state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time." -# # ToDo: set inputs here -# #fmiSetReal(myFMU, InputRef, Value) + # ToDo: set inputs here + #fmiSetReal(myFMU, InputRef, Value) -# t = undual(t) -# x = undual(_x) + t = undual(t) + x = undual(_x) -# # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN -# #c.t = t # this will auto-set time via fx-call! -# c.fmu.default_t = t -# evaluateModel(nfmu, c, x) + # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN + #c.t = t # this will auto-set time via fx-call! + c.next_t = t + evaluateModel(nfmu, c, x) -# out_tmp = zeros(c.fmu.modelDescription.numberOfEventIndicators) -# fmi2GetEventIndicators!(c, out_tmp) + out_tmp = zeros(c.fmu.modelDescription.numberOfEventIndicators) + fmi2GetEventIndicators!(c, out_tmp) -# sense_set!(out, out_tmp) + fd_set!(out, out_tmp) -# c.solution.evals_condition += 1 + c.solution.evals_condition += 1 -# @debug assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) -# return nothing -# end -# function condition!(nfmu::ME_NeuralFMU, c::FMU2Component, out::SubArray{<:ReverseDiff.TrackedReal}, _x, t, integrator) + return nothing +end +function condition(nfmu::ME_NeuralFMU, c::FMU2Component, out::SubArray{<:ReverseDiff.TrackedReal}, _x, t, integrator) -# @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" -# @debug assert_integrator_valid(integrator) + @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" + @debug assert_integrator_valid(integrator) -# #@assert c.state == fmi2ComponentStateContinuousTimeMode "condition!(...): Must be called in mode continuous time." + #@assert c.state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time." -# # ToDo: set inputs here -# #fmiSetReal(myFMU, InputRef, Value) + # ToDo: set inputs here + #fmiSetReal(myFMU, InputRef, Value) -# t = untrack(t) -# x = untrack(_x) + t = untrack(t) + x = untrack(_x) -# # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN -# c.fmu.default_t = t -# evaluateModel(nfmu, c, x) + # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN + c.next_t = t + evaluateModel(nfmu, c, x) -# out_tmp = zeros(c.fmu.modelDescription.numberOfEventIndicators) -# fmi2GetEventIndicators!(c, out_tmp) - -# sense_set!(out, out_tmp) - -# @debug assert_integrator_valid(integrator) + out_tmp = zeros(c.fmu.modelDescription.numberOfEventIndicators) + fmi2GetEventIndicators!(c, out_tmp) -# c.solution.evals_condition += 1 + rd_set!(out, out_tmp) -# return nothing -# end + @debug assert_integrator_valid(integrator) -# [ToDo] for now, ReverseDiff (together with the rrule) seems to have a problem with the SubArray here (when `collect` it accesses array elements that are #undef), -# so I added an additional (single allocating) dispatch... -# Type is ReverseDiff.TrackedReal{Float64, Float64, ReverseDiff.TrackedArray{Float64, Float64, 1, Vector{Float64}, Vector{Float64}}}[#undef, #undef, #undef, ...] -function condition!(nfmu::ME_NeuralFMU, c::FMU2Component, out::AbstractArray{<:ReverseDiff.TrackedReal}, x, t, integrator, handleEventIndicators) - - if !isassigned(out, 1) #isnothing(nfmu.rd_condition_buffer) - logWarning(nfmu.fmu, "There is currently an issue with the condition buffer pre-allocation, the buffer can't be overwritten by the generated rrule.") - #nfmu.rd_condition_buffer = collect(ReverseDiff.TrackedReal{typeof(r.value), typeof(r.deriv), typeof(r.origin)}(0.0, r.deriv, r.tape, r.index, r.origin) for r in out) # copy(out) - #@info "$(out)" - out[:] = zeros(fmi2Real, length(out)) - #nfmu.rd_condition_buffer = true - end - - #condition!(nfmu, c, out_tmp, x, t, integrator) - #out_tmp = zeros(fmi2Real, length(out)) - invoke(condition!, Tuple{ME_NeuralFMU, FMU2Component, Any, Any, Any, Any, Any}, nfmu, c, out, x, t, integrator, handleEventIndicators) - #out[:] = out_tmp + c.solution.evals_condition += 1 - #out[:] = nfmu.rd_condition_buffer - return nothing end +function condition(nfmu::ME_NeuralFMU, c::FMU2Component, out, _x, t, integrator) # Event when event_f(u,t) == 0 -function condition!(nfmu::ME_NeuralFMU, c::FMU2Component, out, x, t, integrator, handleEventIndicators) - - #assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" - @assert c.state == fmi2ComponentStateContinuousTimeMode "condition!(...): Must be called in mode continuous time." + @debug @assert c.state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time." + + #@debug "State condition..." # ToDo: set inputs here #fmiSetReal(myFMU, InputRef, Value) - #t = unsense(t) - #x = unsense(_x) - - # [ToDo] Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN. - # Basically only the layers from very top to FMU need to be evaluated here. - c.fmu.default_t = t - c.fmu.default_ec = out - c.fmu.default_ec_idcs = handleEventIndicators - evaluateModel(nfmu, c, x, nfmu.p) - c.fmu.default_t = -1.0 - c.fmu.default_ec = c.fmu.empty_fmi2Real - c.fmu.default_ec_idcs = c.fmu.empty_fmi2ValueReference - - # write back to condition buffer - out[:] = c.eval_output.ec + t = unsense(t) + x = unsense(_x) + + # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN + c.next_t = t + evaluateModel(nfmu, c, x) # evaluate NeuralFMU (set new states) - #assert_integrator_valid(integrator) + fmi2GetEventIndicators!(c, out) - c.solution.evals_condition += 1 + @debug assert_integrator_valid(integrator) - @debug "condition!(...) -> [typeof=$(typeof(out))]\n$(unsense(out))" + c.solution.evals_condition += 1 return nothing end @@ -360,7 +327,7 @@ end global lastIndicator = nothing global lastIndicatorX = nothing global lastIndicatorT = nothing -function conditionSingle(nfmu::ME_NeuralFMU, c::FMU2Component, index, x, t, integrator) +function conditionSingle(nfmu::ME_NeuralFMU, c::FMU2Component, index, _x, t, integrator) @assert c.state == fmi2ComponentStateContinuousTimeMode "condition(...): Must be called in mode continuous time." @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" @@ -368,10 +335,13 @@ function conditionSingle(nfmu::ME_NeuralFMU, c::FMU2Component, index, x, t, inte # ToDo: set inputs here #fmiSetReal(myFMU, InputRef, Value) - if c.fmu.handleEventIndicators != nothing && index ∉ c.fmu.handleEventIndicators + if c.fmu.executionConfig.handleEventIndicators != nothing && index ∉ c.fmu.executionConfig.handleEventIndicators return 1.0 end + t = unsense(t) + x = unsense(_x) + global lastIndicator # , lastIndicatorX, lastIndicatorT if lastIndicator == nothing || length(lastIndicator) != c.fmu.modelDescription.numberOfEventIndicators @@ -381,12 +351,11 @@ function conditionSingle(nfmu::ME_NeuralFMU, c::FMU2Component, index, x, t, inte # ToDo: Input Function # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN - c.fmu.default_t = t - c.fmu.default_ec = lastIndicator - evaluateModel(nfmu, c, x, nfmu.p) - c.fmu.default_t = -1.0 - c.fmu.default_ec = c.fmu.empty_fmi2Real + c.next_t = t + evaluateModel(nfmu, c, x) # evaluate NeuralFMU (set new states) + fmi2GetEventIndicators!(c, lastIndicator) + c.solution.evals_condition += 1 return lastIndicator[index] @@ -394,40 +363,30 @@ end function f_optim(x, nfmu::ME_NeuralFMU, c::FMU2Component, right_x_fmu) # , idx, direction::Real) # propagete the new state-guess `x` through the NeuralFMU - evaluateModel(nfmu, c, x, nfmu.p) + evaluateModel(nfmu, c, x) #indicators = fmi2GetEventIndicators(c) return Flux.Losses.mse(right_x_fmu, fmi2GetContinuousStates(c)) # - min(-direction*indicators[idx], 0.0) end -# Handles the upcoming event +# Handles the upcoming events. function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) - @debug "affectFMU!" @assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" - assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) - @assert c.state == fmi2ComponentStateContinuousTimeMode "affectFMU!(...): Must be in continuous time mode!" + @debug @assert c.state == fmi2ComponentStateContinuousTimeMode "affectFMU!(...): Must be in continuous time mode!" - # [NOTE] Here unsensing is OK, because we just want to reset the FMU to the correct state! - # The values come directly from the integrator and are NOT function arguments! - t = integrator.t # unsense(integrator.t) - x = integrator.u # unsense(integrator.u) - - if c.x != x - # capture status of `force` - mode = c.force - c.force = true - - # there are fx-evaluations before the event is handled, reset the FMU state to the current integrator step - c.fmu.default_t = t - evaluateModel(nfmu, c, x, nfmu.p) # evaluate NeuralFMU (set new states) - # [NOTE] No need to reset time here, because we did pass a event instance! - # c.fmu.default_t = -1.0 - - c.force = mode - end + t = unsense(integrator.t) + x = unsense(integrator.u) - #@info "$(c.x) - $(integrator.u) - $(x)\n$(fmi2GetEventIndicators(c))" + # there are fx-evaluations before the event is handled, reset the FMU state to the current integrator step + mode = c.force + c.force = true + + c.next_t = t + evaluateModel(nfmu, c, x) # evaluate NeuralFMU (set new states) + + c.force = mode # if inputFunction !== nothing # fmi2SetReal(c, inputValues, inputFunction(integrator.t)) @@ -460,8 +419,7 @@ function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) if c.eventInfo.valuesOfContinuousStatesChanged == fmi2True - left_x = unsense(x) - #right_x = similar(left_x) + left_x = x right_x_fmu = fmi2GetContinuousStates(c) # the new FMU state after handled events @@ -469,38 +427,31 @@ function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) #@debug "affectFMU!(_, _, $idx): NeuralFMU state event from $(left_x) (fmu: $(left_x_fmu)). Indicator [$idx]: $(indicators[idx]). Optimizing new state ..." end - # ToDo: use gradient-based optimization here? + # ToDo: Problem-related parameterization of optimize-call + #result = optimize(x_seek -> f_optim(x_seek, nfmu, right_x_fmu), left_x, LBFGS(); autodiff = :forward) + #result = Optim.optimize(x_seek -> f_optim(x_seek, nfmu, right_x_fmu, idx, sign(indicators[idx])), left_x, Optim.NelderMead()) + # if there is an ANN above the FMU, propaget FMU state through top ANN: - if nfmu.modifiedState + if nfmu.modifiedState == true result = Optim.optimize(x_seek -> f_optim(x_seek, nfmu, c, right_x_fmu), left_x, Optim.NelderMead()) right_x = Optim.minimizer(result) else # if there is no ANN above, then: right_x = right_x_fmu end - # if isdual(integrator.u) - # T, V, N = fd_eltypes(integrator.u) + if isdual(integrator.u) + T, V, N = fd_eltypes(integrator.u) - # new_x = collect(ForwardDiff.Dual{T, V, N}(V(right_x[i]), ForwardDiff.partials(integrator.u[i])) for i in 1:length(integrator.u)) - # #set_u!(integrator, new_x) - # integrator.u .= new_x + new_x = collect(ForwardDiff.Dual{T, V, N}(V(right_x[i]), ForwardDiff.partials(integrator.u[i])) for i in 1:length(integrator.u)) + #set_u!(integrator, new_x) + integrator.u .= new_x - # @debug "affectFMU!(_, _, $idx): NeuralFMU event with state change at $t. Indicator [$idx]. (ForwardDiff) " - # else - # #set_u!(integrator, right_x) - # integrator.u .= right_x - - # @debug "affectFMU!(_, _, $idx): NeuralFMU event with state change at $t. Indicator [$idx]." - # end - for i in 1:length(left_x) - if left_x[i] != 0.0 - scale = right_x[i] ./ left_x[i] - integrator.u[i] *= scale - else # integrator state zero can't be scaled, need to add (but no sensitivities in this case!) - shift = right_x[i] - left_x[i] - integrator.u[i] += shift - logWarning(c.fmu, "Probably wrong sensitivities for ∂x^+ / ∂x^-\nCan't scale from zero state (state #$(i)=0.0)") - end + @debug "affectFMU!(_, _, $idx): NeuralFMU event with state change at $t. Indicator [$idx]. (ForwardDiff) " + else + #set_u!(integrator, right_x) + integrator.u .= right_x + + @debug "affectFMU!(_, _, $idx): NeuralFMU event with state change at $t. Indicator [$idx]." end u_modified!(integrator, true) @@ -514,13 +465,12 @@ function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) end if c.eventInfo.nominalsOfContinuousStatesChanged == fmi2True - # ToDo: Do something with that information, e.g. pass to solver x_nom = fmi2GetNominalsOfContinuousStates(c) end ignore_derivatives() do if idx != -1 - e = FMU2Event(unsense(t), UInt64(idx), unsense(left_x), unsense(right_x)) + e = FMU2Event(t, UInt64(idx), left_x, right_x) push!(c.solution.events, e) end @@ -537,7 +487,7 @@ function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) ratio = ne / pt if ne >= 100 && ratio > c.fmu.executionConfig.maxStateEventsPerSecond - logError(c.fmu, "Event chattering detected $(round(Integer, ratio)) events/s, aborting at t=$(t) (rel. t=$(pt)) at event $(ne):") + logError(c.fmu, "Event jittering detected $(round(Integer, ratio)) events/s, aborting at t=$(t) (rel. t=$(pt)) at event $(ne):") for i in 0:c.fmu.modelDescription.numberOfEventIndicators num = 0 for e in c.solution.events @@ -556,14 +506,14 @@ function affectFMU!(nfmu::ME_NeuralFMU, c::FMU2Component, integrator, idx) c.solution.evals_affect += 1 - assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) end # Does one step in the simulation. function stepCompleted(nfmu::ME_NeuralFMU, c::FMU2Component, x, t, integrator, tStart, tStop) #@assert getCurrentComponent(nfmu.fmu) == c "Thread `$(Threads.threadid())` wants to evaluate wrong component!" - assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) c.solution.evals_stepcompleted += 1 @@ -592,7 +542,7 @@ function stepCompleted(nfmu::ME_NeuralFMU, c::FMU2Component, x, t, integrator, t #@debug "Step completed at $(ForwardDiff.value(t)) with $(collect(ForwardDiff.value(xs) for xs in x))" end - assert_integrator_valid(integrator) + @debug assert_integrator_valid(integrator) end # save FMU values @@ -604,8 +554,8 @@ function saveValues(nfmu::ME_NeuralFMU, c::FMU2Component, recordValues, _x, t, i c.solution.evals_savevalues += 1 # ToDo: Evaluate on light-weight model (sub-model) without fmi2GetXXX or similar and the bottom ANN - c.fmu.default_t = t - evaluateModel(nfmu, c, x, nfmu.p) # evaluate NeuralFMU (set new states) + c.next_t = t + evaluateModel(nfmu, c, x) # evaluate NeuralFMU (set new states) # Todo set inputs @@ -614,28 +564,30 @@ end # TODO import DifferentiableEigen -function saveEigenvalues(nfmu::ME_NeuralFMU, c::FMU2Component, _x, _t, integrator, sensitivity) +function saveEigenvalues(nfmu::ME_NeuralFMU, c::FMU2Component, _x, p, _t, integrator, sensitivity) #@assert c.state == fmi2ComponentStateContinuousTimeMode "saveEigenvalues(...): Must be in continuous time mode." c.solution.evals_saveeigenvalues += 1 - c.fmu.default_t = t + t = unsense(_t) + + c.next_t = t A = nothing #if sensitivity == :ForwardDiff - A = ForwardDiff.jacobian(x -> evaluateModel(nfmu, c, x, nfmu.p), _x) # TODO: chunk_size! + A = ForwardDiff.jacobian(x -> evaluateReModel(nfmu, c, x, p), _x) # TODO: chunk_size! # elseif sensitivity == :ReverseDiff - # A = ReverseDiff.jacobian(x -> evaluateModel(nfmu, c, x, nfmu.p), _x) + # A = ReverseDiff.jacobian(x -> evaluateReModel(nfmu, c, x, p), _x) # elseif sensitivity == :Zygote - # A = Zygote.jacobian(x -> evaluateModel(nfmu, c, x, nfmu.p), _x)[1] + # A = Zygote.jacobian(x -> evaluateReModel(nfmu, c, x, p), _x)[1] # elseif sensitivity == :none - # A = ForwardDiff.jacobian(x -> evaluateModel(nfmu, c, x, nfmu.p), unsense(_x)) + # A = ForwardDiff.jacobian(x -> evaluateReModel(nfmu, c, x, p), unsense(_x)) # end eigs, _ = DifferentiableEigen.eigen(A) # x = unsense(_x) - # c.fmu.default_t = t + # c.next_t = t # evaluateModel(nfmu, c, x) return (eigs...,) @@ -653,32 +605,31 @@ function fx(nfmu::ME_NeuralFMU, #@assert nanx && nanu "NaN in start fx nanx = $nanx nanu = $nanu @ $(t)." - if isnothing(c) + if c === nothing # this should never happen! - @warn "fx() called without allocated FMU instance!" return zeros(length(x)) end ignore_derivatives() do - #t = unsense(t) - c.fmu.default_t = t + t = unsense(t) + c.next_t = t end ############ - evaluateModel(nfmu, c, dx, x, p) + evaluateReModel(nfmu, c, dx, x, p) # if isdual(dx) - # dx_tmp = evaluateModel(nfmu, c, x, p) + # dx_tmp = evaluateReModel(nfmu, c, x, p) # fd_set!(dx, dx_tmp) # elseif istracked(dx) - # dx_tmp = evaluateModel(nfmu, c, x, p) + # dx_tmp = evaluateReModel(nfmu, c, x, p) # rd_set!(dx, dx_tmp) # else # #@info "dx: $(dx)" # #@info "dx_tmp: $(dx_tmp)" - # evaluateModel(nfmu, c, dx, x, p) + # evaluateReModel(nfmu, c, dx, x, p) # end ignore_derivatives() do @@ -703,11 +654,11 @@ function fx(nfmu::ME_NeuralFMU, ignore_derivatives() do c.solution.evals_fx_outofplace += 1 - #t = unsense(t) - c.fmu.default_t = t + t = unsense(t) + c.next_t = t end - return evaluateModel(nfmu, c, x, p) + return evaluateReModel(nfmu, c, x, p) end @@ -741,7 +692,7 @@ function ME_NeuralFMU(fmu::FMU2, end p, re = Flux.destructure(model) - nfmu = ME_NeuralFMU{typeof(model), typeof(re)}(model, p, re) + nfmu = ME_NeuralFMU{typeof(model), typeof(p), typeof(re)}(model, p, re) ###### @@ -827,7 +778,7 @@ Evaluates the ME_NeuralFMU in the timespan given during construction or in a cus function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, tspan::Tuple{Float64, Float64} = nfmu.tspan; showProgress::Bool = false, - progressDescr::String=DEFAULT_PROGRESS_DESCR, + progressDescr::String="Simulating ME-NeuralFMU ...", tolerance::Union{Real, Nothing} = nothing, parameters::Union{Dict{<:Any, <:Any}, Nothing} = nothing, setup::Union{Bool, Nothing} = nothing, @@ -856,13 +807,13 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, recordValues = prepareValueReference(nfmu.fmu, recordValues) saving = (length(recordValues) > 0) - + inPlace = nfmu.fmu.executionConfig.inPlace + t_start = tspan[1] t_stop = tspan[end] nfmu.tspan = tspan nfmu.x0 = x_start - nfmu.p = p #unsense(p) ignore_derivatives() do @debug "ME_NeuralFMU..." @@ -872,8 +823,8 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, nfmu.tolerance = tolerance if isnothing(parameters) - if !isnothing(nfmu.fmu.default_p_refs) - nfmu.parameters = Dict(nfmu.fmu.default_p_refs .=> unsense(nfmu.fmu.default_p)) + if !isnothing(nfmu.fmu.optim_p_refs) + nfmu.parameters = Dict(nfmu.fmu.optim_p_refs .=> unsense(nfmu.fmu.optim_p)) end else nfmu.parameters = parameters @@ -919,23 +870,11 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, # state event callback if c.fmu.hasStateEvents && c.fmu.executionConfig.handleStateEvents - - handleIndicators = nothing - - # if we want a specific subset - if !isnothing(c.fmu.handleEventIndicators) - handleIndicators = c.fmu.handleEventIndicators - else # handle all - handleIndicators = collect(UInt32(i) for i in 1:c.fmu.modelDescription.numberOfEventIndicators) - end - - numEvents = length(handleIndicators) - if c.fmu.executionConfig.useVectorCallbacks - eventCb = VectorContinuousCallback((out, x, t, integrator) -> condition!(nfmu, c, out, x, t, integrator, handleIndicators), + eventCb = VectorContinuousCallback((out, x, t, integrator) -> condition(nfmu, c, out, x, t, integrator), (integrator, idx) -> affectFMU!(nfmu, c, integrator, idx), - numEvents; + Int64(c.fmu.modelDescription.numberOfEventIndicators); rootfind=RightRootFind, save_positions=(saveEventPositions, saveEventPositions), interp_points=c.fmu.executionConfig.rootSearchInterpolationPoints) @@ -1012,10 +951,10 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, savingCB = nothing if saveat === nothing - savingCB = SavingCallback((u,t,integrator) -> saveEigenvalues(nfmu, c, u, t, integrator, recordEigenvaluesSensitivity), + savingCB = SavingCallback((u,t,integrator) -> saveEigenvalues(nfmu, c, u, p, t, integrator, recordEigenvaluesSensitivity), c.solution.eigenvalues) else - savingCB = SavingCallback((u,t,integrator) -> saveEigenvalues(nfmu, c, u, t, integrator, recordEigenvaluesSensitivity), + savingCB = SavingCallback((u,t,integrator) -> saveEigenvalues(nfmu, c, u, p, t, integrator, recordEigenvaluesSensitivity), c.solution.eigenvalues, saveat=saveat) end @@ -1026,9 +965,15 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, prob = nothing - ff = ODEFunction{true}((dx, x, p, t) -> fx(nfmu, c, dx, x, p, t), + if inPlace + ff = ODEFunction{true}((dx, x, p, t) -> fx(nfmu, c, dx, x, p, t), tgrad=nothing) - prob = ODEProblem{true}(ff, nfmu.x0, nfmu.tspan, p) + prob = ODEProblem{true}(ff, nfmu.x0, nfmu.tspan, p) + else + ff = ODEFunction{false}((x, p, t) -> fx(nfmu, c, x, p, t), + tgrad=nothing) # zero_tgrad) + prob = ODEProblem{false}(ff, nfmu.x0, nfmu.tspan, p) + end # if (length(callbacks) == 2) # only start and stop callback, so the system is pure continuous # startCallback(nfmu, nfmu.tspan[1]) @@ -1039,11 +984,11 @@ function (nfmu::ME_NeuralFMU)(x_start::Union{Array{<:Real}, Nothing} = nfmu.x0, if isnothing(sensealg) # when using state events, we (currently) need AD-through-Solver - # if c.fmu.hasStateEvents && c.fmu.executionConfig.handleStateEvents - sensealg = ReverseDiffAdjoint() # Support for multi-state-event simulations, but a little bit slower than QuadratureAdjoint - # else # otherwise, we can use the faster Adjoint-over-Solver (does this work? FMUs seem not be solvable in reverse time - in general) - # sensealg = QuadratureAdjoint(; autojacvec=ReverseDiffVJP(true)) # Faster than ReverseDiffAdjoint - # end + if c.fmu.hasStateEvents && c.fmu.executionConfig.handleStateEvents + sensealg = ReverseDiffAdjoint() # Support for multi-state-event simulations, but a little bit slower than QuadratureAdjoint + else # otherwise, we can use the faster Adjoint-over-Solver + sensealg = QuadratureAdjoint(; autojacvec=ReverseDiffVJP()) # Faster than ReverseDiffAdjoint + end end args = Vector{Any}() @@ -1266,7 +1211,7 @@ function computeGradient(loss, params, gradient, chunk_size, multiObjective::Boo elseif chunk_size == :auto_fmiflux - chunk_size = DEFAULT_CHUNK_SIZE + chunk_size = 32 if multiObjective conf = ForwardDiff.JacobianConfig(loss, params, ForwardDiff.Chunk{min(chunk_size, length(params))}()); @@ -1319,36 +1264,29 @@ function computeGradient(loss, params, gradient, chunk_size, multiObjective::Boo end # WIP -function trainStep(loss, params, gradient, chunk_size, optim::Optim.AbstractOptimizer, printStep, proceed_on_assert, cb, state::Union{Optim.AbstractOptimizerState, Nothing}=nothing, d::Union{Optim.OnceDifferentiable, Nothing}=nothing) - - @assert length(params) == 1 "Currently only length(params)==1 is supported!" - - options = Optim.Options(iterations=1) - autodiff = :finite - inplace = true +function trainStep(loss, params, gradient, chunk_size, optim::Optim.AbstractOptimizer, printStep, proceed_on_assert, cb; state=nothing) - function g!(G, _params) - grad = computeGradient(loss, _params, gradient, chunk_size) - @assert !isnothing(grad) "Gradient nothing!" - G[:] = grad - - if printStep - @info "Grad: Min = $(min(abs.(grad)...)) Max = $(max(abs.(grad)...))" + try + if isnothing(state) + state = initial_state(optim, options, d, initial_x) end - end + + for j in 1:length(params) - if isnothing(d) - d = Optim.promote_objtype(optimizer, params[1], autodiff, inplace, loss) - end + grad = computeGradient(loss, params[j], gradient, chunk_size) + + @assert !isnothing(grad) "Gradient nothing!" - if isnothing(state) - state = Optim.initial_state(optimizer, options, d, params[1]) - end + update_state!(d, state, optim) + step = Flux.Optimise.apply!(optim, params[j], grad) + params[j] .-= step - try - res = optimize(d, params[1], optimizer, options, state) - params[1][:] = Optim.minimizer(res) + if printStep + @info "Grad: Min = $(min(abs.(grad)...)) Max = $(max(abs.(grad)...))" + @info "Step: Min = $(min(abs.(step)...)) Max = $(max(abs.(step)...))" + end + end catch e @@ -1386,7 +1324,7 @@ function trainStep(loss, params, gradient, chunk_size, optim::Flux.Optimise.Abst has_nothing = any(collect(any(isnothing.(grad)) for grad in grads)) || any(isnothing.(grads)) if gradient != :ForwardDiff && (has_nan || has_nothing) - @warn "Gradient determination with $(gradient) failed, because gradient contains `NaNs` and/or `nothing`.\nThis might be because the FMU is throwing redundant events, which is currently not supported.\nTrying ForwardDiff as back-up.\nIf this message gets printed (almost) every step, consider using keyword `gradient=:ForwardDiff` to fix ForwardDiff as sensitivity system." + @warn "Gradient determination with $(gradient) failed, because gradient contains `NaNs` and/or `nothing`.\nTrying ForwardDiff as back-up.\nIf this message gets printed (almost) every step, consider using keyword `gradient=:ForwardDiff` to fix ForwardDiff as sensitivity system." gradient = :ForwardDiff grads = computeGradient(loss, params[j], gradient, chunk_size, multiObjective) end @@ -1484,14 +1422,14 @@ end # checks gradient determination for all available sensitivity configurations, see: # https://docs.sciml.ai/SciMLSensitivity/stable/manual/differential_equation_sensitivities/ -using FMISensitivity.SciMLSensitivity +using FMIImport.SciMLSensitivity function checkSensalgs!(loss, neuralFMU::Union{ME_NeuralFMU, CS_NeuralFMU}; - gradients=(:ReverseDiff, :Zygote, :ForwardDiff), # :FiniteDiff is slow ... - max_msg_len=192, chunk_size=DEFAULT_CHUNK_SIZE, - OtD_autojacvecs=(false, true, TrackerVJP(), ZygoteVJP(), ReverseDiffVJP(false), ReverseDiffVJP(true)), # EnzymeVJP() deadlocks in the current release xD + gradients=(:ForwardDiff, :ReverseDiff, :Zygote), # :FiniteDiff is slow ... + max_msg_len=192, chunk_size=32, + OtD_autojacvecs=(false, true, TrackerVJP(), ZygoteVJP(), ReverseDiffVJP()), # EnzymeVJP() deadlocks in the current release xD OtD_sensealgs=(BacksolveAdjoint, InterpolatingAdjoint, QuadratureAdjoint), OtD_checkpointings=(true, false), - DtO_sensealgs=(ReverseDiffAdjoint, ForwardDiffSensitivity, TrackerAdjoint), # TrackerAdjoint, ZygoteAdjoint freeze the REPL + DtO_sensealgs=(ReverseDiffAdjoint, ZygoteAdjoint, TrackerAdjoint, ForwardDiffSensitivity), multiObjective::Bool=false, bestof::Int=2, timeout_seconds::Real=60.0, diff --git a/test/batching.jl b/test/batching.jl index b8883f7b..df176fb2 100644 --- a/test/batching.jl +++ b/test/batching.jl @@ -24,6 +24,7 @@ fmu = fmiLoad("SpringPendulum1D", EXPORTINGTOOL, EXPORTINGVERSION; type=fmi2Type using FMI.FMIImport using FMI.FMIImport.FMICore +import FMI.FMIImport: unsense # loss function for training function losssum_single(p) diff --git a/test/fmu_params.jl b/test/fmu_params.jl index 95236491..732bd4e1 100644 --- a/test/fmu_params.jl +++ b/test/fmu_params.jl @@ -16,15 +16,15 @@ t_stop = 5.0 tData = t_start:t_step:t_stop # generate training data -velData = sin.(tData.*4.0) +velData = sin.(tData.*3.0) x0 = [0.5, 0.0] # load FMU for NeuralFMU -# [TODO] Replace by a suitable discontinuous FMU -fmu = fmiLoad("SpringPendulum1D", EXPORTINGTOOL, EXPORTINGVERSION; type=:ME) +fmu = fmiLoad("SpringFrictionPendulum1D", EXPORTINGTOOL, EXPORTINGVERSION; type=:ME) using FMI.FMIImport using FMI.FMIImport.FMICore +import FMI.FMIImport: unsense c = fmi2Instantiate!(fmu) fmi2SetupExperiment(c, 0.0, 1.0) @@ -36,17 +36,17 @@ p = fmi2GetReal(c, p_refs) # loss function for training function losssum(p) - #@info "$p" - global problem, x0, posData, solution + global problem, x0, posData solution = problem(x0; p=p, showProgress=true) if !solution.success return Inf end + # posNet = fmi2GetSolutionState(solution, 1; isIndex=true) velNet = fmi2GetSolutionState(solution, 2; isIndex=true) - return Flux.Losses.mse(velNet, velData) + return Flux.Losses.mse(velNet, velData) # Flux.Losses.mse(posNet, posData) end # callback function for training @@ -60,23 +60,24 @@ function callb(p) loss = losssum(p[1]) @info "[$(iterCB)] Loss: $loss" @test loss < lastLoss + #@test p[1][1] == fmu.optim_p[1] + #@info "$(fmu.optim_p[1])" + #@info "$(p)" + #@info "$(problem.parameters)" lastLoss = loss end end numStates = fmiGetNumberOfStates(fmu) -dx = zeros(fmi2Real, numStates) - # the "Chain" for training net = Chain(FMUParameterRegistrator(fmu, p_refs, p), - x -> fmu(x=x, dx=dx)) # , fmuLayer(p)) + x -> fmu(;x=x)) # , fmuLayer(p)) -optim = Adam(1e-4) +optim = Adam(1e-3) solver = Tsit5() problem = ME_NeuralFMU(fmu, net, (t_start, t_stop), solver; saveat=tData) -problem.modifiedState = false @test problem != nothing solutionBefore = problem(x0) @@ -88,26 +89,14 @@ solutionBefore = problem(x0) # train it ... p_net = Flux.params(problem) @test length(p_net) == 1 -@test length(p_net[1]) == 3 +@test length(p_net[1]) == 12 iterCB = 0 lastLoss = losssum(p_net[1]) @info "Start-Loss for net: $lastLoss" -# [ToDo] Check pure gradients -j_fin = FiniteDiff.finite_difference_gradient(losssum, p_net[1]) -plot!(collect(unsense(t) for t in solution.states.t), collect(unsense(u[2]) for u in solution.states.u); label="FD") -@test length(solution.events) == 6 - -j_fwd = ForwardDiff.gradient(losssum, p_net[1]) -plot!(collect(unsense(t) for t in solution.states.t), collect(unsense(u[2]) for u in solution.states.u); label="FWD") -@test length(solution.events) == 6 - -j_rwd = ReverseDiff.gradient(losssum, p_net[1]) -plot!(collect(unsense(t) for t in solution.states.t), collect(unsense(u[2]) for u in solution.states.u); label="RD") -@test length(solution.events) == 6 - -FMIFlux.train!(losssum, p_net, Iterators.repeated((), NUMSTEPS), optim; cb=()->callb(p_net), gradient=GRADIENT) +@warn "FMU parameter tests disabled." +# FMIFlux.train!(losssum, p_net, Iterators.repeated((), NUMSTEPS), optim; cb=()->callb(p_net), gradient=GRADIENT) # check results solutionAfter = problem(x0) diff --git a/test/hybrid_ME.jl b/test/hybrid_ME.jl index 50668bdb..70c2fa65 100644 --- a/test/hybrid_ME.jl +++ b/test/hybrid_ME.jl @@ -67,18 +67,16 @@ c2 = CacheRetrieveLayer(c1) c3 = CacheLayer() c4 = CacheRetrieveLayer(c3) -dx = zeros(fmi2Real, numStates) - # 1. default ME-NeuralFMU (learn dynamics and states, almost-neutral setup, parameter count << 100) net = Chain(Dense(numStates, numStates, identity; init=Flux.identity_init), - x -> myFMU(x=x, dx=dx), + x -> myFMU(;x=x), c3, Dense(numStates, numStates, identity; init=Flux.identity_init), x -> c4([1], x[2], [])) push!(nets, net) # 2. default ME-NeuralFMU (learn dynamics) -net = Chain(x -> myFMU(x=x, dx=dx), +net = Chain(x -> myFMU(;x=x), x -> c1(x), Dense(numStates, 16, identity; init=Flux.identity_init), Dense(16, 16, identity; init=Flux.identity_init), @@ -92,7 +90,7 @@ net = Chain(x -> c1(x), Dense(16, 16, identity, init=Flux.identity_init), Dense(16, numStates, identity, init=Flux.identity_init), x -> c2([1], x[2], []), - x -> myFMU(x=x, dx=dx)) + x -> myFMU(;x=x)) push!(nets, net) # 4. default ME-NeuralFMU (learn dynamics and states) @@ -100,7 +98,7 @@ net = Chain(x -> c1(x), Dense(numStates, 16, identity; init=Flux.identity_init), Dense(16, numStates, identity; init=Flux.identity_init), x -> c2([1], x[2], []), - x -> myFMU(x=x, dx=dx), + x -> myFMU(;x=x), x -> c3(x), Dense(numStates, 16, identity, init=Flux.identity_init), Dense(16, 16, identity, init=Flux.identity_init), @@ -109,7 +107,7 @@ net = Chain(x -> c1(x), push!(nets, net) # 5. NeuralFMU with hard setting time to 0.0 -net = Chain(x -> myFMU(x=x, dx=dx, t=0.0), +net = Chain(states -> myFMU(;x=states, t=0.0), x -> c1(x), Dense(numStates, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -120,8 +118,7 @@ push!(nets, net) # 6. NeuralFMU with additional getter getVRs = [fmi2StringToValueReference(myFMU, "mass.s")] numGetVRs = length(getVRs) -y = zeros(fmi2Real, numGetVRs) -net = Chain(x -> myFMU(x=x, dx=dx, y_refs=getVRs, y=y), +net = Chain(x -> myFMU(;x=x, y_refs=getVRs), x -> c1(x), Dense(numStates+numGetVRs, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -132,7 +129,7 @@ push!(nets, net) # 7. NeuralFMU with additional setter setVRs = [fmi2StringToValueReference(myFMU, "mass.m")] numSetVRs = length(setVRs) -net = Chain(x -> myFMU(x=x, dx=dx, u_refs=setVRs, u=[1.1]), +net = Chain(x -> myFMU(;x=x, u_refs=setVRs, u=[1.1]), x -> c1(x), Dense(numStates, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -141,7 +138,7 @@ net = Chain(x -> myFMU(x=x, dx=dx, u_refs=setVRs, u=[1.1]), push!(nets, net) # 8. NeuralFMU with additional setter and getter -net = Chain(x -> myFMU(x=x, dx=dx, u_refs=setVRs, u=[1.1], y_refs=getVRs, y=y), +net = Chain(x -> myFMU(;x=x, u_refs=setVRs, u=[1.1], y_refs=getVRs), x -> c1(x), Dense(numStates+numGetVRs, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -150,7 +147,7 @@ net = Chain(x -> myFMU(x=x, dx=dx, u_refs=setVRs, u=[1.1], y_refs=getVRs, y=y), push!(nets, net) # 9. an empty NeuralFMU (this does only make sense for debugging) -net = Chain(x -> myFMU(x=x, dx=dx)) +net = Chain(x -> myFMU(;x=x)) push!(nets, net) for i in 1:length(nets) diff --git a/test/hybrid_ME_dis.jl b/test/hybrid_ME_dis.jl index 812162f5..22079734 100644 --- a/test/hybrid_ME_dis.jl +++ b/test/hybrid_ME_dis.jl @@ -22,8 +22,6 @@ realSimData = fmiSimulate(fmu, (t_start, t_stop); parameters=pdict, recordValues x0 = collect(realSimData.values.saveval[1]) @test x0 == [0.5, 0.0] -dx = copy(x0) - # setup traing data velData = fmi2GetSolutionValue(realSimData, "mass.v") @@ -75,14 +73,14 @@ c4 = CacheRetrieveLayer(c3) net = Chain(x -> c1(x), Dense(numStates, numStates, identity; init=Flux.identity_init), x -> c2([1], x[2], []), - x -> fmu(;x=x, dx=dx), + x -> fmu(;x=x), x -> c3(x), Dense(numStates, numStates, identity; init=Flux.identity_init), x -> c4([1], x[2], [])) push!(nets, net) # 2. default ME-NeuralFMU (learn dynamics) -net = Chain(x -> fmu(;x=x, dx=dx), +net = Chain(x -> fmu(;x=x), x -> c1(x), Dense(numStates, 16, identity; init=Flux.identity_init), Dense(16, 16, identity; init=Flux.identity_init), @@ -96,7 +94,7 @@ net = Chain(x -> c1(x), Dense(16, 16, identity, init=Flux.identity_init), Dense(16, numStates, identity, init=Flux.identity_init), x -> c2([1], x[2], []), - x -> fmu(;x=x, dx=dx)) + x -> fmu(;x=x)) push!(nets, net) # 4. default ME-NeuralFMU (learn dynamics and states) @@ -104,7 +102,7 @@ net = Chain(x -> c1(x), Dense(numStates, 16, identity; init=Flux.identity_init), Dense(16, numStates, identity; init=Flux.identity_init), x -> c2([1], x[2], []), - x -> fmu(;x=x, dx=dx), + x -> fmu(;x=x), x -> c3(x), Dense(numStates, 16, identity, init=Flux.identity_init), Dense(16, 16, identity, init=Flux.identity_init), @@ -113,7 +111,7 @@ net = Chain(x -> c1(x), push!(nets, net) # 5. NeuralFMU with hard setting time to 0.0 -net = Chain(states -> fmu(;x=states, t=0.0, dx=dx), +net = Chain(states -> fmu(;x=states, t=0.0), x -> c1(x), Dense(numStates, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -123,9 +121,8 @@ push!(nets, net) # 6. NeuralFMU with additional getter getVRs = [fmi2StringToValueReference(fmu, "mass.s")] -y = zeros(fmi2Real, length(getVRs)) numGetVRs = length(getVRs) -net = Chain(x -> fmu(;x=x, y_refs=getVRs, dx=dx, y=y), +net = Chain(x -> fmu(;x=x, y_refs=getVRs), x -> c1(x), Dense(numStates+numGetVRs, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -136,7 +133,7 @@ push!(nets, net) # 7. NeuralFMU with additional setter setVRs = [fmi2StringToValueReference(fmu, "mass.m")] numSetVRs = length(setVRs) -net = Chain(x -> fmu(;x=x, u_refs=setVRs, u=[1.1], dx=dx), +net = Chain(x -> fmu(;x=x, u_refs=setVRs, u=[1.1]), x -> c1(x), Dense(numStates, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), @@ -145,7 +142,7 @@ net = Chain(x -> fmu(;x=x, u_refs=setVRs, u=[1.1], dx=dx), push!(nets, net) # 8. NeuralFMU with additional setter and getter -net = Chain(x -> fmu(;x=x, u_refs=setVRs, u=[1.1], y_refs=getVRs, y=y, dx=dx), +net = Chain(x -> fmu(;x=x, u_refs=setVRs, u=[1.1], y_refs=getVRs), x -> c1(x), Dense(numStates+numGetVRs, 8, identity; init=Flux.identity_init), Dense(8, 16, identity; init=Flux.identity_init), diff --git a/test/runtests.jl b/test/runtests.jl index 977dc560..c2cf6874 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -7,10 +7,8 @@ using FMIFlux using Test using FMIZoo -import FMIFlux.FMISensitivity: FiniteDiff, ForwardDiff, ReverseDiff - using FMIFlux.FMIImport: fmi2StringToValueReference, fmi2ValueReference, prepareSolveFMU -using FMIFlux.FMIImport: FMU2_EXECUTION_CONFIGURATIONS +using FMIFlux.FMIImport: FMU2_EXECUTION_CONFIGURATION_NO_FREEING, FMU2_EXECUTION_CONFIGURATION_NO_RESET, FMU2_EXECUTION_CONFIGURATION_RESET using FMIFlux: fmi2GetSolutionState, fmi2GetSolutionValue, fmi2GetSolutionTime exportingToolsWindows = [("Dymola", "2022x")] @@ -22,8 +20,8 @@ global GRADIENT = nothing global EXPORTINGTOOL = nothing global EXPORTINGVERSION = nothing -# enable assertions for warnings/errors for all default execution configurations - we want strict tests here -for exec in FMU2_EXECUTION_CONFIGURATIONS +# enable assertions for warnings/errors for all default execution configurations +for exec in [FMU2_EXECUTION_CONFIGURATION_NO_FREEING, FMU2_EXECUTION_CONFIGURATION_NO_RESET, FMU2_EXECUTION_CONFIGURATION_RESET] exec.assertOnError = true exec.assertOnWarning = true end @@ -35,21 +33,16 @@ function runtests(exportingTool) @info "Testing FMUs exported from $(EXPORTINGTOOL) ($(EXPORTINGVERSION))" @testset "Testing FMUs exported from $(EXPORTINGTOOL) ($(EXPORTINGVERSION))" begin - for _GRADIENT ∈ (:ReverseDiff, :ForwardDiff, :FiniteDiff) + for _GRADIENT ∈ (:ReverseDiff, :ForwardDiff) global GRADIENT = _GRADIENT @info "Gradient: $(GRADIENT)" @testset "Gradient: $(GRADIENT)" begin - # @info "Layers (layers.jl)" - # @testset "Layers" begin - # include("layers.jl") - # end - - # @info "Solution Gradients (solution_gradients.jl)" - # @testset "Solution Gradients" begin - # include("solution_gradients.jl") - # end + @info "Layers (layers.jl)" + @testset "Layers" begin + include("layers.jl") + end @info "ME-NeuralFMU (Continuous) (hybrid_ME.jl)" @testset "ME-NeuralFMU (Continuous)" begin @@ -61,42 +54,42 @@ function runtests(exportingTool) include("hybrid_ME_dis.jl") end - # @info "NeuralFMU with FMU parameter optimization (fmu_params.jl)" - # @testset "NeuralFMU with FMU parameter optimization" begin - # include("fmu_params.jl") - # end + @info "NeuralFMU with FMU parameter optimization (fmu_params.jl)" + @testset "NeuralFMU with FMU parameter optimization" begin + include("fmu_params.jl") + end - # @info "Training modes (train_modes.jl)" - # @testset "Training modes" begin - # include("train_modes.jl") - # end + @info "Training modes (train_modes.jl)" + @testset "Training modes" begin + include("train_modes.jl") + end # @info "Multi-threading (multi_threading.jl)" # @testset "Multi-threading" begin # include("multi_threading.jl") # end - # @info "CS-NeuralFMU (hybrid_CS.jl)" - # @testset "CS-NeuralFMU" begin - # include("hybrid_CS.jl") - # end + @info "CS-NeuralFMU (hybrid_CS.jl)" + @testset "CS-NeuralFMU" begin + include("hybrid_CS.jl") + end - # @info "Multiple FMUs (multi.jl)" - # @testset "Multiple FMUs" begin - # include("multi.jl") - # end + @info "Multiple FMUs (multi.jl)" + @testset "Multiple FMUs" begin + include("multi.jl") + end - # @info "Batching (batching.jl)" - # @testset "Batching" begin - # include("batching.jl") - # end + @info "Batching (batching.jl)" + @testset "Batching" begin + include("batching.jl") + end end end - # @info "Benchmark: Supported sensitivities (supported_sensitivities.jl)" - # @testset "Benchmark: Supported sensitivities " begin - # include("supported_sensitivities.jl") - # end + @info "Benchmark: Supported sensitivities (supported_sensitivities.jl)" + @testset "Benchmark: Supported sensitivities " begin + include("supported_sensitivities.jl") + end end end diff --git a/test/solution_gradients.jl b/test/solution_gradients.jl deleted file mode 100644 index 92bf0781..00000000 --- a/test/solution_gradients.jl +++ /dev/null @@ -1,277 +0,0 @@ -# -# Copyright (c) 2021 Tobias Thummerer, Lars Mikelsons -# Licensed under the MIT license. See LICENSE file in the project root for details. -# - -using FMI -using Flux -using DifferentialEquations -using FMIFlux, FMIZoo, Test -import FMISensitivity.SciMLSensitivity.SciMLBase: RightRootFind, LeftRootFind -import FMIFlux: unsense -using ForwardDiff, ReverseDiff, FiniteDiff, Zygote -EXPORTINGTOOL = "Dymola" -EXPORTINGVERSION = "2022x" - -import Random -Random.seed!(5678); - -ENERGY_LOSS = 0.7 -RADIUS = 0.0 -GRAVITY = 9.81 -DBL_MIN = 2.2250738585072013830902327173324040642192159804623318306e-308 - -t_start = 0.0 -t_step = 0.05 -t_stop = 2.0 -tData = t_start:t_step:t_stop -posData = ones(length(tData)) - -x0 = [1.0, 0.0] -dx = [0.0, 0.0] -numStates = 2 -solver = Tsit5() - -# setup BouncingBallODE -function fx_bb(x) - dx = [x[2], -GRAVITY] - dx -end -function fx(dx, x, p, t) - # if rand(1:10)%10 == 0 - # @info "$(typeof(x))" - # end - dx[:] = re_bb(p)(x) -end - -net_bb = Chain(fx_bb, - Dense([1.0 0.0; 0.0 1.0], [0.0, 0.0], identity)) - -ff = ODEFunction{true}(fx) -prob_bb = ODEProblem{true}(ff, x0, (t_start, t_stop), ()) - -function condition(out, x, t, integrator) - out[1] = x[1]-RADIUS - #out[2] = x[1]-RADIUS -end - -global events = 0 -function affect_right!(integrator, idx) - s_new = RADIUS + DBL_MIN - v_new = -integrator.u[2]*ENERGY_LOSS - u_new = [s_new, v_new] - - global events - events += 1 - @info "[$(events)] New state at $(integrator.t) is $(u_new)" - - integrator.u .= u_new -end -function affect_left!(integrator, idx) - s_new = integrator.u[1] - v_new = -integrator.u[2]*ENERGY_LOSS - u_new = [s_new, v_new] - - global events - events += 1 - @info "[$(events)] New state at $(integrator.t) is $(u_new)" - - integrator.u .= u_new -end - -rightCb = VectorContinuousCallback(condition, - affect_right!, - 1; - rootfind=RightRootFind, save_positions=(false, false)) -leftCb = VectorContinuousCallback(condition, - affect_left!, - 1; - rootfind=LeftRootFind, save_positions=(false, false)) - -# load FMU for NeuralFMU -#fmu = fmiLoad("BouncingBall1D", EXPORTINGTOOL, EXPORTINGVERSION; type=:ME) -fmu = fmiLoad("BouncingBall", "ModelicaReferenceFMUs", "0.0.25") -fmu.handleEventIndicators = nothing - -net = Chain(x -> fmu(;x=x, dx=dx), - Dense([1.0 0.0; 0.0 1.0], [0.0; 0.0], identity)) - -prob = ME_NeuralFMU(fmu, net, (t_start, t_stop); saveat=tData) -prob.modifiedState = false - -# ANNs -global solution = nothing - -function losssum(p; sensealg=nothing) - global solution, prob, x0, posData, solver - solution = prob(x0; p=p, solver=solver, sensealg=sensealg) - - if !solution.success - @error "Solution failed!" - return Inf - end - - posNet = fmi2GetSolutionState(solution, 1; isIndex=true) - - return Flux.Losses.mae(posNet, posData) -end - -function losssum_bb(p; sensealg=nothing, root=:Right) - - posNet = mysolve(p; sensealg=sensealg, root=root) - - return Flux.Losses.mae(posNet, posData) -end - -function mysolve(p; sensealg=nothing, root=:Right) - global solution, prob_bb, x0, posData, solver, events - events = 0 - - callback = nothing - if root == :Right - callback = CallbackSet(rightCb) - else - callback = CallbackSet(leftCb) - end - solution = solve(prob_bb, solver; p=p, saveat=tData, callback=callback, sensealg=sensealg) - - if !isa(solution, AbstractArray) - if solution.retcode != FMI.ReturnCode.Success - @error "Solution failed!" - return Inf - end - - return solution.u - else - return solution[1,:] - end -end - -p_net = Flux.params(prob)[1] -p_net_bb, re_bb = Flux.destructure(net_bb) - -using SciMLSensitivity, Plots -sensealg = ReverseDiffAdjoint() # QuadratureAdjoint(autojacvec=ReverseDiffVJP()) - -### START CHECK CONDITIONS - -function condition_bb_check(x) - buffer = similar(x, 1) - condition(buffer, x, t_start, nothing) - return buffer -end -function condition_nfmu_check(x) - buffer = similar(x, 1) - FMIFlux.condition!(prob, FMIFlux.getComponent(prob), buffer, x, t_start, nothing, [UInt32(1)]) - return buffer -end -jac_con1 = ForwardDiff.jacobian(condition_bb_check, x0) -jac_con2 = ForwardDiff.jacobian(condition_nfmu_check, x0) - -jac_con1 = ReverseDiff.jacobian(condition_bb_check, x0) -jac_con2 = ReverseDiff.jacobian(condition_nfmu_check, x0) - -### START CHECK AFFECT - -import SciMLSensitivity: u_modified! -import FMI: fmi2SimulateME -function u_modified!(::NamedTuple, ::Any) - return nothing -end -function affect_bb_check(x) - integrator = (t=t_start, u=x) - affect_right!(integrator, 1) - return integrator.u -end -function affect_nfmu_check(x) - integrator = (t=t_start, u=x, opts=(internalnorm=(a,b)->1.0,) ) - #FMIFlux.affectFMU!(prob, FMIFlux.getComponent(prob), integrator, 1) - integrator.u[1] = DBL_MIN - integrator.u[2] = -0.7 * integrator.u[2] - return integrator.u -end -t_first_event_time = 0.451523640985728 -x_first_event_right = [2.2250738585072014e-308, 3.1006128426489954] -#sol = fmi2SimulateME(fmu, (t_start, t_first_event_time); solver=solver) -jac_con1 = ForwardDiff.jacobian(affect_bb_check, x0) -jac_con2 = ForwardDiff.jacobian(affect_nfmu_check, x0) - -jac_con1 = ReverseDiff.jacobian(affect_bb_check, x0) -jac_con2 = ReverseDiff.jacobian(affect_nfmu_check, x0) - -### - -fig = plot(;ylims=(-0.1, 1.1)) - -# Solution (plain) -losssum(p_net; sensealg=sensealg) -length(solution.events) -plot!(fig, tData, collect(u[1] for u in solution.states.u); label="NFMU: Sol") - -losssum_bb(p_net_bb; sensealg=sensealg) -events -plot!(fig, tData, collect(u[1] for u in solution.u); label="NODE: Sol") - -# Solution FWD -grad_fwd1 = ForwardDiff.gradient(p -> losssum(p; sensealg=sensealg), p_net) -length(solution.events) -plot!(fig, tData, collect(unsense(u[1]) for u in solution.states.u); label="NFMU: FWD") - -grad_fwd2 = ForwardDiff.gradient(p -> losssum_bb(p; sensealg=sensealg), p_net_bb) -events -plot!(fig, tData, collect(unsense(u[1]) for u in solution.u); label="NODE: FWD") - -# Solution ReverseDiff -grad_rwd1 = ReverseDiff.gradient(p -> losssum(p; sensealg=sensealg), p_net) -length(solution.events) -plot!(fig, tData, collect(unsense(u[1]) for u in solution.states.u); label="NFMU: RWD") - -grad_rwd2 = ReverseDiff.gradient(p -> losssum_bb(p; sensealg=sensealg), p_net_bb) -events -plot!(fig, tData, collect(unsense(u[1]) for u in solution[1,:]); label="NODE: RWD") - -# Solution Zygote -# grad_zyg1 = Zygote.gradient(p -> losssum(p; sensealg=sensealg), p_net)[1] -# plot!(fig, tData, collect(unsense(u[1]) for u in solution.states.u); label="NFMU: ZYG") - -# grad_zyg2 = Zygote.gradient(p -> losssum_bb(p; sensealg=sensealg), p_net_bb)[1] -# plot!(fig, tData, collect(unsense(u[1]) for u in solution[1,:]); label="NODE: ZYG") - -# Ground Truth -grad_fin1 = FiniteDiff.finite_difference_gradient(p -> losssum_bb(p; sensealg=sensealg), p_net_bb, Val{:central}; absstep=1e-8) -grad_fin2 = FiniteDiff.finite_difference_gradient(p -> losssum(p; sensealg=sensealg), p_net, Val{:central}; absstep=1e-8) - -atol = 1e-5 -@test isapprox(grad_fin1, grad_fwd1; atol=atol) -@test isapprox(grad_fin2, grad_fwd2; atol=atol) - -@test isapprox(grad_fin1, grad_rwd1; atol=atol) -@test isapprox(grad_fin2, grad_rwd2; atol=atol) - -# Jacobina Test - -jac_fwd1 = ForwardDiff.jacobian(p -> mysolve(p; sensealg=sensealg), p_net) -plot(tData, collect(unsense(u[1]) for u in solution.u); label="FWD") -jac_rwd1 = ReverseDiff.jacobian(p -> mysolve(p; sensealg=sensealg), p_net) -plot!(tData, collect(unsense(u[1]) for u in solution[1,:]); label="RWD (Right)") -jac_rwd1 = ReverseDiff.jacobian(p -> mysolve(p; sensealg=sensealg, root=:Left), p_net) -plot!(tData, collect(unsense(u[1]) for u in solution[1,:]); label="RWD (Left)") -jac_fin1 = FiniteDiff.finite_difference_jacobian(p -> jac_bb(p; sensealg=sensealg), p_net) - -### - -atol = 1e-4 -@test isapprox(grad_fin1, grad_fwd1; atol=atol) -@test isapprox(grad_fin2, grad_fwd2; atol=atol) - -@test isapprox(grad_fin1, grad_rwd1; atol=atol) -@test isapprox(grad_fin2, grad_rwd2; atol=atol) - -# atol = 1e-4 -# @test isapprox(collect(u[1] for u in sol.states.u), collect(u[1] for u in sol_fin.states.u); atol=atol) -# @test isapprox(collect(u[1] for u in sol.states.u), collect(u[1] for u in sol_fwd.states.u); atol=atol) -# @test isapprox(collect(u[1] for u in sol.states.u), collect(u[1] for u in sol_rwd.states.u); atol=atol) - -### - -fmiUnload(fmu) diff --git a/test/supported_sensitivities.jl b/test/supported_sensitivities.jl index 844efb71..47b09b3d 100644 --- a/test/supported_sensitivities.jl +++ b/test/supported_sensitivities.jl @@ -9,51 +9,33 @@ using Flux import Random Random.seed!(5678); -# setup BouncingBallODE -function fx_bb(x) - dx = [x[2], -9.81] - dx -end -function fx(dx, x, p, t) - dx[:] = re(p)(x) -end - -net = Chain(fx_bb, - Dense([1.0 0.0; 0.0 1.0], [0.0, 0.0], identity)) - -ff = ODEFunction{true}(fx) # , tgrad=nothing) -bb_prob = ODEProblem{true}(ff, x0, tspan, params) - -function condition(out, x, t, integrator) - out[1] = x[1] -end - -function affect!(integrator, idx) - u_new = [max(integrator.u[1], 0.0), -integrator.u[2]*0.9] - integrator.u .= u_new -end - -eventCb = VectorContinuousCallback(condition, - affect!, - 2; - rootfind=RightRootFind, save_positions=(false, false)) - t_start = 0.0 t_step = 0.1 -t_stop = 5.0 +t_stop = 3.0 tData = t_start:t_step:t_stop -posData = ones(length(tData)) +velData = sin.(tData) # load FMU for NeuralFMU -fmu = fmiLoad("BouncingBall1D", EXPORTINGTOOL, EXPORTINGVERSION; type=:ME) -fmu.handleEventIndicators = [1] +fmu = fmiLoad("SpringFrictionPendulum1D", EXPORTINGTOOL, EXPORTINGVERSION; type=:ME) x0 = [1.0, 0.0] -dx = [0.0, 0.0] numStates = length(x0) -net = Chain(x -> fmu(;x=x, dx=dx), - Dense([1.0 0.0; 0.0 1.0], [0.0; 0.0], identity)) +c1 = CacheLayer() +c2 = CacheRetrieveLayer(c1) +c3 = CacheLayer() +c4 = CacheRetrieveLayer(c3) + +# default ME-NeuralFMU (learn dynamics and states, almost-neutral setup, parameter count << 100) +net = Chain(x -> c1(x), + Dense(numStates, 32, identity; init=Flux.identity_init), + Dense(32, numStates, identity; init=Flux.identity_init), + x -> c2([1], x[2], []), + x -> fmu(;x=x), + x -> c3(x), + Dense(numStates, 32, identity; init=Flux.identity_init), + Dense(32, numStates, identity; init=Flux.identity_init), + x -> c4([1], x[2], [])) # loss function for training function losssum(p) @@ -64,43 +46,12 @@ function losssum(p) return Inf end - posNet = fmi2GetSolutionState(solution, 1; isIndex=true) + velNet = fmi2GetSolutionState(solution, 2; isIndex=true) - return FMIFlux.Losses.mse(posNet, posData) -end - -function losssum_bb(p) - global nfmu, x0, posData - solution = nfmu(x0; p=p) - - if !solution.success - return Inf - end - - posNet = fmi2GetSolutionState(solution, 1; isIndex=true) - - return FMIFlux.Losses.mse(posNet, posData) + return FMIFlux.Losses.mse(velNet, velData) end nfmu = ME_NeuralFMU(fmu, net, (t_start, t_stop); saveat=tData) -nfmu.modifiedState = false - -using SciMLSensitivity -params = Flux.params(nfmu) -fmu.executionConfig.sensealg = ReverseDiffAdjoint() # QuadratureAdjoint(autojacvec=ReverseDiffVJP(false)) -grad_fd = ForwardDiff.gradient(losssum, params[1]) -grad_rd = ReverseDiff.gradient(losssum, params[1]) -abc = 1 - -import ReverseDiff: increment_deriv!, ZeroTangent -function ReverseDiff.increment_deriv!(::ReverseDiff.TrackedReal, ::ZeroTangent) - return nothing -end - -import DifferentialEquations.DiffEqBase: AbstractODEIntegrator -function Base.show(io::IO, ::MIME"text/plain", ::AbstractODEIntegrator) - print(io, "[AbstractODEIntegrator]") -end FMIFlux.checkSensalgs!(losssum, nfmu)