Improve performance of selected operations on SE(n)

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.79"
+version = "0.8.80"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"

src/groups/special_euclidean.jl

@@ -692,3 +692,27 @@ end
 function project!(M::SpecialEuclideanInGeneralLinear, Y, p, X)
     return copyto!(Y, project(M, p, X))
 end
+
+### Special methods for better performance of selected operations
+
+function exp(M::SpecialEuclidean, p::ArrayPartition, X::ArrayPartition)
+    M1, M2 = M.manifold.manifolds
+    return ArrayPartition(
+        exp(M1.manifold, p.x[1], X.x[1]),
+        exp(M2.manifold, p.x[2], X.x[2]),
+    )
+end
+function log(M::SpecialEuclidean, p::ArrayPartition, q::ArrayPartition)
+    M1, M2 = M.manifold.manifolds
+    return ArrayPartition(
+        log(M1.manifold, p.x[1], q.x[1]),
+        log(M2.manifold, p.x[2], q.x[2]),
+    )
+end
+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 compose(::SpecialEuclidean, p::ArrayPartition, q::ArrayPartition)
+    return ArrayPartition(p.x[2] * q.x[1] + p.x[1], p.x[2] * q.x[2])
+end

src/manifolds/GeneralUnitaryMatrices.jl

@@ -220,6 +220,18 @@ end
 function exp(M::GeneralUnitaryMatrices{2,ℝ}, p::SMatrix, X::SMatrix, t::Real)
     return exp(M, p, t * X)
 end
+function exp(M::GeneralUnitaryMatrices{3,ℝ}, p::SMatrix, X::SMatrix)
+    θ = norm(M, p, X) / sqrt(2)
+    if θ ≈ 0
+        a = 1 - θ^2 / 6
+        b = θ / 2
+    else
+        a = sin(θ) / θ
+        b = (1 - cos(θ)) / θ^2
+    end
+    pinvq = I + a .* X .+ b .* (X^2)
+    return p * pinvq
+end
 function exp!(M::GeneralUnitaryMatrices{2,ℝ}, q, p, X)
     @assert size(q) == (2, 2)
     θ = get_coordinates(M, p, X, DefaultOrthogonalBasis())[1]
@@ -405,6 +417,9 @@ end
 function get_vector_orthogonal(::GeneralUnitaryMatrices{2,ℝ}, p::SMatrix, Xⁱ, ::RealNumbers)
     return @SMatrix [0 -Xⁱ[]; Xⁱ[] 0]
 end
+function get_vector_orthogonal(::GeneralUnitaryMatrices{3,ℝ}, p::SMatrix, Xⁱ, ::RealNumbers)
+    return @SMatrix [0 -Xⁱ[3] Xⁱ[2]; Xⁱ[3] 0 -Xⁱ[1]; -Xⁱ[2] Xⁱ[1] 0]
+end
 
 function get_vector_orthogonal!(::GeneralUnitaryMatrices{1,ℝ}, X, p, Xⁱ, N::RealNumbers)
     return X .= 0

test/groups/special_euclidean.jl

@@ -346,4 +346,29 @@ Random.seed!(10)
             @test isapprox(G, pts[1], hat(G, pts[1], fXp.data), fXp2)
         end
     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
+        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
+        end
+    end
 end