add tgrad to ODESystem

Author: ChrisRackauckas

src/systems/diffeqs/diffeqsystem.jl

@@ -65,6 +65,11 @@ struct ODESystem <: AbstractSystem
     """Parameter variables."""
     ps::Vector{Variable}
     """
+    Time-derivative matrix. Note: this field will not be defined until
+    [`calculate_tgrad`](@ref) is called on the system.
+    """
+    tgrad::RefValue{Vector{Expression}}
+    """
     Jacobian matrix. Note: this field will not be defined until
     [`calculate_jacobian`](@ref) is called on the system.
     """
@@ -99,10 +104,11 @@ function ODESystem(eqs)
 end
 
 function ODESystem(deqs::AbstractVector{DiffEq}, iv, dvs, ps)
+    tgrad = RefValue(Vector{Expression}(undef, 0))
     jac = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact   = RefValue(Matrix{Expression}(undef, 0, 0))
     Wfact_t = RefValue(Matrix{Expression}(undef, 0, 0))
-    ODESystem(deqs, iv, dvs, ps, jac, Wfact, Wfact_t)
+    ODESystem(deqs, iv, dvs, ps, tgrad, jac, Wfact, Wfact_t)
 end
 
 function ODESystem(deqs::AbstractVector{<:Equation}, iv, dvs, ps)
@@ -133,6 +139,15 @@ independent_variables(sys::ODESystem) = Set{Variable}([sys.iv])
 dependent_variables(sys::ODESystem) = Set{Variable}(sys.dvs)
 parameters(sys::ODESystem) = Set{Variable}(sys.ps)
 
+function calculate_tgrad(sys::ODESystem)
+  isempty(sys.tgrad[]) || return sys.tgrad[]  # use cached tgrad, if possible
+  rhs = [detime_dvs(eq.rhs) for eq ∈ sys.eqs]
+  iv = sys.iv()
+  notime_tgrad = [expand_derivatives(ModelingToolkit.Differential(iv)(r)) for r in rhs]
+  tgrad = retime_dvs.(notime_tgrad,(sys.dvs,),iv)
+  sys.tgrad[] = tgrad
+  return tgrad
+end
 
 function calculate_jacobian(sys::ODESystem)
     isempty(sys.jac[]) || return sys.jac[]  # use cached Jacobian, if possible
@@ -160,6 +175,11 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
+function generate_tgrad(sys::ODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)
+    tgrad = calculate_tgrad(sys)
+    return build_function(tgrad, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
+end
+
 function generate_jacobian(sys::ODESystem, dvs = sys.dvs, ps = sys.ps, expression = Val{true}; kwargs...)
     jac = calculate_jacobian(sys)
     return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
@@ -218,13 +238,21 @@ are used to set the order of the dependent variable and parameter vectors,
 respectively.
 """
 function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs = sys.dvs, ps = sys.ps;
-                                     version = nothing,
+                                     version = nothing, tgrad=false,
                                      jac = false, Wfact = false) where {iip}
     f_oop,f_iip = generate_function(sys, dvs, ps, Val{false})
 
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
 
+    if tgrad
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false})
+        _tgrad(u,p,t) = tgrad_oop(u,p,t)
+        _tgrad(J,u,p,t) = tgrad_iip(J,u,p,t)
+    else
+        _tgrad = nothing
+    end
+
     if jac
         jac_oop,jac_iip = generate_jacobian(sys, dvs, ps, Val{false})
         _jac(u,p,t) = jac_oop(u,p,t)
@@ -246,6 +274,7 @@ function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs = sys.dvs, ps = sys.ps;
     end
 
     ODEFunction{iip}(f,jac=_jac,
+                      tgrad = tgrad,
                       Wfact = _Wfact,
                       Wfact_t = _Wfact_t,
                       syms = string.(sys.dvs))

src/utils.jl

@@ -1,6 +1,5 @@
 using MacroTools
 
-
 function Base.convert(::Type{Expression}, ex::Expr)
     ex.head === :if && (ex = Expr(:call, ifelse, ex.args...))
     ex.head === :call || throw(ArgumentError("internal representation does not support non-call Expr"))
@@ -13,6 +12,9 @@ end
 Base.convert(::Type{Expression}, x::Expression) = x
 Base.convert(::Type{Expression}, x::Number) = Constant(x)
 Base.convert(::Type{Expression}, x::Bool) = Constant(x)
+Base.convert(::Type{Expression}, x::Variable) = convert(Operation,x)
+Base.convert(::Type{Expression}, x::Operation) = x
+Base.convert(::Type{Expression}, x::Symbol) = Base.convert(Expression,eval(x))
 Expression(x::Bool) = Constant(x)
 
 function build_expr(head::Symbol, args)
@@ -33,6 +35,24 @@ function flatten_expr!(x)
     x
 end
 
+function detime_dvs(op::Operation)
+  if op.op isa Variable
+    Operation(Variable(op.op.name),Expression[])
+  else
+    Operation(op.op,detime_dvs.(op.args))
+  end
+end
+detime_dvs(op::Constant) = op
+
+function retime_dvs(op::Operation,dvs,iv)
+  if op.op isa Variable && op.op ∈ dvs
+    Operation(Variable(op.op.name),Expression[iv])
+  else
+    Operation(op.op,retime_dvs.(op.args,(dvs,),iv))
+  end
+end
+retime_dvs(op::Constant,dvs,iv) = op
+
 function build_function(rhss, vs, ps = (), args = (), conv = simplified_expr, expression = Val{true};
                         checkbounds = false, constructor=nothing, linenumbers = true)
     _vs = map(x-> x isa Operation ? x.op : x, vs)