Add support for Riemanian objective through JuMP (#448)

* Add support for Riemanian objective through JuMP
* Fix no objective case
* Add entry to Changelog
* Use zero_vector
* Fix Changelog
* Stop using scaled objective
* Fix doc

---------

Co-authored-by: Ronny Bergmann <git@ronnybergmann.net>

Author: blegat

Changelog.md

@@ -6,6 +6,12 @@ The file was started with Version `0.4`.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.5.13] unreleased
+
+### Added
+
+* Allow setting `AbstractManifoldObjective` through JuMP
+
 ## [0.5.12] April 13, 2025
 
 ### Added

docs/src/extensions.md

@@ -109,7 +109,7 @@ MOI.get(model::Manopt.JuMP_Optimizer, ::MOI.NumberOfVariables)
 MOI.supports(::Manopt.JuMP_Optimizer, ::MOI.VariablePrimalStart, ::Type{MOI.VariableIndex})
 MOI.set(::Manopt.JuMP_Optimizer, ::MOI.VariablePrimalStart, ::MOI.VariableIndex, ::Union{Real,Nothing})
 MOI.set(::Manopt.JuMP_Optimizer, ::MOI.ObjectiveSense, ::MOI.OptimizationSense)
-MOI.set(::Manopt.JuMP_Optimizer, ::MOI.ObjectiveFunction{F}, ::F) where {F}
+MOI.set(::Manopt.JuMP_Optimizer, ::MOI.ObjectiveFunction, func::MOI.AbstractScalarFunction)
 MOI.supports(::Manopt.JuMP_Optimizer, ::Union{MOI.ObjectiveSense,MOI.ObjectiveFunction})
 JuMP.build_variable(::Function, ::Any, ::Manopt.AbstractManifold)
 MOI.get(::Manopt.JuMP_Optimizer, ::MOI.ResultCount)

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
 
@@ -270,7 +274,7 @@ function MOI.set(model::Optimizer, ::MOI.ObjectiveSense, sense::MOI.Optimization
 end
 
 """
-    MOI.set(model::Optimizer, ::MOI.ObjectiveSense, sense::MOI.OptimizationSense)
+    MOI.get(model::Optimizer, ::MOI.ObjectiveSense)
 
 Return the objective sense, defaults to `MOI.FEASIBILITY_SENSE` if no sense has
 already been set.
@@ -282,9 +286,38 @@ 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}
+function MOI.set(model::Optimizer, attr::MOI.ObjectiveFunction, func::MOI.AbstractScalarFunction)
+    backend = MOI.Nonlinear.SparseReverseMode()
+    vars = [MOI.VariableIndex(i) for i in eachindex(model.variable_primal_start)]
+    nlp_model = MOI.Nonlinear.Model()
     nl = convert(MOI.ScalarNonlinearFunction, func)
-    MOI.Nonlinear.set_objective(model.nlp_model, nl)
+    MOI.Nonlinear.set_objective(nlp_model, nl)
+    evaluator = MOI.Nonlinear.Evaluator(nlp_model, backend, vars)
+    MOI.initialize(evaluator, [:Grad])
+    function eval_f_cb(M, x)
+        val = MOI.eval_objective(evaluator, JuMP.vectorize(x, _shape(model.manifold)))
+        if model.sense == MOI.MAX_SENSE
+            val = -val
+        end
+        return val
+    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)
+        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)
+    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 +337,11 @@ 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
-        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 = model.objective
+    if model.sense == MOI.FEASIBILITY_SENSE
+        objective = Manopt.ManifoldGradientObjective((_, _) -> 0.0, ManifoldsBase.zero_vector)
     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]