Skip to content

Commit

Permalink
feat: add FMUComponent with support for v2 ME FMUs
Browse files Browse the repository at this point in the history
  • Loading branch information
AayushSabharwal committed Dec 20, 2024
1 parent 4792360 commit fa246ed
Show file tree
Hide file tree
Showing 3 changed files with 220 additions and 1 deletion.
5 changes: 4 additions & 1 deletion Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -64,6 +64,7 @@ Unitful = "1986cc42-f94f-5a68-af5c-568840ba703d"
BifurcationKit = "0f109fa4-8a5d-4b75-95aa-f515264e7665"
ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
DeepDiffs = "ab62b9b5-e342-54a8-a765-a90f495de1a6"
FMI = "14a09403-18e3-468f-ad8a-74f8dda2d9ac"
HomotopyContinuation = "f213a82b-91d6-5c5d-acf7-10f1c761b327"
InfiniteOpt = "20393b10-9daf-11e9-18c9-8db751c92c57"
LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
Expand All @@ -72,9 +73,10 @@ LabelledArrays = "2ee39098-c373-598a-b85f-a56591580800"
MTKBifurcationKitExt = "BifurcationKit"
MTKChainRulesCoreExt = "ChainRulesCore"
MTKDeepDiffsExt = "DeepDiffs"
MTKFMIExt = "FMI"
MTKHomotopyContinuationExt = "HomotopyContinuation"
MTKLabelledArraysExt = "LabelledArrays"
MTKInfiniteOptExt = "InfiniteOpt"
MTKLabelledArraysExt = "LabelledArrays"

[compat]
AbstractTrees = "0.3, 0.4"
Expand Down Expand Up @@ -105,6 +107,7 @@ FindFirstFunctions = "1"
ForwardDiff = "0.10.3"
FunctionWrappers = "1.1"
FunctionWrappersWrappers = "0.1"
FMI = "0.14"
Graphs = "1.5.2"
HomotopyContinuation = "2.11"
InfiniteOpt = "0.5"
Expand Down
214 changes: 214 additions & 0 deletions ext/MTKFMIExt.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,214 @@
module MTKFMIExt

using ModelingToolkit
using ModelingToolkit: t_nounits as t, D_nounits as D
import ModelingToolkit as MTK
import FMI

macro statuscheck(expr)
@assert Meta.isexpr(expr, :call)
fn = expr.args[1]
@assert Meta.isexpr(fn, :.)
@assert fn.args[1] == :FMI
fnname = fn.args[2]

instance = expr.args[2]

return quote
status = $expr
fnname = $fnname
if (status isa Tuple && status[1] == FMI.fmi2True) ||
(!(status isa Tuple) && status != FMI.fmi2StatusOK &&
status != FMI.fmi2StatusWarning)
if status != FMI.fmi2StatusFatal
FMI.fmi2Terminate(wrapper.instance)
end
FMI.fmi2FreeInstance!(wrapper.instance)
wrapper.instance = nothing
error("FMU Error: status $status")
end
end |> esc
end

function MTK.FMIComponent(::Val{2}, ::Val{:ME}; fmu = nothing, tolerance = 1e-6, name)
value_references = Dict()
defs = Dict()
states = []
diffvars = []
observed = Equation[]
stateT = Float64
for (valRef, snames) in FMI.getStateValueReferencesAndNames(fmu)
stateT = FMI.dataTypeForValueReference(fmu, valRef)
snames = map(parseFMIVariableName, snames)
vars = [MTK.unwrap(only(@variables $sname(t)::stateT)) for sname in snames]
for i in eachindex(vars)
if i == 1
push!(diffvars, vars[i])
else
push!(observed, vars[i] ~ vars[1])
end
value_references[vars[i]] = valRef
end
append!(states, vars)
end

inputs = []
for (valRef, snames) in FMI.getInputValueReferencesAndNames(fmu)
snames = map(parseFMIVariableName, snames)
vars = [MTK.unwrap(only(@variables $sname(t)::stateT)) for sname in snames]
for i in eachindex(vars)
if i == 1
push!(inputs, vars[i])
else
push!(observed, vars[i] ~ vars[1])
end
value_references[vars[i]] = valRef
end
append!(states, vars)
end

outputs = []
for (valRef, snames) in FMI.getOutputValueReferencesAndNames(fmu)
snames = map(parseFMIVariableName, snames)
vars = [MTK.unwrap(only(@variables $sname(t)::stateT)) for sname in snames]
for i in eachindex(vars)
if i == 1
push!(outputs, vars[i])
else
push!(observed, vars[i] ~ vars[1])
end
value_references[vars[i]] = valRef
end
append!(states, vars)
end

params = []
parameter_dependencies = Equation[]
for (valRef, pnames) in FMI.getParameterValueReferencesAndNames(fmu)
defval = FMI.getStartValue(fmu, valRef)
T = FMI.dataTypeForValueReference(fmu, valRef)
pnames = map(parseFMIVariableName, pnames)
vars = [MTK.unwrap(only(@parameters $pname::T)) for pname in pnames]
for i in eachindex(vars)
if i == 1
push!(params, vars[i])
else
push!(parameter_dependencies, vars[i] ~ vars[1])
end
value_references[vars[i]] = valRef
end
defs[vars[1]] = defval
end

input_value_references = UInt32[value_references[var] for var in inputs]
param_value_references = UInt32[value_references[var] for var in params]
@parameters wrapper::FMI2InstanceWrapper = FMI2InstanceWrapper(
fmu, param_value_references, input_value_references, tolerance)

output_value_references = UInt32[value_references[var] for var in outputs]
buffer_length = length(diffvars) + length(outputs)
_functor = FMI2MEFunctor(zeros(buffer_length), output_value_references)
@parameters (functor::(typeof(_functor)))(..)[1:buffer_length] = _functor
call_expr = functor(wrapper, copy(diffvars), copy(inputs), copy(params), t)

diffeqs = Equation[]
for (i, var) in enumerate([D.(diffvars); outputs])
push!(diffeqs, var ~ call_expr[i])
end

finalize_affect = MTK.FunctionalAffect(fmi2MEFinalize!, [], [wrapper], [])
step_affect = MTK.FunctionalAffect(fmi2MEStep!, [], [wrapper], [])
instance_management_callback = MTK.SymbolicDiscreteCallback(
(t != t - 1), step_affect; finalize = finalize_affect)

push!(params, wrapper, functor)
eqs = [observed; diffeqs]
return ODESystem(eqs, t, states, params; parameter_dependencies, defaults = defs,
discrete_events = [instance_management_callback], name)
end

function parseFMIVariableName(name::AbstractString)
return Symbol(replace(name, "." => "__"))
end

mutable struct FMI2InstanceWrapper
const fmu::FMI.FMU2
const param_value_references::Vector{UInt32}
const input_value_references::Vector{UInt32}
const tolerance::FMI.fmi2Real
instance::Union{FMI.FMU2Component, Nothing}
end

function FMI2InstanceWrapper(fmu, params, inputs, tolerance)
FMI2InstanceWrapper(fmu, params, inputs, tolerance, nothing)
end

function get_instance!(wrapper::FMI2InstanceWrapper, states, inputs, params, t)
if wrapper.instance === nothing
wrapper.instance = FMI.fmi2Instantiate!(wrapper.fmu)::FMI.FMU2Component
if !isempty(params)
@statuscheck FMI.fmi2SetReal(wrapper.instance, wrapper.param_value_references,
Csize_t(length(wrapper.param_value_references)), params)
end
@statuscheck FMI.fmi2SetupExperiment(
wrapper.instance, FMI.fmi2True, wrapper.tolerance, t, FMI.fmi2False, t)
@statuscheck FMI.fmi2EnterInitializationMode(wrapper.instance)
if !isempty(inputs)
@statuscheck FMI.fmi2SetReal(wrapper.instance, wrapper.input_value_references,
Csize_t(length(wrapper.param_value_references)), inputs)
end
@statuscheck FMI.fmi2ExitInitializationMode(wrapper.instance)
eventInfo = FMI.fmi2NewDiscreteStates(wrapper.instance)
@assert eventInfo.newDiscreteStatesNeeded == FMI.fmi2False
# TODO: Support FMU events
@statuscheck FMI.fmi2EnterContinuousTimeMode(wrapper.instance)
end
instance = wrapper.instance
@statuscheck FMI.fmi2SetTime(instance, t)
@statuscheck FMI.fmi2SetContinuousStates(instance, states)
if !isempty(inputs)
@statuscheck FMI.fmi2SetReal(instance, wrapper.input_value_references,
Csize_t(length(wrapper.param_value_references)), inputs)
end

return instance
end

function complete_step!(wrapper::FMI2InstanceWrapper)
wrapper.instance === nothing && return
@statuscheck FMI.fmi2CompletedIntegratorStep(wrapper.instance, FMI.fmi2True)
end

function reset_instance!(wrapper::FMI2InstanceWrapper)
wrapper.instance === nothing && return
@statuscheck FMI.fmi2Reset(wrapper.instance)
end

struct FMI2MEFunctor{T}
return_buffer::Vector{T}
output_value_references::Vector{UInt32}
end

function (fn::FMI2MEFunctor)(wrapper::FMI2InstanceWrapper, states, inputs, params, t)
instance = get_instance!(wrapper, states, inputs, params, t)

states_buffer = zeros(length(states))
@statuscheck FMI.fmi2GetDerivatives!(instance, states_buffer)
outputs_buffer = zeros(length(fn.output_value_references))
FMI.fmi2GetReal!(instance, fn.output_value_references, outputs_buffer)
return [states_buffer; outputs_buffer]
end

function fmi2MEStep!(integrator, u, p, ctx)
wrapper_idx = p[1]
wrapper = integrator.ps[wrapper_idx]
complete_step!(wrapper)
end

function fmi2MEFinalize!(integrator, u, p, ctx)
wrapper_idx = p[1]
wrapper = integrator.ps[wrapper_idx]
reset_instance!(wrapper)
end

end # module
2 changes: 2 additions & 0 deletions src/ModelingToolkit.jl
Original file line number Diff line number Diff line change
Expand Up @@ -291,4 +291,6 @@ export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunab

export HomotopyContinuationProblem

function FMIComponent end

end # module

0 comments on commit fa246ed

Please sign in to comment.