optimization-test.jl

using BenchmarkTools
using StaticArrays
using Optimization, OptimizationOptimJL
using SimpleOptimization
using Plots

# Define vector types
struct Vec{T} <: FieldVector{2, T}
    x::T
    y::T
end

mutable struct MutVec{T} <: FieldVector{2, T}
    x::T
    y::T
end

StaticArrays.similar_type(::Type{<:Vec}, ::Type{T}, s::Size{(2,)}) where T = Vec{T}
Base.similar(::Type{<:Vec}, ::Type{T}) where T = Vec(zero(T), zero(T))

StaticArrays.similar_type(::Type{<:MutVec}, ::Type{T}, s::Size{(2,)}) where T = MutVec{T}
Base.similar(::Type{<:MutVec}, ::Type{T}) where T = MutVec(zero(T), zero(T))

function mutableidentity(::MutVec{T}) where T
    MMatrix{2, 2, T}(1, 0, 0, 1)
end

# Optimization problem
function rosenbrock(v, p)
    (p[1] - v.x)^2 + p[2] * (v.y - v.x^2)^2
end

# Test on problem
fn = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff())
p = (0.1, 100.)
v₀ = Vec(0., 0.)
prob = OptimizationProblem(fn, v₀, p; lb = Vec(-1., -1.), ub = Vec(1., 0.5 * p[1]))
maxiters = 100_000

@btime solve($prob, $(SimpleLBFGS()); maxiters = $maxiters); 
sol = solve(prob, SimpleLBFGS(); maxiters = maxiters)

mv₀ = MutVec(0., 0.)
mutprob = OptimizationProblem(fn, mv₀, p; lb = MutVec(-1., -1.), ub = MutVec(1., 0.5 * p[1]))

@btime solve($mutprob, $(LBFGS()))
mutsol = solve(mutprob, LBFGS())

@assert maximum(abs2, mutsol.u .- sol.u) < 1e-5

# Test on grid
n = 100
p₁space = range(0, 1; length = n)
p₂space = range(50, 150; length = n + 1)
pgrid = Iterators.product(p₁space, p₂space)

failedoptimisation = fill(0, size(pgrid))
for (i, p) in enumerate(pgrid)
    prob = OptimizationProblem(fn, Vec(0., 0.), p; lb = Vec(-1., -1.), ub = Vec(1., 0.5 * p[1]))

    sol = solve(prob, SimpleLBFGS(); maxiters = maxiters)

    if !SciMLBase.successful_retcode(sol)
        failedoptimisation[i] = 1
        @warn "Error at $p: $(sol.retcode)"
    end
end

sum(failedoptimisation) / length(failedoptimisation)
heatmap(p₁space, p₂space, failedoptimisation'; xlabel = "p1", ylabel = "p2", c = :grays)

# Timing
@benchmark optimovergrid!($mutmaximiser, $pgrid)
@benchmark simpleoptimovergrid!($maximiser, $pgrid)

status.txt

Status `/private/tmp/optimisation/Project.toml`
  [2b5f629d] DiffEqBase v6.183.1
  [7f7a1694] Optimization v4.5.0
  [36348300] OptimizationOptimJL v0.4.3
  [295af30f] Revise v3.9.0
  [657a9c50] SimpleOptimization v0.1.0 `https://github.com/SciML/SimpleOptimization.jl#main`
  [90137ffa] StaticArrays v1.9.14