handle time dependent parameters -- not just names

Author: shashi

src/ModelingToolkit.jl

@@ -19,7 +19,7 @@ import IfElse
 using RecursiveArrayTools
 
 import SymbolicUtils
-import SymbolicUtils: Term, Sym, to_symbolic, FnType, @rule, Rewriters, substitute
+import SymbolicUtils: Term, Sym, to_symbolic, FnType, @rule, Rewriters, substitute, similarterm
 
 using LinearAlgebra: LU, BlasInt
 
@@ -50,7 +50,7 @@ end
 const show_numwrap = Ref(false)
 
 Num(x::Num) = x # ideally this should never be called
-(n::Num)(args...) = value(n)(map(value,args)...)
+(n::Num)(args...) = Num(value(n)(map(value,args)...))
 value(x) = x
 value(x::Num) = x.val
 
@@ -105,7 +105,6 @@ end
 @num_method Base.isless isless(value(a), value(b))
 @num_method Base.isequal isequal(value(a), value(b)) (Number, Symbolic)
 @num_method Base.:(==) value(a) == value(b) (Number,)
-Base.real(x::Num) = Num(real(value(x)))
 
 Base.hash(x::Num, h::UInt) = hash(value(x), h)
 

src/differentials.jl

@@ -137,8 +137,6 @@ function expand_derivatives(n::Num, simplify=true; occurances=nothing)
     Num(expand_derivatives(value(n), simplify; occurances=occurances))
 end
 
-_iszero(x::Number) = iszero(x)
-_isone(x::Number) = isone(x)
 _iszero(x) = false
 _isone(x) = false
 

src/systems/diffeqs/abstractodesystem.jl

@@ -1,16 +1,15 @@
 function calculate_tgrad(sys::AbstractODESystem;
                          simplify=true)
   isempty(sys.tgrad[]) || return sys.tgrad[]  # use cached tgrad, if possible
+
+  # We need to remove explicit time dependence on the state because when we
+  # have `u(t) * t` we want to have the tgrad to be `u(t)` instead of `u'(t) *
+  # t + u(t)`.
   rhs = [detime_dvs(eq.rhs) for eq ∈ equations(sys)]
   iv = sys.iv
-  for r in rhs
-      @show r
-      @show expand_derivatives(Differential(iv)(r))
-  end
   notime_tgrad = [expand_derivatives(ModelingToolkit.Differential(iv)(r)) for r in rhs]
-  tgrad = retime_dvs.(notime_tgrad,(states(sys),),iv)
   if simplify
-      tgrad = ModelingToolkit.simplify.(tgrad)
+      tgrad = ModelingToolkit.simplify.(notime_tgrad)
   end
   sys.tgrad[] = tgrad
   return tgrad

src/systems/diffeqs/odesystem.jl

@@ -102,7 +102,7 @@ function ODESystem(eqs, iv=nothing; kwargs...)
     # NOTE: this assumes that the order of algebric equations doesn't matter
     diffvars = OrderedSet()
     allstates = OrderedSet()
-    ps = OrderedSet{Sym}()
+    ps = OrderedSet()
     # reorder equations such that it is in the form of `diffeq, algeeq`
     diffeq = Equation[]
     algeeq = Equation[]

src/utils.jl

@@ -33,21 +33,21 @@ end
 
 function detime_dvs(op::Term)
   if op.op isa Sym
-      op.op
+      Sym{Number}(nameof(op.op))
   else
     Term(op.op,detime_dvs.(op.args))
   end
 end
 detime_dvs(op) = op
 
-function retime_dvs(op::Operation,dvs,iv)
-  if op.op isa Variable && op.op ∈ dvs
-    Operation(Variable{vartype(op.op)}(op.op.name),Expression[iv])
-  else
-    Operation(op.op,retime_dvs.(op.args,(dvs,),iv))
-  end
+function retime_dvs(op::Sym,dvs,iv)
+    Sym{FnType{Tuple{symtype(iv)}, Number}}(nameof(op))(iv)
+end
+
+function retime_dvs(op::Term, dvs, iv)
+    similarterm(op, op.op, retime_dvs.(op.args,(dvs,),(iv,)))
 end
-retime_dvs(op::Constant,dvs,iv) = op
+retime_dvs(op,dvs,iv) = op
 
 is_constant(::Constant) = true
 is_constant(::Any) = false

test/odesystem.jl

@@ -58,6 +58,7 @@ de = ODESystem(eqs) # This is broken
 ModelingToolkit.calculate_tgrad(de)
 
 tgrad_oop, tgrad_iip = eval.(ModelingToolkit.generate_tgrad(de))
+
 @test tgrad_oop(u,p,t) == [0.0,-u[2],0.0]
 du = zeros(3)
 tgrad_iip(du,u,p,t)
@@ -82,7 +83,7 @@ tgrad_iip(du,u,p,t)
            D(y) ~ x*(ρ-z)-y,
            D(z) ~ x*y - β*z]
     de = ODESystem(eqs)
-    test_diffeq_inference("single internal iv-varying", de, t, (x, y, z), (σ, ρ, β))
+    test_diffeq_inference("single internal iv-varying", de, t, (x, y, z), (σ(t-1), ρ, β))
     @test begin
         f = eval(generate_function(de, [x,y,z], [σ,ρ,β])[2])
         du = [0.0,0.0,0.0]
@@ -92,7 +93,7 @@ tgrad_iip(du,u,p,t)
 
     eqs = [D(x) ~ x + 10σ(t-1) + 100σ(t-2) + 1000σ(t^2)]
     de = ODESystem(eqs)
-    test_diffeq_inference("many internal iv-varying", de, t, (x,), (σ,))
+    test_diffeq_inference("many internal iv-varying", de, t, (x,), (σ(t-2),σ(t^2), σ(t-1)))
     @test begin
         f = eval(generate_function(de, [x], [σ])[2])
         du = [0.0]