[WIP] feat: add FMUComponent

Project.toml

@@ -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"
@@ -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"
@@ -106,6 +108,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"
@@ -141,7 +144,7 @@ SpecialFunctions = "0.7, 0.8, 0.9, 0.10, 1.0, 2"
 StaticArrays = "0.10, 0.11, 0.12, 1.0"
 StochasticDiffEq = "6.72.1"
 StochasticDelayDiffEq = "1.8.1"
-SymbolicIndexingInterface = "0.3.36"
+SymbolicIndexingInterface = "0.3.37"
 SymbolicUtils = "3.7"
 Symbolics = "6.19"
 URIs = "1"

src/ModelingToolkit.jl

@@ -291,4 +291,6 @@ export MTKParameters, reorder_dimension_by_tunables!, reorder_dimension_by_tunab
 
 export HomotopyContinuationProblem
 
+function FMIComponent end
+
 end # module

src/parameters.jl

@@ -33,7 +33,12 @@ end
 
 function getcalledparameter(x)
     x = unwrap(x)
-    return getmetadata(x, CallWithParent)
+    # `parent` is a `CallWithMetadata` with the correct metadata,
+    # but no namespacing. `operation(x)` has the correct namespacing,
+    # but is not a `CallWithMetadata` and doesn't have any metadata.
+    # This approach combines both.
+    parent = getmetadata(x, CallWithParent)
+    return CallWithMetadata(operation(x), metadata(parent))
 end
 
 """

src/structural_transformation/symbolics_tearing.jl

@@ -593,7 +593,7 @@ function tearing_reassemble(state::TearingState, var_eq_matching,
     @set! sys.unknowns = unknowns
 
     obs, subeqs, deps = cse_and_array_hacks(
-        obs, subeqs, unknowns, neweqs; cse = cse_hack, array = array_hack)
+        sys, obs, subeqs, unknowns, neweqs; cse = cse_hack, array = array_hack)
 
     @set! sys.eqs = neweqs
     @set! sys.observed = obs
@@ -627,7 +627,7 @@ if all `p[i]` are present and the unscalarized form is used in any equation (obs
 not) we first count the number of times the scalarized form of each observed variable
 occurs in observed equations (and unknowns if it's split).
 """
-function cse_and_array_hacks(obs, subeqs, unknowns, neweqs; cse = true, array = true)
+function cse_and_array_hacks(sys, obs, subeqs, unknowns, neweqs; cse = true, array = true)
     # HACK 1
     # mapping of rhs to temporary CSE variable
     # `f(...) => tmpvar` in above example
@@ -696,6 +696,7 @@ function cse_and_array_hacks(obs, subeqs, unknowns, neweqs; cse = true, array =
                     tempvar; T = Symbolics.symtype(rhs_arr)))
                 tempvar = setmetadata(
                     tempvar, Symbolics.ArrayShapeCtx, Symbolics.shape(rhs_arr))
+                vars!(all_vars, rhs_arr)
                 tempeq = tempvar ~ rhs_arr
                 rhs_to_tempvar[rhs_arr] = tempvar
                 push!(obs, tempeq)
@@ -718,12 +719,16 @@ function cse_and_array_hacks(obs, subeqs, unknowns, neweqs; cse = true, array =
         Symbolics.shape(sym) != Symbolics.Unknown() || continue
         arg1 = arguments(sym)[1]
         cnt = get(arr_obs_occurrences, arg1, 0)
-        cnt == 0 && continue
         arr_obs_occurrences[arg1] = cnt + 1
     end
     for eq in neweqs
         vars!(all_vars, eq.rhs)
     end
+
+    # also count unscalarized variables used in callbacks
+    for ev in Iterators.flatten((continuous_events(sys), discrete_events(sys)))
+        vars!(all_vars, ev)
+    end
     obs_arr_eqs = Equation[]
     for (arrvar, cnt) in arr_obs_occurrences
         cnt == length(arrvar) || continue
@@ -737,7 +742,9 @@ function cse_and_array_hacks(obs, subeqs, unknowns, neweqs; cse = true, array =
         # try to `create_array(OffsetArray{...}, ...)` which errors.
         # `term(Origin(firstind), scal)` doesn't retain the `symtype` and `size`
         # of `scal`.
-        push!(obs_arr_eqs, arrvar ~ change_origin(Origin(firstind), scal))
+        rhs = scal
+        rhs = change_origin(firstind, rhs)
+        push!(obs_arr_eqs, arrvar ~ rhs)
     end
     append!(obs, obs_arr_eqs)
     append!(subeqs, obs_arr_eqs)
@@ -764,10 +771,10 @@ getindex_wrapper(x, i) = x[i...]
 
 # PART OF HACK 2
 function change_origin(origin, arr)
-    return origin(arr)
+    return Origin(origin)(arr)
 end
 
-@register_array_symbolic change_origin(origin::Origin, arr::AbstractArray) begin
+@register_array_symbolic change_origin(origin::Any, arr::AbstractArray) begin
     size = size(arr)
     eltype = eltype(arr)
     ndims = ndims(arr)

src/structural_transformation/utils.jl

@@ -218,11 +218,18 @@ function find_eq_solvables!(state::TearingState, ieq, to_rm = Int[], coeffs = no
     all_int_vars = true
     coeffs === nothing || empty!(coeffs)
     empty!(to_rm)
+
+    vars_buffer = Set()
     for j in 𝑠neighbors(graph, ieq)
         var = fullvars[j]
         isirreducible(var) && (all_int_vars = false; continue)
         a, b, islinear = linear_expansion(term, var)
         a, b = unwrap(a), unwrap(b)
+        vars!(vars_buffer, b)
+        if islinear && isequal(a, 0) && var in vars_buffer
+            islinear = false
+        end
+        empty!(vars_buffer)
         islinear || (all_int_vars = false; continue)
         a = ModelingToolkit.fold_constants(a)
         b = ModelingToolkit.fold_constants(b)

src/systems/callbacks.jl

@@ -77,6 +77,13 @@ function has_functional_affect(cb)
     (affects(cb) isa FunctionalAffect || affects(cb) isa ImperativeAffect)
 end
 
+function vars!(vars, aff::FunctionalAffect; op = Differential)
+    for var in Iterators.flatten((unknowns(aff), parameters(aff), discretes(aff)))
+        vars!(vars, var)
+    end
+    return vars
+end
+
 #################################### continuous events #####################################
 
 const NULL_AFFECT = Equation[]
@@ -333,6 +340,22 @@ function continuous_events(sys::AbstractSystem)
     filter(!isempty, cbs)
 end
 
+function vars!(vars, cb::SymbolicContinuousCallback; op = Differential)
+    for eq in equations(cb)
+        vars!(vars, eq; op)
+    end
+    for aff in (affects(cb), affect_negs(cb), initialize_affects(cb), finalize_affects(cb))
+        if aff isa Vector{Equation}
+            for eq in aff
+                vars!(vars, eq; op)
+            end
+        elseif aff !== nothing
+            vars!(vars, aff; op)
+        end
+    end
+    return vars
+end
+
 #################################### discrete events #####################################
 
 struct SymbolicDiscreteCallback
@@ -469,6 +492,28 @@ function discrete_events(sys::AbstractSystem)
     cbs
 end
 
+function vars!(vars, cb::SymbolicDiscreteCallback; op = Differential)
+    if symbolic_type(cb.condition) == NotSymbolic
+        if cb.condition isa AbstractArray
+            for eq in cb.condition
+                vars!(vars, eq; op)
+            end
+        end
+    else
+        vars!(vars, cb.condition; op)
+    end
+    for aff in (cb.affects, cb.initialize, cb.finalize)
+        if aff isa Vector{Equation}
+            for eq in aff
+                vars!(vars, eq; op)
+            end
+        elseif aff !== nothing
+            vars!(vars, aff; op)
+        end
+    end
+    return vars
+end
+
 ################################# compilation functions ####################################
 
 # handles ensuring that affect! functions work with integrator arguments

src/systems/diffeqs/abstractodesystem.jl

@@ -1369,11 +1369,13 @@ function InitializationProblem{iip, specialize}(sys::AbstractSystem,
     for sym in unknowns(isys)
         haskey(fullmap, sym) || continue
         symbolic_type(fullmap[sym]) == NotSymbolic() || continue
+        is_array_of_symbolics(fullmap[sym]) && continue
         u0T = promote_type(u0T, typeof(fullmap[sym]))
     end
     for eq in observed(isys)
         haskey(fullmap, eq.lhs) || continue
         symbolic_type(fullmap[eq.lhs]) == NotSymbolic() || continue
+        is_array_of_symbolics(fullmap[eq.lhs]) && continue
         u0T = promote_type(u0T, typeof(fullmap[eq.lhs]))
     end
     if u0T != Union{}

src/systems/imperative_affect.jl

@@ -38,7 +38,7 @@ in the returned tuple, in which case the associated field will not be updated.
     skip_checks::Bool
 end
 
-function ImperativeAffect(f::Function;
+function ImperativeAffect(f;
         observed::NamedTuple = NamedTuple{()}(()),
         modified::NamedTuple = NamedTuple{()}(()),
         ctx = nothing,
@@ -48,18 +48,18 @@ function ImperativeAffect(f::Function;
         collect(values(modified)), collect(keys(modified)),
         ctx, skip_checks)
 end
-function ImperativeAffect(f::Function, modified::NamedTuple;
+function ImperativeAffect(f, modified::NamedTuple;
         observed::NamedTuple = NamedTuple{()}(()), ctx = nothing, skip_checks = false)
     ImperativeAffect(
         f, observed = observed, modified = modified, ctx = ctx, skip_checks = skip_checks)
 end
 function ImperativeAffect(
-        f::Function, modified::NamedTuple, observed::NamedTuple; ctx = nothing, skip_checks = false)
+        f, modified::NamedTuple, observed::NamedTuple; ctx = nothing, skip_checks = false)
     ImperativeAffect(
         f, observed = observed, modified = modified, ctx = ctx, skip_checks = skip_checks)
 end
 function ImperativeAffect(
-        f::Function, modified::NamedTuple, observed::NamedTuple, ctx; skip_checks = false)
+        f, modified::NamedTuple, observed::NamedTuple, ctx; skip_checks = false)
     ImperativeAffect(
         f, observed = observed, modified = modified, ctx = ctx, skip_checks = skip_checks)
 end
@@ -216,3 +216,20 @@ function compile_user_affect(affect::ImperativeAffect, cb, sys, dvs, ps; kwargs.
 end
 
 scalarize_affects(affects::ImperativeAffect) = affects
+
+function vars!(vars, aff::ImperativeAffect; op = Differential)
+    for var in Iterators.flatten((observed(aff), modified(aff)))
+        if symbolic_type(var) == NotSymbolic()
+            if var isa AbstractArray
+                for v in var
+                    v = unwrap(v)
+                    vars!(vars, v)
+                end
+            end
+        else
+            var = unwrap(var)
+            vars!(vars, var)
+        end
+    end
+    return vars
+end

src/systems/index_cache.jl

@@ -299,6 +299,9 @@ function IndexCache(sys::AbstractSystem)
                 for v in vs
                     if (idx = get(disc_idxs, v, nothing)) !== nothing
                         push!(timeseries, idx.clock_idx)
+                    elseif iscall(v) && operation(v) === getindex &&
+                           (idx = get(disc_idxs, arguments(v)[1], nothing)) !== nothing
+                        push!(timeseries, idx.clock_idx)
                     elseif haskey(observed_syms_to_timeseries, v)
                         union!(timeseries, observed_syms_to_timeseries[v])
                     elseif haskey(dependent_pars_to_timeseries, v)

src/systems/nonlinear/initializesystem.jl

@@ -186,6 +186,14 @@ function generate_initializesystem(sys::AbstractSystem;
         defs[eq.lhs] = eq.rhs
     end
 
+    # even if `p => tovar(p)` is in `paramsubs`, `isparameter(p[1]) === true` after substitution
+    # so add scalarized versions as well
+    for k in collect(keys(paramsubs))
+        symbolic_type(k) == ArraySymbolic() || continue
+        for i in eachindex(k)
+            paramsubs[k[i]] = paramsubs[k][i]
+        end
+    end
     eqs_ics = Symbolics.substitute.([eqs_ics; trueobs], (paramsubs,))
     vars = [vars; collect(values(paramsubs))]
     for k in keys(defs)

src/systems/problem_utils.jl

@@ -88,6 +88,18 @@ function symbols_to_symbolics!(sys::AbstractSystem, varmap::AbstractDict)
     end
 end
 
+"""
+    $(TYPEDSIGNATURES)
+
+Utility function to get the value `val` corresponding to key `var` in `varmap`, and
+return `getindex(val, idx)` if it exists or `nothing` otherwise.
+"""
+function get_and_getindex(varmap, var, idx)
+    val = get(varmap, var, nothing)
+    val === nothing && return nothing
+    return val[idx]
+end
+
 """
     $(TYPEDSIGNATURES)
 
@@ -115,8 +127,9 @@ function add_fallbacks!(
             val = map(eachindex(var)) do idx
                 # @something is lazy and saves from writing a massive if-elseif-else
                 @something(get(varmap, var[idx], nothing),
-                    get(varmap, ttvar[idx], nothing), get(fallbacks, var, nothing)[idx],
-                    get(fallbacks, ttvar, nothing)[idx], get(fallbacks, var[idx], nothing),
+                    get(varmap, ttvar[idx], nothing), get_and_getindex(fallbacks, var, idx),
+                    get_and_getindex(fallbacks, ttvar, idx), get(
+                        fallbacks, var[idx], nothing),
                     get(fallbacks, ttvar[idx], nothing), Some(nothing))
             end
             # only push the missing entries
@@ -561,6 +574,10 @@ function maybe_build_initialization_problem(
             p = unwrap(p)
             stype = symtype(p)
             op[p] = get_temporary_value(p)
+            if iscall(p) && operation(p) === getindex
+                arrp = arguments(p)[1]
+                op[arrp] = collect(arrp)
+            end
         end
 
         for v in missing_unknowns

src/systems/systems.jl

@@ -57,6 +57,7 @@ function structural_simplify(
         ks = collect(keys(defs))  # take copy to avoid mutating defs while iterating.
         for k in ks
             if Symbolics.isarraysymbolic(k) && Symbolics.shape(k) !== Symbolics.Unknown()
+                defs[k] === missing && continue
                 for i in eachindex(k)
                     defs[k[i]] = defs[k][i]
                 end