Allow the stochastic cost to be a vector of costs. (#259)

* Allow the stochastic cost to be a vector of costs.
* Let's remove Julia precomp cache for now.
* We only need IJulia in the tutorials/quarto env.
* Fix admonition.
* add elementwise access for vectprial costs.
* restrict ManifoldDiff Compat for now.
* bump version.

.github/workflows/documenter.yml

@@ -19,8 +19,6 @@ jobs:
           version: 1.9
       - name: Julia Cache
         uses: julia-actions/cache@v1
-        with:
-          cache-compiled: "true"
       - name: Cache Quarto
         id: cache-quarto
         uses: actions/cache@v3

Project.toml

@@ -1,7 +1,7 @@
 name = "Manopt"
 uuid = "0fc0a36d-df90-57f3-8f93-d78a9fc72bb5"
 authors = ["Ronny Bergmann <[email protected]>"]
-version = "0.4.22"
+version = "0.4.23"
 
 [deps]
 ColorSchemes = "35d6a980-a343-548e-a6ea-1d62b119f2f4"
@@ -38,7 +38,7 @@ ColorTypes = "0.9.1, 0.10, 0.11"
 Colors = "0.11.2, 0.12"
 DataStructures = "0.17, 0.18"
 LRUCache = "1.4"
-ManifoldDiff = "0.2, 0.3"
+ManifoldDiff = "0.2, 0.3 - 0.3.2"
 Manifolds = "0.8.57"
 ManifoldsBase = "0.14.4"
 Requires = "0.5, 1"

docs/Project.toml

@@ -6,7 +6,6 @@ CondaPkg = "992eb4ea-22a4-4c89-a5bb-47a3300528ab"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 FiniteDifferences = "26cc04aa-876d-5657-8c51-4c34ba976000"
-IJulia = "7073ff75-c697-5162-941a-fcdaad2a7d2a"
 Images = "916415d5-f1e6-5110-898d-aaa5f9f070e0"
 JLD2 = "033835bb-8acc-5ee8-8aae-3f567f8a3819"
 LRUCache = "8ac3fa9e-de4c-5943-b1dc-09c6b5f20637"
@@ -21,8 +20,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 
 [compat]
 BenchmarkTools = "1.3"
-Documenter = "0.27"
-IJulia = "1"
 CondaPkg = "0.2"
+Documenter = "0.27"
 Manifolds = "0.8.27"
 ManifoldsBase = "0.13, 0.14"

docs/make.jl

@@ -11,9 +11,6 @@ if Base.active_project() != joinpath(@__DIR__, "Project.toml")
     Pkg.develop(PackageSpec(; path=(@__DIR__) * "/../"))
     Pkg.resolve()
     Pkg.instantiate()
-    if "--quarto" ∈ ARGS
-        Pkg.build("IJulia") # to activate the right kernel
-    end
 end
 
 # (b) Did someone say render? Then we render!
@@ -26,6 +23,7 @@ if "--quarto" ∈ ARGS
         Pkg.activate(tutorials_folder)
         Pkg.resolve()
         Pkg.instantiate()
+        Pkg.build("IJulia") # build IJulia to the right version.
         Pkg.activate(@__DIR__) # but return to the docs one before
         run(`quarto render $(tutorials_folder)`)
     end

src/plans/stochastic_gradient_plan.jl

@@ -24,9 +24,12 @@ respectively.
         cost=Missing(), evaluation=AllocatingEvaluation()
     )
 
-Create a Stochastic gradient problem with an optional `cost` and the gradient either as one
+Create a Stochastic gradient problem with the gradient either as one
 function (returning an array of tangent vectors) or a vector of functions (each returning one tangent vector).
 
+The optional cost can also be given as either a single function (returning a number)
+pr a vector of functions, each returning a value.
+
 # Used with
 [`stochastic_gradient_descent`](@ref)
 
@@ -39,25 +42,40 @@ struct ManifoldStochasticGradientObjective{T<:AbstractEvaluationType,TCost,TGrad
     gradient!!::TGradient
 end
 function ManifoldStochasticGradientObjective(
-    grad_f!!;
-    cost::Union{Function,Missing}=Missing(),
-    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
-)
-    return ManifoldStochasticGradientObjective{
-        typeof(evaluation),typeof(cost),typeof(grad_f!!)
-    }(
-        cost, grad_f!!
-    )
+    grad_f!!::G; cost::C=Missing(), evaluation::E=AllocatingEvaluation()
+) where {
+    E<:AbstractEvaluationType,
+    G<:Union{Function,AbstractVector{<:Function}},
+    C<:Union{Function,AbstractVector{<:Function},Missing},
+}
+    return ManifoldStochasticGradientObjective{E,C,G}(cost, grad_f!!)
 end
-function ManifoldStochasticGradientObjective(
-    grad_f!!::AbstractVector{<:Function};
-    cost::Union{Function,Missing}=Missing(),
-    evaluation::AbstractEvaluationType=AllocatingEvaluation(),
-)
-    return ManifoldStochasticGradientObjective{
-        typeof(evaluation),typeof(cost),typeof(grad_f!!)
-    }(
-        cost, grad_f!!
+
+function get_cost(
+    M::AbstractManifold, sgo::ManifoldStochasticGradientObjective{E,C}, p
+) where {E<:AbstractEvaluationType,C<:AbstractVector{<:Function}}
+    return sum(f(M, p) for f in sgo.cost)
+end
+
+@doc raw"""
+    get_cost(M::AbstractManifold, sgo::ManifoldStochasticGradientObjective, p, i)
+
+Evaluate the `i`th summand of the cost.
+
+If you use a single function for the stochastic cost, then only the index `ì=1`` is available
+to evaluate the whole cost.
+"""
+function get_cost(
+    M::AbstractManifold, sgo::ManifoldStochasticGradientObjective{E,C}, p, i
+) where {E<:AbstractEvaluationType,C<:AbstractVector{<:Function}}
+    return sgo.cost[i](M, p)
+end
+function get_cost(
+    M::AbstractManifold, sgo::ManifoldStochasticGradientObjective{E,C}, p, i
+) where {E<:AbstractEvaluationType,C<:Function}
+    (i == 1) && return sgo.cost(M, p)
+    return error(
+        "The cost is implemented as a single function and can not be accessed element wise at $i since the index is larger than 1.",
     )
 end
 

test/plans/test_stochastic_gradient_plan.jl

@@ -11,16 +11,30 @@ include("../utils/dummy_types.jl")
         X in [zeros(3), [s, 0.0, 0.0], [-s, 0.0, 0.0], [0.0, s, 0.0], [0.0, -s, 0.0]]
     ]
     f(M, y) = 1 / 2 * sum([distance(M, y, x)^2 for x in pts])
+    f2 = [(M, y) -> 1 / 2 * distance(M, y, x) for x in pts]
     sgrad_f1(M, y) = [-log(M, y, x) for x in pts]
     sgrad_f2 = [((M, y) -> -log(M, y, x)) for x in pts]
-    msgo1 = ManifoldStochasticGradientObjective(sgrad_f1; cost=f)
-    msgo2 = ManifoldStochasticGradientObjective(sgrad_f2; cost=f)
+    msgo_ff = ManifoldStochasticGradientObjective(sgrad_f1; cost=f)
+    msgo_vf = ManifoldStochasticGradientObjective(sgrad_f2; cost=f)
+    msgo_fv = ManifoldStochasticGradientObjective(sgrad_f1; cost=f2)
+    msgo_vv = ManifoldStochasticGradientObjective(sgrad_f2; cost=f2)
+    @testset "Elementwide Cost access" begin
+        for msgo in [msgo_ff, msgo_vf]
+            @test get_cost(M, msgo, p) == get_cost(M, msgo, p, 1)
+            @test_throws ErrorException get_cost(M, msgo, p, 2)
+        end
+        for msgo in [msgo_fv, msgo_vv]
+            for i in 1:length(f2)
+                @test get_cost(M, msgo, p, i) == f2[i](M, p)
+            end
+        end
+    end
     @testset "Objetive Decorator passthrough" begin
         X = zero_vector(M, p)
         Y = zero_vector(M, p)
         Xa = [zero_vector(M, p) for p in pts]
         Ya = [zero_vector(M, p) for p in pts]
-        for obj in [msgo1, msgo2]
+        for obj in [msgo_ff, msgo_vf, msgo_fv, msgo_vv]
             ddo = DummyDecoratedObjective(obj)
             @test get_gradients(M, obj, p) == get_gradients(M, ddo, p)
             get_gradients!(M, Xa, obj, p)
@@ -39,7 +53,7 @@ include("../utils/dummy_types.jl")
         Y = zero_vector(M, p)
         Xa = [zero_vector(M, p) for p in pts]
         Ya = [zero_vector(M, p) for p in pts]
-        for obj in [msgo1, msgo2]
+        for obj in [msgo_ff, msgo_vf, msgo_fv, msgo_vv]
             ddo = ManifoldCountObjective(
                 M, obj, [:StochasticGradient, :StochasticGradients]
             )
@@ -62,7 +76,7 @@ include("../utils/dummy_types.jl")
         Y = zero_vector(M, p)
         Xa = [zero_vector(M, p) for p in pts]
         Ya = [zero_vector(M, p) for p in pts]
-        for obj in [msgo1, msgo2]
+        for obj in [msgo_ff, msgo_vf, msgo_fv, msgo_vv]
             ddo = ManifoldCountObjective(
                 M, obj, [:StochasticGradient, :StochasticGradients]
             )

tutorials/ImplementASolver.qmd

@@ -13,9 +13,10 @@ Afterwards, we will show how to implement the algorithm.
 Finally, we will discuss how to make the algorithm both nice for the user as well as
 initialized in a way, that it can benefit from features already available in `Manopt.jl`.
 
-| !!! note
-|     If you have implemented your own solver, we would be very happy to have that within `Manopt.jl` as well, so maybe consider [opening a Pull Request](https://github.com/JuliaManifolds/Manopt.jl)
-
+```{=commonmark}
+!!! note
+    If you have implemented your own solver, we would be very happy to have that within `Manopt.jl` as well, so maybe consider [opening a Pull Request](https://github.com/JuliaManifolds/Manopt.jl)
+```
 
 ```{julia}
 #| echo: false