Merge pull request #36 from SciML/u/fix

Append solve time in solution interface

Author: utkarsh530

README.md

@@ -11,28 +11,19 @@ Supports generic Julia's SciML interface
 
 using PSOGPU, StaticArrays, CUDA
 
-lb = @SArray [-1.0f0, -1.0f0]
-ub = @SArray [1.0f0, 1.0f0]
+lb = @SArray [-1.0f0, -1.0f0, -1.0f0]
+ub = @SArray [10.0f0, 10.0f0, 10.0f0]
 
 function rosenbrock(x, p)
     sum(p[2] * (x[i + 1] - x[i]^2)^2 + (p[1] - x[i])^2 for i in 1:(length(x) - 1))
 end
 
-x0 = @SArray zeros(Float32, 2)
-p = @SArray Float32[2.0, 100.0]
+x0 = @SArray zeros(Float32, 3)
+p = @SArray Float32[1.0, 100.0]
 
 prob = OptimizationProblem(rosenbrock, x0, p; lb = lb, ub = ub)
 
-## select the no. of particles
-n_particles = 100
-
-## initialize particles on the grid and the best current position
-gbest, particles = PSOGPU.init_particles(prob, n_particles)
-
-## Offload particles to GPU
-particles = cu(particles)
-
-sol = PSOGPU.pso_solve_gpu!(prob, gbest, particles)
-
-sol[].position
+sol = solve(prob,
+    ParallelSyncPSOKernel(1000, backend = CUDA.CUDABackend()),
+    maxiters = 500)
 ```

src/ode_pso.jl

@@ -59,7 +59,9 @@ function parameter_estim_ode!(prob::ODEProblem,
     backend = get_backend(gpu_particles)
     update_states! = PSOGPU._update_particle_states!(backend)
 
-    losses = KernelAbstractions.ones(backend, 1, length(gpu_particles))
+    losses = KernelAbstractions.allocate(backend,
+        typeof(prob.u0),
+        (1, length(gpu_particles)))
     update_costs! = PSOGPU._update_particle_costs!(backend)
 
     improb = make_prob_compatible(prob)

src/solve.jl

@@ -11,10 +11,11 @@ function SciMLBase.__solve(prob::OptimizationProblem,
     lb, ub = check_init_bounds(prob)
     prob = remake(prob; lb = lb, ub = ub)
 
-    gbest, particles = pso_solve(prob, opt, args...; maxiters, kwargs...)
+    gbest, particles, solve_time = 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)
+        gbest.position, prob.f(gbest.position, prob.p), original = particles_positions,
+        stats = Optimization.OptimizationStats(; time = solve_time))
 end
 
 function pso_solve(prob::OptimizationProblem,
@@ -33,34 +34,40 @@ function pso_solve(prob::OptimizationProblem,
     gpu_init_gbest = KernelAbstractions.allocate(backend, typeof(init_gbest), (1,))
     copyto!(gpu_init_gbest, [init_gbest])
 
+    t0 = time()
     gbest, particles = vectorized_solve!(prob,
         gpu_init_gbest,
         gpu_particles,
         opt,
         Val(opt.global_update),
         args...;
         kwargs...)
-    gbest, particles
+    t1 = time()
+    gbest, particles, t1 - t0
 end
 
 function pso_solve(prob::OptimizationProblem,
         opt::ParallelPSOArray,
         args...;
         kwargs...)
     init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
+    t0 = time()
     gbest, particles = vectorized_solve!(prob,
         init_gbest,
         particles,
         opt,
         args...;
         kwargs...)
-    gbest, particles
+    t1 = time()
+    gbest, particles, t1 - t0
 end
 
 function pso_solve(prob::OptimizationProblem, opt::SerialPSO, args...; kwargs...)
     init_gbest, particles = init_particles(prob, opt, typeof(prob.u0))
+    t0 = time()
     gbest, particles = vectorized_solve!(prob, init_gbest, particles, opt; kwargs...)
-    gbest, particles
+    t1 = time()
+    gbest, particles, t1 - t0
 end
 
 function pso_solve(prob::OptimizationProblem,
@@ -74,6 +81,13 @@ function pso_solve(prob::OptimizationProblem,
     copyto!(gpu_particles, particles)
     init_gbest = init_gbest
 
-    gbest = vectorized_solve!(prob, init_gbest, gpu_particles, opt, args...; kwargs...)
-    gbest
+    t0 = time()
+    gbest, particles = vectorized_solve!(prob,
+        init_gbest,
+        gpu_particles,
+        opt,
+        args...;
+        kwargs...)
+    t1 = time()
+    gbest, particles, t1 - t0
 end

test/lbfgs.jl

@@ -57,7 +57,7 @@ l0 = rosenbrock(x0, p)
 
 @time sol = Optimization.solve(prob,
     PSOGPU.HybridPSO(; backend = CUDABackend()),
-    locaL_maxiters = 30)
+    local_maxiters = 30)
 @show sol.objective
 
 @time sol = Optimization.solve(prob,