fix reduction

Author: shashi

src/equations.jl

@@ -15,6 +15,8 @@ end
 Base.:(==)(a::Equation, b::Equation) = all(isequal.((a.lhs, a.rhs), (b.lhs, b.rhs)))
 Base.hash(a::Equation, salt::UInt) = hash(a.lhs, hash(a.rhs, salt))
 
+SymbolicUtils.simplify(x::Equation; kw...) = simplify(x.lhs; kw...) ~ simplify(x.rhs; kw...)
+
 """
 $(TYPEDSIGNATURES)
 
@@ -41,7 +43,6 @@ Base.:~(lhs::Number    , rhs::Num) = Equation(value(lhs), value(rhs))
 Base.:~(lhs::Symbolic, rhs::Symbolic) = Equation(value(lhs), value(rhs))
 Base.:~(lhs::Symbolic, rhs::Any    ) = Equation(value(lhs), value(rhs))
 Base.:~(lhs::Any, rhs::Symbolic    ) = Equation(value(lhs), value(rhs))
-Base.:~(lhs::Number    , rhs::Num) = Equation(value(lhs), value(rhs))
 
 struct ConstrainedEquation
   constraints

src/solve.jl

@@ -72,7 +72,7 @@ end
 function _solve(A, b)
     A = SymbolicUtils.simplify.(to_symbolic.(A), polynorm=true)
     b = SymbolicUtils.simplify.(to_symbolic.(b), polynorm=true)
-    map(to_mtk, SymbolicUtils.simplify.(ldiv(sym_lu(A), b)))
+    SymbolicUtils.simplify.(ldiv(sym_lu(A), b))
 end
 
 LinearAlgebra.:(\)(A::AbstractMatrix{<:Expression}, b::AbstractVector{<:Expression}) = _solve(A, b)

src/systems/abstractsystem.jl

@@ -117,6 +117,9 @@ Generate a function to evaluate the system's equations.
 """
 function generate_function end
 
+getname(x::Sym) = nameof(x)
+getname(t::Term) = t.op isa Sym ? getname(t.op) : error("Cannot get name of $t")
+
 function Base.getproperty(sys::AbstractSystem, name::Symbol)
 
     if name ∈ fieldnames(typeof(sys))
@@ -128,27 +131,23 @@ function Base.getproperty(sys::AbstractSystem, name::Symbol)
         end
     end
 
-    i = findfirst(x->x.name==name,sys.states)
+    i = findfirst(x->getname(x) == name, sys.states)
+
     if i !== nothing
-        x = rename(sys.states[i],renamespace(sys.name,name))
-        if :iv ∈ fieldnames(typeof(sys))
-            return x(getfield(sys,:iv)())
-        else
-            return x()
-        end
+        return rename(sys.states[i],renamespace(sys.name,name))
     end
 
     if :ps ∈ fieldnames(typeof(sys))
-        i = findfirst(x->x.name==name,sys.ps)
+        i = findfirst(x->getname(x) == name,sys.ps)
         if i !== nothing
-            return rename(sys.ps[i],renamespace(sys.name,name))()
+            return rename(sys.ps[i],renamespace(sys.name,name))
         end
     end
 
     if :observed ∈ fieldnames(typeof(sys))
-        i = findfirst(x->convert(Variable,x.lhs).name==name,sys.observed)
+        i = findfirst(x->getname(x.lhs)==name,sys.observed)
         if i !== nothing
-            return rename(convert(Variable,sys.observed[i].lhs),renamespace(sys.name,name))(getfield(sys,:iv)())
+            return rename(sys.observed[i].lhs,renamespace(sys.name,name))
         end
     end
 
@@ -172,19 +171,23 @@ end
 namespace_equations(sys::AbstractSystem) = namespace_equation.(equations(sys),sys.name,sys.iv.name)
 
 function namespace_equation(eq::Equation,name,ivname)
-    _lhs = namespace_operation(eq.lhs,name,ivname)
-    _rhs = namespace_operation(eq.rhs,name,ivname)
+    _lhs = namespace_expr(eq.lhs,name,ivname)
+    _rhs = namespace_expr(eq.rhs,name,ivname)
     _lhs ~ _rhs
 end
 
-function namespace_operation(O::Operation,name,ivname)
-    if O.op isa Sym && O.op.name != ivname
-        Operation(rename(O.op,renamespace(name,O.op.name)),namespace_operation.(O.args,name,ivname))
+function namespace_expr(O::Sym,name,ivname)
+    O.name == ivname ? O : rename(O,renamespace(name,O.name))
+end
+
+function namespace_expr(O::Term,name,ivname)
+    if O.op isa Sym
+        Term(rename(O.op,renamespace(name,O.op.name)),namespace_expr.(O.args,name,ivname))
     else
-        Operation(O.op,namespace_operation.(O.args,name,ivname))
+        Term(O.op,namespace_expr.(O.args,name,ivname))
     end
 end
-namespace_operation(O::Constant,name,ivname) = O
+namespace_expr(O,name,ivname) = O
 
 independent_variable(sys::AbstractSystem) = sys.iv
 states(sys::AbstractSystem) = unique(isempty(sys.systems) ? setdiff(sys.states, convert.(Variable,sys.pins)) : [sys.states;reduce(vcat,namespace_variables.(sys.systems))])

src/systems/diffeqs/odesystem.jl

@@ -31,7 +31,7 @@ struct ODESystem <: AbstractODESystem
     states::Vector
     """Parameter variables."""
     ps::Vector
-    pins::Vector{Variable}
+    pins::Vector{Num}
     observed::Vector{Equation}
     """
     Time-derivative matrix. Note: this field will not be defined until
@@ -114,8 +114,6 @@ function ODESystem(eqs, iv=nothing; kwargs...)
                 break
             end
         end
-    else
-        iv = convert(Variable, iv)
     end
     iv === nothing && throw(ArgumentError("Please pass in independent variables."))
     for eq in eqs

src/systems/reduction.jl

@@ -34,46 +34,55 @@ function make_lhs_0(eq)
         0 ~ eq.lhs - eq.rhs
     end
 end
+isvar(s::Sym) = !isparameter(s)
+isvar(s::Term) = isvar(s.op)
+isvar(s::Any) = false
+
+function filterexpr(f, s)
+    vs = []
+    Rewriters.Prewalk(Rewriters.Chain([@rule((~x::f) => push!(vs, ~x))]))(s)
+    vs
+end
 
 function alias_elimination(sys::ODESystem)
     eqs = vcat(equations(sys), observed(sys))
 
     # make all algebraic equations have 0 on LHS
     eqs = map(eqs) do eq
-        if eq.lhs isa Operation && eq.lhs.op isa Differential
+        if eq.lhs isa Term && eq.lhs.op isa Differential
             eq
         else
             make_lhs_0(eq)
         end
     end
 
-    new_stateops = map(eqs) do eq
-        if eq.lhs isa Operation && eq.lhs.op isa Differential
-            get_variables(eq.lhs)
+    newstates = map(eqs) do eq
+        if eq.lhs isa Term && eq.lhs.op isa Differential
+            filterexpr(isvar, eq.lhs)
         else
             []
         end
     end |> Iterators.flatten |> collect |> unique
 
+
     all_vars = map(eqs) do eq
-        filter(x->!isparameter(x.op), get_variables(eq.rhs))
+        @show eq.rhs
+        @show filterexpr(isvar, eq.rhs)
+        filterexpr(isvar, eq.rhs)
     end |> Iterators.flatten |> collect |> unique
 
-    newstates = convert.(Variable, new_stateops)
-
-
-    alg_idxs = findall(x->x.lhs isa Constant && iszero(x.lhs), eqs)
-
-    eliminate = setdiff(convert.(Variable, all_vars), newstates)
+    alg_idxs = findall(x->!(x.lhs isa Term) && iszero(x.lhs), eqs)
+    @show all_vars, newstates
 
-    vars = map(x->x(sys.iv()), eliminate)
+    eliminate = setdiff(all_vars, newstates)
+    @show eliminate
 
-    outputs = solve_for(eqs[alg_idxs], vars)
+    outputs = solve_for(eqs[alg_idxs], eliminate)
 
     diffeqs = eqs[setdiff(1:length(eqs), alg_idxs)]
 
-    diffeqs′ = substitute_aliases(diffeqs, Dict(vars .=> outputs))
+    diffeqs′ = substitute_aliases(diffeqs, Dict(eliminate .=> outputs))
 
-    ODESystem(diffeqs′, sys.iv(), new_stateops, parameters(sys), observed=vars .~ outputs)
+    ODESystem(diffeqs′, sys.iv, newstates, parameters(sys), observed=eliminate .~ outputs)
 end
 

src/variables.jl

@@ -63,6 +63,7 @@ function Variable(name, indices...)
 end
 
 rename(x::Sym{T},name) where T = Sym{T}(name)
+rename(x::Term, name) where T = x.op isa Sym ? rename(x.op, name)(x.args...) : error("can't rename $x to $name")
 
 """
 $(TYPEDEF)

test/reduction.jl

@@ -61,7 +61,7 @@ flattened_system = ModelingToolkit.flatten(connected)
 
 aliased_flattened_system = alias_elimination(flattened_system)
 
-@test states(aliased_flattened_system) == convert.(Variable, [
+@test isequal(states(aliased_flattened_system), [
         lorenz1.x
         lorenz1.y
         lorenz1.z
@@ -70,15 +70,15 @@ aliased_flattened_system = alias_elimination(flattened_system)
         lorenz2.z
        ])
 
-@test setdiff(parameters(aliased_flattened_system), convert.(Variable, [
+@test setdiff(parameters(aliased_flattened_system), [
         lorenz1.σ
         lorenz1.ρ
         lorenz1.β
         lorenz1.F
         lorenz2.F
         lorenz2.ρ
         lorenz2.β
-       ])) |> isempty
+       ]) |> isempty
 
 test_equal.(equations(aliased_flattened_system), [
        D(lorenz1.x) ~ lorenz1.σ*(lorenz1.y-lorenz1.x) + lorenz2.x - lorenz2.y - lorenz2.z,