Restart a few thoughts on non-array-types.

Author: kellertuer

Project.toml

@@ -31,6 +31,7 @@ RipQP = "1e40b3f8-35eb-4cd8-8edd-3e515bb9de08"
 
 [extensions]
 ManoptJuMPExt = "JuMP"
+ManoptJuMPManifoldsExt = ["JuMP", "Manifolds"]
 ManoptLRUCacheExt = "LRUCache"
 ManoptLineSearchesExt = "LineSearches"
 ManoptManifoldsExt = "Manifolds"

ext/ManoptJuMPExt.jl

@@ -194,10 +194,11 @@ end
 Return a `Bool` indicating whether `attr.name` is a valid option name
 for `Manopt`.
 """
-function MOI.supports(::ManoptOptimizer, ::MOI.RawOptimizerAttribute)
+function MOI.supports(::ManoptOptimizer, attr::MOI.RawOptimizerAttribute)
     # FIXME Ideally, this should only return `true` if it is a valid keyword argument for
     #       one of the `...DescentState()` constructors. Is there an easy way to check this ?
     #       Does it depend on the different solvers ?
+    @info attr.name
     return true
 end
 
@@ -503,6 +504,10 @@ function MOI.optimize!(model::ManoptOptimizer)
     return nothing
 end
 
+#
+#
+# A wrapper for points that are just array shaped
+
 @doc """
     ManifoldPointArrayShape{N} <: JuMP.AbstractShape
 
@@ -588,29 +593,65 @@ At the moment, we only support manifolds for which the shape is a `Array`.
 function _shape(m::ManifoldsBase.AbstractManifold)
     return ManifoldPointArrayShape(ManifoldsBase.representation_size(m))
 end
-
-_in(mime::MIME"text/plain") = "in"
-_in(mime::MIME"text/latex") = "\\in"
-
-function JuMP.in_set_string(mime, set::ManifoldsBase.AbstractManifold)
-    return _in(mime) * " " * string(set)
+function _shape(m::ManifoldsBase.AbstractManifold, ::Nothing)
+    return _shape(m)
+end
+function _shape(m::ManifoldsBase.AbstractManifold, ::P) where {P}
+    error("TODO")
 end
 
 """
-    JuMP.build_variable(::Function, func, m::ManifoldsBase.AbstractManifold)
+    JuMP.build_variable(::Function, func, m::ManifoldsBase.AbstractManifold, ::Type{P}=Nothing)
 
 Build a `JuMP.VariablesConstrainedOnCreation` object containing variables
 and the [`ManifoldSet`](@ref) in which they should belong as well as the
 `shape` that can be used to go from the vectorized MOI representation to the
 shape of the manifold, that is, [`ManifoldPointArrayShape`](@ref).
+
+The optional parameter `P` can be used to indicate that the point is different from an integer
 """
-function JuMP.build_variable(::Function, func, m::ManifoldsBase.AbstractManifold)
+function JuMP.build_variable(::Function, point, m::ManifoldsBase.AbstractManifold)
+    @info point
     shape = _shape(m)
     return JuMP.VariablesConstrainedOnCreation(
-        JuMP.vectorize(func, shape), ManifoldSet(m), shape
+        JuMP.vectorize(point, shape), ManifoldSet(m), shape
     )
 end
 
+#
+#
+# Non-Array points, either general structs or those that require wrappers
+# Test examples:
+# * HyperboloidPoint (the same as array on Hyperbolic, so basically an array)
+# * SVDMPoint on FixedRank, which actualy has 3 fields
+
+@doc """
+    ManifoldPointArrayShape{M, P} <: JuMP.AbstractShape
+
+A wrapper for points on a manifold of type `M` of type `P`
+
+# Fields
+
+* manifold::M
+"""
+struct ManifoldPointShape{M <: AbstractManifold, P} <: JuMP.AbstractShape
+    manifold::M
+end
+
+# Functions we need
+
+
+#
+#
+# Generic further functions
+
+_in(mime::MIME"text/plain") = "in"
+_in(mime::MIME"text/latex") = "\\in"
+
+function JuMP.in_set_string(mime, set::ManifoldsBase.AbstractManifold)
+    return _in(mime) * " " * string(set)
+end
+
 """
     MOI.get(model::ManoptOptimizer, ::MOI.ResultCount)
 

ext/ManoptJuMPManifoldsExt.jl

@@ -0,0 +1,49 @@
+module ManoptJuMPManifoldsExt
+
+using Manopt
+using ManifoldsBase
+using Manifolds
+using LinearAlgebra
+using JuMP: JuMP
+const MOI = JuMP.MOI
+# the order of extensions is as such that
+# ManoptJuMPExt is loaded after Manopt & JuMP but before this one.
+const MJE = Base.get_extension(Manopt, :ManoptJuMPExt)
+#
+#
+# define further conversions between certain Manifolds point/vector types
+# and its vectorized representations for JuMP.
+#
+# Since Manopt & JuMP are already loaded we can asumme that the
+# `Manopt.JuMPManifoldPointShape` and `Manopt.JuMPTangentVectorShape` are defined.
+#
+# #TODO: 1. check that we can use the types
+# #TODO: 2. implement at least the hyperbolic conversions.
+
+# TODO: These are just proof of concept functions to extend conversion to further types.
+# vector -> point
+function MJE._reshape_vector!(
+        v::Vector{T},
+        p::Manifolds.HyperboloidPoint,
+        ::MJE.ManifoldPointShape{M, Manifolds.HyperboloidPoint},
+    ) where {T, M <: ManifoldsBase.AbstractManifold}
+    v .= p.value
+    return v
+end
+function JuMP.reshape_vector(
+        v::Vector{T}, shape::MJE.ManifoldPointShape{M, Manifolds.HyperboloidPoint}
+    ) where {T, M <: ManifoldsBase.AbstractManifold}
+    return HyperboloidPoint(v)
+end
+# point -> vector
+function JuMP.vectorize(
+        p::Manifolds.HyperboloidPoint, ::MJE.ManifoldPointShape{M, Manifolds.HyperboloidPoint}
+    ) where {M <: ManifoldsBase.AbstractManifold}
+    return p.value # is a vector already
+end
+
+function MJE._shape(m::M, ::Manifolds.HyperboloidPoint) where {M}
+    return MJE.ManifoldPointShape{M, Manifolds.HyperboloidPoint}(m)
+end
+
+end # module ManoptJuMPManifoldsExt

test/MOI_wrapper.jl

@@ -75,10 +75,10 @@ function test_sphere(descent_state_type; kws...)
         sprint(show, model),
         "Vector{VariableRef} in ManoptJuMPExt.ManifoldSet{Sphere{ℝ, ManifoldsBase.TypeParameter{Tuple{2}}}}: 1",
     )
-    @test contains(sprint(print, model), "[x[1], x[2], x[3]] in Sphere(2, ℝ)")
+    @test contains(sprint(print, model), "[x[1], x[2], x[3]] in $(Sphere(2))")
     @test contains(
         JuMP.model_string(MIME("text/latex"), model),
-        "[x_{1}, x_{2}, x_{3}] \\in Sphere(2, ℝ)",
+        "[x_{1}, x_{2}, x_{3}] \\in Sphere(2)",
     )
 
     @objective(model, Min, sum(xi^4 for xi in x))