Update tests and fix bugs

Author: utkarsh530

Project.toml

@@ -23,12 +23,4 @@ SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 
 [compat]
-julia = "1.6"
-
-[extras]
-LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
-StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
-
-[targets]
-test = ["Test", "StaticArrays", "LinearAlgebra", "Optimization", "ForwardDiff"]
+julia = "1.6"

src/solve.jl

@@ -13,7 +13,8 @@ struct HybridPSOCache{TPc, TSp, TAlg}
     alg::TAlg
 end
 
-function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
+function __init(prob::OptimizationProblem,
+        opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
         args...; kwargs...) where {T <: QuasiMonteCarlo.SamplingAlgorithm}
     backend = opt.backend
 
@@ -30,16 +31,13 @@ function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
         particles, qmc_samples, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
 
     best_particle = minimum(particles)
-    _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
+    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)
+    return particles, init_gbest
 end
 
-function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
+function __init(prob::OptimizationProblem,
+        opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
         args...; kwargs...) where {T <: GPUSamplingAlgorithm}
     backend = opt.backend
 
@@ -51,13 +49,9 @@ function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
 
     best_particle = minimum(particles)
 
-    _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
+    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)
+    particles, init_gbest
 end
 
 function SciMLBase.init(
@@ -69,37 +63,66 @@ function SciMLBase.init(
     lb, ub = check_init_bounds(prob)
     prob = remake(prob; lb = lb, ub = ub)
 
-    if lb === nothing || ub === nothing || (all(isinf, lb) && all(isinf, ub))
+    particles, _init_gbest = if lb === nothing || ub === nothing ||
+                                (all(isinf, lb) && all(isinf, ub))
         __init(prob, opt, GPUUnboundedSampler(), args...; kwargs...)
     else
         __init(prob, opt, sampler, args...; kwargs...)
     end
+
+    init_gbest = KernelAbstractions.allocate(opt.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
+        prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
     @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)
+    particles, init_gbest = if lb === nothing || ub === nothing ||
+                               (all(isinf, lb) && all(isinf, ub))
+        __init(prob, opt, GPUUnboundedSampler(), args...; kwargs...)
+    else
+        __init(prob, opt, sampler, args...; kwargs...)
+    end
 
     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
+#     @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)
+
+#     return PSOCache{
+#         typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
+#         prob, opt, particles, init_gbest)
+# end
+
 function SciMLBase.reinit!(cache::Union{PSOCache, HybridPSOCache}; kwargs...)
     reinit_cache!(cache, cache.alg)
 end
@@ -170,6 +193,7 @@ function SciMLBase.solve!(
         cache.particles,
         opt,
         args...;
+        maxiters,
         kwargs...)
     t1 = time()
 
@@ -182,7 +206,7 @@ end
 function SciMLBase.solve(prob::OptimizationProblem,
         opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel, HybridPSO},
         args...; maxiters = 100, kwargs...)
-    solve!(init(prob, opt, args...; maxiters, kwargs...), opt)
+    solve!(init(prob, opt, args...; kwargs...), opt, args...; maxiters, kwargs...)
 end
 
 function SciMLBase.__solve(prob::OptimizationProblem,

src/utils.jl

@@ -97,7 +97,6 @@ end
 
     @inbounds particles[i] = SPSOParticle(
         position, velocity, cost, best_position, best_cost)
-
 end
 
 function init_particles!(particles, prob, opt, ::Type{T}) where {T <: SArray}

test/Project.toml

@@ -3,6 +3,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+QuasiMonteCarlo = "8a4e6c94-4038-4cdc-81c3-7e6ffdb2a71b"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"

test/regression.jl

@@ -1,4 +1,5 @@
 using PSOGPU, StaticArrays, SciMLBase, Test, LinearAlgebra, Random, KernelAbstractions
+using QuasiMonteCarlo
 
 @testset "Rosenbrock test dimension = $(N)" for N in 2:3
 
@@ -36,12 +37,26 @@ using PSOGPU, StaticArrays, SciMLBase, Test, LinearAlgebra, Random, KernelAbstra
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelSyncPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelSyncPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
@@ -66,12 +81,26 @@ using PSOGPU, StaticArrays, SciMLBase, Test, LinearAlgebra, Random, KernelAbstra
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelSyncPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelSyncPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
@@ -93,12 +122,24 @@ using PSOGPU, StaticArrays, SciMLBase, Test, LinearAlgebra, Random, KernelAbstra
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
 
     @test sol.objective < 1e-4
 
+    sol = solve!(
+        init(prob, ParallelSyncPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 500)
+
     sol = solve(prob,
         ParallelSyncPSOKernel(n_particles; backend = CPU()),
         maxiters = 500)
@@ -138,12 +179,26 @@ end
 
     @test sol.objective < 2e-3
 
+    sol = solve!(
+        init(prob, ParallelPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 2000)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend = CPU()),
         maxiters = 2000)
 
     @test sol.objective < 2e-2
 
+    sol = solve!(
+        init(prob, ParallelSyncPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 2000)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelSyncPSOKernel(n_particles; backend = CPU()),
         maxiters = 2000)
@@ -171,12 +226,26 @@ end
 
     @test sol.objective < 2e-3
 
+    sol = solve!(
+        init(prob, ParallelPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 1000)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelPSOKernel(n_particles; backend = CPU()),
         maxiters = 1000)
 
     @test sol.objective < 2e-3
 
+    sol = solve!(
+        init(prob, ParallelSyncPSOKernel(n_particles; backend = CPU());
+            sampler = LatinHypercubeSample()),
+        maxiters = 2000)
+
+    @test sol.retcode == ReturnCode.Default
+
     sol = solve(prob,
         ParallelSyncPSOKernel(n_particles; backend = CPU()),
         maxiters = 2000)