Merge pull request #2102 from SciML/fb/linear_opt

Optimize `linearize` by avoiding function generation

Author: baggepinnen

docs/src/systems/DiscreteSystem.md

@@ -11,6 +11,7 @@ DiscreteSystem
   - `get_eqs(sys)` or `equations(sys)`: The equations that define the Discrete System.
   - `get_delay_val(sys)`: The delay of the Discrete System.
   - `get_iv(sys)`: The independent variable of the Discrete System.
+  - `get_u0_p(sys, u0map, parammap)` Numeric arrays for the initial condition and parameters given `var => value` maps.
 
 ## Transformations
 

docs/src/systems/NonlinearSystem.md

@@ -11,6 +11,7 @@ NonlinearSystem
   - `get_eqs(sys)` or `equations(sys)`: The equations that define the nonlinear system.
   - `get_states(sys)` or `states(sys)`: The set of states in the nonlinear system.
   - `get_ps(sys)` or `parameters(sys)`: The parameters of the nonlinear system.
+  - `get_u0_p(sys, u0map, parammap)` Numeric arrays for the initial condition and parameters given `var => value` maps.
 
 ## Transformations
 

docs/src/systems/ODESystem.md

@@ -12,6 +12,7 @@ ODESystem
   - `get_states(sys)` or `states(sys)`: The set of states in the ODE.
   - `get_ps(sys)` or `parameters(sys)`: The parameters of the ODE.
   - `get_iv(sys)`: The independent variable of the ODE.
+  - `get_u0_p(sys, u0map, parammap)` Numeric arrays for the initial condition and parameters given `var => value` maps.
 
 ## Transformations
 

src/systems/abstractsystem.jl

@@ -1386,9 +1386,9 @@ lsys = ModelingToolkit.reorder_states(lsys, states(ssys), desired_order)
 function linearize(sys, lin_fun; t = 0.0, op = Dict(), allow_input_derivatives = false,
                    p = DiffEqBase.NullParameters())
     x0 = merge(defaults(sys), op)
-    f, u0, p = process_DEProblem(ODEFunction{true}, sys, x0, p)
+    u0, p2, _ = get_u0_p(sys, x0, p; use_union = false, tofloat = true)
 
-    linres = lin_fun(u0, p, t)
+    linres = lin_fun(u0, p2, t)
     f_x, f_z, g_x, g_z, f_u, g_u, h_x, h_z, h_u = linres
 
     nx, nu = size(f_u)

src/systems/diffeqs/abstractodesystem.jl

@@ -570,6 +570,26 @@ function ODEFunctionExpr{iip}(sys::AbstractODESystem, dvs = states(sys),
     !linenumbers ? striplines(ex) : ex
 end
 
+"""
+    u0, p, defs = get_u0_p(sys, u0map, parammap; use_union=false, tofloat=!use_union)
+
+Take dictionaries with initial conditions and parameters and convert them to numeric arrays `u0` and `p`. Also return the merged dictionary `defs` containing the entire operating point. 
+"""
+function get_u0_p(sys, u0map, parammap; use_union = false, tofloat = !use_union)
+    eqs = equations(sys)
+    dvs = states(sys)
+    ps = parameters(sys)
+
+    defs = defaults(sys)
+    defs = mergedefaults(defs, parammap, ps)
+    defs = mergedefaults(defs, u0map, dvs)
+
+    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
+    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
+    p = p === nothing ? SciMLBase.NullParameters() : p
+    u0, p, defs
+end
+
 function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
                            implicit_dae = false, du0map = nothing,
                            version = nothing, tgrad = false,
@@ -586,13 +606,7 @@ function process_DEProblem(constructor, sys::AbstractODESystem, u0map, parammap;
     ps = parameters(sys)
     iv = get_iv(sys)
 
-    defs = defaults(sys)
-    defs = mergedefaults(defs, parammap, ps)
-    defs = mergedefaults(defs, u0map, dvs)
-
-    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
-    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat, use_union)
-    p = p === nothing ? SciMLBase.NullParameters() : p
+    u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
 
     if implicit_dae && du0map !== nothing
         ddvs = map(Differential(iv), dvs)

src/systems/discrete_system/discrete_system.jl

@@ -418,17 +418,13 @@ function process_DiscreteProblem(constructor, sys::DiscreteSystem, u0map, paramm
                                  linenumbers = true, parallel = SerialForm(),
                                  eval_expression = true,
                                  use_union = false,
+                                 tofloat = !use_union,
                                  kwargs...)
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
 
-    defs = defaults(sys)
-    defs = mergedefaults(defs, parammap, ps)
-    defs = mergedefaults(defs, u0map, dvs)
-
-    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
-    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = !use_union, use_union)
+    u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)
 

src/systems/nonlinear/nonlinearsystem.jl

@@ -320,17 +320,13 @@ function process_NonlinearProblem(constructor, sys::NonlinearSystem, u0map, para
                                   linenumbers = true, parallel = SerialForm(),
                                   eval_expression = true,
                                   use_union = false,
+                                  tofloat = !use_union,
                                   kwargs...)
     eqs = equations(sys)
     dvs = states(sys)
     ps = parameters(sys)
 
-    defs = defaults(sys)
-    defs = mergedefaults(defs, parammap, ps)
-    defs = mergedefaults(defs, u0map, dvs)
-
-    u0 = varmap_to_vars(u0map, dvs; defaults = defs, tofloat = true)
-    p = varmap_to_vars(parammap, ps; defaults = defs, tofloat = !use_union, use_union)
+    u0, p, defs = get_u0_p(sys, u0map, parammap; tofloat, use_union)
 
     check_eqs_u0(eqs, dvs, u0; kwargs...)