OptimizationODE package for ODE solvers

.github/workflows/CI.yml

@@ -29,6 +29,7 @@ jobs:
           - OptimizationMultistartOptimization
           - OptimizationNLopt
           - OptimizationNOMAD
+          - OptimizationODE
           - OptimizationOptimJL
           - OptimizationOptimisers
           - OptimizationPRIMA

lib/OptimizationODE/Project.toml

@@ -0,0 +1,17 @@
+name = "OptimizationODE"
+uuid = "dfa73e59-e644-4d8a-bf84-188d7ecb34e4"
+authors = ["Paras Puneet Singh <[email protected]>"]
+version = "0.1.0"
+
+[deps]
+ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
+Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba"
+OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
+Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+
+[compat]
+Optimization = "3"
+Reexport = "1"
+julia = "1.9"

lib/OptimizationODE/src/OptimizationODE.jl

@@ -0,0 +1,45 @@
+module OptimizationODE
+
+using Reexport
+@reexport using Optimization, Optimization.SciMLBase
+using OrdinaryDiffEq
+
+export ODEOptimizer, ODEGradientDescent, RKChebyshevDescent, RKAccelerated, PRKChebyshevDescent
+
+# ODEOptimizer is a simple wrapper
+struct ODEOptimizer{Solver} end
+
+# Define solver aliases using stable methods
+const ODEGradientDescent    = ODEOptimizer{Euler}()
+const RKChebyshevDescent    = ODEOptimizer{ROCK2}()
+const RKAccelerated         = ODEOptimizer{BS3}()
+const PRKChebyshevDescent   = ODEOptimizer{ROCK4}()
+
+function SciMLBase.supports_opt_cache_interface(::ODEOptimizer)
+    return true
+end
+
+function SciMLBase.requiresgradient(::ODEOptimizer)
+    return true
+end
+
+function SciMLBase.__solve(
+    cache::OptimizationCache{F, RC, LB, UB, LC, UC, S, ODEOptimizer{Solver}, D, P, C}
+) where {F, RC, LB, UB, LC, UC, S, Solver, D, P, C}
+
+    prob = ODEProblem((du, u, p, t) -> begin
+        cache.f.grad(du, u, cache.p)
+        du .*= -1
+    end, cache.u0, (0.0, cache.solver_args[:maxiters] * cache.solver_args[:dt]), cache.p)
+
+    sol = solve(prob, Solver(); dt = cache.solver_args[:dt])
+
+    final_u = sol.u[end]
+    final_f = cache.f(final_u, cache.p)[1]
+    stats = Optimization.OptimizationStats(; time = sol.t[end])
+
+    return SciMLBase.build_solution(cache, ODEOptimizer{Solver}(), final_u, final_f;
+        original = sol, retcode = ReturnCode.Success, stats)
+end
+
+end

lib/OptimizationODE/test/runtests.jl

@@ -0,0 +1,67 @@
+
+using Test
+using Optimization
+using Optimization.SciMLBase
+using Optimization.ADTypes
+using OptimizationODE
+
+function quad(u, p)
+    return sum(u .^ 2) + p[1] * u[2]^2
+end
+
+function quad_grad!(g, u, p, data)
+    g[1] = 2u[1]
+    g[2] = 2p[1]*u[2]
+    return g
+end
+
+function rosenbrock(u, p)
+    return (p[1] - u[1])^2 + p[2]*(u[2] - u[1]^2)^2
+end
+
+function rosenbrock_grad!(g, u, p, data)
+    g[1] = -2(p[1] - u[1]) - 4p[2]*u[1]*(u[2] - u[1]^2)
+    g[2] = 2p[2]*(u[2] - u[1]^2)
+    return g
+end
+
+@testset "OptimizationODE Solvers" begin
+    u0q = [2.0, -3.0]
+    pq = [5.0]
+    fq = OptimizationFunction(quad, SciMLBase.NoAD(); grad = quad_grad!)
+    probQ = OptimizationProblem(fq, u0q, pq)
+
+    @testset "ODEGradientDescent on Quadratic" begin
+        sol = solve(probQ, ODEGradientDescent; dt = 0.1, maxiters = 1000)
+        @test isapprox(sol.u, zeros(length(sol.u)); atol=1e-2)
+    end
+
+    @testset "RKChebyshevDescent on Quadratic" begin
+        sol = solve(probQ, RKChebyshevDescent; dt = 0.1, maxiters = 1000)
+        @test isapprox(sol.u, zeros(length(sol.u)); atol=1e-2)
+    end
+
+    @testset "RKAccelerated on Quadratic" begin
+        sol = solve(probQ, RKAccelerated; dt = 0.1, maxiters = 1000)
+        @test isapprox(sol.u, zeros(length(sol.u)); atol=1e-2)
+    end
+
+    @testset "PRKChebyshevDescent on Quadratic" begin
+        sol = solve(probQ, PRKChebyshevDescent; dt = 0.1, maxiters = 1000)
+        @test isapprox(sol.u, zeros(length(sol.u)); atol=1e-2)
+    end
+
+    u0r = [-1.2, 1.0]
+    pr  = [1.0, 100.0]
+
+    for (mode, desc) in [(SciMLBase.NoAD(), "NoAD"),
+                         (ADTypes.AutoForwardDiff(), "AutoForwardDiff")]
+        @testset "ODEGradientDescent on Rosenbrock ($desc)" begin
+            fr = OptimizationFunction(rosenbrock, mode; grad = rosenbrock_grad!)
+            probR = OptimizationProblem(fr, u0r, pr)
+            sol = solve(probR, ODEGradientDescent; dt = 0.001, maxiters = 5000)
+            @test isapprox(sol.u[1], 1.0; atol=1e-1)
+            @test isapprox(sol.u[2], 1.0; atol=2e-1)
+        end
+    end
+end