Merge pull request #298 from SciML/build_function

Generalize build_function

Author: ChrisRackauckas

README.md

@@ -84,18 +84,18 @@ connected = ODESystem(connections,t,[α],[γ],systems=[lorenz1,lorenz2])
 u0 = [lorenz1.x => 1.0,
       lorenz1.y => 0.0,
       lorenz1.z => 0.0,
-	  lorenz2.x => 0.0,
-	  lorenz2.y => 1.0,
-	  lorenz2.z => 0.0,
-	  α => 2.0]
+      lorenz2.x => 0.0,
+      lorenz2.y => 1.0,
+      lorenz2.z => 0.0,
+      α => 2.0]
 
 p  = [lorenz1.σ => 10.0,
       lorenz1.ρ => 28.0,
       lorenz1.β => 8/3,
-	  lorenz2.σ => 10.0,
-	  lorenz2.ρ => 28.0,
-	  lorenz2.β => 8/3,
-	  γ => 2.0]
+      lorenz2.σ => 10.0,
+      lorenz2.ρ => 28.0,
+      lorenz2.β => 8/3,
+      γ => 2.0]
 
 tspan = (0.0,100.0)
 prob = ODEProblem(connected,u0,tspan,p)

src/build_function.jl

@@ -16,10 +16,10 @@ i.e. f(u,p,args...) for the out-of-place and scalar functions and
 `f!(du,u,p,args..)` for the in-place version.
 
 ```julia
-build_function(ex vs, ps = (), args = (),
-               conv = simplified_expr, expression = Val{true};
+build_function(ex, args...;
+               conv = simplified_expr, expression = Val{true},
                checkbounds = false, constructor=nothing,
-               linenumbers = true, target = JuliaTarget())
+               linenumbers = false, target = JuliaTarget())
 ```
 
 Arguments:
@@ -57,23 +57,23 @@ function build_function(args...;target = JuliaTarget(),kwargs...)
 end
 
 # Scalar output
-function _build_function(target::JuliaTarget, op::Operation, vs, ps = (), args = (),
-                         conv = simplified_expr, expression = Val{true};
+function _build_function(target::JuliaTarget, op::Operation, args...;
+                         conv = simplified_expr, expression = Val{true},
                          checkbounds = false, constructor=nothing,
                          linenumbers = true)
-    _vs = convert.(Variable,vs)
-    _ps = convert.(Variable,ps)
-    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
-    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(_ps)]
-    (ls, rs) = zip(var_pairs..., param_pairs...)
-    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
+
+    argnames = [gensym(:MTKArg) for i in 1:length(args)]
+    arg_pairs = map(vars_to_pairs,zip(argnames,args))
+    ls = reduce(vcat,first.(arg_pairs))
+    rs = reduce(vcat,last.(arg_pairs))
+    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, rs))
 
     fname = gensym(:ModelingToolkitFunction)
     out_expr = conv(op)
     let_expr = Expr(:let, var_eqs, out_expr)
     bounds_block = checkbounds ? let_expr : :(@inbounds begin $let_expr end)
 
-    fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))
+    fargs = Expr(:tuple,argnames...)
 
     oop_ex = :(
         ($(fargs.args...),) -> begin
@@ -92,19 +92,19 @@ function _build_function(target::JuliaTarget, op::Operation, vs, ps = (), args =
     end
 end
 
-function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
-                         conv = simplified_expr, expression = Val{true};
+function _build_function(target::JuliaTarget, rhss, args...;
+                         conv = simplified_expr, expression = Val{true},
                          checkbounds = false, constructor=nothing,
-                         linenumbers = true, multithread=false)
-    _vs = convert.(Variable,vs)
-    _ps = convert.(Variable,ps)
-    var_pairs   = [(u.name, :(u[$i])) for (i, u) ∈ enumerate(_vs)]
-    param_pairs = [(p.name, :(p[$i])) for (i, p) ∈ enumerate(_ps)]
-    (ls, rs) = zip(var_pairs..., param_pairs...)
+                         linenumbers = false, multithread=false)
 
-    var_eqs = Expr(:(=), build_expr(:tuple, ls), build_expr(:tuple, rs))
+    argnames = [gensym(:MTKArg) for i in 1:length(args)]
+    arg_pairs = map(vars_to_pairs,zip(argnames,args))
+    ls = reduce(vcat,first.(arg_pairs))
+    rs = reduce(vcat,last.(arg_pairs))
+    var_eqs = Expr(:(=), ModelingToolkit.build_expr(:tuple, ls), ModelingToolkit.build_expr(:tuple, rs))
 
     fname = gensym(:ModelingToolkitFunction)
+    fargs = Expr(:tuple,argnames...)
 
     X = gensym(:MTIIPVar)
     if rhss isa SparseMatrixCSC
@@ -135,20 +135,20 @@ function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
     if rhss isa Matrix
         arr_sys_expr = build_expr(:vcat, [build_expr(:row,[conv(rhs) for rhs ∈ rhss[i,:]]) for i in 1:size(rhss,1)])
         # : x because ??? what to do in the general case?
-        _constructor = constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} :  x->(out = similar(typeof(u),$(size(rhss)...)); out .= x)) : constructor
+        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} :  x->(out = similar(typeof($(fargs.args[1])),$(size(rhss)...)); out .= x)) : constructor
     elseif typeof(rhss) <: Array && !(typeof(rhss) <: Vector)
         vector_form = build_expr(:vect, [conv(rhs) for rhs ∈ rhss])
         arr_sys_expr = :(reshape($vector_form,$(size(rhss)...)))
-        _constructor = constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SArray{$(size(rhss)...)} :  x->(out = similar(typeof(u),$(size(rhss)...)); out .= x)) : constructor
+        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SArray{$(size(rhss)...)} :  x->(out = similar(typeof($(fargs.args[1])),$(size(rhss)...)); out .= x)) : constructor
     elseif rhss isa SparseMatrixCSC
         vector_form = build_expr(:vect, [conv(rhs) for rhs ∈ nonzeros(rhss)])
-        arr_sys_expr = :(SparseMatrixCSC{eltype(u),Int}($(size(rhss)...), $(rhss.colptr), $(rhss.rowval), $vector_form))
+        arr_sys_expr = :(SparseMatrixCSC{eltype($(first(argnames))),Int}($(size(rhss)...), $(rhss.colptr), $(rhss.rowval), $vector_form))
         # Static and sparse? Probably not a combo that will actually be hit, but give a default anyways
-        _constructor = constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} : x->x) : constructor
+        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.SMatrix{$(size(rhss)...)} : x->x) : constructor
     else # Vector
         arr_sys_expr = build_expr(:vect, [conv(rhs) for rhs ∈ rhss])
         # Handle vector constructor separately using `typeof(u)` to support things like LabelledArrays
-        _constructor = constructor === nothing ? :(u isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof(u), eltype(X)) : x->convert(typeof(u),x)) : constructor
+        _constructor = constructor === nothing ? :($(first(argnames)) isa ModelingToolkit.StaticArrays.StaticArray ? ModelingToolkit.StaticArrays.similar_type(typeof($(fargs.args[1])), eltype(X)) : x->convert(typeof($(fargs.args[1])),x)) : constructor
     end
 
     let_expr = Expr(:let, var_eqs, tuple_sys_expr)
@@ -157,8 +157,6 @@ function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
     arr_bounds_block = checkbounds ? arr_let_expr : :(@inbounds begin $arr_let_expr end)
     ip_bounds_block = checkbounds ? ip_let_expr : :(@inbounds begin $ip_let_expr end)
 
-    fargs = ps == () ? :(u,$(args...)) : :(u,p,$(args...))
-
     oop_ex = :(
         ($(fargs.args...),) -> begin
             # If u is a weird non-StaticArray type and we want a sparse matrix, just do the optimized sparse anyways
@@ -191,6 +189,21 @@ function _build_function(target::JuliaTarget, rhss, vs, ps = (), args = (),
     end
 end
 
+vars_to_pairs(args) = vars_to_pairs(args[1],args[2])
+function vars_to_pairs(name,vs::AbstractArray)
+	_vs = convert.(Variable,vs)
+	names = [Symbol(u) for u ∈ _vs]
+	exs = [:($name[$i]) for (i, u) ∈ enumerate(_vs)]
+	names,exs
+end
+
+function vars_to_pairs(name,vs)
+	_vs = convert(Variable,vs)
+	names = [Symbol(_vs)]
+	exs = [name]
+	names,exs
+end
+
 get_varnumber(varop::Operation,vars::Vector{Operation}) =  findfirst(x->isequal(x,varop),vars)
 get_varnumber(varop::Operation,vars::Vector{<:Variable})  =  findfirst(x->isequal(x,varop.op),vars)
 
@@ -251,8 +264,8 @@ function _build_function(target::StanTarget, eqs, vs, ps, iv,
     """
 end
 
-function _build_function(target::CTarget, eqs, vs, ps, iv,
-                         conv = simplified_expr, expression = Val{true};
+function _build_function(target::CTarget, eqs, vs, ps, iv;
+                         conv = simplified_expr, expression = Val{true},
                          fname = :diffeqf, derivname=:internal_var___du,
                          varname=:internal_var___u,paramname=:internal_var___p)
     differential_equation = string(join([numbered_expr(eq,vs,ps,derivname=derivname,
@@ -265,8 +278,8 @@ function _build_function(target::CTarget, eqs, vs, ps, iv,
     """
 end
 
-function _build_function(target::MATLABTarget, eqs, vs, ps, iv,
-                         conv = simplified_expr, expression = Val{true};
+function _build_function(target::MATLABTarget, eqs, vs, ps, iv;
+                         conv = simplified_expr, expression = Val{true},
                          fname = :diffeqf, derivname=:internal_var___du,
                          varname=:internal_var___u,paramname=:internal_var___p)
     matstr = join([numbered_expr(eq.rhs,vs,ps,derivname=derivname,

src/systems/diffeqs/abstractodesystem.jl

@@ -36,24 +36,27 @@ function (f::ODEToExpr)(O::Operation)
 end
 (f::ODEToExpr)(x) = convert(Expr, x)
 
-function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_tgrad(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
     tgrad = calculate_tgrad(sys)
-    return build_function(tgrad, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
+    return build_function(tgrad, dvs, ps, sys.iv;
+                          conv = ODEToExpr(sys), kwargs...)
 end
 
-function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; sparse = false, kwargs...)
+function generate_jacobian(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); sparse = false, kwargs...)
     jac = calculate_jacobian(sys)
     if sparse
         jac = SparseArrays.sparse(jac)
     end
-    return build_function(jac, dvs, ps, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
+    return build_function(jac, dvs, ps, sys.iv;
+                          conv = ODEToExpr(sys), kwargs...)
 end
 
-function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = parameters(sys); kwargs...)
     rhss = [deq.rhs for deq ∈ equations(sys)]
     dvs′ = convert.(Variable,dvs)
     ps′ = convert.(Variable,ps)
-    return build_function(rhss, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
+    return build_function(rhss, dvs′, ps′, sys.iv;
+                          conv = ODEToExpr(sys),kwargs...)
 end
 
 function calculate_factorized_W(sys::AbstractODESystem, simplify=true)
@@ -88,8 +91,10 @@ function generate_factorized_W(sys::AbstractODESystem, vs = states(sys), ps = pa
                         StaticArrays.LU(LowerTriangular( SMatrix{$siz...}(UnitLowerTriangular(A)) ), UpperTriangular(A), SVector(ntuple(n->n, max($siz...))))
                     end)
 
-    Wfact_func   = build_function(Wfact  , vs, ps, (:__MTKWgamma,:t), ODEToExpr(sys), expression;constructor=constructor,kwargs...)
-    Wfact_t_func = build_function(Wfact_t, vs, ps, (:__MTKWgamma,:t), ODEToExpr(sys), expression;constructor=constructor,kwargs...)
+    Wfact_func   = build_function(Wfact  , vs, ps, Variable(:__MTKWgamma), sys.iv;
+                                  conv = ODEToExpr(sys), expression = expression, constructor=constructor,kwargs...)
+    Wfact_t_func = build_function(Wfact_t, vs, ps, Variable(:__MTKWgamma), sys.iv;
+                                  conv = ODEToExpr(sys), expression = expression, constructor=constructor,kwargs...)
 
     return (Wfact_func, Wfact_t_func)
 end
@@ -136,29 +141,29 @@ function DiffEqBase.ODEFunction{iip}(sys::AbstractODESystem, dvs = states(sys),
                                      jac = false, Wfact = false,
                                      sparse = false,
                                      kwargs...) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false}; kwargs...)
+    f_oop,f_iip = generate_function(sys, dvs, ps; expression=Val{false}, kwargs...)
 
     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}; kwargs...)
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps; expression=Val{false}, kwargs...)
         _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}; sparse = sparse, kwargs...)
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps; sparse = sparse, expression=Val{false}, kwargs...)
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
         _jac = nothing
     end
 
     if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false}; kwargs...)
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps; expression=Val{false}, kwargs...)
         Wfact_oop, Wfact_iip = tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
         _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)

src/systems/diffeqs/sdesystem.jl

@@ -80,10 +80,11 @@ function SDESystem(deqs::AbstractVector{<:Equation}, neqs, iv, dvs, ps;
     SDESystem(deqs, neqs, iv′, dvs′, ps′, tgrad, jac, Wfact, Wfact_t, name, systems)
 end
 
-function generate_diffusion_function(sys::SDESystem, dvs = sys.states, ps = sys.ps, expression = Val{true}; kwargs...)
+function generate_diffusion_function(sys::SDESystem, dvs = sys.states, ps = sys.ps; kwargs...)
     dvs′ = convert.(Variable,dvs)
     ps′ = convert.(Variable,ps)
-    return build_function(sys.noiseeqs, dvs′, ps′, (sys.iv.name,), ODEToExpr(sys), expression; kwargs...)
+    return build_function(sys.noiseeqs, dvs′, ps′, sys.iv;
+                          conv = ODEToExpr(sys),kwargs...)
 end
 
 """
@@ -100,32 +101,32 @@ respectively.
 function DiffEqBase.SDEFunction{iip}(sys::SDESystem, dvs = sys.states, ps = sys.ps;
                                      version = nothing, tgrad=false, sparse = false,
                                      jac = false, Wfact = false, kwargs...) where {iip}
-    f_oop,f_iip = generate_function(sys, dvs, ps, Val{false}; kwargs...)
-    g_oop,g_iip = generate_diffusion_function(sys, dvs, ps, Val{false}; kwargs...)
+    f_oop,f_iip = generate_function(sys, dvs, ps; expression=Val{false}, kwargs...)
+    g_oop,g_iip = generate_diffusion_function(sys, dvs, ps; expression=Val{false}, kwargs...)
 
     f(u,p,t) = f_oop(u,p,t)
     f(du,u,p,t) = f_iip(du,u,p,t)
     g(u,p,t) = g_oop(u,p,t)
     g(du,u,p,t) = g_iip(du,u,p,t)
 
     if tgrad
-        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps, Val{false}; kwargs...)
+        tgrad_oop,tgrad_iip = generate_tgrad(sys, dvs, ps; expression=Val{false}, kwargs...)
         _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}; sparse=sparse, kwargs...)
+        jac_oop,jac_iip = generate_jacobian(sys, dvs, ps; expression=Val{false}, sparse=sparse, kwargs...)
         _jac(u,p,t) = jac_oop(u,p,t)
         _jac(J,u,p,t) = jac_iip(J,u,p,t)
     else
         _jac = nothing
     end
 
     if Wfact
-        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true, Val{false}; kwargs...)
+        tmp_Wfact,tmp_Wfact_t = generate_factorized_W(sys, dvs, ps, true; expression=Val{false}, kwargs...)
         Wfact_oop, Wfact_iip = tmp_Wfact
         Wfact_oop_t, Wfact_iip_t = tmp_Wfact_t
         _Wfact(u,p,dtgamma,t) = Wfact_oop(u,p,dtgamma,t)

src/systems/nonlinear/nonlinearsystem.jl

@@ -47,20 +47,22 @@ function calculate_jacobian(sys::NonlinearSystem)
     return jac
 end
 
-function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys), expression = Val{true};
+function generate_jacobian(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys);
                            sparse = false, kwargs...)
     jac = calculate_jacobian(sys)
     if sparse
         jac = SparseArrays.sparse(jac)
     end
-    return build_function(jac, convert.(Variable,vs), convert.(Variable,ps), (), AbstractSysToExpr(sys))
+    return build_function(jac, convert.(Variable,vs), convert.(Variable,ps),
+                          conv = AbstractSysToExpr(sys))
 end
 
-function generate_function(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys), expression = Val{true}; kwargs...)
+function generate_function(sys::NonlinearSystem, vs = states(sys), ps = parameters(sys); kwargs...)
     rhss = [eq.rhs for eq ∈ sys.eqs]
     vs′ = convert.(Variable,vs)
     ps′ = convert.(Variable,ps)
-    return build_function(rhss, vs′, ps′, (), AbstractSysToExpr(sys), expression; kwargs...)
+    return build_function(rhss, vs′, ps′;
+                          conv = AbstractSysToExpr(sys), kwargs...)
 end
 
 """
@@ -86,7 +88,7 @@ function DiffEqBase.NonlinearProblem{iip}(sys::NonlinearSystem,u0map,tspan,
     ps = parameters(sys)
 
     f = generate_function(sys;checkbounds=checkbounds,linenumbers=linenumbers,
-                              multithread=multithread,sparse=sparse)
+                              multithread=multithread,sparse=sparse,expression=Val{false})
     u0 = varmap_to_vars(u0map,dvs)
     p = varmap_to_vars(parammap,ps)
     NonlinearProblem(f,u0,tspan,p;kwargs...)