Merge pull request #751 from SciML/ms/observed_lowering

Convert observed equations as assignments when lowering

Author: YingboMa

src/build_function.jl

@@ -159,14 +159,16 @@ function Base.getproperty(sys::AbstractSystem, name::Symbol)
     throw(error("Variable $name does not exist"))
 end
 
-renamespace(namespace,name) = Symbol(namespace,:₊,name)
-
-function renamespace(namespace, x::Sym)
-    Sym{symtype(x)}(renamespace(namespace,x.name))
-end
-
-function renamespace(namespace, x::Term)
-    renamespace(namespace, operation(x))(arguments(x)...)
+function renamespace(namespace, x)
+    if x isa Num
+        renamespace(namespace, value(x))
+    elseif istree(x)
+        renamespace(namespace, operation(x))(arguments(x)...)
+    elseif x isa Sym
+        Sym{symtype(x)}(renamespace(namespace,nameof(x)))
+    else
+        Symbol(namespace,:₊,x)
+    end
 end
 
 function namespace_variables(sys::AbstractSystem)
@@ -213,7 +215,7 @@ end
 independent_variable(sys::AbstractSystem) = sys.iv
 function states(sys::AbstractSystem)
     unique(isempty(sys.systems) ?
-           setdiff(sys.states, value.(sys.pins)) :
+           sys.states :
            [sys.states;reduce(vcat,namespace_variables.(sys.systems))])
 end
 parameters(sys::AbstractSystem) = isempty(sys.systems) ? sys.ps : [sys.ps;reduce(vcat,namespace_parameters.(sys.systems))]

src/systems/abstractsystem.jl

@@ -69,13 +69,20 @@ end
 
 function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
     # optimization
-    dvs′ = makesym.(value.(dvs), states=dvs)
-    ps′ = makesym.(value.(ps), states=dvs)
+    obsvars = map(eq->eq.lhs, observed(sys))
+    fulldvs = [dvs; obsvars]
+    fulldvs′ = makesym.(value.(fulldvs))
 
-    sub = Dict(dvs .=> dvs′)
+    sub = Dict(fulldvs .=> fulldvs′)
     # substitute x(t) by just x
     rhss = [substitute(deq.rhs, sub) for deq ∈ equations(sys)]
-    return build_function(rhss, dvs′, ps′, sys.iv;
+    obss = [makesym(value(eq.lhs)) ~ substitute(eq.rhs, sub) for eq ∈ observed(sys)]
+
+    dvs′ = fulldvs′[1:length(dvs)]
+    ps′ = makesym.(value.(ps), states=())
+
+    # TODO: add an optional check on the ordering of observed equations
+    return build_function(Let(obss, rhss), dvs′, ps′, sys.iv;
                           conv = ODEToExpr(sys),kwargs...)
 end
 

src/systems/diffeqs/abstractodesystem.jl

@@ -168,7 +168,7 @@ end
 
 function collect_var!(states, parameters, var, iv)
     isequal(var, iv) && return nothing
-    if isparameter(var) || isparameter(operation(var))
+    if isparameter(var) || (istree(var) && isparameter(operation(var)))
         push!(parameters, var)
     else
         push!(states, var)

src/systems/diffeqs/odesystem.jl

@@ -8,6 +8,7 @@ function flatten(sys::ODESystem)
                          independent_variable(sys),
                          states(sys),
                          parameters(sys),
+                         pins=pins(sys),
                          observed=observed(sys))
     end
 end
@@ -56,8 +57,6 @@ end
 
 function alias_elimination(sys::ODESystem)
     eqs = vcat(equations(sys), observed(sys))
-    neweqs = Equation[]; sizehint!(neweqs, length(eqs))
-    subs = Pair[]
     diff_vars = filter(!isnothing, map(eqs) do eq
             if isdiffeq(eq)
                 arguments(eq.lhs)[1]
@@ -67,6 +66,9 @@ function alias_elimination(sys::ODESystem)
         end) |> Set
 
     deps = Set()
+    subs = Pair[]
+    neweqs = Equation[]; sizehint!(neweqs, length(eqs))
+
     for (i, eq) in enumerate(eqs)
         # only substitute when the variable is algebraic
         if isdiffeq(eq)
@@ -107,15 +109,12 @@ function alias_elimination(sys::ODESystem)
         end
     end
 
-    eqs′ = substitute_aliases(neweqs, Dict(subs))
-
     alias_vars = first.(subs)
     sys_states = states(sys)
-    alias_eqs = alias_vars .~ last.(subs)
-    #alias_eqs = topsort_equations(alias_eqs, sys_states)
+    alias_eqs = topsort_equations(alias_vars .~ last.(subs), sys_states)
 
     newstates = setdiff(sys_states, alias_vars)
-    ODESystem(eqs′, sys.iv, newstates, parameters(sys), observed=alias_eqs)
+    ODESystem(neweqs, sys.iv, newstates, parameters(sys), observed=alias_eqs)
 end
 
 """

src/systems/reduction.jl

@@ -222,7 +222,7 @@ Symbol("z⦗t⦘")
 """
 function tosymbol(t::Term; states=nothing, escape=true)
     if operation(t) isa Sym
-        if states !== nothing && !(any(isequal(t), states))
+        if states !== nothing && !(t in states)
             return nameof(operation(t))
         end
         op = nameof(operation(t))

src/utils.jl

@@ -29,99 +29,113 @@ D = Differential(t)
 
 test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynorm=true))
 
-eqs = [D(x) ~ σ*(y-x),
-       D(y) ~ x*(ρ-z)-y + β,
-       0 ~ sin(z) - x + y,
-       sin(u) ~ x + y,
-       2β ~ 2,
-       x ~ a,
+eqs = [
+       D(x) ~ σ*(y-x)
+       D(y) ~ x*(ρ-z)-y + β
+       0 ~ sin(z) - x + y
+       sin(u) ~ x + y
+       x ~ a
       ]
 
-lorenz1 = ODESystem(eqs,t,[u,x,y,z,a],[σ,ρ,β],name=:lorenz1)
+lorenz1 = ODESystem(eqs,t,name=:lorenz1)
 
 lorenz1_aliased = alias_elimination(lorenz1)
 reduced_eqs = [
                D(x) ~ σ * (y - x),
-               D(y) ~ x*(ρ-z)-y + 1,
+               D(y) ~ x*(ρ-z)-y + β,
                0 ~ sin(z) - x + y,
                0 ~ x + y - sin(u),
               ]
 test_equal.(equations(lorenz1_aliased), reduced_eqs)
-test_equal.(states(lorenz1_aliased), [u, x, y, z])
+@test isempty(setdiff(states(lorenz1_aliased), [u, x, y, z]))
 test_equal.(observed(lorenz1_aliased), [
-    β ~ 1,
     a ~ x,
 ])
 
 # Multi-System Reduction
 
+@variables s
 eqs1 = [
         D(x) ~ σ*(y-x) + F,
         D(y) ~ x*(ρ-z)-u,
         D(z) ~ x*y - β*z,
         u ~ x + y - z,
        ]
 
-lorenz1 = ODESystem(eqs1,pins=[F],name=:lorenz1)
-
-eqs2 = [
-        D(x) ~ F,
-        D(y) ~ x*(ρ-z)-x,
-        D(z) ~ x*y - β*z,
-        u ~ x - y - z
-       ]
+lorenz = name -> ODESystem(eqs1,t,pins=[F],name=name)
+lorenz1 = lorenz(:lorenz1)
+lorenz2 = lorenz(:lorenz2)
 
-lorenz2 = ODESystem(eqs2,pins=[F],name=:lorenz2)
-
-connected = ODESystem([lorenz2.y ~ a + lorenz1.x,
-                       lorenz1.F ~ lorenz2.u,
-                       lorenz2.F ~ lorenz1.u],t,[a],[],systems=[lorenz1,lorenz2])
+connected = ODESystem([s ~ a + lorenz1.x
+                       lorenz2.y ~ s
+                       lorenz1.F ~ lorenz2.u
+                       lorenz2.F ~ lorenz1.u],t,systems=[lorenz1,lorenz2])
 
 # Reduced Flattened System
 
 flattened_system = ModelingToolkit.flatten(connected)
 
 aliased_flattened_system = alias_elimination(flattened_system)
 
-@test isequal(states(aliased_flattened_system), [
+@test setdiff(states(aliased_flattened_system), [
         a
         lorenz1.x
         lorenz1.y
         lorenz1.z
         lorenz2.x
         lorenz2.y
         lorenz2.z
-       ])
+       ]) |> isempty
 
 @test setdiff(parameters(aliased_flattened_system), [
         lorenz1.σ
         lorenz1.ρ
         lorenz1.β
-        lorenz1.F
-        lorenz2.F
+        lorenz2.σ
         lorenz2.ρ
         lorenz2.β
        ]) |> isempty
 
 reduced_eqs = [
-               0 ~ a + lorenz1.x - lorenz2.y, # irreducible by alias elimination
-               D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - lorenz2.y - lorenz2.z,
-               D(lorenz1.y) ~ lorenz1.x*(lorenz1.ρ-lorenz1.z)-(lorenz1.x + lorenz1.y - lorenz1.z),
-               D(lorenz1.z) ~ lorenz1.x*lorenz1.y - lorenz1.β*lorenz1.z,
-               D(lorenz2.x) ~ lorenz1.x + lorenz1.y - lorenz1.z,
-               D(lorenz2.y) ~ lorenz2.x*(lorenz2.ρ-lorenz2.z)-lorenz2.x,
-               D(lorenz2.z) ~ lorenz2.x*lorenz2.y - lorenz2.β*lorenz2.z
+               0 ~ s - lorenz2.y
+               D(lorenz1.x) ~ lorenz1.F + lorenz1.σ*(lorenz1.y + -1lorenz1.x)
+               D(lorenz1.y) ~ -1lorenz1.u + lorenz1.x*(lorenz1.ρ + -1lorenz1.z)
+               D(lorenz1.z) ~ lorenz1.x*lorenz1.y + -1lorenz1.β*lorenz1.z
+               D(lorenz2.x) ~ lorenz2.F + lorenz2.σ*(lorenz2.y + -1lorenz2.x)
+               D(lorenz2.y) ~ -1lorenz2.u + lorenz2.x*(lorenz2.ρ + -1lorenz2.z)
+               D(lorenz2.z) ~ lorenz2.x*lorenz2.y + -1lorenz2.β*lorenz2.z
               ]
 test_equal.(equations(aliased_flattened_system), reduced_eqs)
 
 observed_eqs = [
-                lorenz1.F ~ lorenz2.u,
-                lorenz2.F ~ lorenz1.u,
-                lorenz1.u ~ lorenz1.x + lorenz1.y - lorenz1.z,
-                lorenz2.u ~ lorenz2.x - lorenz2.y - lorenz2.z,
+                s ~ a + lorenz1.x
+                lorenz1.u ~ lorenz1.x + lorenz1.y - lorenz1.z
+                lorenz2.u ~ lorenz2.x + lorenz2.y - lorenz2.z
+                lorenz2.F ~ lorenz1.u
+                lorenz1.F ~ lorenz2.u
                ]
 test_equal.(observed(aliased_flattened_system), observed_eqs)
 
+pp = [
+      lorenz1.σ => 10
+      lorenz1.ρ => 28
+      lorenz1.β => 8/3
+      lorenz2.σ => 10
+      lorenz2.ρ => 28
+      lorenz2.β => 8/3
+     ]
+u0 = [
+      a         => 1.0
+      lorenz1.x => 1.0
+      lorenz1.y => 0.0
+      lorenz1.z => 0.0
+      lorenz2.x => 1.0
+      lorenz2.y => 0.0
+      lorenz2.z => 0.0
+     ]
+prob1 = ODEProblem(aliased_flattened_system, u0, (0.0, 100.0), pp)
+solve(prob1, Rodas5())
+
 # issue #578
 
 let
@@ -131,10 +145,10 @@ let
         D(x) ~ x + y
         x ~ y
     ]
-    sys = ODESystem(eqs, t, [x], [])
+    sys = ODESystem(eqs, t)
     asys = alias_elimination(flatten(sys))
 
-    test_equal.(asys.eqs, [D(x) ~ 2x])
+    test_equal.(asys.eqs, [D(x) ~ x + y])
     test_equal.(asys.observed, [y ~ x])
 end
 
@@ -149,17 +163,17 @@ let
     @parameters k_P
     pc = ODESystem(Equation[u_c ~ k_P * y_c], t, pins=[y_c], name=:pc)
     connections = [
-        ol.u ~ pc.u_c,
+        ol.u ~ pc.u_c
         pc.y_c ~ ol.y
     ]
     connected = ODESystem(connections, t, systems=[ol, pc])
     @test equations(connected) isa Vector{Equation}
     sys = flatten(connected)
     reduced_sys = alias_elimination(sys)
     ref_eqs = [
-               D(ol.x) ~ ol.a*ol.x + ol.b*pc.u_c
-               0 ~ ol.c*ol.x + ol.d*pc.u_c - ol.y
-               0 ~ pc.k_P*ol.y - pc.u_c
+               D(ol.x) ~ ol.a*ol.x + ol.b*ol.u
+               0 ~ ol.c*ol.x + ol.d*ol.u + -1ol.y
+               0 ~ pc.k_P*pc.y_c + -1pc.u_c
               ]
     @test ref_eqs == equations(reduced_sys)
 end