[WIP] Product manifold

.gitignore

@@ -1,2 +1,4 @@
 docs/build
 Manifest.toml
+benchmark/tune.json
+benchmark/results.md

Project.toml

@@ -14,6 +14,7 @@ OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
+UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
 [extras]
 DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
@@ -22,4 +23,4 @@ ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
-test = ["Test"]
+test = ["Test", "DoubleFloats", "ForwardDiff", "ReverseDiff"]

benchmark/benchmarks.jl

@@ -0,0 +1,131 @@
+using Manifolds
+using BenchmarkTools
+using StaticArrays
+using LinearAlgebra
+using Random
+
+# Define a parent BenchmarkGroup to contain our suite
+SUITE = BenchmarkGroup()
+
+Random.seed!(12334)
+
+function add_manifold(M::Manifold, pts, name;
+    test_tangent_vector_broadcasting = true,
+    retraction_methods = [],
+    inverse_retraction_methods = [],
+    point_distributions = [],
+    tvector_distributions = [])
+
+    SUITE["manifolds"][name] = BenchmarkGroup()
+    tv = log(M, pts[1], pts[2])
+    tv1 = log(M, pts[1], pts[2])
+    tv2 = log(M, pts[1], pts[3])
+    p = similar(pts[1])
+    SUITE["manifolds"][name]["similar"] = @benchmarkable similar($(pts[1]))
+    SUITE["manifolds"][name]["log"] = @benchmarkable log($M, $(pts[1]), $(pts[2]))
+    SUITE["manifolds"][name]["log!"] = @benchmarkable log!($M, $tv, $(pts[1]), $(pts[2]))
+    for iretr ∈ inverse_retraction_methods
+        SUITE["manifolds"][name]["inverse_retract: "*string(iretr)] = @benchmarkable inverse_retract($M, $(pts[1]), $(pts[2]), $iretr)
+        SUITE["manifolds"][name]["inverse_retract!: "*string(iretr)] = @benchmarkable inverse_retract!($M, $tv, $(pts[1]), $(pts[2]), $iretr)
+    end
+    SUITE["manifolds"][name]["exp"] = @benchmarkable exp($M, $(pts[1]), $tv1)
+    SUITE["manifolds"][name]["exp!"] = @benchmarkable exp!($M, $p, $(pts[1]), $tv1)
+    for retr ∈ retraction_methods
+        SUITE["manifolds"][name]["retract: "*string(retr)] = @benchmarkable retract($M, $(pts[1]), $tv1, $retr)
+        SUITE["manifolds"][name]["retract!: "*string(retr)] = @benchmarkable retract!($M, $p, $(pts[1]), $tv1, $retr)
+    end
+    SUITE["manifolds"][name]["norm"] = @benchmarkable norm($M, $(pts[1]), $tv1)
+    SUITE["manifolds"][name]["inner"] = @benchmarkable inner($M, $(pts[1]), $tv1, $tv2)
+    SUITE["manifolds"][name]["distance"] = @benchmarkable distance($M, $(pts[1]), $(pts[2]))
+    SUITE["manifolds"][name]["isapprox (pt)"] = @benchmarkable isapprox($M, $(pts[1]), $(pts[2]))
+    SUITE["manifolds"][name]["isapprox (tv)"] = @benchmarkable isapprox($M, $(pts[1]), $tv1, $tv2)
+    SUITE["manifolds"][name]["2 * tv1 + 3 * tv2"] = @benchmarkable 2 * $tv1 + 3 * $tv2
+    SUITE["manifolds"][name]["tv = 2 * tv1 + 3 * tv2"] = @benchmarkable $tv = 2 * $tv1 + 3 * $tv2
+    if test_tangent_vector_broadcasting
+        SUITE["manifolds"][name]["tv = 2 .* tv1 .+ 3 .* tv2"] = @benchmarkable $tv = 2 .* $tv1 .+ 3 .* $tv2
+        SUITE["manifolds"][name]["tv .= 2 .* tv1 .+ 3 .* tv2"] = @benchmarkable $tv .= 2 .* $tv1 .+ 3 .* $tv2
+    end
+    for pd ∈ point_distributions
+        distr_name = string(pd)
+        distr_name = distr_name[1:min(length(distr_name), 50)]
+        SUITE["manifolds"][name]["point distribution "*distr_name] = @benchmarkable rand($pd)
+    end
+    for tvd ∈ point_distributions
+        distr_name = string(tvd)
+        distr_name = distr_name[1:min(length(distr_name), 50)]
+        SUITE["manifolds"][name]["tangent vector distribution "*distr_name] = @benchmarkable rand($tvd)
+    end
+end
+
+# General manifold benchmarks
+function add_manifold_benchmarks()
+
+    SUITE["manifolds"] = BenchmarkGroup()
+
+    s2 = Manifolds.Sphere(2)
+    array_s2 = ArrayManifold(s2)
+    r2 = Manifolds.Euclidean(2)
+
+    pts_r2 = [Size(2)([1.0, 1.0]),
+              Size(2)([-2.0, 3.0]),
+              Size(2)([3.0, -2.0])]
+
+    add_manifold(r2,
+                 pts_r2,
+                 "Euclidean{2} -- SizedArray")
+
+    add_manifold(r2,
+                 [MVector{2,Float64}([1.0, 1.0]),
+                  MVector{2,Float64}([-2.0, 3.0]),
+                  MVector{2,Float64}([3.0, -2.0])],
+                  "Euclidean{2} -- MVector")
+
+    ud_sphere = Manifolds.uniform_distribution(s2, Size(3)([1.0, 0.0, 0.0]))
+    gtd_sphere = Manifolds.normal_tvector_distribution(s2, Size(3)([1.0, 0.0, 0.0]), 1.0)
+
+    pts_s2 = [Size(3)([1.0, 0.0, 0.0]),
+              Size(3)([0.0, 1.0, 0.0]),
+              Size(3)([0.0, 0.0, 1.0])]
+
+    add_manifold(s2,
+                 pts_s2,
+                 "Sphere{2} -- SizedArray";
+                 point_distributions = [ud_sphere],
+                 tvector_distributions = [gtd_sphere])
+
+    add_manifold(array_s2,
+                 [Size(3)([1.0, 0.0, 0.0]),
+                  Size(3)([0.0, 1.0, 0.0]),
+                  Size(3)([0.0, 0.0, 1.0])],
+                 "ArrayManifold{Sphere{2}} -- SizedArray";
+                 test_tangent_vector_broadcasting = false)
+
+    retraction_methods_rot = [Manifolds.PolarRetraction(),
+                              Manifolds.QRRetraction()]
+
+    inverse_retraction_methods_rot = [Manifolds.PolarInverseRetraction(),
+                                      Manifolds.QRInverseRetraction()]
+
+    so2 = Manifolds.Rotations(2)
+    angles = (0.0, π/2, 2π/3)
+    add_manifold(so2,
+                 [Size(2, 2)([cos(ϕ) sin(ϕ); -sin(ϕ) cos(ϕ)]) for ϕ in angles],
+                 "Rotations(2) -- SizedArray",
+                 retraction_methods = retraction_methods_rot,
+                 inverse_retraction_methods = inverse_retraction_methods_rot)
+
+    m_prod = Manifolds.ProductManifold(s2, r2)
+    shape_s2r2 = Manifolds.ShapeSpecification(s2, r2)
+
+    pts_prd_base = [[1.0, 0.0, 0.0, 0.0, 0.0],
+                    [0.0, 1.0, 0.0, 1.0, 0.0],
+                    [0.0, 0.0, 1.0, 0.0, 0.1]]
+    pts_prod = map(p -> Manifolds.ProductArray(shape_s2r2, p), pts_prd_base)
+    add_manifold(m_prod, pts_prod, "ProductManifold with ProductArray")
+
+    pts_prod_mpoints = [Manifolds.ProductMPoint(p[1], p[2]) for p in zip(pts_s2, pts_r2)]
+    add_manifold(m_prod, pts_prod_mpoints, "ProductManifold with MPoint";
+        test_tangent_vector_broadcasting = false)
+end
+
+add_manifold_benchmarks()

benchmark/runbenchmarks.jl

@@ -0,0 +1,9 @@
+using PkgBenchmark
+
+#--track-allocation=all
+config = BenchmarkConfig(id = nothing,
+    juliacmd = `julia -O3`,
+    env = Dict("JULIA_NUM_THREADS" => 4))
+
+results = benchmarkpkg("Manifolds", config)
+export_markdown("benchmark/results.md", results)

docs/Project.toml

@@ -0,0 +1,5 @@
+[deps]
+Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
+
+[compat]
+Documenter = "0.22"

docs/make.jl

@@ -8,9 +8,22 @@ makedocs(
     pages = [
         "Home" => "index.md",
         "Manifolds" => [
-            "Sphere" => "manifolds/sphere.md",
-            "Rotations" => "manifolds/rotations.md",
+            "Basic manifolds" => [
+                "Euclidean" => "manifolds/euclidean.md",
+                "Rotations" => "manifolds/rotations.md",
+                "Sphere" => "manifolds/sphere.md"
+            ],
+            "Combined manifolds" => [
+                "Product manifold" => "manifolds/product.md"
+            ],
+            "Manifold decorators" => [
+                "Array manifold" => "manifolds/array.md",
+                "Metric manifold" => "manifolds/metric.md"
+            ]
         ],
-        "Distributions" => "distributions.md"
+        "Distributions" => "distributions.md",
+        "Library" => [
+            "Internals" => "lib/internals.md"
+        ]
     ]
 )

docs/src/distributions.md

@@ -7,3 +7,7 @@ Modules = [Manifolds]
 Pages = ["DistributionsBase.jl"]
 Order = [:type, :function]
 ```
+```@docs
+Manifolds.ProjectedPointDistribution
+Manifolds.ProjectedTVectorDistribution
+```

docs/src/index.md

@@ -12,17 +12,8 @@ Order = [:type, :function]
 ```
 
 ## Decorators
-### ArrayManifold
-```@autodocs
-Modules = [Manifolds]
-Pages = ["ArrayManifold.jl"]
-Order = [:type, :function]
-```
-### Lie Group
 
-### Metric
-```@autodocs
-Modules = [Manifolds]
-Pages = ["Metric.jl"]
-Order = [:type, :function]
-```
+* [Array manifold](@ref)
+* [Metric manifold](@ref)
+
+### Lie Group

docs/src/lib/internals.md

@@ -0,0 +1,9 @@
+# Internal Documentation
+
+Documentation for `Manifolds.jl`'s internal methods.
+
+```@docs
+Manifolds.SizedAbstractArray
+Manifolds.ziptuples
+Manifolds.size_to_tuple
+```

docs/src/manifolds/array.md

@@ -0,0 +1,7 @@
+# Array manifold
+
+```@autodocs
+Modules = [Manifolds]
+Pages = ["ArrayManifold.jl"]
+Order = [:type, :function]
+```

docs/src/manifolds/euclidean.md

@@ -0,0 +1,7 @@
+# Sphere
+
+```@autodocs
+Modules = [Manifolds]
+Pages = ["Euclidean.jl"]
+Order = [:type, :function]
+```

docs/src/manifolds/metric.md

@@ -0,0 +1,7 @@
+# Metric manifold
+
+```@autodocs
+Modules = [Manifolds]
+Pages = ["Metric.jl"]
+Order = [:type, :function]
+```

docs/src/manifolds/product.md

@@ -0,0 +1,9 @@
+# Product Manifold
+
+Product manifold $M = M_1 \times M_2 \times \dots M_n$ of manifolds $M_1, M_2, \dots, M_n$. Points on the product manifold can be constructed using [`Manifolds.prod_point`](@ref) with canonical projections $\Pi_i \colon M \to M_i$ for $i \in 1, 2, \dots, n$ provided by [`Manifolds.proj_product`](@ref).
+
+```@autodocs
+Modules = [Manifolds]
+Pages = ["ProductManifold.jl"]
+Order = [:type, :function]
+```

src/Euclidean.jl

@@ -19,6 +19,14 @@ struct Euclidean{T<:Tuple} <: Manifold where {T} end
 Euclidean(n::Int) = Euclidean{Tuple{n}}()
 Euclidean(m::Int, n::Int) = Euclidean{Tuple{m,n}}()
 
+function representation_size(::Euclidean{Tuple{n}}, ::Type{T}) where {n,T<:Union{MPoint, TVector, CoTVector}}
+    return (n,)
+end
+
+function representation_size(::Euclidean{Tuple{m,n}}, ::Type{T}) where {m,n,T<:Union{MPoint, TVector, CoTVector}}
+    return (m,n)
+end
+
 @generated manifold_dimension(::Euclidean{T}) where {T} = *(T.parameters...)
 
 struct EuclideanMetric <: RiemannianMetric end
@@ -37,12 +45,50 @@ det_local_metric(M::MetricManifold{<:Manifold,EuclideanMetric}, x) = one(eltype(
 
 log_local_metric_density(M::MetricManifold{<:Manifold,EuclideanMetric}, x) = zero(eltype(x))
 
-inner(::Euclidean, x, v, w) = dot(v, w)
-inner(::MetricManifold{<:Manifold,EuclideanMetric}, x, v, w) = dot(v, w)
+@inline inner(::Euclidean, x, v, w) = dot(v, w)
+@inline inner(::MetricManifold{<:Manifold,EuclideanMetric}, x, v, w) = dot(v, w)
 
+distance(::Euclidean, x, y) = norm(x-y)
 norm(::Euclidean, x, v) = norm(v)
 norm(::MetricManifold{<:Manifold,EuclideanMetric}, x, v) = norm(v)
 
 exp!(M::Euclidean, y, x, v) = (y .= x + v)
 
 log!(M::Euclidean, v, x, y) = (v .= y - x)
+
+function zero_tangent_vector!(M::Euclidean, v, x)
+    fill!(v, 0)
+    return v
+end
+
+project_point!(M::Euclidean, x) = x
+
+function project_tangent!(M::Euclidean, w, x, v)
+    w .= v
+    return w
+end
+
+"""
+    projected_distribution(M::Euclidean, d, [x])
+
+Wraps standard distribution `d` into a manifold-valued distribution. Generated
+points will be of similar type to `x`. By default, the type is not changed.
+"""
+function projected_distribution(M::Euclidean, d, x)
+    return ProjectedPointDistribution(M, d, project_point!, x)
+end
+
+function projected_distribution(M::Euclidean, d)
+    return ProjectedPointDistribution(M, d, project_point!, rand(d))
+end
+
+"""
+    normal_tvector_distribution(S::Euclidean, x, σ)
+
+Normal distribution in ambient space with standard deviation `σ`
+projected to tangent space at `x`.
+"""
+function normal_tvector_distribution(M::Euclidean{Tuple{N}}, x, σ) where N
+    d = Distributions.MvNormal(zero(x), σ)
+    return ProjectedTVectorDistribution(M, x, d, project_tangent!, x)
+end