Add support for Riemanian objective through JuMP

Changelog.md

@@ -14,6 +14,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
   especially turn maximisation problems into minimisation ones using a scaling of `-1`.
 * Introduce a `ManifoldConstrainedSetObjective`
 * Introduce a `projected_gradient_method`
+* Allow setting `AbstractManifoldObjective` through JuMP
 
 
 ## [0.5.11] April 8, 2025

ext/ManoptJuMPExt.jl

@@ -29,6 +29,10 @@ function MOI.dimension(set::VectorizedManifold)
     return prod(ManifoldsBase.representation_size(set.manifold))
 end
 
+struct RiemannianFunction{MO<:Manopt.AbstractManifoldObjective} <: MOI.AbstractScalarFunction
+    func::MO
+end
+
 mutable struct Optimizer <: MOI.AbstractOptimizer
     # Manifold in which all the decision variables leave
     manifold::Union{Nothing,ManifoldsBase.AbstractManifold}
@@ -40,8 +44,8 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
     variable_primal_start::Vector{Union{Nothing,Float64}}
     # Sense of the optimization, that is whether it is for example min, max or no objective
     sense::MOI.OptimizationSense
-    # Model used to compute gradient of the objective function with AD
-    nlp_model::MOI.Nonlinear.Model
+    # Objective function of the optimization
+    objective::Union{Nothing,Manopt.AbstractManifoldObjective}
     # Solver parameters set with `MOI.RawOptimizerAttribute`
     options::Dict{String,Any}
     function Optimizer()
@@ -51,7 +55,7 @@ mutable struct Optimizer <: MOI.AbstractOptimizer
             nothing,
             Union{Nothing,Float64}[],
             MOI.FEASIBILITY_SENSE,
-            MOI.Nonlinear.Model(),
+            nothing,
             Dict{String,Any}(DESCENT_STATE_TYPE => Manopt.GradientDescentState),
         )
     end
@@ -68,7 +72,7 @@ MOI.get(::Optimizer, ::MOI.SolverVersion) = "0.4.37"
 function MOI.is_empty(model::Optimizer)
     return isnothing(model.manifold) &&
            isempty(model.variable_primal_start) &&
-           MOI.is_empty(model.nlp_model) &&
+           isnothing(model.objective) &&
            model.sense == MOI.FEASIBILITY_SENSE
 end
 
@@ -83,7 +87,7 @@ function MOI.empty!(model::Optimizer)
     model.state = nothing
     empty!(model.variable_primal_start)
     model.sense = MOI.FEASIBILITY_SENSE
-    MOI.empty!(model.nlp_model)
+    model.objective = nothing
     return nothing
 end
 
@@ -282,9 +286,29 @@ MOI.get(model::Optimizer, ::MOI.ObjectiveSense) = model.sense
 
 Set the objective function as `func` for `model`.
 """
-function MOI.set(model::Optimizer, ::MOI.ObjectiveFunction{F}, func::F) where {F}
-    nl = convert(MOI.ScalarNonlinearFunction, func)
-    MOI.Nonlinear.set_objective(model.nlp_model, nl)
+function MOI.set(model::Optimizer, attr::MOI.ObjectiveFunction{F}, func::F) where {F<:MOI.AbstractScalarFunction}
+    backend = MOI.Nonlinear.SparseReverseMode()
+    vars = [MOI.VariableIndex(i) for i in eachindex(model.variable_primal_start)]
+    nlp_model = MOI.Nonlinear.Model()
+    evaluator = MOI.Nonlinear.Evaluator(nlp_model, backend, vars)
+    MOI.initialize(evaluator, [:Grad])
+    function eval_f_cb(M, x)
+        return MOI.eval_objective(evaluator, JuMP.vectorize(x, _shape(model.manifold)))
+    end
+    function eval_grad_f_cb(M, X)
+        x = JuMP.vectorize(X, _shape(model.manifold))
+        grad_f = zeros(length(x))
+        MOI.eval_objective_gradient(evaluator, grad_f, x)
+        reshaped_grad_f = JuMP.reshape_vector(grad_f, _shape(model.manifold))
+        return ManifoldDiff.riemannian_gradient(model.manifold, X, reshaped_grad_f)
+    end
+    objective = RiemannianFunction(Manopt.ManifoldGradientObjective(eval_f_cb, eval_grad_f_cb))
+    MOI.set(model, MOI.ObjectiveFunction{typeof(objective)}(), objective)
+    return nothing
+end
+
+function MOI.set(model::Optimizer, ::MOI.ObjectiveFunction, func::RiemannianFunction)
+    model.objective = func.func
     model.problem = nothing
     model.state = nothing
     return nothing
@@ -304,36 +328,20 @@ function MOI.optimize!(model::Optimizer)
             model.variable_primal_start[i]
         end for i in eachindex(model.variable_primal_start)
     ]
-    backend = MOI.Nonlinear.SparseReverseMode()
-    vars = [MOI.VariableIndex(i) for i in eachindex(model.variable_primal_start)]
-    evaluator = MOI.Nonlinear.Evaluator(model.nlp_model, backend, vars)
-    MOI.initialize(evaluator, [:Grad])
-    function eval_f_cb(M, x)
-        if model.sense == MOI.FEASIBILITY_SENSE
-            return 0.0
-        end
-        obj = MOI.eval_objective(evaluator, JuMP.vectorize(x, _shape(model.manifold)))
-        if model.sense == MOI.MAX_SENSE
-            obj = -obj
+    objective = model.objective
+    if isnothing(objective) || model.sense == MOI.FEASIBILITY_SENSE
+        eval_f_cb(_, _) = 0.0
+        function eval_grad_f_cb(M, X)
+            x = JuMP.vectorize(X, _shape(model.manifold))
+            grad_f = zeros(length(x))
+            reshaped_grad_f = JuMP.reshape_vector(grad_f, _shape(model.manifold))
+            return ManifoldDiff.riemannian_gradient(model.manifold, X, reshaped_grad_f)
         end
-        return obj
-    end
-    function eval_grad_f_cb(M, X)
-        x = JuMP.vectorize(X, _shape(model.manifold))
-        grad_f = zeros(length(x))
-        if model.sense == MOI.FEASIBILITY_SENSE
-            grad_f .= zero(eltype(grad_f))
-        else
-            MOI.eval_objective_gradient(evaluator, grad_f, x)
-        end
-        if model.sense == MOI.MAX_SENSE
-            LinearAlgebra.rmul!(grad_f, -1)
-        end
-        reshaped_grad_f = JuMP.reshape_vector(grad_f, _shape(model.manifold))
-        return ManifoldDiff.riemannian_gradient(model.manifold, X, reshaped_grad_f)
+        objective = Manopt.ManifoldGradientObjective(eval_f_cb, eval_grad_f_cb)
+    elseif model.sense == MOI.MAX_SENSE
+        objective = -objective
     end
-    mgo = Manopt.ManifoldGradientObjective(eval_f_cb, eval_grad_f_cb)
-    dmgo = decorate_objective!(model.manifold, mgo)
+    dmgo = decorate_objective!(model.manifold, objective)
     model.problem = DefaultManoptProblem(model.manifold, dmgo)
     reshaped_start = JuMP.reshape_vector(start, _shape(model.manifold))
     descent_state_type = model.options[DESCENT_STATE_TYPE]