Skip to content

Commit 94eadef

Browse files
utkarsh530utkarsh530
committed
Update sampling to allow qmc methods
1 parent 12dfa7e commit 94eadef

File tree

3 files changed

+193
-10
lines changed

3 files changed

+193
-10
lines changed

‎src/algorithms.jl‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,14 @@
11

22
abstract type PSOAlgorithm end
33
abstract type HybridPSOAlgorithm{LocalOpt} end
4+
abstract type GPUSamplingAlgorithm end
5+
6+
struct GPUUniformSampler <: GPUSamplingAlgorithm
7+
end
8+
9+
struct GPUUnboundedSampler <: GPUSamplingAlgorithm
10+
end
11+
412
"""
513
```julia
614
ParallelPSOKernel(num_particles = 100)

‎src/solve.jl‎

Lines changed: 45 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -13,23 +13,44 @@ struct HybridPSOCache{TPc, TSp, TAlg}
1313
alg::TAlg
1414
end
1515

16-
function SciMLBase.init(
17-
prob::OptimizationProblem, opt::ParallelPSOKernel, args...; kwargs...)
16+
function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
17+
args...; kwargs...) where {T <: QuasiMonteCarlo.SamplingAlgorithm}
1818
backend = opt.backend
19-
@assert prob.u0 isa SArray
2019

21-
## Bounds check
22-
lb, ub = check_init_bounds(prob)
23-
lb, ub = check_init_bounds(prob)
24-
prob = remake(prob; lb = lb, ub = ub)
20+
particles = KernelAbstractions.allocate(
21+
backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
22+
23+
qmc_samples = QuasiMonteCarlo.sample(opt.num_particles, prob.lb, prob.ub, sampler)
24+
25+
qmc_samples = adapt(backend, qmc_samples)
26+
27+
kernel! = gpu_init_particles!(backend)
28+
29+
kernel!(
30+
particles, qmc_samples, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
31+
32+
best_particle = minimum(particles)
33+
_init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
34+
35+
init_gbest = KernelAbstractions.allocate(backend, typeof(_init_gbest), (1,))
36+
copyto!(init_gbest, [_init_gbest])
37+
return PSOCache{
38+
typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
39+
prob, opt, particles, init_gbest)
40+
end
41+
42+
function __init(prob::OptimizationProblem, opt::ParallelPSOKernel, sampler::T,
43+
args...; kwargs...) where {T <: GPUSamplingAlgorithm}
44+
backend = opt.backend
2545

2646
particles = KernelAbstractions.allocate(
2747
backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
2848
kernel! = gpu_init_particles!(backend)
2949

30-
kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
50+
kernel!(particles, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
3151

3252
best_particle = minimum(particles)
53+
3354
_init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
3455

3556
init_gbest = KernelAbstractions.allocate(backend, typeof(_init_gbest), (1,))
@@ -39,6 +60,22 @@ function SciMLBase.init(
3960
prob, opt, particles, init_gbest)
4061
end
4162

63+
function SciMLBase.init(
64+
prob::OptimizationProblem, opt::ParallelPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
65+
@assert prob.u0 isa SArray
66+
67+
## Bounds check
68+
lb, ub = check_init_bounds(prob)
69+
lb, ub = check_init_bounds(prob)
70+
prob = remake(prob; lb = lb, ub = ub)
71+
72+
if lb === nothing || ub === nothing || (all(isinf, lb) && all(isinf, ub))
73+
__init(prob, opt, GPUUnboundedSampler(), args...; kwargs...)
74+
else
75+
__init(prob, opt, sampler, args...; kwargs...)
76+
end
77+
end
78+
4279
function SciMLBase.init(
4380
prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; kwargs...)
4481
backend = opt.backend

‎src/utils.jl‎

Lines changed: 140 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -11,15 +11,70 @@ function uniform(dim::Int, lb::AbstractArray{T}, ub::AbstractArray{T}) where {T}
1111
return arr
1212
end
1313

14-
@kernel function gpu_init_particles!(particles, prob, opt, ::Type{T}) where {T <: SArray}
14+
function _gen_sampling_kernel(prob, sampling::QuasiMonteCarlo.SamplingAlgorithm)
15+
dim = length(prob.u0)
16+
lb = prob.lb
17+
ub = prob.ub
18+
19+
positions = QuasiMonteCarlo.sample(num_particles, lb, ub, sampler)
20+
end
21+
22+
@kernel function gpu_init_particles!(
23+
particles, prob, opt, ::Type{T}, ::Type{GPUUnboundedSampler}) where {T <: SArray}
24+
i = @index(Global, Linear)
25+
26+
dim = length(prob.u0)
27+
lb = prob.lb
28+
ub = prob.ub
29+
cost_func = prob.f
30+
p = prob.p
31+
32+
# gbest_position = StaticArrays.sacollect(T,
33+
# ifelse(
34+
# abs(prob.u0[i]) > 0, prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
35+
# rand(eltype(prob.u0))) for i in 1:dim)
36+
37+
# gbest_position = convert(T, gbest_position)
38+
# gbest_cost = cost_func(gbest_position, p)
39+
# if !isnothing(prob.f.cons)
40+
# penalty = calc_penalty(gbest_position, prob, 1, opt.θ, opt.γ, opt.h)
41+
# gbest_cost = cost_func(gbest_position, p) + penalty
42+
# else
43+
# gbest_cost = cost_func(gbest_position, p)
44+
# end
45+
# gbest_cost = cost_func(gbest_position, p)
46+
47+
position = StaticArrays.sacollect(T,
48+
ifelse(abs(prob.u0[i]) > 0,
49+
prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
50+
rand(eltype(prob.u0))) for i in 1:dim)
51+
52+
# KernelAbstractions.@print "$(position[1]) \n"
53+
54+
velocity = zero(T)
55+
56+
if !isnothing(prob.f.cons)
57+
penalty = calc_penalty(position, prob, 1, opt.θ, opt.γ, opt.h)
58+
cost = cost_func(position, p) + penalty
59+
else
60+
cost = cost_func(position, p)
61+
end
62+
63+
best_position = position
64+
best_cost = cost
65+
@inbounds particles[i] = SPSOParticle(
66+
position, velocity, cost, best_position, best_cost)
67+
end
68+
69+
@kernel function gpu_init_particles!(
70+
particles, prob, opt, ::Type{T}, ::Type{GPUUniformSampler}) where {T <: SArray}
1571
i = @index(Global, Linear)
1672

1773
dim = length(prob.u0)
1874
lb = prob.lb
1975
ub = prob.ub
2076
cost_func = prob.f
2177
p = prob.p
22-
num_particles = opt.num_particles
2378

2479
if lb === nothing || (all(isinf, lb) && all(isinf, ub))
2580
gbest_position = StaticArrays.sacollect(T,
@@ -64,6 +119,89 @@ end
64119
position, velocity, cost, best_position, best_cost)
65120
end
66121

122+
@kernel function gpu_init_particles!(particles, qmc_particles, prob, opt, ::Type{T},
123+
::Type{T1}) where {T <: SArray, T1 <: QuasiMonteCarlo.SamplingAlgorithm}
124+
i = @index(Global, Linear)
125+
126+
dim = length(prob.u0)
127+
lb = prob.lb
128+
ub = prob.ub
129+
cost_func = prob.f
130+
p = prob.p
131+
num_particles = opt.num_particles
132+
133+
@inbounds position = StaticArrays.sacollect(T, qmc_particles[j, i] for j in 1:dim)
134+
velocity = zero(T)
135+
136+
if !isnothing(prob.f.cons)
137+
penalty = calc_penalty(position, prob, 1, opt.θ, opt.γ, opt.h)
138+
cost = cost_func(position, p) + penalty
139+
else
140+
cost = cost_func(position, p)
141+
end
142+
143+
best_position = position
144+
best_cost = cost
145+
146+
@inbounds particles[i] = SPSOParticle(
147+
position, velocity, cost, best_position, best_cost)
148+
# push!(particles, SPSOParticle(position, velocity, cost, best_position, best_cost))
149+
150+
end
151+
152+
# @kernel function gpu_init_particles!(gpu_particles, prob, opt, ::Type{T}) where {T <: SArray}
153+
# i = @index(Global, Linear)
154+
155+
# dim = length(prob.u0)
156+
# lb = prob.lb
157+
# ub = prob.ub
158+
# cost_func = prob.f
159+
# p = prob.p
160+
# num_particles = opt.num_particles
161+
162+
# if lb === nothing || (all(isinf, lb) && all(isinf, ub))
163+
# gbest_position = StaticArrays.sacollect(T,
164+
# ifelse(
165+
# abs(prob.u0[i]) > 0, prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
166+
# rand(eltype(prob.u0))) for i in 1:dim)
167+
# else
168+
# gbest_position = StaticArrays.sacollect(T, uniform_itr(dim, lb, ub))
169+
# end
170+
171+
# gbest_position = convert(T, gbest_position)
172+
# gbest_cost = cost_func(gbest_position, p)
173+
# if !isnothing(prob.f.cons)
174+
# penalty = calc_penalty(gbest_position, prob, 1, opt.θ, opt.γ, opt.h)
175+
# gbest_cost = cost_func(gbest_position, p) + penalty
176+
# else
177+
# gbest_cost = cost_func(gbest_position, p)
178+
# end
179+
# gbest_cost = cost_func(gbest_position, p)
180+
181+
# if lb === nothing || (all(isinf, lb) && all(isinf, ub))
182+
# position = StaticArrays.sacollect(T,
183+
# ifelse(abs(prob.u0[i]) > 0,
184+
# prob.u0[i] + rand(eltype(prob.u0)) * abs(prob.u0[i]),
185+
# rand(eltype(prob.u0))) for i in 1:dim)
186+
# else
187+
# position = StaticArrays.sacollect(T, uniform_itr(dim, lb, ub))
188+
# end
189+
190+
# velocity = zero(T)
191+
192+
# if !isnothing(prob.f.cons)
193+
# penalty = calc_penalty(position, prob, 1, opt.θ, opt.γ, opt.h)
194+
# cost = cost_func(position, p) + penalty
195+
# else
196+
# cost = cost_func(position, p)
197+
# end
198+
199+
# best_position = position
200+
# best_cost = cost
201+
# @inbounds particles[i] = SPSOParticle(
202+
# position, velocity, cost, best_position, best_cost)
203+
# end
204+
67205
function init_particles!(particles, prob, opt, ::Type{T}) where {T <: SArray}
68206
dim = length(prob.u0)
69207
lb = prob.lb

0 commit comments

Comments
 (0)