Update caching in hybrid pso

Author: utkarsh530

src/hybrid.jl

@@ -5,27 +5,83 @@
     result[i] = sol.u
 end
 
-function SciMLBase.__solve(prob::SciMLBase.OptimizationProblem,
-        opt::HybridPSO{Backend, LocalOpt},
-        args...;
-        abstol = nothing,
-        reltol = nothing,
-        maxiters = 100,
-        local_maxiters = 10,
+struct HybridPSOCache{TPc, TSp, TAlg}
+    pso_cache::TPc
+    start_points::TSp
+    alg::TAlg
+end
+
+function SciMLBase.init(
+        prob::OptimizationProblem, opt::HybridPSO{Backend, LocalOpt}, args...;
         kwargs...) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
-    t0 = time()
     psoalg = opt.pso
-    local_opt = opt.local_opt
     backend = opt.backend
 
-    sol_pso = solve(prob, psoalg, args...; maxiters, kwargs...)
+    pso_cache = init(prob, psoalg)
+
+    start_points = KernelAbstractions.allocate(
+        backend, typeof(prob.u0), opt.pso.num_particles)
+
+    return HybridPSOCache{
+        typeof(pso_cache), typeof(start_points), typeof(opt)}(pso_cache, start_points, opt)
+end
+
+function reinit_cache!(cache::HybridPSOCache,
+        opt::HybridPSO{Backend, LocalOpt}) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
+    reinit!(cache.pso_cache)
+    fill!(cache.start_points, zero(eltype(cache.start_points)))
+    # prob = cache.prob
+    # backend = opt.backend
+    # particles = cache.particles
+
+    # kernel! = PSOGPU.gpu_init_particles!(backend)
+    # kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
+
+    # best_particle = minimum(particles)
+    # _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
+
+    # copyto!(cache.gbest, [_init_gbest])
+
+    return nothing
+end
+
+function Base.getproperty(cache::HybridPSOCache, name::Symbol)
+    if name ∈ (:start_points, :pso_cache, :alg)
+        return getfield(cache, name)
+    else
+        return getproperty(cache.pso_cache, name)
+    end
+end
+
+function Base.setproperty!(cache::HybridPSOCache, name::Symbol, val)
+    if name ∈ (:start_points, :pso_cache, :alg)
+        return setfield!(cache, name, val)
+    else
+        return setproperty!(cache.pso_cache, name, val)
+    end
+end
+
+function SciMLBase.solve!(
+        cache::HybridPSOCache, opt::HybridPSO{Backend, LocalOpt}, args...;
+        abstol = nothing,
+        reltol = nothing,
+        maxiters = 100, local_maxiters = 10, kwargs...) where {
+        Backend, LocalOpt <: Union{LBFGS, BFGS}}
+
+    pso_cache = cache.pso_cache
+
+    sol_pso = solve!(pso_cache)
     x0s = sol_pso.original
-    prob = remake(prob, lb = nothing, ub = nothing)
+
+    backend = opt.backend
+
+    prob = remake(cache.prob, lb = nothing, ub = nothing)
     f = Base.Fix2(prob.f.f, prob.p)
     ∇f = instantiate_gradient(f, prob.f.adtype)
 
     kernel = simplebfgs_run!(backend)
-    result = KernelAbstractions.allocate(backend, typeof(prob.u0), length(x0s))
+    result = cache.start_points
+    copyto!(result, x0s)
     nlprob = NonlinearProblem{false}(∇f, prob.u0)
 
     nlalg = LocalOpt isa LBFGS ?

src/solve.jl

@@ -35,32 +35,6 @@ function SciMLBase.init(
         prob, opt, particles, init_gbest)
 end
 
-function SciMLBase.init(
-        prob::OptimizationProblem, opt::ParallelPSOKernel, args...; kwargs...)
-    backend = opt.backend
-    @assert prob.u0 isa SArray
-
-    ## Bounds check
-    lb, ub = check_init_bounds(prob)
-    lb, ub = check_init_bounds(prob)
-    prob = remake(prob; lb = lb, ub = ub)
-
-    particles = KernelAbstractions.allocate(
-        backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
-    kernel! = gpu_init_particles!(backend)
-
-    kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
-
-    best_particle = minimum(particles)
-    _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
-
-    init_gbest = KernelAbstractions.allocate(backend, typeof(_init_gbest), (1,))
-    copyto!(init_gbest, [_init_gbest])
-    return PSOCache{
-        typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
-        prob, opt, particles, init_gbest)
-end
-
 function SciMLBase.init(
         prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; kwargs...)
     backend = opt.backend
@@ -85,7 +59,7 @@ function SciMLBase.init(
         prob, opt, particles, init_gbest)
 end
 
-function SciMLBase.reinit!(cache::PSOCache; kwargs...)
+function SciMLBase.reinit!(cache::Union{PSOCache, HybridPSOCache}; kwargs...)
     reinit_cache!(cache, cache.alg)
 end
 
@@ -122,7 +96,8 @@ function reinit_cache!(cache::PSOCache, opt::ParallelSyncPSOKernel)
     return nothing
 end
 
-function SciMLBase.solve!(cache, args...; maxiters = 100, kwargs...)
+function SciMLBase.solve!(
+        cache::Union{PSOCache, HybridPSOCache}, args...; maxiters = 100, kwargs...)
     solve!(cache, cache.alg, args...; maxiters, kwargs...)
 end
 
@@ -150,8 +125,8 @@ function SciMLBase.solve!(
     prob = cache.prob
     t0 = time()
     gbest, particles = vectorized_solve!(prob,
-        init_gbest,
-        gpu_particles,
+        cache.gbest,
+        cache.particles,
         opt,
         args...;
         kwargs...)
@@ -163,7 +138,8 @@ function SciMLBase.solve!(
         stats = Optimization.OptimizationStats(; time = t1 - t0))
 end
 
-function SciMLBase.solve(prob::OptimizationProblem, opt::ParallelPSOKernel,
+function SciMLBase.solve(prob::OptimizationProblem,
+        opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel, HybridPSO},
         args...; maxiters = 100, kwargs...)
     solve!(init(prob, opt, args...; maxiters, kwargs...), opt)
 end