Merge pull request #3940 from AayushSabharwal/as/fix-optprob

fix: fix sparse cost hessian

Author: AayushSabharwal

src/systems/codegen.jl

@@ -667,10 +667,9 @@ function calculate_cost_hessian(sys::System; sparse = false, simplify = false)
     obj = cost(sys)
     dvs = unknowns(sys)
     if sparse
-        exprs = Symbolics.sparsehessian(obj, dvs; simplify)::AbstractSparseArray
-        sparsity = similar(exprs, Float64)
+        return Symbolics.sparsehessian(obj, dvs; simplify)::AbstractSparseArray
     else
-        exprs = Symbolics.hessian(obj, dvs; simplify)
+        return Symbolics.hessian(obj, dvs; simplify)
     end
 end
 

test/initializationsystem.jl

@@ -5,6 +5,7 @@ using SymbolicIndexingInterface, SciMLStructures
 using SciMLStructures: Tunable
 using ModelingToolkit: t_nounits as t, D_nounits as D, observed
 using DynamicQuantities
+using DiffEqBase: BrownFullBasicInit
 
 @parameters g
 @variables x(t) y(t) [state_priority = 10] λ(t)

test/odesystem.jl

@@ -565,33 +565,32 @@ let
     eqs = [D(x[1]) ~ x[2]
            D(x[2]) ~ -x[1] - 0.5 * x[2] + k
            y ~ 0.9 * x[1] + x[2]]
-    @named sys = System(eqs, t, vcat(x, [y]), [k], defaults = Dict(x .=> 0))
+    @named sys = System(eqs, t, vcat(x, [y]), [k])
     sys = mtkcompile(sys)
 
     u0 = x .=> [0.5, 0]
     du0 = D.(x) .=> 0.0
-    prob = DAEProblem(sys, [du0; u0], (0, 50))
-    @test prob[x] ≈ [0.5, 0.0]
+    prob = DAEProblem(sys, du0, (0, 50); guesses = u0)
+    @test prob[x] ≈ [0.5, 1.0]
     @test prob.du0 ≈ [0.0, 0.0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
     sol = solve(prob, IDA())
     @test sol[y] ≈ 0.9 * sol[x[1]] + sol[x[2]]
     @test isapprox(sol[x[1]][end], 1, atol = 1e-3)
 
-    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5],
-        (0, 50))
+    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0], (0, 50); guesses = u0)
 
-    @test prob[x] ≈ [0.5, 0]
+    @test prob[x] ≈ [0.5, 1]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 1
     sol = solve(prob, IDA())
     @test isapprox(sol[x[1]][end], 1, atol = 1e-3)
 
-    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, x[1] => 0.5, k => 2],
-        (0, 50))
-    @test prob[x] ≈ [0.5, 0]
+    prob = DAEProblem(sys, [D(y) => 0, D(x[1]) => 0, D(x[2]) => 0, k => 2],
+        (0, 50); guesses = u0)
+    @test prob[x] ≈ [0.5, 3]
     @test prob.du0 ≈ [0, 0]
     @test prob.p isa MTKParameters
     @test prob.ps[k] ≈ 2
@@ -600,7 +599,7 @@ let
 
     # no initial conditions for D(x[1]) and D(x[2]) provided
     @test_throws ModelingToolkit.MissingVariablesError prob=DAEProblem(
-        sys, Pair[], (0, 50))
+        sys, Pair[], (0, 50); guesses = u0)
 
     prob = ODEProblem(sys, Pair[x[1] => 0], (0, 50))
     sol = solve(prob, Rosenbrock23())

test/optimizationsystem.jl

@@ -390,3 +390,21 @@ end
     obj = myeigvals_1(m)
     @test_nowarn OptimizationSystem(obj, p_free, []; name = :osys)
 end
+
+@testset "Test sparse hessian" begin
+    rosenbrock(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
+    @variables x[1:2]
+    @named sys = OptimizationSystem(rosenbrock(x))
+    sys = complete(sys)
+    prob = OptimizationProblem(sys, [x => [42.0, 12.37]]; hess = true, sparse = true)
+
+    symbolic_hess = Symbolics.hessian(cost(sys), x)
+    symbolic_hess_value = Symbolics.fast_substitute(symbolic_hess, Dict(x[1] => prob[x[1]], x[2] => prob[x[2]]))
+
+    oop_hess = prob.f.hess(prob.u0, prob.p)
+    @test oop_hess ≈ symbolic_hess_value
+
+    iip_hess = similar(prob.f.hess_prototype)
+    prob.f.hess(iip_hess, prob.u0, prob.p)
+    @test iip_hess ≈ symbolic_hess_value
+end