Use EnsembleProblem instead of custom struct

src/Optimization.jl

@@ -23,6 +23,7 @@ include("utils.jl")
 include("function.jl")
 include("adtypes.jl")
 include("cache.jl")
+include("ensemble.jl")
 
 @static if !isdefined(Base, :get_extension)
     function __init__()

src/ensemble.jl

@@ -1,13 +1,20 @@
-struct EnsembleOptimizationProblem{T1} <: SciMLBase.AbstractEnsembleProblem
-    prob::OptimizationProblem{iip, F, T1} where {iip, F}
-    u0s::Vector{T1}
+function SciMLBase.EnsembleProblem(prob::OptimizationProblem, u0s::Vector{Vector{T}}; kwargs...) where {T}
+    prob_func = (prob, i, repeat = nothing) -> remake(prob, u0 = u0s[i])
+    return SciMLBase.EnsembleProblem(prob; prob_func, kwargs...)
 end
 
-function SciMLBase.__init(prob::EnsembleOptimizationProblem{T}, args...; kwargs...) where {T <: OptimizationProblem}
-    probs = [remake(prob.prob, u0=u0; kwargs...) for u0 in prob.u0s]
-    return [SciMLBase.__init(prob, args...; kwargs...) for prob in probs]
+function SciMLBase.init(prob::EnsembleProblem{T}, args...; kwargs...) where {T <: OptimizationProblem}
+    SciMLBase.__init(prob, args...; kwargs...)
 end
 
-function SciMLBase.__solve(caches::Vector{OptimizationCache}, args...; kwargs...)
+function SciMLBase.__init(prob::EnsembleProblem{T}, args...; trajectories,  kwargs...) where {T <: OptimizationProblem}
+    return [SciMLBase.__init(prob.prob_func(prob.prob, i), args...; kwargs...) for i in 1:trajectories]
+end
+
+function SciMLBase.solve!(cache::Vector{<:OptimizationCache}; kwargs...)
+    return [SciMLBase.solve!(cache[i]; kwargs...) for i in eachindex(cache)]
+end
+
+function SciMLBase.__solve(caches::Vector{<:OptimizationCache}, args...; kwargs...)
     return [SciMLBase.__solve(cache, args...; kwargs...) for cache in caches]
 end

test/ensemble.jl

@@ -0,0 +1,15 @@
+using Optimization, OptimizationOptimJL, ForwardDiff
+
+x0 = zeros(2)
+rosenbrock(x, p = nothing) = (1 - x[1])^2 + 100 * (x[2] - x[1]^2)^2
+l1 = rosenbrock(x0)
+
+optf = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())
+prob = OptimizationProblem(optf, x0)
+sol1 = Optimization.solve(prob, OptimizationOptimJL.BFGS(), maxiters = 5)
+
+
+ensembleprob = Optimization.EnsembleProblem(prob, [x0, x0 .+ rand(2), x0 .+ rand(2), x0 .+ rand(2)])
+sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(), trajectories = 4, maxiters = 5)
+
+@test findmin(i -> sol[i].objective, 1:4) < sol1.objective