Merge pull request #234 from JuliaDynamics/hw/cse_observed

Author: hexaeder

Project.toml

@@ -17,6 +17,7 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Mixers = "2a8e4939-dab8-5edc-8f64-72a8776f13de"
 NNlib = "872c559c-99b0-510c-b3b7-b6c96a88d5cd"
 NonlinearSolve = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
+OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d"
 Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PreallocationTools = "d236fae5-4411-538c-8e31-a6e3d9e00b46"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
@@ -37,11 +38,12 @@ Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
+SymbolicUtils = "d1185830-fcd6-423d-90d6-eec64667417b"
 Symbolics = "0c5d862f-8b57-4792-8d23-62f2024744c7"
 
 [extensions]
 NetworkDynamicsCUDAExt = ["CUDA", "Adapt"]
-NetworkDynamicsMTKExt = ["ModelingToolkit", "Symbolics"]
+NetworkDynamicsMTKExt = ["ModelingToolkit", "SymbolicUtils", "Symbolics"]
 NetworkDynamicsSymbolicsExt = ["Symbolics", "MacroTools"]
 
 [compat]
@@ -62,6 +64,7 @@ Mixers = "0.1.2"
 ModelingToolkit = "9.67"
 NNlib = "0.9.13"
 NonlinearSolve = "4"
+OrderedCollections = "1.8.0"
 Polyester = "0.7.12"
 PreallocationTools = "0.4.23"
 PrecompileTools = "1.2.1"
@@ -75,6 +78,7 @@ StaticArrays = "1.9.4"
 SteadyStateDiffEq = "2.2.0"
 StyledStrings = "1.0.3"
 SymbolicIndexingInterface = "0.3.27"
+SymbolicUtils = "3.24"
 Symbolics = "6.19.0"
 TimerOutputs = "0.5.23"
 julia = "1.10"

ext/NetworkDynamicsMTKExt.jl

@@ -2,12 +2,14 @@ module NetworkDynamicsMTKExt
 
 using ModelingToolkit: Symbolic, iscall, operation, arguments, build_function
 using ModelingToolkit: ModelingToolkit, Equation, ODESystem, Differential
-using ModelingToolkit: full_equations, get_variables, structural_simplify, getname, unwrap
+using ModelingToolkit: equations, full_equations, get_variables, structural_simplify, getname, unwrap
 using ModelingToolkit: parameters, unknowns, independent_variables, observed, defaults
 using Symbolics: Symbolics, fixpoint_sub, substitute
 using RecursiveArrayTools: RecursiveArrayTools
 using ArgCheck: @argcheck
 using LinearAlgebra: Diagonal, I
+using SymbolicUtils.Code: Let, Assignment
+using OrderedCollections: OrderedDict
 
 using NetworkDynamics: NetworkDynamics, set_metadata!,
                        PureFeedForward, FeedForward, NoFeedForward, PureStateMap
@@ -299,8 +301,10 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
     params = setdiff(allparams, Set(allinputs))
 
     # extract the main equations and observed equations
-    eqs::Vector{Equation} = ModelingToolkit.subs_constants(full_equations(sys))
+    eqs::Vector{Equation} = ModelingToolkit.subs_constants(equations(sys))
+    obseqs_sorted::Vector{Equation} = ModelingToolkit.subs_constants(observed(sys))
     fix_metadata!(eqs, sys);
+    fix_metadata!(obseqs_sorted, sys);
 
     # assert the ordering of states and equations
     explicit_states = Symbolic[eq_type(eq)[2] for eq in eqs if !isnothing(eq_type(eq)[2])]
@@ -311,8 +315,8 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
     end
 
     # check hat there are no rhs differentials in the equations
-    if !isempty(rhs_differentials(eqs))
-        diffs = rhs_differentials(eqs)
+    if !isempty(rhs_differentials(vcat(eqs, obseqs_sorted)))
+        diffs = rhs_differentials(vcat(eqs, obseqs_sorted))
         buf = IOBuffer()
         println(buf, "Equations contain differentials in their rhs: ", diffs)
         # for (i, eqs) in enumerate(eqs)
@@ -323,23 +327,17 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
         throw(ArgumentError(String(take!(buf))))
     end
 
-    # extract observed equations. They might depend on eachother so resolve them
-    obs_subs = Dict(eq.lhs => eq.rhs for eq in observed(sys))
-    obseqs = map(observed(sys)) do eq
-        expanded_rhs = fixpoint_sub(eq.rhs, obs_subs)
-        eq.lhs ~ ModelingToolkit.subs_constants(expanded_rhs)
-    end
-    fix_metadata!(obseqs, sys);
     # obs can only depend on parameters (including allinputs) or states
-    obs_deps = _all_rhs_symbols(obseqs)
+    obs_subs = OrderedDict(eq.lhs => eq.rhs for eq in obseqs_sorted)
+    obs_deps = _all_rhs_symbols(fixpoint_sub(obseqs_sorted, obs_subs))
     if !(obs_deps ⊆ Set(allparams) ∪ Set(states) ∪ independent_variables(sys))
         @warn "obs_deps !⊆ parameters ∪ unknowns. Difference: $(setdiff(obs_deps, Set(allparams) ∪ Set(states)))"
     end
 
     # if some states shadow outputs (out ~ state in observed)
     # switch their names. I.e. prioritize use of name `out`
     renamings = Dict()
-    for eq in obseqs
+    for eq in obseqs_sorted
         if eq.lhs ∈ Set(alloutputs) && iscall(eq.rhs) &&
             operation(eq.rhs) isa Symbolics.BasicSymbolic && eq.rhs ∈ Set(states)
             verbose && @info "Encountered trivial equation $eq. Swap out $(eq.lhs) <=> $(eq.rhs) everywhere."
@@ -349,36 +347,37 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
     end
     if !isempty(renamings)
         eqs = map(eq -> substitute(eq, renamings), eqs)
-        obseqs = map(eq -> substitute(eq, renamings), obseqs)
+        obseqs_sorted = map(eq -> substitute(eq, renamings), obseqs_sorted)
+        obs_subs = OrderedDict(eq.lhs => eq.rhs for eq in obseqs_sorted)
         states = map(s -> substitute(s, renamings), states)
         verbose && @info "New States:" states
     end
 
-    # find the output equations, this might remove them from obseqs!
+    # find the output equations, this might remove them from obseqs_sorted (obs_subs stays intact)
     outeqs = Equation[]
     for out in Iterators.flatten(outputss)
         if out ∈ Set(states)
             push!(outeqs, out ~ out)
-        else
-            idx = findfirst(eq -> isequal(eq.lhs, out), obseqs)
-            if isnothing(idx)
-                throw(ArgumentError("Output $out was neither foundin states nor in observed equations."))
-            end
-            eq = obseqs[idx]
-            if !isempty(rhs_differentials(eq))
-                println(obs_subs[out])
-                throw(ArgumentError("Algebraic FF equation for output $out contains differentials in the RHS: $(rhs_differentials(eq))"))
-            end
-            deleteat!(obseqs, idx)
-
-            if ff_to_constraint && !isempty(get_variables(eq.rhs) ∩ allinputs)
+        elseif out ∈ keys(obs_subs)
+            # if its a observed, we need to check for ff behavior
+            fulleq = out ~ fixpoint_sub(obs_subs[out], obs_subs)
+            if ff_to_constraint && !isempty(get_variables(fulleq.rhs) ∩ allinputs)
                 verbose && @info "Output $out would lead to FF in g, promote to state instead."
-                push!(eqs, 0 ~ eq.lhs - eq.rhs)
-                push!(states, eq.lhs)
-                push!(outeqs, eq.lhs ~ eq.lhs)
-            else
-                push!(outeqs, eq)
+                # not observed anymore, delete from observed and put in equations
+                push!(eqs, 0 ~ out - obs_subs[out])
+                push!(states, out)
+                push!(outeqs, out ~ out)
+
+                deleteat!(obseqs_sorted, findfirst(eq -> isequal(eq.lhs, out), obseqs_sorted))
+                delete!(obs_subs, out)
+            else # "normal" observed state
+                push!(outeqs, out ~ obs_subs[out])
+                # delete from obs equations but *not* from obs_subs (otherwise can't be reference)
+                # in equations
+                deleteat!(obseqs_sorted, findfirst(eq -> isequal(eq.lhs, out), obseqs_sorted))
             end
+        else
+            throw(ArgumentError("Output $out was neither foundin states nor in observed equations."))
         end
     end
 
@@ -399,26 +398,28 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
     end
 
     iv = only(independent_variables(sys))
-    out_deps = _all_rhs_symbols(outeqs)
+
+    out_deps = _all_rhs_symbols(fixpoint_sub(outeqs, obs_subs))
     fftype = _determine_fftype(out_deps, states, allinputs, params, iv)
 
     # filter out unnecessary parameters
-    var_deps = _all_rhs_symbols(eqs)
+    var_deps = _all_rhs_symbols(fixpoint_sub(eqs, obs_subs))
     unused_params = Set(setdiff(params, (var_deps ∪ out_deps))) # do not exclud obs_deps
     if verbose && !isempty(unused_params)
         @info "Parameters $(unused_params) do not appear in equations of f and g and will be marked as unused."
     end
 
     # TODO: explore Symbolcs/SymbolicUtils CSE
     # now generate the actual functions
-    formulas = [eq.rhs for eq in eqs]
-    if !isempty(formulas)
+    if !isempty(eqs)
+        formulas = _get_formulas(eqs, obs_subs)
         _, f_ip = build_function(formulas, states, inputss..., params, iv; cse=false, expression)
     else
         f_ip = nothing
     end
 
-    gformulas = [eq.rhs for eq in outeqs]
+    # find all observable assigments necessary for outeqs
+    gformulas = _get_formulas(outeqs, obs_subs)
     gformargs = if fftype isa PureFeedForward
         (inputss..., params, iv)
     elseif fftype isa FeedForward
@@ -436,10 +437,13 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
         MultipleOutputWrapper{fftype, length(outputss), typeof(_g_ip)}(_g_ip)
     end
 
-    # and the observed functions
-    obsstates = [eq.lhs for eq in obseqs]
-    obsformulas = [eq.rhs for eq in obseqs]
-    _, obsf_ip = build_function(obsformulas, states, inputss..., params, iv; cse=false, expression)
+    obsstates = [eq.lhs for eq in obseqs_sorted]
+    if !isempty(obsstates)
+        obsformulas = _get_formulas([s ~ s for s in obsstates], obs_subs)
+        _, obsf_ip = build_function(obsformulas, states, inputss..., params, iv; cse=false, expression)
+    else
+        obsf_ip = nothing
+    end
 
     return (;
         f=f_ip, g=g_ip,
@@ -450,15 +454,60 @@ function generate_io_function(_sys, inputss::Tuple, outputss::Tuple;
         obsstates,
         fftype,
         obsf = obsf_ip,
-        equations=formulas,
-        outputeqs=Dict(Iterators.flatten(outputss) .=> gformulas),
-        observed=Dict(getname.(obsstates) .=> obsformulas),
+        equations=eqs,
+        outputeqs=outeqs,
+        observed=obseqs_sorted,
         odesystem=sys,
         params,
         unused_params
     )
 end
 
+function _get_formulas(eqs, obs_subs)
+    # Bit hacky, were building a function like this,
+    # where all (necessary) obs and eqs are contained in the bgin block of the first output
+    # out[1] = begin
+    #     obs1   = ...
+    #     obs2   = ...
+    #     ...
+    #     state1 = ...
+    #     state2 = ...
+    #     ...
+    #     state1 # ens up in out[1]
+    # end
+    # out[2] = state2
+    # ...
+    isempty(eqs) && return []
+    obsdeps = _collect_deps_on_obs([eq.rhs for eq in eqs], obs_subs)
+    obs_assignments = [Assignment(k, v) for (k,v) in obs_subs if k ∈ obsdeps]
+
+    # implicit equations are not use via assigments, so we filter for e
+    eqs_assignments = [Assignment(eq.lhs, eq.rhs) for eq in eqs
+                          if !isequal(eq.lhs, eq.rhs) && !isequal(eq.lhs, 0)]
+    # since implicit eqs did not end up in assighmets, we use the rhs
+    out = [isequal(eq.lhs, 0) ? eq.rhs : eq.lhs for eq in eqs]
+
+    [Let(vcat(obs_assignments, eqs_assignments), out[1], false), out[2:end]...]
+end
+function _collect_deps_on_obs(terms, obs_subs)
+    deps = Set{Symbolic}()
+    for term in terms
+        _collect_deps_on_obs!(deps, obs_subs, term)
+    end
+    deps
+end
+function _collect_deps_on_obs!(deps, obs_subs, term)
+    termdeps = get_variables(term)
+    for sym in termdeps
+        if haskey(obs_subs, sym)
+            # check recursively whether the observed depends on other observed
+            _collect_deps_on_obs!(deps, obs_subs, obs_subs[sym])
+            push!(deps, sym)
+        end
+    end
+    deps
+end
+
 function _determine_fftype(deps, states, allinputs, params, t)
     if isempty(allinputs ∩ deps) # no ff path
         if isempty(params ∩ deps) && !(t ∈ deps) # no p nor t