faster hat

JuliaManifolds · Oct 6, 2023 · 8bcaf87 · 8bcaf87
1 parent 44cb08c
commit 8bcaf87
Show file tree

Hide file tree

Showing 2 changed files with 57 additions and 21 deletions.
diff --git a/src/groups/special_euclidean.jl b/src/groups/special_euclidean.jl
@@ -713,6 +713,20 @@ 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)
+    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)
+    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 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
@@ -348,27 +348,49 @@ Random.seed!(10)
     end
 
     @testset "performance of selected operations" begin
-        SE3 = SpecialEuclidean(3)
-        R3 = Rotations(3)
-
-        t = SVector{3}.([1:3, 2:4, 4:6])
-        p = SMatrix{3,3}(I)
-        ω = [SA[1.0, 2.0, 3.0], SA[3.0, 2.0, 1.0], SA[1.0, 3.0, 2.0]]
-        pts = [ArrayPartition(ti, exp(R3, p, hat(R3, p, ωi))) for (ti, ωi) in zip(t, ω)]
-        Xs = [
-            ArrayPartition(SA[-1.0, 2.0, 1.0], hat(R3, p, SA[1.0, 0.5, -0.5])),
-            ArrayPartition(SA[-2.0, 1.0, 0.5], hat(R3, p, SA[-1.0, -0.5, 1.1])),
-        ]
-        exp(SE3, pts[1], Xs[1])
-        compose(SE3, pts[1], pts[2])
-        log(SE3, pts[1], pts[2])
-        vee(SE3, pts[1], Xs[2])
-        # @btime shows 0 but `@allocations` is inaccurate
-        @static if VERSION >= v"1.9-DEV"
-            @test (@allocations exp(SE3, pts[1], Xs[1])) <= 4
-            @test (@allocations compose(SE3, pts[1], pts[2])) <= 4
-            @test (@allocations log(SE3, pts[1], pts[2])) <= 12
-            @test (@allocations vee(SE3, pts[1], Xs[2])) <= 13
+        for n in [2, 3]
+            SEn = SpecialEuclidean(n)
+            Rn = Rotations(n)
+
+            p = SMatrix{n,n}(I)
+
+            if n == 2
+                t = SVector{2}.([1:2, 2:3, 3:4])
+                ω = [[1.0], [2.0], [1.0]]
+                pts = [
+                    ArrayPartition(ti, exp(Rn, p, hat(Rn, p, ωi))) for (ti, ωi) in zip(t, ω)
+                ]
+                Xs = [
+                    ArrayPartition(SA[-1.0, 2.0], hat(Rn, p, SA[1.0])),
+                    ArrayPartition(SA[1.0, -2.0], hat(Rn, p, SA[0.5])),
+                ]
+            elseif n == 3
+                t = SVector{3}.([1:3, 2:4, 4:6])
+                ω = [SA[1.0, 2.0, 3.0], SA[3.0, 2.0, 1.0], SA[1.0, 3.0, 2.0]]
+                pts = [
+                    ArrayPartition(ti, exp(Rn, p, hat(Rn, p, ωi))) for (ti, ωi) in zip(t, ω)
+                ]
+                Xs = [
+                    ArrayPartition(SA[-1.0, 2.0, 1.0], hat(Rn, p, SA[1.0, 0.5, -0.5])),
+                    ArrayPartition(SA[-2.0, 1.0, 0.5], hat(Rn, p, SA[-1.0, -0.5, 1.1])),
+                ]
+            end
+            exp(SEn, pts[1], Xs[1])
+            compose(SEn, pts[1], pts[2])
+            log(SEn, pts[1], pts[2])
+            @test isapprox(SEn, log(SEn, pts[1], pts[1]), 0 .* Xs[1])
+            @test isapprox(SEn, exp(SEn, pts[1], 0 .* Xs[1]), pts[1])
+            vee(SEn, pts[1], Xs[2])
+            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)
+            # @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])) <= 12
+                @test (@allocations vee(SEn, pts[1], Xs[2])) <= 13
+                @test (@allocations hat(SEn, pts[1], csen)) <= 13
+            end
         end
     end
 end