Merge pull request #38 from SciML/u/bfgs

Update Hybrid API and misc cleanups

Author: Vaibhavdixit02

.github/workflows/CI.yml

@@ -20,7 +20,7 @@ jobs:
       matrix:
         version:
           - '1.6'
-          - '1.9'
+          - '1'
         os:
           - ubuntu-latest
         arch:

Project.toml

@@ -5,7 +5,8 @@ version = "1.0.0-DEV"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
-CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+DiffEqGPU = "071ae1c0-96b5-11e9-1965-c90190d839ea"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 MLDatasets = "eb30cadb-4394-5ae3-aed4-317e484a6458"
@@ -18,7 +19,6 @@ Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
 SimpleChains = "de6bee2f-e2f4-4ec7-b6ed-219cc6f6e9e5"
 SimpleNonlinearSolve = "727e6d20-b764-4bd8-a329-72de5adea6c7"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
-Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f"
 
 [compat]
 julia = "1.6"
@@ -29,4 +29,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test", "StaticArrays", "LinearAlgebra"]
+test = ["Test", "StaticArrays", "LinearAlgebra", "Optimization", "ForwardDiff", ]

src/PSOGPU.jl

@@ -3,8 +3,8 @@ module PSOGPU
 using SciMLBase, StaticArrays, Setfield, KernelAbstractions
 using QuasiMonteCarlo, Optimization, SimpleNonlinearSolve, ForwardDiff
 import Adapt
-
-# import DiffEqGPU: GPUTsit5, vectorized_asolve, make_prob_compatible
+import Enzyme: autodiff_deferred, Active, Reverse
+import DiffEqGPU: GPUTsit5, vectorized_asolve, make_prob_compatible
 
 ## Use lb and ub either as StaticArray or pass them separately as CuArrays
 ## Passing as CuArrays makes more sense, or maybe SArray? The based on no. of dimension
@@ -48,11 +48,11 @@ end
 
 include("./algorithms.jl")
 include("./utils.jl")
-# include("./ode_pso.jl")
+include("./ode_pso.jl")
 include("./kernels.jl")
 include("./lowerlevel_solve.jl")
 include("./solve.jl")
-include("./lbfgs.jl")
+include("./bfgs.jl")
 include("./hybrid.jl")
 
 export ParallelPSOKernel,

src/algorithms.jl

@@ -1,6 +1,6 @@
 
 abstract type PSOAlgorithm end
-
+abstract type HybridPSOAlgorithm{LocalOpt} end
 """
 ```julia
 ParallelPSOKernel(num_particles = 100)
@@ -118,3 +118,28 @@ end
 
 SciMLBase.allowsbounds(::PSOAlgorithm) = true
 SciMLBase.allowsconstraints(::PSOAlgorithm) = true
+
+struct LBFGS
+    threshold::Int
+end
+
+function LBFGS(; threshold = 10)
+    LBFGS(threshold)
+end
+
+struct BFGS end
+
+struct HybridPSO{Backend, LocalOpt} <: HybridPSOAlgorithm{LocalOpt}
+    pso::PSOAlgorithm
+    local_opt::LocalOpt
+    backend::Backend
+end
+
+function HybridPSO(;
+        backend = CPU(),
+        pso = PSOGPU.ParallelPSOKernel(100; global_update = false, backend),
+        local_opt = LBFGS())
+    HybridPSO(pso, local_opt, backend)
+end
+
+SciMLBase.allowsbounds(::HybridPSOAlgorithm{LocalOpt}) where {LocalOpt} = true

src/bfgs.jl

@@ -0,0 +1,52 @@
+function SciMLBase.__solve(prob::SciMLBase.OptimizationProblem,
+        opt::LBFGS,
+        args...;
+        abstol = nothing,
+        reltol = nothing,
+        maxiters = 1000,
+        kwargs...)
+    f = Base.Fix2(prob.f.f, prob.p)
+    function _g(θ, _p = nothing)
+        return ForwardDiff.gradient(f, θ)
+    end
+    t0 = time()
+    nlprob = NonlinearProblem{false}(_g, prob.u0)
+    nlsol = solve(nlprob,
+        SimpleLimitedMemoryBroyden(; threshold = opt.threshold, linesearch = Val(true));
+        maxiters,
+        abstol,
+        reltol)
+    θ = nlsol.u
+    t1 = time()
+
+    SciMLBase.build_solution(SciMLBase.DefaultOptimizationCache(prob.f, prob.p),
+        opt,
+        θ,
+        prob.f(θ, prob.p))
+end
+
+function SciMLBase.__solve(prob::SciMLBase.OptimizationProblem,
+        opt::BFGS,
+        args...;
+        abstol = nothing,
+        reltol = nothing,
+        maxiters = 1000,
+        kwargs...)
+    f = Base.Fix2(prob.f.f, prob.p)
+    ∇f = instantiate_gradient(f, prob.f.adtype)
+
+    t0 = time()
+    nlprob = NonlinearProblem{false}(∇f, prob.u0)
+    nlsol = solve(nlprob,
+        SimpleBroyden(; linesearch = Val(true));
+        maxiters,
+        abstol,
+        reltol)
+    θ = nlsol.u
+    t1 = time()
+
+    SciMLBase.build_solution(SciMLBase.DefaultOptimizationCache(prob.f, prob.p),
+        opt,
+        θ,
+        prob.f(θ, prob.p))
+end

src/hybrid.jl

@@ -1,88 +1,53 @@
-struct HybridPSOLBFGS
-    pso::PSOAlgorithm
-    lbfgs::LBFGS
-end
-
-function HybridPSOLBFGS(; pso = PSOGPU.ParallelPSOKernel(100 ; global_update = false), lbfgs = LBFGS())
-    HybridPSOLBFGS(pso, lbfgs)
-end
-
-@kernel function lbfgs_run!(nlcaches, x0s, result)
-    i = @index(Global, Linear)
-    # nlcache = reinit!(nlcaches[i], x0s[i])
-    # @show nlcache.u
-    res = solve!(nlcaches[i])
-    # @show res
-    # @show res.resid
-    result[i] = res
-end
-
-@kernel function simplelbfgs_run!(nlprob, x0s, result)
+@kernel function simplebfgs_run!(nlprob, x0s, result, opt, maxiters, abstol, reltol)
     i = @index(Global, Linear)
     nlcache = remake(nlprob; u0 = x0s[i])
-    result[i] = solve(nlcache, SimpleLimitedMemoryBroyden(; threshold = 10))
-end
-
-@kernel function simplebfgs_run!(nlprob, x0s, result, opt, maxiters)
-    i = @index(Global, Linear)
-    nlcache = remake(nlprob; u0 = x0s[i])
-    sol = SciMLBase.__solve(nlcache, opt, maxiters = maxiters, abstol = 1f-6, reltol = 1f-6)
+    sol = solve(nlcache, opt; maxiters, abstol, reltol)
     result[i] = sol.u
 end
 
-function SciMLBase.__solve(prob::SciMLBase.OptimizationProblem, opt::HybridPSOLBFGS, args...; maxiters = 1000, kwargs...)
+function SciMLBase.__solve(prob::SciMLBase.OptimizationProblem,
+        opt::HybridPSO{Backend, LocalOpt},
+        args...;
+        abstol = nothing,
+        reltol = nothing,
+        maxiters = 100,
+        local_maxiters = 10,
+        kwargs...) where {Backend, LocalOpt <: Union{LBFGS, BFGS}}
     t0 = time()
     psoalg = opt.pso
-    lbfgsalg = opt.lbfgs
+    local_opt = opt.local_opt
+    backend = opt.backend
 
     sol_pso = solve(prob, psoalg, args...; maxiters, kwargs...)
-
-    # @show sol_pso.u
-
     x0s = sol_pso.original
-    # @show prob.u0
-    # @show x0s
     prob = remake(prob, lb = nothing, ub = nothing)
-    # @show length(x0s)
-    # f = Optimization.instantiate_function(prob.f, prob.u0, prob.f.adtype, prob.p, 0)
     f = Base.Fix2(prob.f.f, prob.p)
-    function _g(θ, _p = nothing)
-        return ForwardDiff.gradient(f , θ)
-    end
-    # @show prob.u0
-    # nlcaches = [init(NonlinearProblem(NonlinearFunction(_g), x0), LimitedMemoryBroyden(; threshold = lbfgsalg.m, linesearch = LiFukushimaLineSearch()))
-    #     for x0 in x0s
-    # ]
-    # @show nlcaches[1]
-    # @show ismutable(nlcaches[1])
-    backend = lbfgsalg.backend
-    # kernel = lbfgs_run!(backend)
-    # result = KernelAbstractions.allocate(lbfgsalg.backend, SciMLBase.NonlinearSolution, length(x0s))
-
-    # kernel(nlcaches, x0s, result; ndrange = length(x0s))
-
-    # kernel = simplelbfgs_run!(backend)
-    # result = KernelAbstractions.allocate(backend, SciMLBase.NonlinearSolution, length(x0s))
-    # nlprob = NonlinearProblem(NonlinearFunction(_g), prob.u0)
-    # kernel(nlprob, x0s, result; ndrange = length(x0s))
+    ∇f = instantiate_gradient(f, prob.f.adtype)
 
     kernel = simplebfgs_run!(backend)
     result = KernelAbstractions.allocate(backend, typeof(prob.u0), length(x0s))
-    nlprob = NonlinearProblem{false}(_g, prob.u0)
+    nlprob = NonlinearProblem{false}(∇f, prob.u0)
+
+    nlalg = LocalOpt isa LBFGS ?
+            SimpleLimitedMemoryBroyden(;
+        threshold = local_opt.threshold,
+        linesearch = Val(true)) : SimpleBroyden(; linesearch = Val(true))
+
+    kernel(nlprob,
+        x0s,
+        result,
+        nlalg,
+        local_maxiters,
+        abstol,
+        reltol;
+        ndrange = length(x0s))
 
-    kernel(nlprob, x0s, result, SimpleLimitedMemoryBroyden(; threshold = lbfgsalg.m, linesearch = Val(true)), maxiters; ndrange = length(x0s))
-
-    # @show result
     t1 = time()
     sol_bfgs = (x -> prob.f(x, prob.p)).(result)
-    # sol_bfgs = [prob.f(θ, prob.p) for θ in result]
-    # @show typeof(sol_bfgs)
-
-    sol_bfgs = (x -> isnan(x) ? Inf32 : x).(sol_bfgs)
-    # @show minimum(sol_bfgs)
+    sol_bfgs = (x -> isnan(x) ? convert(eltype(prob.u0), Inf) : x).(sol_bfgs)
 
     minobj, ind = findmin(sol_bfgs)
 
     SciMLBase.build_solution(SciMLBase.DefaultOptimizationCache(prob.f, prob.p), opt,
-     view(result, ind) , minobj)
+        view(result, ind), minobj)
 end

src/lbfgs.jl

@@ -2,12 +2,16 @@ function get_pos(particle)
     return particle.position
 end
 
-function SciMLBase.__solve(prob::OptimizationProblem, opt::PSOAlgorithm, args...; kwargs...)
+function SciMLBase.__solve(prob::OptimizationProblem,
+        opt::PSOAlgorithm,
+        args...;
+        maxiters = 100,
+        kwargs...)
     lb, ub = check_init_bounds(prob)
     lb, ub = check_init_bounds(prob)
     prob = remake(prob; lb = lb, ub = ub)
 
-    gbest, particles = pso_solve(prob, opt)
+    gbest, particles = pso_solve(prob, opt, args...; maxiters, kwargs...)
     particles_positions = get_pos.(particles)
     SciMLBase.build_solution(SciMLBase.DefaultOptimizationCache(prob.f, prob.p), opt,
         gbest.position, prob.f(gbest.position, prob.p), original = particles_positions)
@@ -44,7 +48,12 @@ function pso_solve(prob::OptimizationProblem,
         args...;
         kwargs...)
     init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
-    gbest, particles = vectorized_solve!(prob, init_gbest, particles, opt, args...; kwargs...)
+    gbest, particles = vectorized_solve!(prob,
+        init_gbest,
+        particles,
+        opt,
+        args...;
+        kwargs...)
     gbest, particles
 end
 
@@ -65,6 +74,6 @@ function pso_solve(prob::OptimizationProblem,
     copyto!(gpu_particles, particles)
     init_gbest = init_gbest
 
-    gbest = vectorized_solve!(prob, init_gbest, gpu_particles, opt; kwargs...)
+    gbest = vectorized_solve!(prob, init_gbest, gpu_particles, opt, args...; kwargs...)
     gbest
 end

src/solve.jl

@@ -193,3 +193,11 @@ Based on the paper: Particle swarm optimization method for constrained optimizat
     penalty = h(T(iter)) * penalty
     penalty
 end
+
+@inline function instantiate_gradient(f, adtype::AutoForwardDiff)
+    (θ, p) -> ForwardDiff.gradient(f, θ)
+end
+
+@inline function instantiate_gradient(f, adtype::AutoEnzyme)
+    (θ, p) -> autodiff_deferred(Reverse, f, Active, Active(θ))[1][1]
+end

src/utils.jl

@@ -1,6 +1,7 @@
 [deps]
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SafeTestsets = "1bc83da4-3b8d-516f-aca4-4fe02f6d838f"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"