[WIP] Power manifold

Project.toml

@@ -7,6 +7,7 @@ version = "0.1.0"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Einsum = "b7d42ee7-0b51-5a75-98ca-779d3107e4c0"
+HybridArrays = "1baab800-613f-4b0a-84e4-9cd3431bfbb9"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Markdown = "d6f4376e-aef5-505a-96c1-9c027394607a"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
@@ -15,6 +16,11 @@ SimpleTraits = "699a6c99-e7fa-54fc-8d76-47d257e15c1d"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UnsafeArrays = "c4a57d5a-5b31-53a6-b365-19f8c011fbd6"
 
+[compat]
+DoubleFloats = ">= 0.9.2"
+HybridArrays = ">= 0.3"
+julia = "1.0"
+
 [extras]
 DoubleFloats = "497a8b3b-efae-58df-a0af-a86822472b78"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -24,7 +30,3 @@ Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [targets]
 test = ["Test", "DoubleFloats", "ForwardDiff", "OrdinaryDiffEq", "ReverseDiff"]
-
-[compat]
-julia = "1.0"
-DoubleFloats = ">= 0.9.2"

benchmark/benchmarks.jl

@@ -1,6 +1,7 @@
 using Manifolds
 using BenchmarkTools
 using StaticArrays
+using HybridArrays
 using LinearAlgebra
 using Random
 
@@ -152,6 +153,69 @@ function add_manifold_benchmarks()
             test_tangent_vector_broadcasting = false)
     end
 
+    # power manifolds
+    begin
+        Ms = Sphere(2)
+        Ms1 = PowerManifold(Ms, 5)
+        Ms2 = PowerManifold(Ms, 5, 7)
+        Mr = Manifolds.Rotations(3)
+        Mr1 = PowerManifold(Mr, 5)
+        Mr2 = PowerManifold(Mr, 5, 7)
+
+        types_s1 = [Array{Float64,2},
+                    HybridArray{Tuple{3,StaticArrays.Dynamic()}, Float64, 2}]
+        types_s2 = [Array{Float64,3},
+                    HybridArray{Tuple{3,StaticArrays.Dynamic(),StaticArrays.Dynamic()}, Float64, 3}]
+
+        types_r1 = [Array{Float64,3},
+                    HybridArray{Tuple{3,3,StaticArrays.Dynamic()}, Float64, 3}]
+        types_r2 = [Array{Float64,4},
+                    HybridArray{Tuple{3,3,StaticArrays.Dynamic(),StaticArrays.Dynamic()}, Float64, 4}]
+
+        retraction_methods = [Manifolds.PowerRetraction(Manifolds.ExponentialRetraction())]
+        inverse_retraction_methods = [Manifolds.InversePowerRetraction(Manifolds.LogarithmicInverseRetraction())]
+
+        sphere_dist = Manifolds.uniform_distribution(Ms, @SVector [1.0, 0.0, 0.0])
+        power_s1_pt_dist = Manifolds.PowerPointDistribution(Ms1, sphere_dist, randn(Float64, 3, 5))
+        power_s2_pt_dist = Manifolds.PowerPointDistribution(Ms2, sphere_dist, randn(Float64, 3, 5, 7))
+        sphere_tv_dist = Manifolds.normal_tvector_distribution(Ms, (@MVector [1.0, 0.0, 0.0]), 1.0)
+        power_s1_tv_dist = Manifolds.PowerFVectorDistribution(TangentBundleFibers(Ms1), rand(power_s1_pt_dist), sphere_tv_dist)
+        power_s2_tv_dist = Manifolds.PowerFVectorDistribution(TangentBundleFibers(Ms2), rand(power_s2_pt_dist), sphere_tv_dist)
+
+        id_rot = @SMatrix [1.0 0.0 0.0; 0.0 1.0 0.0; 0.0 0.0 1.0]
+        rotations_dist = Manifolds.normal_rotation_distribution(Mr, id_rot, 1.0)
+        power_r1_pt_dist = Manifolds.PowerPointDistribution(Mr1, rotations_dist, randn(Float64, 3, 3, 5))
+        power_r2_pt_dist = Manifolds.PowerPointDistribution(Mr2, rotations_dist, randn(Float64, 3, 3, 5, 7))
+        rotations_tv_dist = Manifolds.normal_tvector_distribution(Mr, MMatrix(id_rot), 1.0)
+        power_r1_tv_dist = Manifolds.PowerFVectorDistribution(TangentBundleFibers(Mr1), rand(power_r1_pt_dist), rotations_tv_dist)
+        power_r2_tv_dist = Manifolds.PowerFVectorDistribution(TangentBundleFibers(Mr2), rand(power_r2_pt_dist), rotations_tv_dist)
+
+        trim(s::String) = s[1:min(length(s), 20)]
+
+        for T in types_s1
+            pts1 = [convert(T, rand(power_s1_pt_dist)) for _ in 1:3]
+            add_manifold(Ms1, pts1, "power manifold S²^(5,), type $(trim(string(T)))";
+                test_tangent_vector_broadcasting = true)
+        end
+        for T in types_s2
+            pts2 = [convert(T, rand(power_s2_pt_dist)) for _ in 1:3]
+            add_manifold(Ms2, pts2, "power manifold S²^(5,7), type $(trim(string(T)))";
+                test_tangent_vector_broadcasting = true)
+        end
+
+        for T in types_r1
+            pts1 = [convert(T, rand(power_r1_pt_dist)) for _ in 1:3]
+            add_manifold(Mr1, pts1, "power manifold SO(3)^(5,), type $(trim(string(T)))";
+                test_tangent_vector_broadcasting = true)
+        end
+        for T in types_r2
+            pts2 = [convert(T, rand(power_r2_pt_dist)) for _ in 1:3]
+            add_manifold(Mr2, pts2, "power manifold SO(3)^(5,7), type $(trim(string(T)))";
+                test_tangent_vector_broadcasting = true)
+        end
+
+    end
+
 end
 
 add_manifold_benchmarks()

docs/make.jl

@@ -14,6 +14,7 @@ makedocs(
                 "Sphere" => "manifolds/sphere.md"
             ],
             "Combined manifolds" => [
+                "Power manifold" => "manifolds/power.md",
                 "Product manifold" => "manifolds/product.md",
                 "Vector bundle" => "manifolds/vector_bundle.md"
             ],

docs/src/lib/public.md

@@ -8,4 +8,7 @@ submanifold_component
 Manifolds.ProductArray
 ProductRepr
 Manifolds.prod_point
+Manifolds.StaticReshaper
+Manifolds.ArrayReshaper
+Manifolds.make_reshape
 ```

docs/src/manifolds/power.md

@@ -0,0 +1,9 @@
+# Product Manifold
+
+Power manifold $P = M^{N_1 \times N_2 \times \dots \times N_n}$ of manifold $N$ with shape $N_1 \times N_2 \times \dots \times N_n$. Canonical projections may depend on multiple indices $i_1, i_2, \dots, i_n$. For example, manifold-valued time series require a single index ($n=1$) and manifold-valued images require two indices ($n=2$).
+
+```@autodocs
+Modules = [Manifolds]
+Pages = ["PowerManifold.jl"]
+Order = [:type, :function]
+```

src/Manifolds.jl

@@ -9,10 +9,13 @@ import Base: isapprox,
     getindex,
     setindex!,
     size,
+    length,
     copy,
     copyto!,
     convert,
     dataids,
+    axes,
+    promote_rule,
     +,
     -,
     *
@@ -25,6 +28,7 @@ import LinearAlgebra: dot,
     Diagonal
 using Requires
 using StaticArrays
+using HybridArrays
 import Markdown: @doc_str
 import Distributions: _rand!, support
 import Random: rand
@@ -517,7 +521,7 @@ vectors shorter than $d$ (has a left inverse).
 injectivity_radius(M::Manifold, x, ::AbstractRetractionMethod) = injectivity_radius(M, x)
 
 """
-    injectivity_radius(M::Manifold, x)
+    injectivity_radius(M::Manifold)
 
 Infimum of the injectivity radii of all manifold points.
 """
@@ -625,6 +629,7 @@ is_tangent_vector(M::Manifold, x, v; kwargs...) = true
 is_tangent_vector(M::Manifold, x::MPoint, v::TVector) = error("A validation for a $(typeof(v)) in the tangent space of a $(typeof(x)) on $(typeof(M)) not implemented.")
 
 include("utils.jl")
+include("SizedAbstractArray.jl")
 
 include("ProductRepresentations.jl")
 include("ArrayManifold.jl")
@@ -634,6 +639,7 @@ include("DistributionsBase.jl")
 include("Metric.jl")
 include("Euclidean.jl")
 include("ProductManifold.jl")
+include("PowerManifold.jl")
 include("Rotations.jl")
 include("Sphere.jl")
 include("ProjectedDistribution.jl")
@@ -662,6 +668,7 @@ export Manifold,
     Euclidean,
     Sphere,
     ProductManifold,
+    PowerManifold,
     ProductRepr,
     VectorSpaceType,
     TangentSpace,