diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index ef2e6850e..8cfed5ae4 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -24,11 +24,11 @@ jobs: - OptimizationEvolutionary - OptimizationFlux - OptimizationGCMAES + - OptimizationManopt - OptimizationMetaheuristics - OptimizationMOI - OptimizationMultistartOptimization - OptimizationNLopt - #- OptimizationNonconvex - OptimizationNOMAD - OptimizationOptimJL - OptimizationOptimisers diff --git a/lib/OptimizationManopt/Project.toml b/lib/OptimizationManopt/Project.toml new file mode 100644 index 000000000..f60db4bb0 --- /dev/null +++ b/lib/OptimizationManopt/Project.toml @@ -0,0 +1,23 @@ +name = "OptimizationManopt" +uuid = "e57b7fff-7ee7-4550-b4f0-90e9476e9fb6" +authors = ["Mateusz Baran "] +version = "0.1.0" + +[deps] +LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +ManifoldDiff = "af67fdf4-a580-4b9f-bbec-742ef357defd" +Manifolds = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e" +ManifoldsBase = "3362f125-f0bb-47a3-aa74-596ffd7ef2fb" +Manopt = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5" +Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" +Reexport = "189a3867-3050-52da-a836-e630ba90ab69" + +[extras] +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" +ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[targets] +test = ["Enzyme", "ForwardDiff", "Random", "Test", "Zygote"] diff --git a/lib/OptimizationManopt/src/OptimizationManopt.jl b/lib/OptimizationManopt/src/OptimizationManopt.jl new file mode 100644 index 000000000..6f536d21a --- /dev/null +++ b/lib/OptimizationManopt/src/OptimizationManopt.jl @@ -0,0 +1,311 @@ +module OptimizationManopt + +using Reexport +@reexport using Manopt +using Optimization, Manopt, ManifoldsBase, ManifoldDiff, Optimization.SciMLBase + +""" + abstract type AbstractManoptOptimizer end + +A Manopt solver without things specified by a call to `solve` (stopping criteria) and +internal state. +""" +abstract type AbstractManoptOptimizer end + +SciMLBase.supports_opt_cache_interface(opt::AbstractManoptOptimizer) = true + +function __map_optimizer_args!(cache::OptimizationCache, + opt::AbstractManoptOptimizer; + callback = nothing, + maxiters::Union{Number, Nothing} = nothing, + maxtime::Union{Number, Nothing} = nothing, + abstol::Union{Number, Nothing} = nothing, + reltol::Union{Number, Nothing} = nothing, + kwargs...) + solver_kwargs = (; kwargs...) + + if !isnothing(maxiters) + solver_kwargs = (; + solver_kwargs..., stopping_criterion = [Manopt.StopAfterIteration(maxiters)]) + end + + if !isnothing(maxtime) + if haskey(solver_kwargs, :stopping_criterion) + solver_kwargs = (; solver_kwargs..., + stopping_criterion = push!( + solver_kwargs.stopping_criterion, Manopt.StopAfterTime(maxtime))) + else + solver_kwargs = (; + solver_kwargs..., stopping_criterion = [Manopt.StopAfter(maxtime)]) + end + end + + if !isnothing(abstol) + if haskey(solver_kwargs, :stopping_criterion) + solver_kwargs = (; solver_kwargs..., + stopping_criterion = push!( + solver_kwargs.stopping_criterion, Manopt.StopWhenChangeLess(abstol))) + else + solver_kwargs = (; + solver_kwargs..., stopping_criterion = [Manopt.StopWhenChangeLess(abstol)]) + end + end + + if !isnothing(reltol) + @warn "common reltol is currently not used by $(typeof(opt).super)" + end + return solver_kwargs +end + +## gradient descent +struct GradientDescentOptimizer <: AbstractManoptOptimizer end + +function call_manopt_optimizer( + M::ManifoldsBase.AbstractManifold, opt::GradientDescentOptimizer, + loss, + gradF, + x0; + stopping_criterion::Union{Manopt.StoppingCriterion, Manopt.StoppingCriterionSet}, + evaluation::AbstractEvaluationType = Manopt.AllocatingEvaluation(), + stepsize::Stepsize = ArmijoLinesearch(M), + kwargs...) + opts = gradient_descent(M, + loss, + gradF, + x0; + return_state = true, + evaluation, + stepsize, + stopping_criterion) + # we unwrap DebugOptions here + minimizer = Manopt.get_solver_result(opts) + return (; minimizer = minimizer, minimum = loss(M, minimizer), options = opts) +end + +## Nelder-Mead +struct NelderMeadOptimizer <: AbstractManoptOptimizer end + +function call_manopt_optimizer(M::ManifoldsBase.AbstractManifold, opt::NelderMeadOptimizer, + loss, + gradF, + x0; + stopping_criterion::Union{Manopt.StoppingCriterion, Manopt.StoppingCriterionSet}, + kwargs...) + opts = NelderMead(M, + loss; + return_state = true, + stopping_criterion) + minimizer = Manopt.get_solver_result(opts) + return (; minimizer = minimizer, minimum = loss(M, minimizer), options = opts) +end + +## conjugate gradient descent +struct ConjugateGradientDescentOptimizer <: AbstractManoptOptimizer end + +function call_manopt_optimizer(M::ManifoldsBase.AbstractManifold, + opt::ConjugateGradientDescentOptimizer, + loss, + gradF, + x0; + stopping_criterion::Union{Manopt.StoppingCriterion, Manopt.StoppingCriterionSet}, + evaluation::AbstractEvaluationType = InplaceEvaluation(), + stepsize::Stepsize = ArmijoLinesearch(M), + kwargs...) + opts = conjugate_gradient_descent(M, + loss, + gradF, + x0; + return_state = true, + evaluation, + stepsize, + stopping_criterion) + # we unwrap DebugOptions here + minimizer = Manopt.get_solver_result(opts) + return (; minimizer = minimizer, minimum = loss(M, minimizer), options = opts) +end + +## particle swarm +struct ParticleSwarmOptimizer <: AbstractManoptOptimizer end + +function call_manopt_optimizer(M::ManifoldsBase.AbstractManifold, + opt::ParticleSwarmOptimizer, + loss, + gradF, + x0; + stopping_criterion::Union{Manopt.StoppingCriterion, Manopt.StoppingCriterionSet}, + evaluation::AbstractEvaluationType = InplaceEvaluation(), + population_size::Int = 100, + retraction_method::AbstractRetractionMethod = default_retraction_method(M), + inverse_retraction_method::AbstractInverseRetractionMethod = default_inverse_retraction_method(M), + vector_transport_method::AbstractVectorTransportMethod = default_vector_transport_method(M), + kwargs...) + initial_population = vcat([x0], [rand(M) for _ in 1:(population_size - 1)]) + opts = particle_swarm(M, + loss; + x0 = initial_population, + n = population_size, + return_state = true, + retraction_method, + inverse_retraction_method, + vector_transport_method, + stopping_criterion) + # we unwrap DebugOptions here + minimizer = Manopt.get_solver_result(opts) + return (; minimizer = minimizer, minimum = loss(M, minimizer), options = opts) +end + +## quasi Newton + +struct QuasiNewtonOptimizer <: AbstractManoptOptimizer end + +function call_manopt_optimizer(M::Manopt.AbstractManifold, + opt::QuasiNewtonOptimizer, + loss, + gradF, + x0; + stopping_criterion::Union{Manopt.StoppingCriterion, Manopt.StoppingCriterionSet}, + evaluation::AbstractEvaluationType = InplaceEvaluation(), + retraction_method::AbstractRetractionMethod = default_retraction_method(M), + vector_transport_method::AbstractVectorTransportMethod = default_vector_transport_method(M), + stepsize = WolfePowellLinesearch(M; + retraction_method = retraction_method, + vector_transport_method = vector_transport_method, + linesearch_stopsize = 1e-12), + kwargs... +) + opts = quasi_Newton(M, + loss, + gradF, + x0; + return_state = true, + evaluation, + retraction_method, + vector_transport_method, + stepsize, + stopping_criterion) + # we unwrap DebugOptions here + minimizer = Manopt.get_solver_result(opts) + return (; minimizer = minimizer, minimum = loss(M, minimizer), options = opts) +end + +## Optimization.jl stuff + +function build_loss(f::OptimizationFunction, prob, cb) + function (::AbstractManifold, θ) + x = f.f(θ, prob.p) + cb(x, θ) + __x = first(x) + return prob.sense === Optimization.MaxSense ? -__x : __x + end +end + +function build_gradF(f::OptimizationFunction{true}, cur) + function g(M::AbstractManifold, G, θ) + f.grad(G, θ, cur...) + G .= riemannian_gradient(M, θ, G) + end + function g(M::AbstractManifold, θ) + G = zero(θ) + f.grad(G, θ, cur...) + return riemannian_gradient(M, θ, G) + end +end + +# TODO: +# 1) convert tolerances and other stopping criteria +# 2) return convergence information +# 3) add callbacks to Manopt.jl + +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 <: + AbstractManoptOptimizer, + D, + P, + C +} + local x, cur, state + + manifold = haskey(cache.solver_args, :manifold) ? cache.solver_args[:manifold] : nothing + + if manifold === nothing + throw(ArgumentError("Manifold not specified in the problem for e.g. `OptimizationProblem(f, x, p; manifold = SymmetricPositiveDefinite(5))`.")) + end + + if cache.data !== Optimization.DEFAULT_DATA + maxiters = length(cache.data) + else + maxiters = cache.solver_args.maxiters + end + + cur, state = iterate(cache.data) + + function _cb(x, θ) + opt_state = Optimization.OptimizationState(iter = 0, + u = θ, + objective = x[1]) + cb_call = cache.callback(opt_state, x...) + if !(cb_call isa Bool) + error("The callback should return a boolean `halt` for whether to stop the optimization process.") + end + nx_itr = iterate(cache.data, state) + if isnothing(nx_itr) + true + else + cur, state = nx_itr + cb_call + end + end + solver_kwarg = __map_optimizer_args!(cache, cache.opt, callback = _cb, + maxiters = maxiters, + maxtime = cache.solver_args.maxtime, + abstol = cache.solver_args.abstol, + reltol = cache.solver_args.reltol; + ) + + _loss = build_loss(cache.f, cache, _cb) + + gradF = build_gradF(cache.f, cur) + + if haskey(solver_kwarg, :stopping_criterion) + stopping_criterion = Manopt.StopWhenAny(solver_kwarg.stopping_criterion...) + else + stopping_criterion = Manopt.StopAfterIteration(500) + end + + opt_res = call_manopt_optimizer(manifold, cache.opt, _loss, gradF, cache.u0; + solver_kwarg..., stopping_criterion = stopping_criterion) + + asc = get_active_stopping_criteria(opt_res.options.stop) + + opt_ret = any(Manopt.indicates_convergence, asc) ? ReturnCode.Success : + ReturnCode.Failure + + return SciMLBase.build_solution(cache, + cache.opt, + opt_res.minimizer, + cache.sense === Optimization.MaxSense ? + -opt_res.minimum : opt_res.minimum; + original = opt_res.options, + retcode = opt_ret) +end + +end # module OptimizationManopt diff --git a/lib/OptimizationManopt/test/runtests.jl b/lib/OptimizationManopt/test/runtests.jl new file mode 100644 index 000000000..ca9e5d083 --- /dev/null +++ b/lib/OptimizationManopt/test/runtests.jl @@ -0,0 +1,137 @@ +using OptimizationManopt +using Optimization +using Manifolds +using ForwardDiff, Zygote, Enzyme +using Manopt +using Test +using Optimization.SciMLBase +using LinearAlgebra + +rosenbrock(x, p) = (p[1] - x[1])^2 + p[2] * (x[2] - x[1]^2)^2 + +function rosenbrock_grad!(storage, x, p) + storage[1] = -2.0 * (p[1] - x[1]) - 4.0 * p[2] * (x[2] - x[1]^2) * x[1] + storage[2] = 2.0 * p[2] * (x[2] - x[1]^2) +end + +R2 = Euclidean(2) + +@testset "Error on no or mismatching manifolds" begin + x0 = zeros(2) + p = [1.0, 100.0] + + stepsize = Manopt.ArmijoLinesearch(R2) + opt = OptimizationManopt.GradientDescentOptimizer() + + optprob_forwarddiff = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff()) + prob_forwarddiff = OptimizationProblem(optprob_forwarddiff, x0, p) + @test_throws ArgumentError("Manifold not specified in the problem for e.g. `OptimizationProblem(f, x, p; manifold = SymmetricPositiveDefinite(5))`.") Optimization.solve( + prob_forwarddiff, opt) +end + +@testset "Gradient descent" begin + x0 = zeros(2) + p = [1.0, 100.0] + + stepsize = Manopt.ArmijoLinesearch(R2) + opt = OptimizationManopt.GradientDescentOptimizer() + + optprob_forwarddiff = OptimizationFunction(rosenbrock, Optimization.AutoEnzyme()) + prob_forwarddiff = OptimizationProblem( + optprob_forwarddiff, x0, p; manifold = R2, stepsize = stepsize) + sol = Optimization.solve(prob_forwarddiff, opt) + @test sol.minimum < 0.2 + + optprob_grad = OptimizationFunction(rosenbrock; grad = rosenbrock_grad!) + prob_grad = OptimizationProblem(optprob_grad, x0, p; manifold = R2, stepsize = stepsize) + sol = Optimization.solve(prob_grad, opt) + @test sol.minimum < 0.2 +end + +@testset "Nelder-Mead" begin + x0 = zeros(2) + p = [1.0, 100.0] + + opt = OptimizationManopt.NelderMeadOptimizer() + + optprob = OptimizationFunction(rosenbrock) + prob = OptimizationProblem(optprob, x0, p; manifold = R2) + + sol = Optimization.solve(prob, opt) + @test sol.minimum < 1e-6 +end + +@testset "Conjugate gradient descent" begin + x0 = zeros(2) + p = [1.0, 100.0] + + stepsize = Manopt.ArmijoLinesearch(R2) + opt = OptimizationManopt.ConjugateGradientDescentOptimizer() + + optprob = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff()) + prob = OptimizationProblem(optprob, x0, p; manifold = R2) + + sol = Optimization.solve(prob, opt, stepsize = stepsize) + @test sol.minimum < 0.5 +end + +@testset "Quasi Newton" begin + x0 = zeros(2) + p = [1.0, 100.0] + + opt = OptimizationManopt.QuasiNewtonOptimizer() + function callback(state, l) + println(state.u) + println(l) + return false + end + optprob = OptimizationFunction(rosenbrock, Optimization.AutoForwardDiff()) + prob = OptimizationProblem(optprob, x0, p; manifold = R2) + + sol = Optimization.solve(prob, opt, callback = callback, maxiters = 30) + @test sol.minimum < 1e-14 +end + +@testset "Particle swarm" begin + x0 = zeros(2) + p = [1.0, 100.0] + + opt = OptimizationManopt.ParticleSwarmOptimizer() + + optprob = OptimizationFunction(rosenbrock) + prob = OptimizationProblem(optprob, x0, p; manifold = R2) + + sol = Optimization.solve(prob, opt) + @test sol.minimum < 0.1 +end + +@testset "Custom constraints" begin + cons(res, x, p) = (res .= [x[1]^2 + x[2]^2, x[1] * x[2]]) + + x0 = zeros(2) + p = [1.0, 100.0] + opt = OptimizationManopt.GradientDescentOptimizer() + + optprob_cons = OptimizationFunction(rosenbrock; grad = rosenbrock_grad!, cons = cons) + prob_cons = OptimizationProblem(optprob_cons, x0, p) + @test_throws SciMLBase.IncompatibleOptimizerError Optimization.solve(prob_cons, opt) +end + +@testset "SPD Manifold" begin + M = SymmetricPositiveDefinite(5) + m = 100 + σ = 0.005 + q = Matrix{Float64}(I, 5, 5) .+ 2.0 + data2 = [exp(M, q, σ * rand(M; vector_at = q)) for i in 1:m] + + f(M, x, p = nothing) = sum(distance(M, x, data2[i])^2 for i in 1:m) + f(x, p = nothing) = sum(distance(M, x, data2[i])^2 for i in 1:m) + + optf = OptimizationFunction(f, Optimization.AutoForwardDiff()) + prob = OptimizationProblem(optf, data2[1]; manifold = M, maxiters = 1000) + + opt = OptimizationManopt.GradientDescentOptimizer() + @time sol = Optimization.solve(prob, opt) + + @test sol.u≈q atol=1e-2 +end