Pass hessian prototype to OptimizationFunction (#1606)

Co-authored-by: Christopher Rackauckas <[email protected]>

Author: Vaibhavdixit02

src/systems/optimization/optimizationsystem.jl

@@ -184,10 +184,17 @@ function DiffEqBase.OptimizationProblem{iip}(sys::OptimizationSystem, u0map,
         _hess = nothing
     end
 
+    if sparse
+        hess_prototype = hessian_sparsity(sys)
+    else
+        hess_prototype = nothing
+    end
+
     _f = DiffEqBase.OptimizationFunction{iip}(f,
                                               SciMLBase.NoAD();
                                               grad = _grad,
-                                              hess = _hess)
+                                              hess = _hess,
+                                              hess_prototype = hess_prototype)
 
     defs = defaults(sys)
     defs = mergedefaults(defs, parammap, ps)
@@ -251,6 +258,12 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
         _hess = :nothing
     end
 
+    if sparse
+        hess_prototype = hessian_sparsity(sys)
+    else
+        hess_prototype = nothing
+    end
+
     defs = defaults(sys)
     defs = mergedefaults(defs, parammap, ps)
     defs = mergedefaults(defs, u0map, dvs)
@@ -269,7 +282,8 @@ function OptimizationProblemExpr{iip}(sys::OptimizationSystem, u0,
         ub = $ub
         _f = OptimizationFunction{iip}(f, SciMLBase.NoAD();
                                        grad = grad,
-                                       hess = hess)
+                                       hess = hess,
+                                       hess_prototype = hess_prototype)
         OptimizationProblem{$iip}(_f, u0, p; lb = lb, ub = ub, kwargs...)
     end
 end

test/optimizationsystem.jl

@@ -21,7 +21,10 @@ calculate_hessian(combinedsys)
 generate_function(combinedsys)
 generate_gradient(combinedsys)
 generate_hessian(combinedsys)
-ModelingToolkit.hessian_sparsity(combinedsys)
+hess_sparsity = ModelingToolkit.hessian_sparsity(sys1)
+sparse_prob = OptimizationProblem(sys1, [x, y], [a, b], grad = true, sparse = true)
+@test sparse_prob.f.hess_prototype.rowval == hess_sparsity.rowval
+@test sparse_prob.f.hess_prototype.colptr == hess_sparsity.colptr
 
 u0 = [sys1.x => 1.0
       sys1.y => 2.0