How to record optimisation traces of multistart (ensemble) optimisation runs?

[TorkelE]

So, I'd like to record the trace of objective function evaluations (to evaluate the optimisation run), at least in systems biology this is a common workflow. I think I can achieve this using a callback for single runs. However, I am not sure how to do this for optimisation runs on the ensemble simulations.

[ChrisRackauckas]

push!(losses, l) in a callback. We will improve logging in the near future but for now that should be all you need.

[TorkelE]

So for single run I understand how to do that, but is there a way for multistart (where all the traces do not get combined n a single vector)?

From what I can tell, there is no way for the callback to know from which of the ensemble runs it is being run. if there is it would be good to know (and if not, add that as a suggestion!).

[ChrisRackauckas]

You can make a callback and add it to the prob in the prob_func, and then have a different array for each of those callbacks for the save.

[TorkelE]

Not fully sure, but will see if I can figure something like that out. Right now the callback only have access to an OptimizationState structure (and the loss value):
https://docs.sciml.ai/Optimization/stable/API/optimization_state/. Potentially there is something in the original field, I will explore this, but otherwise I don't think there is anything that allows you to differentiate callback calls from different ensemble runs from each other.

[ChrisRackauckas]

using Optimization, OptimizationOptimJL, Random

Random.seed!(100)

rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2
x0 = zeros(2)

optf = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())
prob = OptimizationProblem(optf, x0, [1.0, 100.0])
@time sol1 = Optimization.solve(prob, OptimizationOptimJL.BFGS(), maxiters = 5)

@show sol1.objective

x0s = [x0, x0 .+ rand(2), x0 .+ rand(2), x0 .+ rand(2)]
trajs = 100
traces = [Float64[] for i in 1:trajs]

function prob_func(prob, i, repeat)
    cb = function (state, loss_val)
     push!(traces[i], loss_val
   end
    remake(prob, u0 = x0s[1], callback = cb)
end

ensembleprob = Optimization.EnsembleProblem(prob; prob_func)
@time sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(),
    EnsembleThreads(), trajectories = 100, maxiters = 5)
@show findmin(i -> sol[i].objective, 1:100)[1]

should do it

[TorkelE]

Yes, that should work, thanks for the help sorting this out!