Fix ODESymstems and add extra indirection

Author: YingboMa

src/build_function.jl

@@ -458,27 +458,12 @@ end
 
 vars_to_pairs(args) = vars_to_pairs(args[1],args[2])
 function vars_to_pairs(name,vs::AbstractArray)
-    vs_names = [term_to_symbol(value(u)) for u ∈ vs]
-	exs = [:($name[$i]) for (i, u) ∈ enumerate(vs)]
-	vs_names,exs
+    vs_names = [tosymbol(u) for u ∈ vs]
+    exs = [:($name[$i]) for (i, u) ∈ enumerate(vs)]
+    vs_names,exs
 end
-
-function term_to_symbol(t::Term)
-    if operation(t) isa Sym
-        s = nameof(operation(t))
-    else
-        error("really?")
-    end
-end
-
-term_to_symbol(s::Sym) = nameof(s)
-
 function vars_to_pairs(name,vs)
-    [term_to_symbol(value(vs))], [name]
-end
-
-function rm_calls_with_iv(expr)
-    Rewriters.Prewalk(Rewriters.Chain([@rule((~f::(x->x isa Sym))(~t::(x->x isa Sym)) => Sym{symtype((~f)(~t))}((term_to_symbol(~f))))]))(value(expr))
+    [tosymbol(vs)], [name]
 end
 
 get_varnumber(varop, vars::Vector) =  findfirst(x->isequal(x,varop),vars)
@@ -494,7 +479,7 @@ function numbered_expr(O::Union{Term,Sym},args...;varordering = args[1],offset =
 		end
 	end
   end
-  return Expr(:call, O isa Sym ? nameof(O) : Symbol(O.op),
+  return Expr(:call, O isa Sym ? nameof(O) : tosymbol(O.op),
          [numbered_expr(x,args...;offset=offset,lhsname=lhsname,
                         rhsnames=rhsnames,varordering=varordering) for x in O.args]...)
 end
@@ -504,7 +489,7 @@ function numbered_expr(de::ModelingToolkit.Equation,args...;varordering = args[1
 
     varordering = value.(args[1])
     var = var_from_nested_derivative(de.lhs)[1]
-    i = findfirst(x->isequal(x isa Sym ? term_to_symbol(x) : term_to_symbol(x.op),term_to_symbol(var)),varordering)
+    i = findfirst(x->isequal(tosymbol(x isa Sym ? x : x.op), tosymbol(var)),varordering)
     :($lhsname[$(i+offset)] = $(numbered_expr(de.rhs,args...;offset=offset,
 											  varordering = varordering,
 											  lhsname = lhsname,

src/systems/diffeqs/abstractodesystem.jl

@@ -11,6 +11,9 @@ function calculate_tgrad(sys::AbstractODESystem;
   if simplify
       tgrad = ModelingToolkit.simplify.(notime_tgrad)
   end
+  xs = states(sys)
+  rule = Dict(map((x, xt) -> x=>xt, detime_dvs.(xs), xs))
+  tgrad = substitute.(tgrad, Ref(rule))
   sys.tgrad[] = tgrad
   return tgrad
 end
@@ -41,7 +44,8 @@ ODEToExpr(@nospecialize(sys)) = ODEToExpr(sys,states(sys))
 (f::ODEToExpr)(O::Num) = f(value(O))
 function (f::ODEToExpr)(O::Term)
     if isa(O.op, Sym)
-        any(isequal(O), f.states) && return O.op.name  # dependent variables
+        any(isequal(O), f.states) && return tosymbol(O)
+        # dependent variables
         return build_expr(:call, Any[O.op.name; f.(O.args)])
     end
     return build_expr(:call, Any[O.op; f.(O.args)])
@@ -64,9 +68,15 @@ function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = param
                           conv = ODEToExpr(sys), kwargs...)
 end
 
-Base.Symbol(t::Num) = Symbol(value(t))
-function Base.Symbol(t::Term)
+Base.Symbol(x::Union{Num,Symbolic}) = tosymbol(x)
+tosymbol(x; kwargs...) = x
+tosymbol(x::Sym; kwargs...) = nameof(x)
+tosymbol(t::Num; kwargs...) = tosymbol(value(t); kwargs...)
+function tosymbol(t::Term; states=nothing)
     if t.op isa Sym
+        if states !== nothing && !(any(isequal(t), states))
+            return nameof(t.op)
+        end
         op = nameof(t.op)
         args = t.args
     elseif t.op isa Differential
@@ -75,7 +85,7 @@ function Base.Symbol(t::Term)
         end
         op = Symbol(nameof(t.args[1].op),
                     :ˍ,
-                    Symbol(t.op.x))
+                    tosymbol(t.op.x))
         args = t.args[1].args
     else
         @goto err
@@ -86,13 +96,13 @@ function Base.Symbol(t::Term)
     error("Cannot convert $t to a symbol")
 end
 
-makesym(t::Symbolic) = Sym{symtype(t)}(Symbol(t))
-makesym(t::Num) = makesym(value(t))
+makesym(t::Symbolic; kwargs...) = Sym{symtype(t)}(tosymbol(t; kwargs...))
+makesym(t::Num; kwargs...) = makesym(value(t); kwargs...)
 
 function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
     # optimization
-    dvs′ = makesym.(value.(dvs))
-    ps′ = makesym.(value.(ps))
+    dvs′ = makesym.(value.(dvs), states=dvs)
+    ps′ = makesym.(value.(ps), states=dvs)
 
     sub = Dict(dvs .=> dvs′)
     # substitute x(t) by just x
@@ -107,7 +117,7 @@ function calculate_massmatrix(sys::AbstractODESystem; simplify=true)
     M = zeros(length(eqs),length(eqs))
     for (i,eq) in enumerate(eqs)
         if eq.lhs isa Term && eq.lhs.op isa Differential
-            j = findfirst(x->isequal(term_to_symbol(x),term_to_symbol(var_from_nested_derivative(eq.lhs)[1])),dvs)
+            j = findfirst(x->isequal(tosymbol(x),tosymbol(var_from_nested_derivative(eq.lhs)[1])),dvs)
             M[i,j] = 1
         else
             eq.lhs == 0 || error("Only semi-explicit constant mass matrices are currently supported. Faulty equation: $eq.")

test/derivatives.jl

@@ -97,7 +97,7 @@ t1 = ModelingToolkit.gradient(tmp, [x1, x2])
 @parameters t k
 @variables x(t)
 @derivatives D'~k
-@test ModelingToolkit.makesym(D(x).val).name === :xˍk
+@test ModelingToolkit.makesym(D(x).val).name === Symbol("xˍk⦗t⦘")
 
 using ModelingToolkit
 @variables t x(t)

test/odesystem.jl

@@ -145,6 +145,7 @@ eqs = [D(x) ~ -A*x,
        D(y) ~ A*x - B*_x]
 de = ODESystem(eqs)
 @test begin
+    local f
     f = eval(generate_function(de, [x,y], [A,B,C])[2])
     du = [0.0,0.0]
     f(du, [1.0,2.0], [1,2,3], 0.0)