[WIP] Ensemble

Project.toml

@@ -13,6 +13,7 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36"
 Pkg = "44cfe95a-1eb2-52ea-b672-e2afdf69b78f"
 Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
+QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
 Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"

src/Optimization.jl

@@ -12,7 +12,7 @@ if !isdefined(Base, :get_extension)
 end
 
 using Logging, ProgressLogging, ConsoleProgressMonitor, TerminalLoggers, LoggingExtras
-using ArrayInterface, Base.Iterators, SparseArrays, LinearAlgebra
+using ArrayInterface, Base.Iterators, SparseArrays, LinearAlgebra, QuasiMonteCarlo
 using Pkg
 
 import SciMLBase: OptimizationProblem, OptimizationFunction, ObjSense,
@@ -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

@@ -0,0 +1,19 @@
+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.EnsembleProblem(prob::OptimizationProblem, trajectories::Int; kwargs...)
+    if prob.lb != nothing && prob.ub != nothing
+        u0s = QuasiMonteCarlo.sample(trajectories, prob.lb, prob.ub, LatinHypercubeSample())
+        prob_func = (prob, i, repeat = nothing) -> remake(prob, u0 = u0s[:, i])
+    else
+        error("EnsembleProblem requires either initial points as second argument or lower and upper bounds to be defined with the trajectories second argument method.")
+    end
+    return SciMLBase.EnsembleProblem(prob; prob_func, kwargs...)
+end
+
+
+function SciMLBase.solve(prob::EnsembleProblem{T}, args...; kwargs...) where {T <: OptimizationProblem}
+    return SciMLBase.__solve(prob, args...; kwargs...)
+end

test/ensemble.jl

@@ -0,0 +1,18 @@
+using Optimization, OptimizationOptimJL, ForwardDiff, Test
+
+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(), EnsembleThreads(), trajectories = 4, maxiters = 5)
+@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective
+
+sol = Optimization.solve(ensembleprob, OptimizationOptimJL.BFGS(), EnsembleDistributed(), trajectories = 4, maxiters = 5)
+@test findmin(i -> sol[i].objective, 1:4)[1] < sol1.objective