Less allocations

Author: YingboMa

src/systems/diffeqs/abstractodesystem.jl

@@ -122,6 +122,7 @@ function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = param
            nothing,
     isdde = false,
     has_difference = false,
+    split_parameters = false,
     kwargs...)
     if isdde
         eqs = delay_to_function(sys)
@@ -151,10 +152,12 @@ function generate_function(sys::AbstractODESystem, dvs = states(sys), ps = param
             build_function(rhss, ddvs, u, p, t; postprocess_fbody = pre,
                 states = sol_states,
                 kwargs...)
+        elseif split_parameters
+            build_function(rhss, u, p..., t; postprocess_fbody = pre, states = sol_states,
+                kwargs...)
         else
-            fun = build_function(rhss, u, p..., t; postprocess_fbody = pre, states = sol_states,
+            build_function(rhss, u, p, t; postprocess_fbody = pre, states = sol_states,
                 kwargs...)
-            fun[1], :((out, u, p, t)->$(fun[2])(out, u, p..., t))
         end
     end
 end
@@ -326,15 +329,23 @@ function DiffEqBase.ODEFunction{iip, specialize}(sys::AbstractODESystem, dvs = s
     checkbounds = false,
     sparsity = false,
     analytic = nothing,
+    split_parameters = false,
     kwargs...) where {iip, specialize}
     f_gen = generate_function(sys, dvs, ps; expression = Val{eval_expression},
-        expression_module = eval_module, checkbounds = checkbounds,
+        expression_module = eval_module, checkbounds = checkbounds, split_parameters,
         kwargs...)
     f_oop, f_iip = eval_expression ?
                    (drop_expr(@RuntimeGeneratedFunction(eval_module, ex)) for ex in f_gen) :
                    f_gen
-    f(u, p, t) = f_oop(u, p, t)
-    f(du, u, p, t) = f_iip(du, u, p, t)
+    if split_parameters
+        g(u, p, t) = f_oop(u, p..., t)
+        g(du, u, p, t) = f_iip(du, u, p..., t)
+        f = g
+    else
+        k(u, p, t) = f_oop(u, p, t)
+        k(du, u, p, t) = f_iip(du, u, p, t)
+        f = k
+    end
 
     if specialize === SciMLBase.FunctionWrapperSpecialize && iip
         if u0 === nothing || p === nothing || t === nothing
@@ -687,6 +698,7 @@ function get_u0_p(sys,
     parammap;
     use_union = false,
     tofloat = !use_union,
+    split_parameters = false,
     symbolic_u0 = false)
     eqs = equations(sys)
     dvs = states(sys)
@@ -699,7 +711,7 @@ function get_u0_p(sys,
     if symbolic_u0
         u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = false, use_union = false)
     else
-        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
+        u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = !split_parameters)
     end
     p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
     p = p === nothing ? SciMLBase.NullParameters() : p
@@ -717,15 +729,24 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     use_union = false,
     tofloat = !use_union,
     symbolic_u0 = false,
+    split_parameters = false,
     kwargs...)
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union, symbolic_u0)
-    split_ps, split_idxs = split_parameters_by_type(p)
-    split_sym_ps = Base.Fix1(getindex, parameters(sys)).(split_idxs)
+    u0, p, defs = get_u0_p(sys,
+        u0map,
+        parammap;
+        tofloat,
+        use_union,
+        symbolic_u0,
+        split_parameters)
+    if split_parameters
+        p, split_idxs = split_parameters_by_type(p)
+        ps = Base.Fix1(getindex, parameters(sys)).(split_idxs)
+    end
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)
@@ -739,12 +760,12 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)
 
-    f = constructor(sys, dvs, split_sym_ps, u0; ddvs = ddvs, tgrad = tgrad, jac = jac,
-        checkbounds = checkbounds, p = split_ps,
+    f = constructor(sys, dvs, ps, u0; ddvs = ddvs, tgrad = tgrad, jac = jac,
+        checkbounds = checkbounds, p = p,
         linenumbers = linenumbers, parallel = parallel, simplify = simplify,
-        sparse = sparse, eval_expression = eval_expression,
+        sparse = sparse, eval_expression = eval_expression, split_parameters,
         kwargs...)
-    implicit_dae ? (f, du0, u0, split_ps) : (f, u0, split_ps)
+    implicit_dae ? (f, du0, u0, p) : (f, u0, p)
 end
 
 function ODEFunctionExpr(sys::AbstractODESystem, args...; kwargs...)