OptimizationODE package for ODE solvers

lib/OptimizationODE/Project.toml

@@ -0,0 +1,16 @@
+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"
+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,129 @@
+module OptimizationODE
+
+using Reexport
+@reexport using Optimization
+using Optimization.SciMLBase
+
+export ODEGradientDescent
+
+# The optimizer “type”
+
+struct ODEGradientDescent end
+
+# capability flags
+SciMLBase.requiresbounds(::ODEGradientDescent)           = false
+SciMLBase.allowsbounds(::ODEGradientDescent)             = false
+SciMLBase.allowscallback(::ODEGradientDescent)           = false
+SciMLBase.supports_opt_cache_interface(::ODEGradientDescent) = true
+SciMLBase.requiresgradient(::ODEGradientDescent)         = true
+SciMLBase.requireshessian(::ODEGradientDescent)          = false
+SciMLBase.requiresconsjac(::ODEGradientDescent)          = false
+SciMLBase.requiresconshess(::ODEGradientDescent)         = false
+
+# Map standard kwargs to our solver’s args
+
+function __map_optimizer_args!(
+    cache::OptimizationCache, opt::ODEGradientDescent;
+    callback    = nothing,
+    maxiters::Union{Number,Nothing}=nothing,
+    maxtime::Union{Number,Nothing}=nothing,
+    abstol::Union{Number,Nothing}=nothing,
+    reltol::Union{Number,Nothing}=nothing,
+    η::Float64 = 0.1,
+    tmax::Float64 = 1.0,
+    dt::Float64 = 0.01,
+    kwargs...
+)
+    # override our defaults
+    cache.solver_args = merge(cache.solver_args, (
+        η = η,
+        tmax = tmax,
+        dt = dt,
+    ))
+    # now apply common options
+    if !(isnothing(maxiters))
+        cache.solver_args.maxiters = maxiters
+    end
+    if !(isnothing(maxtime))
+        cache.solver_args.maxtime = maxtime
+    end
+    return nothing
+end
+
+# 3) Initialize the cache (captures f, u0, bounds, and solver_args)
+
+function SciMLBase.__init(
+    prob::SciMLBase.OptimizationProblem,
+    opt::ODEGradientDescent,
+    data = Optimization.DEFAULT_DATA;
+    η::Float64 = 0.1,
+    tmax::Float64 = 1.0,
+    dt::Float64 = 0.01,
+    callback  = (args...)->false,
+    progress  = false,
+    kwargs...
+)
+    return OptimizationCache(
+        prob, opt, data;
+        η        = η,
+        tmax     = tmax,
+        dt       = dt,
+        callback = callback,
+        progress = progress,
+        maxiters = nothing,
+        maxtime  = nothing,
+        kwargs...
+    )
+end
+
+# 4) The actual solve loop: Euler integration of gradient descent
+
+function SciMLBase.__solve(
+    cache::OptimizationCache{F,RC,LB,UB,LC,UC,S,O,D,P,C}
+) where {F,RC,LB,UB,LC,UC,S,O<:ODEGradientDescent,D,P,C}
+
+    # unpack initial state & parameters
+    u0      = cache.u0
+    η       = get(cache.solver_args, :η, 0.1)
+    tmax    = get(cache.solver_args, :tmax, 1.0)
+    dt      = get(cache.solver_args, :dt, 0.01)
+    maxiter = get(cache.solver_args, :maxiters, nothing)
+
+    # prepare working storage
+    u = copy(u0)
+    G = similar(u)
+
+    t    = 0.0
+    iter = 0
+    # Euler loop
+    while (isnothing(maxiter) || iter < maxiter) && t <= tmax
+        # compute gradient in‐place
+        cache.f.grad(G, u, cache.p)
+        # Euler step
+        u .-= η .* G
+        t    += dt
+        iter += 1
+    end
+
+    # final objective
+    fval = cache.f(u, cache.p)
+
+    # record stats: one final f‐eval, iter gradient‐evals
+    stats = Optimization.OptimizationStats(
+        iterations = iter,
+        time       = 0.0,        # could time() if you like
+        fevals     = 1,
+        gevals     = iter,
+        hevals     = 0
+    )
+
+    return SciMLBase.build_solution(
+        cache, cache.opt,
+        u,
+        fval,
+        retcode = ReturnCode.Success,
+        stats   = stats
+    )
+end
+
+end # module

lib/OptimizationODE/test/runtests.jl

@@ -0,0 +1,45 @@
+using Test
+using Optimization
+using Optimization.SciMLBase
+using OptimizationODE
+
+@testset "ODEGradientDescent Tests" begin
+
+    # Define the Rosenbrock objective and its gradient
+    function rosen(u, p)
+        return (p[1] - u[1])^2 + p[2] * (u[2] - u[1]^2)^2
+    end
+
+    function rosen_grad!(g, u, p, data)
+        g[1] = -2 * (p[1] - u[1]) - 4 * p[2] * u[1] * (u[2] - u[1]^2)
+        g[2] =  2 * p[2] * (u[2] - u[1]^2)
+        return g
+    end
+
+    # Set up the problem
+    u0 = [0.0, 0.0]
+    p  = [1.0, 100.0]
+
+    # Wrap into an OptimizationFunction without AD, providing our gradient
+    f = OptimizationFunction(
+        rosen,
+        Optimization.SciMLBase.NoAD();
+        grad = rosen_grad!
+    )
+
+    prob = OptimizationProblem(f, u0, p)
+
+    # Solve with ODEGradientDescent 
+    sol = solve(
+        prob,
+        ODEGradientDescent();
+        η    = 0.001,
+        tmax = 1_000.0,
+        dt   = 0.01
+    )
+
+    # Assertions 
+    @test isapprox(sol.u[1], 1.0; atol = 1e-2)
+    @test isapprox(sol.u[2], 1.0; atol = 1e-2)
+
+end