Fix alias elimination

Author: YingboMa

src/systems/diffeqs/odesystem.jl

@@ -149,5 +149,3 @@ Base.:(==)(sys1::ODESystem, sys2::ODESystem) =
 function rename(sys::ODESystem,name)
     ODESystem(sys.eqs, sys.iv, sys.states, sys.ps, sys.pins, sys.observed, sys.tgrad, sys.jac, sys.Wfact, sys.Wfact_t, name, sys.systems)
 end
-
-isdiffeq(eq) = eq.lhs isa Term && operation(eq.lhs) isa Differential

src/systems/reduction.jl

@@ -29,17 +29,22 @@ function substitute_aliases(diffeqs, dict)
     lhss(diffeqs) .~ fixpoint_sub.(rhss(diffeqs), (dict,))
 end
 
-isvar(s::Sym; param=false) = param ? true : !isparameter(s)
-isvar(s::Term; param=false) = isvar(s.op; param=param)
-isvar(s::Any;param=false) = false
+# Note that we reduce parameters, too
+# i.e. `2param = 3` will be reduced away
+isvar(s::Sym) = true
+isvar(s::Term) = isvar(operation(s))
+isvar(s::Any) = false
 
 function get_α_x(αx)
-    if isvar(αx, param=true)
+    if isvar(αx)
         return 1, αx
     elseif αx isa Term && operation(αx) === (*)
         args = arguments(αx)
-        nums = filter(!isvar, args)
-        syms = filter(isvar, args)
+        nums = []
+        syms = []
+        for arg in args
+            isvar(arg) ? push!(syms, arg) : push!(nums, arg)
+        end
 
         if length(syms) == 1
             return prod(nums), syms[1]
@@ -51,20 +56,7 @@ end
 
 function alias_elimination(sys::ODESystem)
     eqs = vcat(equations(sys), observed(sys))
-
     subs = Pair[]
-    # Case 1: Right hand side is a constant
-    ii = findall(eqs) do eq
-        !(eq.rhs isa Symbolic)
-    end
-    for eq in eqs[ii]
-        α,x = get_α_x(eq.lhs)
-        push!(subs, x => isone(α) ? eq.rhs : eq.rhs / α)
-    end
-    deleteat!(eqs, ii) # remove them
-
-    # Case 2: One side is a differentiated var, the other is an algebraic var
-    #         substitute the algebraic var with the diff var
     diff_vars = filter(!isnothing, map(eqs) do eq
             if isdiffeq(eq)
                 eq.lhs.args[1]
@@ -73,41 +65,19 @@ function alias_elimination(sys::ODESystem)
             end
         end) |> Set
 
-    del = Int[]
-    for (i, eq) in enumerate(eqs)
-        res_left = get_α_x(eq.lhs)
-        if !isnothing(res_left)
-            α, x = res_left
-            res_right = get_α_x(eq.rhs)
-            if !isnothing(res_right)
-                β, y = res_right
-                if y in diff_vars && !(x in diff_vars)
-                    multiple = β / α
-                    push!(subs, x => isone(multiple) ? y : multiple * y)
-                    push!(del, i)
-                elseif x in diff_vars && !(y in diff_vars)
-                    multiple = α / β
-                    push!(subs, y => isone(multiple) ? x : multiple * x)
-                    push!(del, i)
-                end
-            end
-        end
-    end
-    deleteat!(eqs, del)
-
-    # Case 3: Explicit substitutions
+    # only substitute when the variable is algebraic
     del = Int[]
     for (i, eq) in enumerate(eqs)
         isdiffeq(eq) && continue
         res_left = get_α_x(eq.lhs)
-        if !isnothing(res_left)
+        if !isnothing(res_left) && !(res_left[2] in diff_vars)
             # `α x = rhs` => `x = rhs / α`
             α, x = res_left
             push!(subs, x => _isone(α) ? eq.rhs : eq.rhs / α)
             push!(del, i)
         else
             res_right = get_α_x(eq.rhs)
-            if !isnothing(res_right)
+            if !isnothing(res_right) && !(res_right[2] in diff_vars)
                 # `lhs = β y` => `y = lhs / β`
                 β, y = res_right
                 push!(subs, y => _isone(β) ? eq.lhs : β * eq.lhs)
@@ -117,11 +87,9 @@ function alias_elimination(sys::ODESystem)
     end
     deleteat!(eqs, del)
 
-    diffeqs = filter(eq -> eq.lhs isa Term && eq.lhs.op isa Differential, eqs)
-    diffeqs′ = substitute_aliases(diffeqs, Dict(subs))
+    eqs′ = substitute_aliases(eqs, Dict(subs))
+    alias_vars = first.(subs)
 
-    newstates = map(diffeqs) do eq
-        eq.lhs.args[1]
-    end
-    ODESystem(diffeqs′, sys.iv, newstates, parameters(sys), observed=first.(subs) .~ last.(subs))
+    newstates = setdiff(states(sys), alias_vars)
+    ODESystem(eqs′, sys.iv, newstates, parameters(sys), observed=alias_vars .~ last.(subs))
 end

test/reduction.jl

@@ -7,55 +7,52 @@ using ModelingToolkit, OrdinaryDiffEq, Test
 test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynorm=true))
 
 eqs = [D(x) ~ σ*(y-x),
-       D(y) ~ x*(ρ-z)-y,
-       D(z) ~ a*y - β*z,
-       β ~ 2,
-       x ~ a]
+       D(y) ~ x*(ρ-z)-y + β,
+       0 ~ sin(z) - x + y,
+       sin(u) ~ x + y,
+       2β ~ 2,
+       x ~ a,
+      ]
 
-lorenz1 = ODESystem(eqs,t,[x,y,z,a],[σ,ρ,β],name=:lorenz1)
+lorenz1 = ODESystem(eqs,t,[u,x,y,z,a],[σ,ρ,β],name=:lorenz1)
 
 lorenz1_aliased = alias_elimination(lorenz1)
-@test length(equations(lorenz1_aliased)) == 3
-@test length(states(lorenz1_aliased)) == 3
-
-eqs = [D(x) ~ σ*(y-x),
-       D(y) ~ x*(ρ-z)-y,
-       D(z) ~ x*y - 2*z]
-
-# TODO: maybe remove β from ps, or maybe don't allow this example on params
-@test lorenz1_aliased == ODESystem(eqs,t,[x,y,z],[σ,ρ,β],observed=[β ~ 2, a ~ x],name=:lorenz1)
+reduced_eqs = [
+               D(x) ~ σ * (y - x),
+               D(y) ~ x*(ρ-z)-y + 1,
+               0 ~ sin(z) - x + y,
+               sin(u) ~ x + y,
+              ]
+test_equal.(equations(lorenz1_aliased), reduced_eqs)
+test_equal.(states(lorenz1_aliased), [u, x, y, z])
+test_equal.(observed(lorenz1_aliased), [
+    β ~ 1,
+    a ~ x,
+])
 
 # Multi-System Reduction
 
-eqs1 = [D(x) ~ σ*(y-x) + F,
-       D(y) ~ x*(ρ-z)-u,
-       D(z) ~ x*y - β*z]
-
-aliases = [u ~ x + y - z]
+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],observed=aliases,name=:lorenz1)
+lorenz1 = ODESystem(eqs1,pins=[F],name=:lorenz1)
 
-eqs2 = [D(x) ~ F,
-       D(y) ~ x*(ρ-z)-x,
-       D(z) ~ x*y - β*z]
+eqs2 = [
+        D(x) ~ F,
+        D(y) ~ x*(ρ-z)-x,
+        D(z) ~ x*y - β*z,
+        u ~ x - y - z
+       ]
 
-aliases2 = [u ~ x - y - z]
+lorenz2 = ODESystem(eqs2,pins=[F],name=:lorenz2)
 
-lorenz2 = ODESystem(eqs2,pins=[F],observed=aliases2,name=:lorenz2)
-
-connections = [lorenz1.F ~ lorenz2.u,
-               lorenz2.F ~ lorenz1.u]
-
-connected = ODESystem([lorenz2.y ~ a + lorenz1.x],t,[a],[],observed=connections,systems=[lorenz1,lorenz2])
-
-# Reduced Unflattened System
-#=
-
-connections2 = [lorenz1.F ~ lorenz2.u,
-                lorenz2.F ~ lorenz1.u,
-                a ~ -lorenz1.x + lorenz2.y]
-connected = ODESystem(Equation[],t,[],[],observed=connections2,systems=[lorenz1,lorenz2])
-=#
+connected = ODESystem([lorenz2.y ~ a + lorenz1.x,
+                       lorenz1.F ~ lorenz2.u,
+                       lorenz2.F ~ lorenz1.u],t,[a],[],systems=[lorenz1,lorenz2])
 
 # Reduced Flattened System
 
@@ -64,6 +61,7 @@ flattened_system = ModelingToolkit.flatten(connected)
 aliased_flattened_system = alias_elimination(flattened_system)
 
 @test isequal(states(aliased_flattened_system), [
+        a
         lorenz1.x
         lorenz1.y
         lorenz1.z
@@ -83,25 +81,24 @@ aliased_flattened_system = alias_elimination(flattened_system)
        ]) |> isempty
 
 reduced_eqs = [
-               D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - (a + lorenz1.x) - lorenz2.z,
+               lorenz2.y ~ a + lorenz1.x, # 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*(a + lorenz1.x) - lorenz2.β*lorenz2.z
+               D(lorenz2.z) ~ lorenz2.x*lorenz2.y - lorenz2.β*lorenz2.z
               ]
 test_equal.(equations(aliased_flattened_system), reduced_eqs)
 
 observed_eqs = [
-                lorenz2.y ~ a + lorenz1.x,
                 lorenz1.F ~ lorenz2.u,
                 lorenz2.F ~ lorenz1.u,
                 lorenz1.u ~ lorenz1.x + lorenz1.y - lorenz1.z,
                 lorenz2.u ~ lorenz2.x - lorenz2.y - lorenz2.z,
                ]
 test_equal.(observed(aliased_flattened_system), observed_eqs)
 
-
 # issue #578
 
 let