Update OptimizationSystem and NonlinearSystem

Author: AayushSabharwal

src/systems/jumps/jumpsystem.jl

@@ -291,7 +291,8 @@ function DiffEqBase.DiscreteProblem(sys::JumpSystem, u0map, tspan::Union{Tuple,
         end
     end
 
-    df = DiscreteFunction{true, true}(f; syms = Symbol.(states(sys)), paramsyms = Symbol.(ps), sys = sys,
+    df = DiscreteFunction{true, true}(f; syms = Symbol.(states(sys)),
+                                      paramsyms = Symbol.(ps), sys = sys,
                                       observed = observedfun)
     DiscreteProblem(df, u0, tspan, p; kwargs...)
 end
@@ -331,7 +332,8 @@ function DiscreteProblemExpr(sys::JumpSystem, u0map, tspan::Union{Tuple, Nothing
         u0 = $u0
         p = $p
         tspan = $tspan
-        df = DiscreteFunction{true, true}(f, syms = $(Symbol.(states(sys))), paramsyms = Symbol.(ps))
+        df = DiscreteFunction{true, true}(f, syms = $(Symbol.(states(sys))),
+                                          paramsyms = $(Symbol.(parameters(sys))))
         DiscreteProblem(df, u0, tspan, p)
     end
 end

src/systems/nonlinear/nonlinearsystem.jl

@@ -240,7 +240,9 @@ function SciMLBase.NonlinearFunction{iip}(sys::NonlinearSystem, dvs = states(sys
                            jac_prototype = sparse ?
                                            similar(calculate_jacobian(sys, sparse = sparse),
                                                    Float64) : nothing,
-                           syms = Symbol.(states(sys)), observed = observedfun)
+                           syms = Symbol.(states(sys)),
+                           paramsyms = Symbol.(parameters(sys)),
+                           observed = observedfun)
 end
 
 """
@@ -285,7 +287,8 @@ function NonlinearFunctionExpr{iip}(sys::NonlinearSystem, dvs = states(sys),
         NonlinearFunction{$iip}(f,
                                 jac = jac,
                                 jac_prototype = $jp_expr,
-                                syms = $(Symbol.(states(sys))))
+                                syms = $(Symbol.(states(sys))),
+                                paramsyms = $(Symbol.(parameters(sys))))
     end
     !linenumbers ? striplines(ex) : ex
 end
@@ -314,6 +317,7 @@ function process_NonlinearProblem(constructor, sys::NonlinearSystem, u0map, para
 
     f = constructor(sys, dvs, ps, u0; jac = jac, checkbounds = checkbounds,
                     linenumbers = linenumbers, parallel = parallel, simplify = simplify,
+                    syms = Symbol.(dvs), paramsyms = Symbol.(ps),
                     sparse = sparse, eval_expression = eval_expression, kwargs...)
     return f, u0, p
 end

src/systems/optimization/optimizationsystem.jl

@@ -273,6 +273,8 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                                   grad = _grad,
                                                   hess = _hess,
                                                   hess_prototype = hess_prototype,
+                                                  syms = Symbol.(states(sys)),
+                                                  paramsyms = Symbol.(parameters(sys)),
                                                   cons = cons,
                                                   cons_j = cons_j,
                                                   cons_h = cons_h,
@@ -287,6 +289,8 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
                                                   SciMLBase.NoAD();
                                                   grad = _grad,
                                                   hess = _hess,
+                                                  syms = Symbol.(states(sys)),
+                                                  paramsyms = Symbol.(parameters(sys)),
                                                   hess_prototype = hess_prototype,
                                                   expr = obj_expr)
     end
@@ -399,9 +403,13 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
             cons = $cons
             cons_j = $cons_j
             cons_h = $cons_h
+            syms = $(Symbol.(states(sys)))
+            paramsyms = $(Symbol.(parameters(sys)))
             _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
                                            grad = grad,
                                            hess = hess,
+                                           syms = syms,
+                                           paramsyms = paramsyms,
                                            hess_prototype = hess_prototype,
                                            cons = cons,
                                            cons_j = cons_j,
@@ -421,9 +429,13 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
             hess = $_hess
             lb = $lb
             ub = $ub
+            syms = $(Symbol.(states(sys)))
+            paramsyms = $(Symbol.(parameters(sys)))
             _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
                                            grad = grad,
                                            hess = hess,
+                                           syms = syms,
+                                           paramsyms = paramsyms,
                                            hess_prototype = hess_prototype,
                                            expr = obj_expr)
             OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)