Skip to content

Commit be7d98f

Browse files
utkarsh530utkarsh530
committed
Add reinit tests and restructure
1 parent eb4b35b commit be7d98f

File tree

5 files changed

+201
-208
lines changed

5 files changed

+201
-208
lines changed

‎src/PSOGPU.jl‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -67,6 +67,7 @@ include("./utils.jl")
6767
include("./ode_pso.jl")
6868
include("./kernels.jl")
6969
include("./lowerlevel_solve.jl")
70+
include("init.jl")
7071
include("./solve.jl")
7172
include("./bfgs.jl")
7273
include("./hybrid.jl")

‎src/hybrid.jl‎

Lines changed: 0 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -5,56 +5,6 @@
55
result[i] = sol.u
66
end
77

8-
function SciMLBase.init(
9-
prob::OptimizationProblem, opt::HybridPSO{Backend, LocalOpt}, args...;
10-
kwargs...) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
11-
psoalg = opt.pso
12-
backend = opt.backend
13-
14-
pso_cache = init(prob, psoalg)
15-
16-
start_points = KernelAbstractions.allocate(
17-
backend, typeof(prob.u0), opt.pso.num_particles)
18-
19-
return HybridPSOCache{
20-
typeof(pso_cache), typeof(start_points), typeof(opt)}(pso_cache, start_points, opt)
21-
end
22-
23-
function reinit_cache!(cache::HybridPSOCache,
24-
opt::HybridPSO{Backend, LocalOpt}) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
25-
reinit!(cache.pso_cache)
26-
fill!(cache.start_points, zero(eltype(cache.start_points)))
27-
# prob = cache.prob
28-
# backend = opt.backend
29-
# particles = cache.particles
30-
31-
# kernel! = PSOGPU.gpu_init_particles!(backend)
32-
# kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
33-
34-
# best_particle = minimum(particles)
35-
# _init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
36-
37-
# copyto!(cache.gbest, [_init_gbest])
38-
39-
return nothing
40-
end
41-
42-
function Base.getproperty(cache::HybridPSOCache, name::Symbol)
43-
if name ∈ (:start_points, :pso_cache, :alg)
44-
return getfield(cache, name)
45-
else
46-
return getproperty(cache.pso_cache, name)
47-
end
48-
end
49-
50-
function Base.setproperty!(cache::HybridPSOCache, name::Symbol, val)
51-
if name ∈ (:start_points, :pso_cache, :alg)
52-
return setfield!(cache, name, val)
53-
else
54-
return setproperty!(cache.pso_cache, name, val)
55-
end
56-
end
57-
588
function SciMLBase.solve!(
599
cache::HybridPSOCache, opt::HybridPSO{Backend, LocalOpt}, args...;
6010
abstol = nothing,

‎src/init.jl‎

Lines changed: 166 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,166 @@
1+
mutable struct PSOCache{TP, TAlg, TPart, TGbest, TSampler}
2+
prob::TP
3+
alg::TAlg
4+
particles::TPart
5+
gbest::TGbest
6+
sampler::TSampler
7+
end
8+
9+
struct HybridPSOCache{TPc, TSp, TAlg}
10+
pso_cache::TPc
11+
start_points::TSp
12+
alg::TAlg
13+
end
14+
15+
function __init!(particles, prob::OptimizationProblem,
16+
opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
17+
args...; kwargs...) where {T <: QuasiMonteCarlo.SamplingAlgorithm}
18+
backend = opt.backend
19+
20+
qmc_samples = QuasiMonteCarlo.sample(opt.num_particles, prob.lb, prob.ub, sampler)
21+
22+
qmc_samples = adapt(backend, qmc_samples)
23+
24+
kernel! = gpu_init_particles!(backend)
25+
26+
kernel!(
27+
particles, qmc_samples, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
28+
29+
best_particle = minimum(particles)
30+
init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
31+
32+
return particles, init_gbest
33+
end
34+
35+
function __init!(particles, prob::OptimizationProblem,
36+
opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
37+
args...; kwargs...) where {T <: GPUSamplingAlgorithm}
38+
backend = opt.backend
39+
40+
kernel! = gpu_init_particles!(backend)
41+
42+
kernel!(particles, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
43+
44+
best_particle = minimum(particles)
45+
46+
init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
47+
48+
particles, init_gbest
49+
end
50+
51+
function SciMLBase.init(
52+
prob::OptimizationProblem, opt::ParallelPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
53+
@assert prob.u0 isa SArray
54+
55+
## Bounds check
56+
lb, ub = check_init_bounds(prob)
57+
lb, ub = check_init_bounds(prob)
58+
prob = remake(prob; lb = lb, ub = ub)
59+
60+
particles = KernelAbstractions.allocate(
61+
opt.backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
62+
63+
_sampler = if lb === nothing || ub === nothing || (all(isinf, lb) && all(isinf, ub))
64+
GPUUnboundedSampler()
65+
else
66+
sampler
67+
end
68+
69+
particles, _init_gbest = __init!(particles, prob, opt, _sampler, args...; kwargs...)
70+
71+
init_gbest = KernelAbstractions.allocate(opt.backend, typeof(_init_gbest), (1,))
72+
copyto!(init_gbest, [_init_gbest])
73+
74+
return PSOCache{
75+
typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest), typeof(_sampler)}(
76+
prob, opt, particles, init_gbest, _sampler)
77+
end
78+
79+
function SciMLBase.init(
80+
prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
81+
@assert prob.u0 isa SArray
82+
83+
## Bounds check
84+
lb, ub = check_init_bounds(prob)
85+
lb, ub = check_init_bounds(prob)
86+
prob = remake(prob; lb = lb, ub = ub)
87+
88+
particles = KernelAbstractions.allocate(
89+
opt.backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
90+
91+
_sampler = if lb === nothing || ub === nothing || (all(isinf, lb) && all(isinf, ub))
92+
GPUUnboundedSampler()
93+
else
94+
sampler
95+
end
96+
97+
particles, init_gbest = __init!(particles, prob, opt, _sampler, args...; kwargs...)
98+
99+
return PSOCache{
100+
typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest), typeof(_sampler)}(
101+
prob, opt, particles, init_gbest, _sampler)
102+
end
103+
104+
function SciMLBase.reinit!(cache::Union{PSOCache, HybridPSOCache})
105+
reinit_cache!(cache, cache.alg)
106+
end
107+
108+
function reinit_cache!(cache::PSOCache, opt::ParallelPSOKernel)
109+
prob = cache.prob
110+
particles = cache.particles
111+
112+
particles, _init_gbest = __init!(particles, prob, opt, cache.sampler)
113+
114+
copyto!(cache.gbest, [_init_gbest])
115+
116+
return nothing
117+
end
118+
119+
function reinit_cache!(cache::PSOCache, opt::ParallelSyncPSOKernel)
120+
prob = cache.prob
121+
particles = cache.particles
122+
123+
particles, init_gbest = __init!(particles, prob, opt, cache.sampler)
124+
125+
cache.gbest = init_gbest
126+
127+
return nothing
128+
end
129+
130+
function SciMLBase.init(
131+
prob::OptimizationProblem, opt::HybridPSO{Backend, LocalOpt}, args...;
132+
kwargs...) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
133+
psoalg = opt.pso
134+
backend = opt.backend
135+
136+
pso_cache = init(prob, psoalg, args...; kwargs...)
137+
138+
start_points = KernelAbstractions.allocate(
139+
backend, typeof(prob.u0), opt.pso.num_particles)
140+
141+
return HybridPSOCache{
142+
typeof(pso_cache), typeof(start_points), typeof(opt)}(pso_cache, start_points, opt)
143+
end
144+
145+
function reinit_cache!(cache::HybridPSOCache,
146+
opt::HybridPSO{Backend, LocalOpt}) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
147+
reinit!(cache.pso_cache)
148+
fill!(cache.start_points, zero(eltype(cache.start_points)))
149+
return nothing
150+
end
151+
152+
function Base.getproperty(cache::HybridPSOCache, name::Symbol)
153+
if name ∈ (:start_points, :pso_cache, :alg)
154+
return getfield(cache, name)
155+
else
156+
return getproperty(cache.pso_cache, name)
157+
end
158+
end
159+
160+
function Base.setproperty!(cache::HybridPSOCache, name::Symbol, val)
161+
if name ∈ (:start_points, :pso_cache, :alg)
162+
return setfield!(cache, name, val)
163+
else
164+
return setproperty!(cache.pso_cache, name, val)
165+
end
166+
end

‎src/solve.jl‎

Lines changed: 0 additions & 158 deletions
Original file line numberDiff line numberDiff line change
@@ -1,164 +1,6 @@
11
function get_pos(particle)
22
return particle.position
33
end
4-
mutable struct PSOCache{TP, TAlg, TPart, TGbest}
5-
prob::TP
6-
alg::TAlg
7-
particles::TPart
8-
gbest::TGbest
9-
end
10-
struct HybridPSOCache{TPc, TSp, TAlg}
11-
pso_cache::TPc
12-
start_points::TSp
13-
alg::TAlg
14-
end
15-
16-
function __init(prob::OptimizationProblem,
17-
opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
18-
args...; kwargs...) where {T <: QuasiMonteCarlo.SamplingAlgorithm}
19-
backend = opt.backend
20-
21-
particles = KernelAbstractions.allocate(
22-
backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
23-
24-
qmc_samples = QuasiMonteCarlo.sample(opt.num_particles, prob.lb, prob.ub, sampler)
25-
26-
qmc_samples = adapt(backend, qmc_samples)
27-
28-
kernel! = gpu_init_particles!(backend)
29-
30-
kernel!(
31-
particles, qmc_samples, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
32-
33-
best_particle = minimum(particles)
34-
init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
35-
36-
return particles, init_gbest
37-
end
38-
39-
function __init(prob::OptimizationProblem,
40-
opt::Union{ParallelPSOKernel, ParallelSyncPSOKernel}, sampler::T,
41-
args...; kwargs...) where {T <: GPUSamplingAlgorithm}
42-
backend = opt.backend
43-
44-
particles = KernelAbstractions.allocate(
45-
backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
46-
kernel! = gpu_init_particles!(backend)
47-
48-
kernel!(particles, prob, opt, typeof(prob.u0), T; ndrange = opt.num_particles)
49-
50-
best_particle = minimum(particles)
51-
52-
init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
53-
54-
particles, init_gbest
55-
end
56-
57-
function SciMLBase.init(
58-
prob::OptimizationProblem, opt::ParallelPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
59-
@assert prob.u0 isa SArray
60-
61-
## Bounds check
62-
lb, ub = check_init_bounds(prob)
63-
lb, ub = check_init_bounds(prob)
64-
prob = remake(prob; lb = lb, ub = ub)
65-
66-
particles, _init_gbest = if lb === nothing || ub === nothing ||
67-
(all(isinf, lb) && all(isinf, ub))
68-
__init(prob, opt, GPUUnboundedSampler(), args...; kwargs...)
69-
else
70-
__init(prob, opt, sampler, args...; kwargs...)
71-
end
72-
73-
init_gbest = KernelAbstractions.allocate(opt.backend, typeof(_init_gbest), (1,))
74-
copyto!(init_gbest, [_init_gbest])
75-
76-
return PSOCache{
77-
typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
78-
prob, opt, particles, init_gbest)
79-
end
80-
81-
function SciMLBase.init(
82-
prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; sampler = GPUUniformSampler(), kwargs...)
83-
@assert prob.u0 isa SArray
84-
85-
## Bounds check
86-
lb, ub = check_init_bounds(prob)
87-
lb, ub = check_init_bounds(prob)
88-
prob = remake(prob; lb = lb, ub = ub)
89-
90-
particles, init_gbest = if lb === nothing || ub === nothing ||
91-
(all(isinf, lb) && all(isinf, ub))
92-
__init(prob, opt, GPUUnboundedSampler(), args...; kwargs...)
93-
else
94-
__init(prob, opt, sampler, args...; kwargs...)
95-
end
96-
97-
return PSOCache{
98-
typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
99-
prob, opt, particles, init_gbest)
100-
end
101-
102-
# function SciMLBase.init(
103-
# prob::OptimizationProblem, opt::ParallelSyncPSOKernel, args...; kwargs...)
104-
# backend = opt.backend
105-
# @assert prob.u0 isa SArray
106-
107-
# ## Bounds check
108-
# lb, ub = check_init_bounds(prob)
109-
# lb, ub = check_init_bounds(prob)
110-
# prob = remake(prob; lb = lb, ub = ub)
111-
112-
# particles = KernelAbstractions.allocate(
113-
# backend, SPSOParticle{typeof(prob.u0), eltype(typeof(prob.u0))}, opt.num_particles)
114-
# kernel! = gpu_init_particles!(backend)
115-
116-
# kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
117-
118-
# best_particle = minimum(particles)
119-
# init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
120-
121-
# return PSOCache{
122-
# typeof(prob), typeof(opt), typeof(particles), typeof(init_gbest)}(
123-
# prob, opt, particles, init_gbest)
124-
# end
125-
126-
function SciMLBase.reinit!(cache::Union{PSOCache, HybridPSOCache}; kwargs...)
127-
reinit_cache!(cache, cache.alg)
128-
end
129-
130-
function reinit_cache!(cache::PSOCache, opt::ParallelPSOKernel)
131-
prob = cache.prob
132-
backend = opt.backend
133-
particles = cache.particles
134-
135-
kernel! = PSOGPU.gpu_init_particles!(backend)
136-
kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
137-
138-
best_particle = minimum(particles)
139-
_init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
140-
141-
copyto!(cache.gbest, [_init_gbest])
142-
143-
return nothing
144-
end
145-
146-
function reinit_cache!(cache::PSOCache, opt::ParallelSyncPSOKernel)
147-
prob = cache.prob
148-
backend = opt.backend
149-
particles = cache.particles
150-
151-
kernel! = PSOGPU.gpu_init_particles!(backend)
152-
kernel!(particles, prob, opt, typeof(prob.u0); ndrange = opt.num_particles)
153-
154-
best_particle = minimum(particles)
155-
156-
init_gbest = SPSOGBest(best_particle.best_position, best_particle.best_cost)
157-
158-
cache.gbest = init_gbest
159-
160-
return nothing
161-
end
1624

1635
function SciMLBase.solve!(
1646
cache::Union{PSOCache, HybridPSOCache}, args...; maxiters = 100, kwargs...)

0 commit comments

Comments
 (0)