Merge pull request #601 from SciML/Vaibhavdixit02-patch-3

Add callback to MOI
SciML · Oct 5, 2023 · eddf8c4 · eddf8c4
2 parents 7805260 + cb304de
commit eddf8c4
Show file tree

Hide file tree

Showing 4 changed files with 32 additions and 11 deletions.
diff --git a/lib/OptimizationMOI/Project.toml b/lib/OptimizationMOI/Project.toml
@@ -1,7 +1,7 @@
 name = "OptimizationMOI"
 uuid = "fd9f6733-72f4-499f-8506-86b2bdd0dea1"
 authors = ["Vaibhav Dixit <[email protected]> and contributors"]
-version = "0.1.15"
+version = "0.1.16"
 
 [deps]
 Ipopt_jll = "9cc047cb-c261-5740-88fc-0cf96f7bdcc7"

diff --git a/lib/OptimizationMOI/src/OptimizationMOI.jl b/lib/OptimizationMOI/src/OptimizationMOI.jl
@@ -22,10 +22,10 @@ function SciMLBase.allowsconstraints(opt::Union{MOI.AbstractOptimizer,
     MOI.OptimizerWithAttributes})
     true
 end
-function SciMLBase.allowscallback(opt::Union{MOI.AbstractOptimizer,
-    MOI.OptimizerWithAttributes})
-    false
-end
+# function SciMLBase.allowscallback(opt::Union{MOI.AbstractOptimizer,
+#     MOI.OptimizerWithAttributes})
+#     false
+# end
 
 function _create_new_optimizer(opt::MOI.OptimizerWithAttributes)
     return _create_new_optimizer(MOI.instantiate(opt, with_bridge_type = Float64))

diff --git a/lib/OptimizationMOI/src/nlp.jl b/lib/OptimizationMOI/src/nlp.jl
@@ -1,7 +1,7 @@
 mutable struct MOIOptimizationNLPEvaluator{T, F <: OptimizationFunction, RC, LB, UB,
     I,
     JT <: DenseOrSparse{T}, HT <: DenseOrSparse{T},
-    CHT <: DenseOrSparse{T}, S} <:
+    CHT <: DenseOrSparse{T}, S, CB} <:
                MOI.AbstractNLPEvaluator
     f::F
     reinit_cache::RC
@@ -14,6 +14,7 @@ mutable struct MOIOptimizationNLPEvaluator{T, F <: OptimizationFunction, RC, LB,
     J::JT
     H::HT
     cons_H::Vector{CHT}
+    callback::CB
 end
 
 function Base.getproperty(evaluator::MOIOptimizationNLPEvaluator, x::Symbol)
@@ -101,7 +102,7 @@ function SciMLBase.get_paramsyms(sol::SciMLBase.OptimizationSolution{
     sol.cache.evaluator.f.paramsyms
 end
 
-function MOIOptimizationNLPCache(prob::OptimizationProblem, opt; kwargs...)
+function MOIOptimizationNLPCache(prob::OptimizationProblem, opt; callback = nothing, kwargs...)
     reinit_cache = Optimization.ReInitCache(prob.u0, prob.p) # everything that can be changed via `reinit`
 
     num_cons = prob.ucons === nothing ? 0 : length(prob.ucons)
@@ -142,7 +143,8 @@ function MOIOptimizationNLPCache(prob::OptimizationProblem, opt; kwargs...)
         prob.sense,
         J,
         H,
-        cons_H)
+        cons_H,
+        callback)
     return MOIOptimizationNLPCache(evaluator, opt, NamedTuple(kwargs))
 end
 
@@ -169,7 +171,13 @@ function MOI.initialize(evaluator::MOIOptimizationNLPEvaluator,
 end
 
 function MOI.eval_objective(evaluator::MOIOptimizationNLPEvaluator, x)
-    return evaluator.f(x, evaluator.p)
+    if evaluator.callback === nothing
+        return evaluator.f(x, evaluator.p)
+    else
+        l = evaluator.f(x, evaluator.p)
+        evaluator.callback(x, l)
+        return l
+    end
 end
 
 function MOI.eval_constraint(evaluator::MOIOptimizationNLPEvaluator, g, x)
@@ -406,6 +414,11 @@ function SciMLBase.__solve(cache::MOIOptimizationNLPCache)
     MOI.set(opt_setup,
         MOI.NLPBlock(),
         MOI.NLPBlockData(con_bounds, cache.evaluator, true))
+
+    if cache.evaluator.callback !== nothing
+        MOI.set(opt_setup, MOI.Silent(), true)
+    end
+
     MOI.optimize!(opt_setup)
     if MOI.get(opt_setup, MOI.ResultCount()) >= 1
         minimizer = MOI.get(opt_setup, MOI.VariablePrimal(), θ)

diff --git a/lib/OptimizationMOI/test/runtests.jl b/lib/OptimizationMOI/test/runtests.jl
@@ -36,8 +36,16 @@ end
 
     optprob = OptimizationFunction((x, p) -> -rosenbrock(x, p), Optimization.AutoZygote())
     prob = OptimizationProblem(optprob, x0, _p; sense = Optimization.MaxSense)
-
-    sol = solve(prob, Ipopt.Optimizer())
+    global iter = 0
+    callback = function (p, l)
+        global iter
+        iter += 1
+
+        display(l)
+        return false
+    end
+
+    sol = solve(prob, Ipopt.Optimizer(); callback)
     @test 10 * sol.objective < l1
 
     # cache interface