Add callback

Author: tmigot

src/AdaptiveRegularization.jl

@@ -58,6 +58,23 @@ Additional `kwargs` are used for stopping criterion, see `Stopping.jl`.
 # Output
 The returned value is a `GenericExecutionStats`, see `SolverCore.jl`.
 
+# Callback
+The callback is called at each iteration.
+The expected signature of the callback is `callback(nlp, solver, stats)`, and its output is ignored.
+Changing any of the input arguments will affect the subsequent iterations.
+In particular, setting `stats.status = :user` will stop the algorithm.
+All relevant information should be available in `nlp` and `solver`.
+Notably, you can access, and modify, the following:
+- `solver.stp`: stopping object used for the algorithm;
+- `solver.workspace`: additional allocations;
+- `stats`: structure holding the output of the algorithm (`GenericExecutionStats`), which contains, among other things:
+  - `stats.dual_feas`: norm of current gradient;
+  - `stats.iter`: current iteration counter;
+  - `stats.objective`: current objective function value;
+  - `stats.status`: current status of the algorithm. Should be `:unknown` unless the algorithm has attained a stopping criterion. Changing this to anything will stop the algorithm, but you should use `:user` to properly indicate the intention.
+  - `stats.elapsed_time`: elapsed time in seconds.
+
+
 This implementation uses `Stopping.jl`. Therefore, it is also possible to used
         
     TRARC(stp; kwargs...)

src/main.jl

@@ -103,6 +103,7 @@ function SolverCore.solve!(
   stats::GenericExecutionStats{T, S};
   robust::Bool = true,
   verbose::Integer = false,
+  callback = (args...) -> nothing,
   kwargs...,
 ) where {Pb, M, SRC, MStp, LoS, Score, S, T}
   PData = solver.meta
@@ -139,6 +140,7 @@ function SolverCore.solve!(
   set_solution!(stats, nlp_at_x.x)
   set_objective!(stats, nlp_at_x.fx)
   set_dual_residual!(stats, nlp_at_x.current_score)
+  set_iter!(stats, 0)
   set_time!(stats, nlp_at_x.current_time - nlp_stop.meta.start_time)
 
   verbose > 0 && @info log_header(
@@ -147,7 +149,9 @@ function SolverCore.solve!(
   )
   verbose > 0 && @info log_row(Any[iter, ft, norm_∇f, 0.0, "First iteration", α])
 
-  while !OK
+  callback(nlp, solver, stats)
+
+  while !OK && (stats.status != :user)
     preprocess!(nlp_stop, PData, workspace, ∇f, norm_∇f, α)
 
     if ~PData.OK
@@ -232,7 +236,7 @@ function SolverCore.solve!(
     set_dual_residual!(stats, nlp_at_x.current_score)
     set_iter!(stats, nlp_stop.meta.nb_of_stop)
     set_time!(stats, nlp_at_x.current_time - nlp_stop.meta.start_time)
-    # TODO: callback(nlp, solver, stats)
+    callback(nlp, solver, stats)
   end # while !OK
 
   stats

test/callback.jl

@@ -0,0 +1,24 @@
+@testset "Test callback" begin
+  f(x) = (x[1] - 1)^2 + 4 * (x[2] - x[1]^2)^2
+  nlp = ADNLPModel(f, [-1.2; 1.0])
+  function cb(nlp, solver, stats)
+    if stats.iter == 8
+      stats.status = :user
+    end
+  end
+  stats = TRARC(nlp, callback = cb)
+  @test stats.iter == 8
+end
+
+@testset "Test callback for NLS" begin
+  F(x) = [x[1] - 1; 2 * (x[2] - x[1]^2)]
+  nls = ADNLSModel(F, [-1.2; 1.0], 2)
+  function cb(nlp, solver, stats)
+    if stats.iter == 8
+      stats.status = :user
+    end
+  end
+
+  stats = TRARC(nls, callback = cb)
+  @test stats.iter == 8
+end

test/runtests.jl

@@ -7,6 +7,8 @@ using ADNLPModels, NLPModels, OptimizationProblems.ADNLPProblems, SolverCore, So
 # Stopping
 using Stopping
 
+include("callback.jl")
+
 @testset "Testing NLP solvers" begin
   @testset "$name" for name in ALL_solvers
     solver = eval(name)