fixes and first few tests

Author: shashi

src/systems/reduction.jl

@@ -27,9 +27,9 @@ function substitute_aliases(diffeqs, dict)
     lhss(diffeqs) .~ fixpoint_sub.(rhss(diffeqs), (dict,))
 end
 
-isvar(s::Sym) = !isparameter(s)
-isvar(s::Term) = isvar(s.op)
-isvar(s::Any) = false
+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
 
 function filterexpr(f, s)
     vs = []
@@ -84,15 +84,15 @@ function alias_elimination(sys::ODESystem)
 end
 
 function get_α_x(αx)
-    if isvar(αx)
-        return αx, 1
+    if isvar(αx, param=true)
+        return 1, αx
     elseif αx isa Term && operation(αx) === (*)
         args = arguments(αx)
         nums = filter(!isvar, args)
         syms = filter(isvar, args)
 
         if length(syms) == 1
-            return syms[1], prod(nums)
+            return prod(nums), syms[1]
         end
     else
         return nothing
@@ -105,51 +105,68 @@ function alias_elimination2(sys)
     subs = Pair[]
     # Case 1: Right hand side is a constant
     ii = findall(eqs) do eq
-        (eq.lhs isa Sym || (eq.lhs isa Term && !(eq.lhs.op isa Differential))) && !(eq.rhs isa Symbolic)
+        !(eq.rhs isa Symbolic)
     end
     for eq in eqs[ii]
-        substitution_dict[eq.lhs] = eq.rhs
-        push!(subs, eq.lhs => eq.rhs)
+        α,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 = findall(eqs) do eq
-        if eq.lhs isa Term && eq.lhs.op isa Differential
-            eq.lhs.args[1]
-        else
-            nothing
-        end
-    end
+    diff_vars = filter(!isnothing, map(eqs) do eq
+            if eq.lhs isa Term && eq.lhs.op isa Differential
+                eq.lhs.args[1]
+            else
+                nothing
+            end
+        end) |> Set
 
-    for eq in eqs
+    del = Int[]
+    for (i, eq) in enumerate(eqs)
         res_left = get_α_x(eq.lhs)
-        if !isnothing(res)
+        if !isnothing(res_left)
+            α, x = res_left
             res_right = get_α_x(eq.rhs)
-            β, y = res
-            if y in diff_vars && !(x in diff_vars)
-                multiple = β / α
-                push!(subs, x => isone(multiple) ? y : multiple * y)
-            elseif x in diff_vars && !(y in diff_vars)
-                multiple = α / β
-                push!(subs, y => isone(multiple) ? y : multiple * y)
+            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
-    for eq in eqs
+    del = Int[]
+    for (i, eq) in enumerate(eqs)
         res_left = get_α_x(eq.lhs)
-        if !isnothing(res)
+        if !isnothing(res_left)
+            α, x = res_left
             res_right = get_α_x(eq.rhs)
-            β, y = res
-            multiple =  β / α
-            push!(subs, x => isone(multiple) ? x : multiple * x)
+            if !isnothing(res_right)
+                β, y = res_right
+                multiple =  β / α
+                push!(subs, x => _isone(multiple) ? x : multiple * x)
+                push!(del, i)
+            end
         end
     end
+    deleteat!(eqs, del)
 
     diffeqs = filter(eq -> eq.lhs isa Term && eq.lhs.op isa Differential, eqs)
     diffeqs′ = substitute_aliases(diffeqs, Dict(subs))
+
+    newstates = map(diffeqs) do eq
+        eq.lhs.args[1]
+    end
     ODESystem(diffeqs′, sys.iv, newstates, parameters(sys), observed=first.(subs) .~ last.(subs))
 end

test/reduction.jl

@@ -1,4 +1,5 @@
 using ModelingToolkit, OrdinaryDiffEq, Test
+using ModelingToolkit: alias_elimination2
 
 @parameters t σ ρ β
 @variables x(t) y(t) z(t) a(t) u(t) F(t)
@@ -9,19 +10,21 @@ test_equal(a, b) = @test isequal(simplify(a, polynorm=true), simplify(b, polynor
 eqs = [D(x) ~ σ*(y-x),
        D(y) ~ x*(ρ-z)-y,
        D(z) ~ a*y - β*z,
-       0 ~ x - a]
+       β ~ 2,
+       x ~ a]
 
 lorenz1 = ODESystem(eqs,t,[x,y,z,a],[σ,ρ,β],name=:lorenz1)
 
-lorenz1_aliased = alias_elimination(lorenz1)
+lorenz1_aliased = alias_elimination2(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 - β*z]
+       D(z) ~ x*y - 2*z]
 
-@test lorenz1_aliased == ODESystem(eqs,t,[x,y,z],[σ,ρ,β],observed=[a ~ x],name=:lorenz1)
+# 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)
 
 # Multi-System Reduction
 
@@ -44,7 +47,7 @@ lorenz2 = ODESystem(eqs2,pins=[F],observed=aliases2,name=:lorenz2)
 connections = [lorenz1.F ~ lorenz2.u,
                lorenz2.F ~ lorenz1.u]
 
-connected = ODESystem([0 ~ a + lorenz1.x - lorenz2.y],t,[a],[],observed=connections,systems=[lorenz1,lorenz2])
+connected = ODESystem([lorenz2.y ~ a + lorenz1.x ],t,[a],[],observed=connections,systems=[lorenz1,lorenz2])
 
 # Reduced Unflattened System
 #=
@@ -59,7 +62,7 @@ connected = ODESystem(Equation[],t,[],[],observed=connections2,systems=[lorenz1,
 
 flattened_system = ModelingToolkit.flatten(connected)
 
-aliased_flattened_system = alias_elimination(flattened_system)
+aliased_flattened_system = alias_elimination2(flattened_system)
 
 @test isequal(states(aliased_flattened_system), [
         lorenz1.x
@@ -107,7 +110,7 @@ let
         x ~ y
     ];
     sys = ODESystem(eqs, t, [x], []);
-    asys = alias_elimination(ModelingToolkit.flatten(sys))
+    asys = alias_elimination2(ModelingToolkit.flatten(sys))
 
     test_equal.(asys.eqs, [D(x) ~ 2x])
     test_equal.(asys.observed, [y ~ x])