Update benchmarks

Author: utkarsh530

Project.toml

@@ -5,6 +5,7 @@ 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"

benchmarks/Fitzhugh_Nagumo/Manifest.toml

@@ -2,7 +2,7 @@
 
 julia_version = "1.10.0"
 manifest_format = "2.0"
-project_hash = "64f611d6591a3e31f2740258c68d607684b60a0f"
+project_hash = "5385867d5ac9adc9babdf0cd3ced91c0f6a2076f"
 
 [[deps.ADTypes]]
 git-tree-sha1 = "41c37aa88889c171f1300ceac1313c06e891d245"
@@ -150,6 +150,12 @@ git-tree-sha1 = "aebf55e6d7795e02ca500a689d326ac979aaf89e"
 uuid = "9718e550-a3fa-408a-8086-8db961cd8217"
 version = "0.1.1"
 
+[[deps.BenchmarkTools]]
+deps = ["JSON", "Logging", "Printf", "Profile", "Statistics", "UUIDs"]
+git-tree-sha1 = "f1f03a9fa24271160ed7e73051fba3c1a759b53f"
+uuid = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
+version = "1.4.0"
+
 [[deps.BitFlags]]
 git-tree-sha1 = "2dc09997850d68179b69dafb58ae806167a32b1b"
 uuid = "d1d4a3ce-64b1-5f1a-9ba4-7e7e69966f35"
@@ -358,9 +364,9 @@ uuid = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
 version = "4.1.1"
 
 [[deps.DataAPI]]
-git-tree-sha1 = "8da84edb865b0b5b0100c0666a9bc9a0b71c553c"
+git-tree-sha1 = "abe83f3a2f1b857aac70ef8b269080af17764bbe"
 uuid = "9a962f9c-6df0-11e9-0e5d-c546b8b5ee8a"
-version = "1.15.0"
+version = "1.16.0"
 
 [[deps.DataDeps]]
 deps = ["HTTP", "Libdl", "Reexport", "SHA", "Scratch", "p7zip_jll"]
@@ -785,16 +791,22 @@ version = "1.0.0"
 
 [[deps.JLD2]]
 deps = ["FileIO", "MacroTools", "Mmap", "OrderedCollections", "Pkg", "PrecompileTools", "Printf", "Reexport", "Requires", "TranscodingStreams", "UUIDs"]
-git-tree-sha1 = "315b508ec5df53936532097ffe6e5deacbf41861"
+git-tree-sha1 = "7c0008f0b7622c6c0ee5c65cbc667b69f8a65672"
 uuid = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
-version = "0.4.44"
+version = "0.4.45"
 
 [[deps.JLLWrappers]]
 deps = ["Artifacts", "Preferences"]
 git-tree-sha1 = "7e5d6779a1e09a36db2a7b6cff50942a0a7d0fca"
 uuid = "692b3bcd-3c85-4b1f-b108-f13ce0eb3210"
 version = "1.5.0"
 
+[[deps.JSON]]
+deps = ["Dates", "Mmap", "Parsers", "Unicode"]
+git-tree-sha1 = "31e996f0a15c7b280ba9f76636b3ff9e2ae58c9a"
+uuid = "682c06a0-de6a-54ab-a142-c8b1cf79cde6"
+version = "0.21.4"
+
 [[deps.JSON3]]
 deps = ["Dates", "Mmap", "Parsers", "PrecompileTools", "StructTypes", "UUIDs"]
 git-tree-sha1 = "eb3edce0ed4fa32f75a0a11217433c31d56bd48b"
@@ -1215,10 +1227,10 @@ uuid = "ca575930-c2e3-43a9-ace4-1e988b2c1908"
 version = "1.2.0"
 
 [[deps.NonlinearSolve]]
-deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "EnumX", "FastBroadcast", "FastClosures", "FiniteDiff", "ForwardDiff", "LazyArrays", "LineSearches", "LinearAlgebra", "LinearSolve", "MaybeInplace", "PrecompileTools", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SciMLOperators", "SimpleNonlinearSolve", "SparseArrays", "SparseDiffTools", "StaticArrays", "UnPack"]
-git-tree-sha1 = "72b036b728461272ae1b1c3f7096cb4c319d8793"
+deps = ["ADTypes", "ArrayInterface", "ConcreteStructs", "DiffEqBase", "FastBroadcast", "FastClosures", "FiniteDiff", "ForwardDiff", "LazyArrays", "LineSearches", "LinearAlgebra", "LinearSolve", "MaybeInplace", "PrecompileTools", "Preferences", "Printf", "RecursiveArrayTools", "Reexport", "SciMLBase", "SimpleNonlinearSolve", "SparseArrays", "SparseDiffTools", "StaticArraysCore", "TimerOutputs"]
+git-tree-sha1 = "78bdd3a4a62865cf43c53d63783b0cbfddcdbbe6"
 uuid = "8913a72c-1f9b-4ce2-8d82-65094dcecaec"
-version = "3.4.0"
+version = "3.5.0"
 
     [deps.NonlinearSolve.extensions]
     NonlinearSolveBandedMatricesExt = "BandedMatrices"
@@ -1295,9 +1307,9 @@ version = "0.5.5+0"
 
 [[deps.Optimization]]
 deps = ["ADTypes", "ArrayInterface", "ConsoleProgressMonitor", "DocStringExtensions", "LinearAlgebra", "Logging", "LoggingExtras", "Pkg", "Printf", "ProgressLogging", "Reexport", "SciMLBase", "SparseArrays", "SymbolicIndexingInterface", "TerminalLoggers"]
-git-tree-sha1 = "31fe8abda56b2168d262b515ecabdb44d7c36b4d"
+git-tree-sha1 = "e24a89f3f15fd4beff32a12bde4310768f47c5bc"
 uuid = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
-version = "3.21.1"
+version = "3.21.2"
 
     [deps.Optimization.extensions]
     OptimizationEnzymeExt = "Enzyme"
@@ -1447,6 +1459,10 @@ version = "0.5.5"
 deps = ["Unicode"]
 uuid = "de0858da-6303-5e67-8744-51eddeeeb8d7"
 
+[[deps.Profile]]
+deps = ["Printf"]
+uuid = "9abbd945-dff8-562f-b5e8-e1ebf5ef1b79"
+
 [[deps.ProgressLogging]]
 deps = ["Logging", "SHA", "UUIDs"]
 git-tree-sha1 = "80d919dee55b9c50e8d9e2da5eeafff3fe58b539"

benchmarks/Fitzhugh_Nagumo/Project.toml

@@ -1,5 +1,7 @@
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
+BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 PSOGPU = "ab63da0c-63b4-40fa-a3b7-d2cba5be6419"
+StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

benchmarks/Fitzhugh_Nagumo/ode_estimation.jl

@@ -1,25 +1,27 @@
+using Pkg
+Pkg.activate(@__DIR__)
 using PSOGPU, OrdinaryDiffEq, StaticArrays
 
 function f(u, p, t)
-    a,b,τinv,l = p
-    v,w = u
-    dv = v - v^3/3 -w + l
-    dw = τinv*(v +  a - b*w)
+    a, b, τinv, l = p
+    v, w = u
+    dv = v - v^3 / 3 - w + l
+    dw = τinv * (v + a - b * w)
     return SVector{2}(dv, dw)
 end
 
-p = @SArray [0.7f0,0.8f0,0.08f0,0.5f0]              # Parameters used to construct the dataset
+p = @SArray [0.7f0, 0.8f0, 0.08f0, 0.5f0]              # Parameters used to construct the dataset
 r0 = @SArray [1.0f0; 1.0f0]                     # initial value
 tspan = (0.0f0, 30.0f0)                 # sample of 3000 observations over the (0,30) timespan
-prob = ODEProblem(f, r0, tspan,p)
+prob = ODEProblem(f, r0, tspan, p)
 
 tspan2 = (0.0f0, 3.0f0)                 # sample of 300 observations with a timestep of 0.01
-prob_short = ODEProblem(f, r0, tspan2,p)
+prob_short = ODEProblem(f, r0, tspan2, p)
 
-dt = 30.0f0/3000f0
+dt = 30.0f0 / 3000.0f0
 tf = 30.0f0
-tinterval = 0f0:dt:tf
-t  = collect(tinterval)
+tinterval = 0.0f0:dt:tf
+t = collect(tinterval)
 
 h = 0.01f0
 M = 300
@@ -28,20 +30,23 @@ tstop = tstart + M * h
 tinterval_short = 0:h:tstop
 t_short = collect(tinterval_short)
 
-data_sol_short = solve(prob_short,Vern9(),saveat=t_short,reltol=1f-6,abstol=1f-6)
+data_sol_short = solve(prob_short,
+    Vern9(),
+    saveat = t_short,
+    reltol = 1.0f-6,
+    abstol = 1.0f-6)
 data_short = convert(Array, data_sol_short) # This operation produces column major dataset obs as columns, equations as rows
-data_sol = solve(prob,Vern9(),saveat=t,reltol=1f-6,abstol=1f-6)
+data_sol = solve(prob, Vern9(), saveat = t, reltol = 1.0f-6, abstol = 1.0f-6)
 data = convert(Array, data_sol)
 
-using Plots
+# using Plots
 
-plot(data_sol_short)
+# plot(data_sol_short)
 
-plot(data_sol)
+# plot(data_sol)
 
 n_particles = 10_000
 
-
 # obj_short = build_loss_objective(prob_short,Tsit5(),L2Loss(t_short,data_short),tstops=t_short)
 function loss(u, p)
     odeprob, t = p
@@ -50,8 +55,8 @@ function loss(u, p)
     sum(abs2, data_short .- pred)
 end
 
-lb = @SArray fill(0.f0, 4)
-ub = @SArray fill(5.f0, 4)
+lb = @SArray fill(0.0f0, 4)
+ub = @SArray fill(5.0f0, 4)
 
 optprob = OptimizationProblem(loss, prob.p, (prob, t_short); lb = lb, ub = ub)
 
@@ -68,20 +73,30 @@ gpu_particles = cu(particles)
 
 CUDA.allowscalar(false)
 
-function prob_func(prob, gpu_particle)
-    return remake(prob, p = (prob.p[1], gpu_particle.position))
-end
-
 using Adapt
 
 losses = adapt(CUDABackend(), ones(eltype(prob.u0), (1, n_particles)))
 
 solver_cache = (; losses, gpu_particles, gpu_data, gbest)
 
+adaptive = false
+
 @time gsol = PSOGPU.parameter_estim_ode!(prob,
     solver_cache,
     lb,
-    ub;
+    ub, Val(adaptive);
     saveat = t_short,
     dt = 0.1f0,
-    maxiters = 1)
+    maxiters = 100)
+
+using BenchmarkTools
+
+@benchmark PSOGPU.parameter_estim_ode!($prob,
+    $(deepcopy(solver_cache)),
+    $lb,
+    $ub, $Val(adaptive);
+    saveat = t_short,
+    dt = 0.1f0,
+    maxiters = 100)
+
+@show gbest.cost, gsol