Introduce an EmbeddedObjective

src/plans/objective.jl

@@ -47,23 +47,155 @@
 """
 struct InplaceEvaluation <: AbstractEvaluationType end
 
-struct ReturnManifoldObjective{E,P,O<:AbstractManifoldObjective{E}} <:
-       AbstractDecoratedManifoldObjective{E,P}
-    objective::O
+@doc raw"""
+ReturnManifoldObjective{E,O2,O1<:AbstractManifoldObjective{E}} <:
+       AbstractDecoratedManifoldObjective{E,O2}
+
+A wrapper to indicate that `get_solver_result` should return the inner objetcive.
+
+The types are such that one can still dispatch on the undecorated type `O2` of the
+original objective as well.
+"""
+struct ReturnManifoldObjective{E,O2,O1<:AbstractManifoldObjective{E}} <:
+       AbstractDecoratedManifoldObjective{E,O2}
+    objective::O1
 end
 function ReturnManifoldObjective(
     o::O
 ) where {E<:AbstractEvaluationType,O<:AbstractManifoldObjective{E}}
     return ReturnManifoldObjective{E,O,O}(o)
 end
 function ReturnManifoldObjective(
-    o::O
+    o::O1
+) where {
+    E<:AbstractEvaluationType,
+    O2<:AbstractManifoldObjective,
+    O1<:AbstractDecoratedManifoldObjective{E,O2},
+}
+    return ReturnManifoldObjective{E,O2,O1}(o)
+end
+
+@doc raw"""
+    EmbeddedManifoldObjective{P, T, E, O2, O1<:AbstractManifoldObjective{E}} <:
+       AbstractDecoratedManifoldObjective{O2, O1}
+
+Declare an objective to be defined in the embedding.
+This also declares the gradient to be defined in the embedding,
+and especially being the Riesz representer with respect to the metric in the embedding.
+The types can be used to still dispatch on also the undecorated objective type `O2`.
+
+# Fields
+* `objective` – the objective that is defined in the embedding
+* `p`         - (`nothing`) a point in the embedding.
+* `X`         - (`nothing`) a tangent vector in the embedding
+
+When a point in the embedding `p` is provided, `embed!` is used in place of this point to reduce
+memory allocations. Similarly `X` is used when embedding tangent vectors
+
+"""
+struct EmbeddedManifoldObjective{P,T,E,O2,O1<:AbstractManifoldObjective{E}} <:
+       AbstractDecoratedManifoldObjective{E,O2}
+    objective::O1
+    p::P
+    X::T
+end
+function EmbeddedManifoldObjective(
+    o::O, p::P=nothing, X::T=nothing
+) where {P,T,E<:AbstractEvaluationType,O<:AbstractManifoldObjective{E}}
+    return EmbeddedManifoldObjective{P,T,E,O,O}(o, p, X)
+end
+function EmbeddedManifoldObjective(
+    o::O1, p::P=nothing, X::T=nothing
 ) where {
+    P,
+    T,
     E<:AbstractEvaluationType,
-    P<:AbstractManifoldObjective,
-    O<:AbstractDecoratedManifoldObjective{E,P},
+    O2<:AbstractManifoldObjective,
+    O1<:AbstractDecoratedManifoldObjective{E,O2},
 }
-    return ReturnManifoldObjective{E,P,O}(o)
+    return EmbeddedManifoldObjective{P,T,E,P,O}(o, p, X)
+end
+function EmbeddedManifoldObjective(
+    M::AbstractManifold,
+    o::O;
+    q=rand(M),
+    p::P=embed(M, q),
+    X::T=embed(M, q, rand(M; vector_at=q)),
+) where {P,T,O<:AbstractManifoldObjective}
+    return EmbeddedManifoldObjective(o, p, X)
+end
+
+@doc raw"""
+    get_cost(M, emo::EmbeddedManifoldObjective, p)
+
+Evaluate the cost function of an objective defined in the embedding, that is
+call [`embed`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/functions/#ManifoldsBase.embed-Tuple{AbstractManifold,%20Any})
+on the point `p` and call the original cost on this point.
+
+if `emo.p` is not nothing, the embedding is done in place of `emo.p`
+"""
+function get_cost(M, emo::EmbeddedManifoldObjective{Nothing}, p)
+    return get_cost(get_embedding(M), emo.objective, embed(M, p))
+end
+function get_cost(M, emo::EmbeddedManifoldObjective{P}, p) where {P}
+    embed!(M, emo.p, p)
+    return get_cost(get_embedding(M), emo.objective, emo.p)
+end
+
+function get_cost_function(emo::EmbeddedManifoldObjective)
+    return (M, p) -> get_cost(M, emo, p)
+end
+
+@doc raw"""
+    get_gradient(M, emo::EmbeddedManifoldObjective, p)
+    get_gradient(M, X, emo::EmbeddedManifoldObjective, p)
+
+Evaluate the gradient function of an objective defined in the embedding, that is
+call [`embed`](https://juliamanifolds.github.io/ManifoldsBase.jl/stable/functions/#ManifoldsBase.embed-Tuple{AbstractManifold,%20Any})
+on the point `p` and call the original cost on this point.
+And convert the gradient using [`riemannian_gradient`]() on the result.
+
+if `emo.p` is not nothing, the embedding is done in place of `emo.p`
+"""
+function get_gradient(M, emo::EmbeddedManifoldObjective{Nothing,Nothing}, p)
+    return riemannian_gradient(
+        M, p, get_gradient(get_embedding(M), emo.objective, embed(M, p))
+    )
+end
+function get_gradient(M, emo::EmbeddedManifoldObjective{P,Nothing}, p) where {P}
+    embed!(M, emo.p, p)
+    return riemannian_gradient(M, p, get_gradient(get_embedding(M), emo.objective, emo.p))
+end
+function get_gradient(M, emo::EmbeddedManifoldObjective{Nothing,T}, p) where {T}
+    get_gradient!(get_embedding(M), emo.X, emo.objective, embed(M, p))
+    return riemannian_gradient(M, p, emo.X)
+end
+function get_gradient(M, emo::EmbeddedManifoldObjective{P,T}, p) where {P,T}
+    embed!(M, emo.p, p)
+    get_gradient!(get_embedding(M), emo.X, emo.objective, emo.p)
+    return riemannian_gradient(M, p, emo.X)
+end
+function get_gradient!(M, X, emo::EmbeddedManifoldObjective{Nothing,Nothing}, p)
+    riemannian_gradient!(
+        M, X, p, get_gradient(get_embedding(M), emo.objective, embed(M, p))
+    )
+    return X
+end
+function get_gradient!(M, X, emo::EmbeddedManifoldObjective{P,Nothing}, p) where {P}
+    embed!(M, emo.p, p)
+    riemannian_gradient!(M, X, p, get_gradient(get_embedding(M), emo, emo.p))
+    return X
+end
+function get_gradient!(M, X, emo::EmbeddedManifoldObjective{Nothing,T}, p) where {T}
+    get_gradient!(get_embedding(M), emo.X, emo, embed(M, p))
+    riemannian_gradient!(M, X, p, emo.X)
+    return X
+end
+function get_gradient!(M, X, emo::EmbeddedManifoldObjective{P,T}, p) where {P,T}
+    embed!(M, emo.p, p)
+    get_gradient!(get_embedding(M), emo.X, emo, emo.p)
+    riemannian_gradient!(M, X, p, emo.X)
+    return X
 end
 
 """

src/solvers/solver.jl

@@ -64,12 +64,13 @@
 optional arguments provide necessary details on the decorators.
 A specific one is used to activate certain decorators.
 
-* `cache` – (`missing`) specify a cache. Currenlty `:Simple` is supported and `:LRU` if you
+* `cache`     – (`missing`) specify a cache. Currenlty `:Simple` is supported and `:LRU` if you
   load `LRUCache.jl`. For this case a tuple specifying what to cache and how many can be provided,
   i.e. `(:LRU, [:Cost, :Gradient], 10)`, where the number specifies the size of each cache.
   and 10 is the default if one omits the last tuple entry
-* `count` – (`missing`) specify calls to the objective to be called, see [`ManifoldCountObjective`](@ref) for the full list
-
+* `count`     – (`missing`) specify calls to the objective to be called, see [`ManifoldCountObjective`](@ref) for the full list
+* `objective` - (`missing`) specify that an objective is `:Euclidean`, which is equivalent to specifying the default embedding
+                the equivalent statement here is `:Embedded`, which fits better in naming when `M` is an []`
 other keywords are ignored.
 
 # See also
@@ -83,12 +84,25 @@
         Missing,Symbol,Tuple{Symbol,<:AbstractArray},Tuple{Symbol,<:AbstractArray,P}
     }=missing,
     count::Union{Missing,AbstractVector{<:Symbol}}=missing,
+    objective::Union{Missing,Symbol}=missing,
+    _p=ismissing(objective) ? missing : rand(M),
+    embedded_p=ismissing(objective) ? missing : embed(M, _p),
+    embedded_X=ismissing(objective) ? missing : embed(M, _p, rand(M; vector_at=p)),
     return_objective=false,
     kwargs...,
 ) where {O<:AbstractManifoldObjective,P}
-    # Order: First count _then_ cache, so that cache is the outer decorator and
-    # we only count _after_ cache misses
-    deco_o = ismissing(count) ? o : objective_count_factory(M, o, count)
+    # Order:
+    # 1) wrap embedding,
+    # 2) _then_ count
+    # 3) _then_ cache,
+    # count should not be affected by 1) but cache should be on manifold not embedding
+    # => we only count _after_ cache misses
+    # and always last wrapper: ReturnObjective.
+    deco_o = o
+    if !ismissing(objective) && objective ∈ [:Embedding, :Euclidan]
+        deco_o = EmbeddedManifoldObjective(o, embedded_p, embedded_X)
+    end
+    deco_o = ismissing(count) ? o : objective_count_factory(M, deco_o, count)
     deco_o = ismissing(cache) ? deco_o : objective_cache_factory(M, deco_o, cache)
     deco_o = return_objective ? ReturnManifoldObjective(deco_o) : deco_o
     return deco_o