Add more SE2 and SE3 performance dispatches (#660)

* Add more SE2 and SE3 performance dispatches * Specialise only on DefaultOrthogonalBasis * do not create a view on coordinates * Fix formatting * add tests * bump version --------- Co-authored-by: Johannes Terblanche <[email protected]> Co-authored-by: Mateusz Baran <[email protected]>
JuliaManifolds · Oct 17, 2023 · b03cabc · b03cabc · mateuszbaran · Oct 18, 2023
1 parent 106a8f3
commit b03cabc
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 3 deletions.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "Manifolds"
 uuid = "1cead3c2-87b3-11e9-0ccd-23c62b72b94e"
 authors = ["Seth Axen <[email protected]>", "Mateusz Baran <[email protected]>", "Ronny Bergmann <[email protected]>", "Antoine Levitt <[email protected]>"]
-version = "0.8.80"
+version = "0.8.81"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

diff --git a/src/groups/special_euclidean.jl b/src/groups/special_euclidean.jl
@@ -713,20 +713,55 @@ function vee(M::SpecialEuclidean, p::ArrayPartition, X::ArrayPartition)
     M1, M2 = M.manifold.manifolds
     return vcat(vee(M1.manifold, p.x[1], X.x[1]), vee(M2.manifold, p.x[2], X.x[2]))
 end
-function hat(M::SpecialEuclidean{2}, p::ArrayPartition, c::SVector)
+function get_coordinates(
+    M::SpecialEuclidean,
+    p::ArrayPartition,
+    X::ArrayPartition,
+    basis::DefaultOrthogonalBasis,
+)
+    M1, M2 = M.manifold.manifolds
+    return vcat(
+        get_coordinates(M1.manifold, p.x[1], X.x[1], basis),
+        get_coordinates(M2.manifold, p.x[2], X.x[2], basis),
+    )
+end
+function hat(M::SpecialEuclidean{2}, p::ArrayPartition, c::AbstractVector)
     M1, M2 = M.manifold.manifolds
     return ArrayPartition(
         get_vector_orthogonal(M1.manifold, p.x[1], c[SOneTo(2)], ℝ),
         get_vector_orthogonal(M2.manifold, p.x[2], c[SA[3]], ℝ),
     )
 end
-function hat(M::SpecialEuclidean{3}, p::ArrayPartition, c::SVector)
+function get_vector(
+    M::SpecialEuclidean{2},
+    p::ArrayPartition,
+    c::AbstractVector,
+    basis::DefaultOrthogonalBasis,
+)
+    return ArrayPartition(
+        get_vector(M.manifold.manifolds[1].manifold, p.x[1], c[SOneTo(2)], basis),
+        get_vector(M.manifold.manifolds[2].manifold, p.x[2], c[SA[3]], basis),
+    )
+end
+
+function hat(M::SpecialEuclidean{3}, p::ArrayPartition, c::AbstractVector)
     M1, M2 = M.manifold.manifolds
     return ArrayPartition(
         get_vector_orthogonal(M1.manifold, p.x[1], c[SOneTo(3)], ℝ),
         get_vector_orthogonal(M2.manifold, p.x[2], c[SA[4, 5, 6]], ℝ),
     )
 end
+function get_vector(
+    M::SpecialEuclidean{3},
+    p::ArrayPartition,
+    c::AbstractVector,
+    basis::DefaultOrthogonalBasis,
+)
+    return ArrayPartition(
+        get_vector(M.manifold.manifolds[1].manifold, p.x[1], c[SOneTo(3)], basis),
+        get_vector(M.manifold.manifolds[2].manifold, p.x[2], c[SA[4, 5, 6]], basis),
+    )
+end
 function compose(::SpecialEuclidean, p::ArrayPartition, q::ArrayPartition)
     return ArrayPartition(p.x[2] * q.x[1] + p.x[1], p.x[2] * q.x[2])
 end
diff --git a/test/groups/special_euclidean.jl b/test/groups/special_euclidean.jl
@@ -382,15 +382,29 @@ Random.seed!(10)
             @test isapprox(SEn, log(SEn, pts[1], pts[1]), 0 .* Xs[1]; atol=1e-16)
             @test isapprox(SEn, exp(SEn, pts[1], 0 .* Xs[1]), pts[1])
             vee(SEn, pts[1], Xs[2])
+            get_coordinates(SEn, pts[1], Xs[2], DefaultOrthogonalBasis())
             csen = n == 2 ? SA[1.0, 2.0, 3.0] : SA[1.0, 0.0, 2.0, 2.0, -1.0, 1.0]
             hat(SEn, pts[1], csen)
+            get_vector(SEn, pts[1], csen, DefaultOrthogonalBasis())
             # @btime shows 0 but `@allocations` is inaccurate
             @static if VERSION >= v"1.9-DEV"
                 @test (@allocations exp(SEn, pts[1], Xs[1])) <= 4
                 @test (@allocations compose(SEn, pts[1], pts[2])) <= 4
                 @test (@allocations log(SEn, pts[1], pts[2])) <= 28
                 @test (@allocations vee(SEn, pts[1], Xs[2])) <= 13
+                @test (@allocations get_coordinates(
+                    SEn,
+                    pts[1],
+                    Xs[2],
+                    DefaultOrthogonalBasis(),
+                )) <= 13
                 @test (@allocations hat(SEn, pts[1], csen)) <= 13
+                @test (@allocations get_vector(
+                    SEn,
+                    pts[1],
+                    csen,
+                    DefaultOrthogonalBasis(),
+                )) <= 13
             end
         end
     end