Merge pull request #155 from JuliaDiffEq/builder

Function-local safe symbolic Jacobians of numerical functions

Author: ChrisRackauckas

README.md

@@ -118,25 +118,24 @@ The `@variables` and `@parameters` macros support this with the following syntax
 julia> @variables x[1:3];
 julia> x
 3-element Array{Operation,1}:
- x[1]()
- x[2]()
- x[3]()
+ x₁()
+ x₂()
+ x₃()
 
 # support for arbitrary ranges and tensors
 julia> @variables y[2:3,1:5:6];
 julia> y
 2×2 Array{Operation,2}:
- y[2,1]()  y[2,6]()
- y[3,1]()  y[3,6]()
-
+    y₂̒₁() y₂̒₆()
+    y₃̒₁() y₃̒₆()
 
 # also works for dependent variables
 julia> @parameters t; @variables z[1:3](t);
 julia> z
 3-element Array{Operation,1}:
- z[1](t())
- z[2](t())
- z[3](t())
+ z₁(t())
+ z₂(t())
+ z₃(t())
  ```
 
 ## Core Principles
@@ -201,7 +200,7 @@ aliased to the given call, allowing implicit use of dependents for convenience.
 @parameters t α σ(..) β[1:2]
 @variables w(..) x(t) y() z(t, α, x)
 
-expr = β[1] * x + y^α + σ(3) * (z - t) - β[2] * w(t - 1)
+expr = β₁* x + y^α + σ(3) * (z - t) - β₂ * w(t - 1)
 ```
 
 Note that `@parameters` and `@variables` implicitly add `()` to values that

src/ModelingToolkit.jl

@@ -84,11 +84,11 @@ include("variables.jl")
 include("operations.jl")
 include("differentials.jl")
 include("equations.jl")
-include("systems/diffeqs/diffeqsystem.jl")
-include("systems/diffeqs/first_order_transform.jl")
-include("systems/nonlinear/nonlinear_system.jl")
 include("function_registration.jl")
 include("simplify.jl")
 include("utils.jl")
+include("systems/diffeqs/diffeqsystem.jl")
+include("systems/diffeqs/first_order_transform.jl")
+include("systems/nonlinear/nonlinear_system.jl")
 
 end # module

src/systems/diffeqs/diffeqsystem.jl

@@ -209,15 +209,82 @@ Create an `ODEFunction` from the [`ODESystem`](@ref). The arguments `dvs` and `p
 are used to set the order of the dependent variable and parameter vectors,
 respectively.
 """
-function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs, ps;
+function DiffEqBase.ODEFunction{iip}(sys::ODESystem, dvs, ps,
+                                     safe = Val{true};
                                      version = nothing,
-                                     jac = false, Wfact = false) where iip
-    expr = eval(generate_function(sys, dvs, ps))
-    jac_expr = jac ? nothing : eval(generate_jacobian(sys, dvs, ps))
-    Wfact_expr,Wfact_t_expr = Wfact ? (nothing,nothing) : eval.(generate_factorized_W(sys, dvs, ps))
-    ODEFunction{iip}(eval(expr),jac=jac_expr,
-                      Wfact = Wfact_expr, Wfact_t = Wfact_t_expr)
+                                     jac = false, Wfact = false) where {iip}
+    _f = eval(generate_function(sys, dvs, ps))
+    out_f_safe(u,p,t) = ModelingToolkit.fast_invokelatest(_f,typeof(u),u,p,t)
+    out_f_safe(du,u,p,t) = ModelingToolkit.fast_invokelatest(_f,Nothing,du,u,p,t)
+    out_f(u,p,t) = _f(u,p,t)
+    out_f(du,u,p,t) = _f(du,u,p,t)
+
+    if jac
+        @show generate_jacobian(sys, dvs, ps)
+        _jac = eval(generate_jacobian(sys, dvs, ps))
+        jac_f_safe(u,p,t) = ModelingToolkit.fast_invokelatest(_jac,Matrix{eltype(u)},u,p,t)
+        jac_f_safe(J,u,p,t) = ModelingToolkit.fast_invokelatest(_jac,Nothing,J,u,p,t)
+        jac_f(u,p,t) = _jac(u,p,t)
+        jac_f(J,u,p,t) = _jac(J,u,p,t)
+    else
+        jac_f_safe = nothing
+        jac_f = nothing
+    end
+
+    if Wfact
+        _Wfact,_Wfact_t = eval.(generate_factorized_W(sys, dvs, ps))
+        Wfact_f_safe(u,p,t) = ModelingToolkit.fast_invokelatest(_Wfact,Matrix{eltype(u)},u,p,t)
+        Wfact_f_safe(J,u,p,t) = ModelingToolkit.fast_invokelatest(_Wfact,Nothing,J,u,p,t)
+        Wfact_f_t_safe(u,p,t) = ModelingToolkit.fast_invokelatest(_Wfact,Matrix{eltype(u)},u,p,t)
+        Wfact_f_t_safe(J,u,p,t) = ModelingToolkit.fast_invokelatest(_Wfact,Nothing,J,u,p,t)
+        Wfact_f(u,p,t) = _Wfact(u,p,t)
+        Wfact_f(J,u,p,t) = _Wfact(J,u,p,t)
+        Wfact_f_t(u,p,t) = _Wfact_t(u,p,t)
+        Wfact_f_t(J,u,p,t) = _Wfact_t(J,u,p,t)
+    else
+        Wfact_f_safe = nothing
+        Wfact_f_t_safe = nothing
+        Wfact_f = nothing
+        Wfact_t_f = nothing
+    end
+
+    if safe === Val{true}
+        ODEFunction{iip}(out_f_safe,jac=jac_f_safe,
+                          Wfact = Wfact_f_safe,
+                          Wfact_t = Wfact_f_t_safe)
+    else
+        ODEFunction{iip}(out_f,jac=jac_f,
+                          Wfact = Wfact_f,
+                          Wfact_t = Wfact_t_f)
+    end
 end
+
 function DiffEqBase.ODEFunction(sys::ODESystem, args...; kwargs...)
     ODEFunction{true}(sys, args...; kwargs...)
 end
+
+"""
+$(SIGNATURES)
+
+Generate `ODESystem`, dependent variables, and parameters from an `ODEProblem`.
+"""
+function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
+    @parameters t
+    vars = [Variable(:x, i)(t) for i in eachindex(prob.u0)]
+    params = [Variable(:α,i; known = true)() for i in eachindex(prob.p)]
+    @derivatives D'~t
+
+    rhs = [D(var) for var in vars]
+
+    if DiffEqBase.isinplace(prob)
+        lhs = similar(vars, Any)
+        prob.f(lhs, vars, params, t)
+    else
+        lhs = prob.f(vars, params, t)
+    end
+
+    eqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
+    de = ODESystem(eqs)
+
+    de, vars, params
+end

src/utils.jl

@@ -51,10 +51,11 @@ function build_function(rhss, vs, ps, args = (), conv = rhs -> convert(Expr, rhs
     quote
         function $fname($X,u,p,$(args...))
             $ip_let_expr
+            nothing
         end
         function $fname(u,p,$(args...))
             X = $let_expr
-            T = $(constructor === nothing ? :(u isa StaticArrays.StaticArray ? StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, eltype(X)); du .= x)) : constructor)
+            T = $(constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->(du=similar(u, eltype(X)); du .= x)) : constructor)
             T(X)
         end
     end
@@ -91,26 +92,10 @@ function vars!(vars, O)
     return vars
 end
 
-"""
-$(SIGNATURES)
-
-Generate `ODESystem`, dependent variables, and parameters from an `ODEProblem`.
-"""
-function modelingtoolkitize(prob::DiffEqBase.ODEProblem)
-    t, = @parameters t; vars = [Variable(Symbol(:x, i))(t) for i in eachindex(prob.u0)]; params = [Variable(Symbol(:α, i); known = true)() for i in eachindex(prob.p)];
-    D, = @derivatives D'~t
-
-    rhs = [D(var) for var in vars]
-
-    if DiffEqBase.isinplace(prob)
-        lhs = similar(vars, Any)
-        prob.f(lhs, vars, params, t)
-    else
-        lhs = prob.f(vars, params, t)
-    end
-
-    eqs = vcat([rhs[i] ~ lhs[i] for i in eachindex(prob.u0)]...)
-    de = ODESystem(eqs)
-
-    de, vars, params
+@inline @generated function fast_invokelatest(f, ::Type{rt}, args...) where rt
+  tupargs = Expr(:tuple,(a==Nothing ? Int : a for a in args)...)
+  quote
+    _f = $(Expr(:cfunction, Base.CFunction, :f, rt, :((Core.svec)($((a==Nothing ? Int : a for a in args)...))), :(:ccall)))
+    return ccall(_f.ptr,rt,$tupargs,$((:(getindex(args,$i) === nothing ? 0 : getindex(args,$i)) for i in 1:length(args))...))
+  end
 end

src/variables.jl

@@ -1,5 +1,20 @@
 export Variable, @variables, @parameters
 
+const IndexMap = Dict{Char,Char}(
+            '0' => '₀',
+            '1' => '₁',
+            '2' => '₂',
+            '3' => '₃',
+            '4' => '₄',
+            '5' => '₅',
+            '6' => '₆',
+            '7' => '₇',
+            '8' => '₈',
+            '9' => '₉')
+function map_subscripts(indices)
+    str = string(indices)
+    join(IndexMap[c] for c in str)
+end
 
 """
 $(TYPEDEF)
@@ -20,7 +35,7 @@ struct Variable <: Function
     Variable(name; known = false) = new(name, known)
 end
 function Variable(name, indices...; known = false)
-    var_name = Symbol("$name[$(join(indices, ","))]")
+    var_name = Symbol("$(name)$(join(map_subscripts.(indices), "̒"))")
     Variable(var_name; known=known)
 end
 
@@ -113,7 +128,7 @@ end
 function _construct_var(var_name, known, call_args)
     if call_args === nothing
         :(Variable($(Meta.quot(var_name)); known = $known)())
-    elseif call_args[end] == :..
+    elseif !isempty(call_args) && call_args[end] == :..
         :(Variable($(Meta.quot(var_name)); known = $known))
     else
         :(Variable($(Meta.quot(var_name)); known = $known)($(call_args...)))
@@ -123,7 +138,7 @@ end
 function _construct_var(var_name, known, call_args, ind)
     if call_args === nothing
         :(Variable($(Meta.quot(var_name)), $ind...; known = $known)())
-    elseif call_args[end] == :..
+    elseif !isempty(call_args) && call_args[end] == :..
         :(Variable($(Meta.quot(var_name)), $ind...; known = $known))
     else
         :(Variable($(Meta.quot(var_name)), $ind...; known = $known)($(call_args...)))

test/variable_parsing.jl

@@ -56,3 +56,6 @@ s1 = [Variable(:s, 1, 1; known = true)() Variable(:s, 1, 2; known = true)()
 @variables x[1:2](t)
 x1 = [Variable(:x, 1)(t), Variable(:x, 2)(t)]
 @test isequal(x1, x)
+
+@variables a[1:11,1:2]
+@variables a()