Merge pull request #935 from mxpoch/master

added callbacks to pycma wrapper

Author: ChrisRackauckas

docs/src/optimization_packages/pycma.md

@@ -13,7 +13,7 @@ Pkg.add("OptimizationPyCMA")
 
 ## Methods
 
-`PyCMAOpt` supports the usual keyword arguments `maxiters`, `maxtime`, `abstol`, `reltol` in addition to any PyCMA-specific options (passed verbatim via keyword arguments to `solve`).
+`PyCMAOpt` supports the usual keyword arguments `maxiters`, `maxtime`, `abstol`, `reltol`, `callback` in addition to any PyCMA-specific options (passed verbatim via keyword arguments to `solve`).
 
 ## Example
 

lib/OptimizationPyCMA/src/OptimizationPyCMA.jl

@@ -20,7 +20,7 @@ end
 # Defining the SciMLBase interface for PyCMAOpt
 SciMLBase.allowsbounds(::PyCMAOpt) = true
 SciMLBase.supports_opt_cache_interface(opt::PyCMAOpt) = true
-SciMLBase.allowscallback(::PyCMAOpt) = false
+SciMLBase.allowscallback(::PyCMAOpt) = true
 SciMLBase.requiresgradient(::PyCMAOpt) = false
 SciMLBase.requireshessian(::PyCMAOpt) = false
 SciMLBase.requiresconsjac(::PyCMAOpt) = false
@@ -47,7 +47,7 @@ function __map_optimizer_args(prob::OptimizationCache, opt::PyCMAOpt;
 
     # mapping Optimization.jl args
     mapped_args["bounds"] = (prob.lb, prob.ub)
-    
+
     if !("verbose" ∈ keys(mapped_args))
         mapped_args["verbose"] = -1   
     end 
@@ -116,46 +116,57 @@ function SciMLBase.__solve(cache::OptimizationCache{
 }
     local x
 
-    # doing conversions
-    maxiters = Optimization._check_and_convert_maxiters(cache.solver_args.maxiters)
-    maxtime = Optimization._check_and_convert_maxtime(cache.solver_args.maxtime)
-
     # wrapping the objective function
     _loss = function (θ)
         x = cache.f(θ, cache.p)    
         return first(x)
     end
 
+    _cb = function(es)
+        opt_state = Optimization.OptimizationState(; iter = pyconvert(Int, es.countiter),
+            u = pyconvert(Vector{Float64}, es.best.x),
+            objective = pyconvert(Float64, es.best.f),
+            original = es)
+
+        cb_call = cache.callback(opt_state, x...)
+        if !(cb_call isa Bool)
+            error("The callback should return a boolean `halt` for whether to stop the optimization process.")
+        end
+        if cb_call
+            es.opts.set(Dict("termination_callback" => es -> true))
+        end
+    end
+
+    # doing conversions
+    maxiters = Optimization._check_and_convert_maxiters(cache.solver_args.maxiters)
+    maxtime = Optimization._check_and_convert_maxtime(cache.solver_args.maxtime)
+
     # converting the Optimization.jl Args to PyCMA format
     opt_args = __map_optimizer_args(cache, cache.opt; cache.solver_args...,
         maxiters = maxiters,
         maxtime = maxtime)
 
     # init the CMAopt class
     es = get_cma().CMAEvolutionStrategy(cache.u0, 1, pydict(opt_args))
-    logger = es.logger
 
     # running the optimization
     t0 = time()
-    opt_res = es.optimize(_loss)
+    opt_res = es.optimize(_loss, callback = _cb)
     t1 = time()
-
-    # loading logged files from disk
-    logger.load()
-
+    
     # reading the results
     opt_ret_dict = opt_res.stop()
     retcode = __map_pycma_retcode(pyconvert(Dict{String, Any}, opt_ret_dict))
 
     # logging and returning results of the optimization
     stats = Optimization.OptimizationStats(;
-        iterations = length(logger.xmean),
+        iterations = pyconvert(Int, es.countiter),
         time = t1 - t0,
-        fevals = length(logger.xmean))
+        fevals = pyconvert(Int, es.countevals))
 
     SciMLBase.build_solution(cache, cache.opt,
-        pyconvert(Float64, logger.xrecent[-1][-1]),
-        pyconvert(Float64, logger.f[-1][-1]); original = opt_res,
+    pyconvert(Vector{Float64}, opt_res.result.xbest),
+    pyconvert(Float64, opt_res.result.fbest); original = opt_res,
         retcode = retcode,
         stats = stats)
 end

lib/OptimizationPyCMA/test/runtests.jl

@@ -8,6 +8,15 @@ using OptimizationPyCMA, Test
     f = OptimizationFunction(rosenbrock)
     prob = OptimizationProblem(f, x0, _p, lb = [-1.0, -1.0], ub = [0.8, 0.8])
     sol = solve(prob, PyCMAOpt())
-    @test 10 * sol.objective < l1    
-    sol = solve(prob, PyCMAOpt(), maxiters = 100, verbose=-1, seed=42)
+    @test 10 * sol.objective < l1
+
+    # test callback function
+    callback = function (state, l) 
+        if state.iter > 10
+            return true
+        end
+        return false
+    end
+    
+    sol = solve(prob, PyCMAOpt(), callback = callback, maxiters = 25, verbose=-1, seed=42)
 end