Merge pull request #265 from SciML/mtk_sparsity

Finish up AutoModelingToolkit sparsity handling 2

Author: Vaibhavdixit02

Project.toml

@@ -15,6 +15,7 @@ ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 TerminalLoggers = "5d786b92-1e48-4d6f-9151-6b4477ca9bed"
 
 [compat]

lib/OptimizationMOI/Project.toml

@@ -14,9 +14,10 @@ Optimization = "3"
 
 [extras]
 Ipopt = "b6b21f68-93f8-5de0-b562-5493be1d77c9"
+ModelingToolkit = "961ee093-0014-501f-94e3-6117800e7a78"
 NLopt = "76087f3c-5699-56af-9a33-bf431cd00edd"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [targets]
-test = ["Ipopt", "NLopt", "Test", "Zygote"]
+test = ["Ipopt", "ModelingToolkit", "NLopt", "Test", "Zygote"]

lib/OptimizationMOI/src/OptimizationMOI.jl

@@ -1,15 +1,15 @@
 module OptimizationMOI
 
-using MathOptInterface, Optimization, Optimization.SciMLBase
+using MathOptInterface, Optimization, Optimization.SciMLBase, Optimization.SparseArrays
 const MOI = MathOptInterface
 
 struct MOIOptimizationProblem{T,F<:OptimizationFunction,uType,P} <: MOI.AbstractNLPEvaluator
     f::F
     u0::uType
     p::P
-    J::Matrix{T}
-    H::Matrix{T}
-    cons_H::Vector{Matrix{T}}
+    J::Union{Matrix{T},SparseMatrixCSC{T}}
+    H::Union{Matrix{T},SparseMatrixCSC{T}}
+    cons_H::Vector{<:Union{Matrix{T},SparseMatrixCSC{T}}}
 end
 
 function MOIOptimizationProblem(prob::OptimizationProblem)
@@ -21,9 +21,9 @@ function MOIOptimizationProblem(prob::OptimizationProblem)
         f,
         prob.u0,
         prob.p,
-        zeros(T, num_cons, n),
-        zeros(T, n, n),
-        Matrix{T}[zeros(T, n, n) for i in 1:num_cons],
+        isnothing(f.cons_jac_prototype) ? zeros(T, num_cons, n) : convert.(T, f.cons_jac_prototype),
+        isnothing(f.hess_prototype) ? zeros(T, n, n) : convert.(T, f.hess_prototype),
+        isnothing(f.cons_hess_prototype) ? Matrix{T}[zeros(T, n, n) for i in 1:num_cons] : [convert.(T, f.cons_hess_prototype[i]) for i in 1:num_cons],
     )
 end
 

lib/OptimizationMOI/test/runtests.jl

@@ -1,4 +1,4 @@
-using OptimizationMOI, Optimization, Ipopt, NLopt, Zygote
+using OptimizationMOI, Optimization, Ipopt, NLopt, Zygote, ModelingToolkit
 using Test
 
 @testset "OptimizationMOI.jl" begin
@@ -30,4 +30,14 @@ using Test
 
     sol = solve(prob, OptimizationMOI.MOI.OptimizerWithAttributes(NLopt.Optimizer, "algorithm" => :LD_LBFGS))
     @test 10 * sol.minimum < l1
+
+    cons_circ = (x, p) -> [x[1]^2 + x[2]^2]
+    optprob = OptimizationFunction(rosenbrock, Optimization.AutoModelingToolkit(true, true); cons=cons_circ)
+    prob = OptimizationProblem(optprob, x0, _p, ucons=[Inf], lcons=[0.0])
+
+    sol = solve(prob, Ipopt.Optimizer())
+    @test 10 * sol.minimum < l1
+
+    sol = solve(prob, OptimizationMOI.MOI.OptimizerWithAttributes(Ipopt.Optimizer, "max_cpu_time" => 60.0))
+    @test 10 * sol.minimum < l1
 end

lib/OptimizationOptimJL/src/OptimizationOptimJL.jl

@@ -1,6 +1,6 @@
 module OptimizationOptimJL
 
-using Reexport, Optimization, Optimization.SciMLBase
+using Reexport, Optimization, Optimization.SciMLBase, Optimization.SparseArrays
 @reexport using Optim
 decompose_trace(trace::Optim.OptimizationTrace) = last(trace)
 decompose_trace(trace::Optim.OptimizationState) = trace
@@ -133,7 +133,8 @@ function ___solve(prob::OptimizationProblem, opt::Optim.AbstractOptimizer,
                 H .*= false
             end
         end
-        optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0)
+        F = eltype(prob.u0)
+        optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0, real(zero(F)), Optim.NLSolversBase.alloc_DF(prob.u0, real(zero(F))), isnothing(f.hess_prototype) ? Optim.NLSolversBase.alloc_H(prob.u0, real(zero(F))) : convert.(F, f.hess_prototype))
     end
 
     opt_args = __map_optimizer_args(prob, opt, callback=_cb, maxiters=maxiters, maxtime=maxtime, abstol=abstol, reltol=reltol; kwargs...)
@@ -285,7 +286,8 @@ function ___solve(prob::OptimizationProblem, opt::Optim.ConstrainedOptimizer,
             H .*= false
         end
     end
-    optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0)
+    F = eltype(prob.u0)
+    optim_f = Optim.TwiceDifferentiable(_loss, gg, fg!, hh, prob.u0, real(zero(F)), Optim.NLSolversBase.alloc_DF(prob.u0, real(zero(F))), isnothing(f.hess_prototype) ? Optim.NLSolversBase.alloc_H(prob.u0, real(zero(F))) : convert.(F, f.hess_prototype))
 
     cons! = (res, θ) -> res .= f.cons(θ)
 

lib/OptimizationOptimJL/test/runtests.jl

@@ -18,7 +18,7 @@ using Test
     @test 10 * sol.minimum < l1
 
     prob = OptimizationProblem(rosenbrock, x0, _p)
-    sol = solve(prob, Optim.NelderMead(;initial_simplex=Optim.AffineSimplexer(;a = 0.025, b = 0.5)))
+    sol = solve(prob, Optim.NelderMead(; initial_simplex=Optim.AffineSimplexer(; a=0.025, b=0.5)))
     @test 10 * sol.minimum < l1
 
     cons = (x, p) -> [x[1]^2 + x[2]^2]
@@ -96,4 +96,12 @@ using Test
     prob = OptimizationProblem(optprob, x0, _p)
     sol = solve(prob, Optim.BFGS())
     @test 10 * sol.minimum < l1
+
+    optprob = OptimizationFunction(rosenbrock, Optimization.AutoModelingToolkit(true, false))
+    prob = OptimizationProblem(optprob, x0, _p)
+    sol = solve(prob, Optim.Newton())
+    @test 10 * sol.minimum < l1
+
+    sol = solve(prob, Optim.KrylovTrustRegion())
+    @test 10 * sol.minimum < l1
 end

src/Optimization.jl

@@ -9,7 +9,7 @@ using Reexport
 using Requires
 using DiffResults
 using Logging, ProgressLogging, ConsoleProgressMonitor, TerminalLoggers, LoggingExtras
-using ArrayInterfaceCore, Base.Iterators
+using ArrayInterfaceCore, Base.Iterators, SparseArrays
 using Pkg
 
 import SciMLBase: OptimizationProblem, OptimizationFunction, AbstractADType

src/function/function.jl

@@ -5,7 +5,7 @@ This function is used internally by Optimization.jl to construct
 the necessary extra functions (gradients, Hessians, etc.) before
 optimization. Each of the ADType dispatches use the supplied automatic
 differentiation type in order to specify how the construction process
-occurs. 
+occurs.
 
 If no ADType is given, then the default `NoAD` dispatch simply
 defines closures on any supplied gradient function to enclose the
@@ -31,7 +31,7 @@ function instantiate_function(f, x, ::AbstractADType, p, num_cons = 0)
     cons_j = f.cons_j === nothing ? nothing : (res,x)->f.cons_j(res,x,p)
     cons_h = f.cons_h === nothing ? nothing : (res,x)->f.cons_h(res,x,p)
 
-    return OptimizationFunction{true}(f.f, SciMLBase.NoAD(); grad=grad, hess=hess, hv=hv, 
+    return OptimizationFunction{true}(f.f, SciMLBase.NoAD(); grad=grad, hess=hess, hv=hv,
         cons=cons, cons_j=cons_j, cons_h=cons_h,
-        hess_prototype=nothing, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
+        hess_prototype=f.hess_prototype, cons_jac_prototype=f.cons_jac_prototype, cons_hess_prototype=f.cons_hess_prototype)
 end

src/function/mtk.jl

@@ -1,50 +1,3 @@
-"""
-AutoModelingToolkit <: AbstractADType
-
-An AbstractADType choice for use in OptimizationFunction for automatically
-generating the unspecified derivative functions. Usage:
-
-```julia
-OptimizationFunction(f,AutoModelingToolkit();kwargs...)
-```
-
-This uses the [ModelingToolkit.jl](https://github.com/SciML/ModelingToolkit.jl)
-symbolic system for automatically converting the `f` function into
-a symbolic equation and uses symbolic differentiation in order to generate
-a fast derivative code. Note that this will also compile a new version
-of your `f` function that is automatically optimized. Because of the
-required symbolic analysis, the state and parameters are required in
-the function definition, i.e.:
-
-Summary:
-
-- Not compatible with GPUs
-- Compatible with Hessian-based optimization
-- Not compatible with Hv-based optimization
-- Not compatible with constraint functions
-
-## Constructor
-
-```julia
-OptimizationFunction(f,AutoModelingToolkit(),x0,p,
-                     grad = false, hess = false, sparse = false,
-                     checkbounds = false,
-                     linenumbers = true,
-                     parallel=SerialForm(),
-                     kwargs...)
-```
-
-The special keyword arguments are as follows:
-
-- `grad`: whether to symbolically generate the gradient function.
-- `hess`: whether to symbolically generate the Hessian function.
-- `sparse`: whether to use sparsity detection in the Hessian.
-- `checkbounds`: whether to perform bounds checks in the generated code.
-- `linenumbers`: whether to include line numbers in the generated code.
-- `parallel`: whether to automatically parallelize the calculations.
-
-For more information, see the [ModelingToolkit.jl `OptimizationSystem` documentation](https://mtk.sciml.ai/dev/systems/OptimizationSystem/)
-"""
 struct AutoModelingToolkit <: AbstractADType
     obj_sparse::Bool
     cons_sparse::Bool
@@ -56,6 +9,8 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons=0)
     p = isnothing(p) ? SciMLBase.NullParameters() : p
     sys = ModelingToolkit.modelingtoolkitize(OptimizationProblem(f, x, p))
 
+    hess_prototype, cons_jac_prototype, cons_hess_prototype = nothing, nothing, nothing
+
     if f.grad === nothing
         grad_oop, grad_iip = ModelingToolkit.generate_gradient(sys, expression=Val{false})
         grad(J, u) = (grad_iip(J, u, p); J)
@@ -72,7 +27,7 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons=0)
 
     if f.hv === nothing
         hv = function (H, θ, v, args...)
-            res = ArrayInterfaceCore.zeromatrix(θ)
+            res = adtype.obj_sparse ? hess_prototype : ArrayInterfaceCore.zeromatrix(θ)
             hess(res, θ, args...)
             H .= res * v
         end
@@ -105,7 +60,19 @@ function instantiate_function(f, x, adtype::AutoModelingToolkit, p, num_cons=0)
         cons_h = f.cons_h
     end
 
-    return OptimizationFunction{true}(f.f, adtype; grad=grad, hess=hess, hv=hv, 
+    if adtype.obj_sparse
+        _hess_prototype = ModelingToolkit.hessian_sparsity(sys)
+        hess_prototype = convert.(eltype(x), _hess_prototype)
+    end
+
+    if adtype.cons_sparse
+        _cons_jac_prototype = ModelingToolkit.jacobian_sparsity(cons_sys)
+        cons_jac_prototype = convert.(eltype(x), _cons_jac_prototype)
+        _cons_hess_prototype = ModelingToolkit.hessian_sparsity(cons_sys)
+        cons_hess_prototype = [convert.(eltype(x), _cons_hess_prototype[i]) for i in 1:num_cons]
+    end
+
+    return OptimizationFunction{true}(f.f, adtype; grad=grad, hess=hess, hv=hv,
         cons=cons, cons_j=cons_j, cons_h=cons_h,
-        hess_prototype=nothing, cons_jac_prototype=nothing, cons_hess_prototype=nothing)
+        hess_prototype=hess_prototype, cons_jac_prototype=cons_jac_prototype, cons_hess_prototype=cons_hess_prototype)
 end

test/ADtests.jl

@@ -1,6 +1,7 @@
 using Optimization, OptimizationOptimJL, OptimizationOptimisers, Test
 using ForwardDiff, Zygote, ReverseDiff, FiniteDiff, Tracker
 using ModelingToolkit
+
 x0 = zeros(2)
 rosenbrock(x, p=nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
 l1 = rosenbrock(x0)
@@ -61,13 +62,16 @@ optf = OptimizationFunction(rosenbrock, Optimization.AutoModelingToolkit(true, t
 optprob = Optimization.instantiate_function(optf, x0, Optimization.AutoModelingToolkit(true, true), nothing, 2)
 using SparseArrays
 sH = sparse([1, 1, 2, 2], [1, 2, 1, 2], zeros(4))
+@test findnz(sH)[1:2] == findnz(optprob.hess_prototype)[1:2]
 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))
+@test findnz(sJ)[1:2] == findnz(optprob.cons_jac_prototype)[1:2]
 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))]
+sH3 = [sparse([1, 2], [1, 2], zeros(2)), sparse([1, 1, 2], [1, 2, 1], zeros(3))]
+@test getindex.(findnz.(sH3), Ref([1,2])) == getindex.(findnz.(optprob.cons_hess_prototype), Ref([1,2]))
 optprob.cons_h(sH3, x0)
 @test Array.(sH3) == [[2.0 0.0; 0.0 2.0], [-0.0 1.0; 1.0 0.0]]