Merge pull request #241 from SciML/Vaibhavdixit02-patch-1

Add `obj_sparse` and `cons_sparse` fields to `AutoModelingToolkit` and pass to derivative functions

Author: Vaibhavdixit02

src/function/mtk.jl

@@ -1,6 +1,11 @@
-struct AutoModelingToolkit <: AbstractADType end
+struct AutoModelingToolkit <: AbstractADType
+    obj_sparse::Bool
+    cons_sparse::Bool
+end
+
+AutoModelingToolkit() = AutoModelingToolkit(false, false)
 
-function instantiate_function(f, x, ::AutoModelingToolkit, p, num_cons=0)
+function instantiate_function(f, x, ad::AutoModelingToolkit, p, num_cons=0)
     p = isnothing(p) ? SciMLBase.NullParameters() : p
     sys = ModelingToolkit.modelingtoolkitize(OptimizationProblem(f, x, p))
 
@@ -12,8 +17,8 @@ function instantiate_function(f, x, ::AutoModelingToolkit, p, num_cons=0)
     end
 
     if f.hess === nothing
-        hess_oop, hess_iip = ModelingToolkit.generate_hessian(sys, expression=Val{false})
-        hess(J, u) = (hess_iip(J, u, p); J)
+        hess_oop, hess_iip = ModelingToolkit.generate_hessian(sys, expression=Val{false}, sparse = ad.obj_sparse)
+        hess(H, u) = (hess_iip(H, u, p); H)
     else
         hess = f.hess
     end
@@ -36,17 +41,17 @@ function instantiate_function(f, x, ::AutoModelingToolkit, p, num_cons=0)
     end
 
     if f.cons !== nothing && f.cons_j === nothing
+        jac_oop, jac_iip = ModelingToolkit.generate_jacobian(cons_sys, expression=Val{false}, sparse=ad.cons_sparse)
         cons_j = function (J, θ)
-            jac_oop, jac_iip = ModelingToolkit.generate_jacobian(cons_sys, expression=Val{false})
             jac_iip(J, θ, p)
         end
     else
         cons_j = f.cons_j
     end
 
     if f.cons !== nothing && f.cons_h === nothing
+        cons_hess_oop, cons_hess_iip = ModelingToolkit.generate_hessian(cons_sys, expression=Val{false}, sparse=ad.cons_sparse)
         cons_h = function (res, θ)
-            cons_hess_oop, cons_hess_iip = ModelingToolkit.generate_hessian(cons_sys, expression=Val{false})
             cons_hess_iip(res, θ, p)
         end
     else

test/ADtests.jl

@@ -57,6 +57,20 @@ H3 = [Array{Float64}(undef, 2, 2), Array{Float64}(undef, 2, 2)]
 optprob.cons_h(H3, x0)
 @test H3 == [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]
 
+optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoModelingToolkit(true, true), cons=con2_c)
+optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoModelingToolkit(true, true), nothing, 2)
+using SparseArrays
+sH = sparse([1, 1, 2, 2], [1, 2, 1, 2], zeros(4))
+optprob.hess(sH, x0)
+@test sH == H2
+@test optprob.cons(x0) == [0.0, 0.0]
+sJ = sparse([1, 1, 2, 2], [1, 2, 1, 2], zeros(4))
+optprob.cons_j(sJ, [5.0, 3.0])
+@test all(isapprox(sJ, [10.0 6.0; -0.149013 -0.958924]; rtol=1e-3))
+sH3 = [sparse([1,2], [1, 2], zeros(2)), sparse([1, 1, 2], [1, 2, 1], zeros(3))]
+optprob.cons_h(sH3, x0)
+@test Array.(sH3) == [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]
+
 optf = OptimizationFunction(rosenbrock, GalacticOptim.AutoForwardDiff())
 optprob = GalacticOptim.instantiate_function(optf, x0, GalacticOptim.AutoForwardDiff(), nothing)
 optprob.grad(G2, x0)

test/Project.toml

@@ -12,6 +12,7 @@ Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Tracker = "9f7883ad-71c0-57eb-9f7f-b5c9e6d3789c"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"