Move ManifoldDiff extension here (#623)

* Move ManifoldDiff extension here

* fix

* bumps

* fix ManifoldDiff imports

* test for some mutating methods

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.68"
+version = "0.8.69"
 
 [deps]
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
@@ -48,7 +48,7 @@ Einsum = "0.4"
 Graphs = "1.4"
 HybridArrays = "0.4"
 Kronecker = "0.4, 0.5"
-ManifoldDiff = "0.2.1, 0.3"
+ManifoldDiff = "0.3.3"
 ManifoldsBase = "0.14.1"
 MatrixEquations = "2.2"
 OrdinaryDiffEq = "6.31"

src/Manifolds.jl

@@ -291,18 +291,26 @@ using ManifoldDiff:
     jacobian,
     _jacobian,
     _jacobian!,
-    riemannian_gradient,
-    riemannian_gradient!,
     set_default_differential_backend!
 using ManifoldDiff:
     AbstractDiffBackend,
     AbstractRiemannianDiffBackend,
+    CoprojectorOntoVector,
     ExplicitEmbeddedBackend,
+    IdentityProjector,
     NoneDiffBackend,
+    ProjectorOntoVector,
     RiemannianProjectionBackend,
     TangentDiffBackend
 
-import ManifoldDiff: riemannian_gradient, riemannian_gradient!
+import ManifoldDiff:
+    adjoint_Jacobi_field,
+    adjoint_Jacobi_field!,
+    diagonalizing_projectors,
+    jacobi_field,
+    jacobi_field!,
+    riemannian_gradient,
+    riemannian_gradient!
 
 using Markdown: @doc_str
 using MatrixEquations: lyapc, sylvc

src/groups/group.jl

@@ -392,6 +392,29 @@ function adjoint_action!(
     return Y
 end
 
+function ManifoldDiff.differential_exp_argument_lie_approx!(
+    M::AbstractManifold,
+    Z,
+    p,
+    X,
+    Y;
+    n=20,
+)
+    tmp = copy(M, p, Y)
+    a = -1.0
+    zero_vector!(M, Z, p)
+    for k in 0:n
+        a *= -1 // (k + 1)
+        Z .+= a .* tmp
+        if k < n
+            copyto!(tmp, lie_bracket(M, X, tmp))
+        end
+    end
+    q = exp(M, p, X)
+    translate_diff!(M, Z, q, Identity(M), Z)
+    return Z
+end
+
 @doc raw"""
     inv(G::AbstractDecoratorManifold, p)
 

src/manifolds/Circle.jl

@@ -16,6 +16,10 @@ struct Circle{𝔽} <: AbstractManifold{𝔽} end
 
 Circle(𝔽::AbstractNumbers=ℝ) = Circle{𝔽}()
 
+function adjoint_Jacobi_field(::Circle{ℝ}, p, q, t, X, β::Tβ) where {Tβ}
+    return X
+end
+
 @doc raw"""
     check_point(M::Circle, p)
 
@@ -92,6 +96,12 @@ Compute the inner product of two (complex) numbers with in the complex plane.
 complex_dot(a, b) = dot(map(real, a), map(real, b)) + dot(map(imag, a), map(imag, b))
 complex_dot(a::Number, b::Number) = (real(a) * real(b) + imag(a) * imag(b))
 
+function diagonalizing_projectors(M::Circle{ℝ}, p, X)
+    sbv = sign(X[])
+    proj = ProjectorOntoVector(M, p, @SVector [sbv == 0 ? one(sbv) : sbv])
+    return ((zero(number_eltype(p)), proj),)
+end
+
 @doc raw"""
     distance(M::Circle, p, q)
 
@@ -266,6 +276,10 @@ Return true. [`Circle`](@ref) is a flat manifold.
 """
 is_flat(M::Circle) = true
 
+function jacobi_field(::Circle{ℝ}, p, q, t, X, β::Tβ) where {Tβ}
+    return X
+end
+
 @doc raw"""
     log(M::Circle, p, q)
 

src/manifolds/Euclidean.jl

@@ -38,6 +38,10 @@ function active_traits(f, ::Euclidean, args...)
     )
 end
 
+function adjoint_Jacobi_field(::Euclidean{Tuple{}}, p, q, t, X, β::Tβ) where {Tβ}
+    return X
+end
+
 Base.:^(𝔽::AbstractNumbers, n) = Euclidean(n...; field=𝔽)
 
 Base.:^(M::Euclidean, n::Int) = ^(M, (n,))
@@ -93,6 +97,10 @@ function det_local_metric(
     return one(eltype(p))
 end
 
+function diagonalizing_projectors(::Euclidean, p, X)
+    return ((zero(number_eltype(p)), IdentityProjector()),)
+end
+
 """
     distance(M::Euclidean, p, q)
 
@@ -390,6 +398,10 @@ Return true. [`Euclidean`](@ref) is a flat manifold.
 """
 is_flat(M::Euclidean) = true
 
+function jacobi_field(::Euclidean{Tuple{}}, p, q, t, X, β::Tβ) where {Tβ}
+    return X
+end
+
 function local_metric(
     ::MetricManifold{𝔽,<:AbstractManifold,EuclideanMetric},
     p,

src/manifolds/Hyperbolic.jl

@@ -186,6 +186,15 @@ for (P, T) in zip(_HyperbolicPointTypes, _HyperbolicTangentTypes)
     end
 end
 
+function diagonalizing_projectors(M::Hyperbolic, p, X)
+    X_norm = norm(M, p, X)
+    X_normed = X / X_norm
+    return (
+        (zero(number_eltype(p)), ProjectorOntoVector(M, p, X_normed)),
+        (-one(number_eltype(p)), CoprojectorOntoVector(M, p, X_normed)),
+    )
+end
+
 get_embedding(::Hyperbolic{N}) where {N} = Lorentz(N + 1, MinkowskiMetric())
 
 embed(::Hyperbolic, p::AbstractArray) = p

src/manifolds/ProductManifold.jl

@@ -143,6 +143,40 @@ function active_traits(f, ::ProductManifold, args...)
     return merge_traits(IsDefaultMetric(ProductMetric()))
 end
 
+function adjoint_Jacobi_field(
+    M::ProductManifold,
+    p::ArrayPartition,
+    q::ArrayPartition,
+    t,
+    X::ArrayPartition,
+    β::Tβ,
+) where {Tβ}
+    return ArrayPartition(
+        map(
+            adjoint_Jacobi_field,
+            M.manifolds,
+            submanifold_components(M, p),
+            submanifold_components(M, q),
+            ntuple(_ -> t, length(M.manifolds)),
+            submanifold_components(M, X),
+            ntuple(_ -> β, length(M.manifolds)),
+        )...,
+    )
+end
+function adjoint_Jacobi_field!(M::ProductManifold, Y, p, q, t, X, β::Tβ) where {Tβ}
+    map(
+        adjoint_Jacobi_field!,
+        M.manifolds,
+        submanifold_components(M, Y),
+        submanifold_components(M, p),
+        submanifold_components(M, q),
+        ntuple(_ -> t, length(M.manifolds)),
+        submanifold_components(M, X),
+        ntuple(_ -> β, length(M.manifolds)),
+    )
+    return Y
+end
+
 function allocate_coordinates(M::AbstractManifold, p::ArrayPartition, T, n::Int)
     return allocate_coordinates(M, p.x[1], T, n)
 end
@@ -880,6 +914,40 @@ function Base.log(M::ProductManifold, p::ArrayPartition, q::ArrayPartition)
     )
 end
 
+function jacobi_field(
+    M::ProductManifold,
+    p::ArrayPartition,
+    q::ArrayPartition,
+    t,
+    X::ArrayPartition,
+    β::Tβ,
+) where {Tβ}
+    return ArrayPartition(
+        map(
+            jacobi_field,
+            M.manifolds,
+            submanifold_components(M, p),
+            submanifold_components(M, q),
+            ntuple(_ -> t, length(M.manifolds)),
+            submanifold_components(M, X),
+            ntuple(_ -> β, length(M.manifolds)),
+        )...,
+    )
+end
+function jacobi_field!(M::ProductManifold, Y, p, q, t, X, β::Tβ) where {Tβ}
+    map(
+        jacobi_field!,
+        M.manifolds,
+        submanifold_components(M, Y),
+        submanifold_components(M, p),
+        submanifold_components(M, q),
+        ntuple(_ -> t, length(M.manifolds)),
+        submanifold_components(M, X),
+        ntuple(_ -> β, length(M.manifolds)),
+    )
+    return Y
+end
+
 function log!(M::ProductManifold, X, p, q)
     map(
         log!,

src/manifolds/Sphere.jl

@@ -128,6 +128,15 @@ function check_vector(M::AbstractSphere, p, X; kwargs...)
     return nothing
 end
 
+function diagonalizing_projectors(M::AbstractSphere{ℝ}, p, X)
+    X_norm = norm(M, p, X)
+    X_normed = X / X_norm
+    return (
+        (zero(number_eltype(p)), ProjectorOntoVector(M, p, X_normed)),
+        (one(number_eltype(p)), CoprojectorOntoVector(M, p, X_normed)),
+    )
+end
+
 @doc raw"""
     distance(M::AbstractSphere, p, q)
 

test/groups/special_orthogonal.jl

@@ -181,4 +181,39 @@ include("group_utils.jl")
             @test !is_vector(G, gI, ones(n)) # wrong size
         end
     end
+
+    @testset "differentials" begin
+        G = SpecialOrthogonal(3)
+        p = Matrix(I, 3, 3)
+
+        ω = [[1.0, 2.0, 3.0], [3.0, 2.0, 1.0], [1.0, 3.0, 2.0]]
+        pts = [exp(G, p, hat(G, p, ωi)) for ωi in ω]
+        Xpts = [hat(G, p, [-1.0, 2.0, 0.5]), hat(G, p, [1.0, 0.0, 0.5])]
+
+        q2 = exp(G, pts[1], Xpts[2])
+        @test isapprox(
+            G,
+            q2,
+            ManifoldDiff.differential_exp_argument_lie_approx(
+                G,
+                pts[1],
+                Xpts[1],
+                Xpts[2];
+                n=0,
+            ),
+            Xpts[2],
+        )
+        diff_ref = [
+            0.0 -0.7482721017619345 -0.508151233069837
+            0.7482721017619345 0.0 -0.10783358474129323
+            0.508151233069837 0.10783358474129323 0.0
+        ]
+        @test isapprox(
+            G,
+            q2,
+            ManifoldDiff.differential_exp_argument_lie_approx(G, pts[1], Xpts[1], Xpts[2]),
+            diff_ref;
+            atol=1e-12,
+        )
+    end
 end

test/manifolds/circle.jl

@@ -95,6 +95,26 @@ using Manifolds: TFVector, CoTFVector
         # isapprox for 1-element vectors
         @test isapprox(M, [1.0], [1.0])
         @test isapprox(M, [0.0], [1.0], [1.0])
+
+        # ManifoldDiff
+        @test ManifoldDiff.adjoint_Jacobi_field(
+            M,
+            0.0,
+            1.0,
+            0.5,
+            2.0,
+            ManifoldDiff.βdifferential_shortest_geodesic_startpoint,
+        ) === 2.0
+        @test ManifoldDiff.diagonalizing_projectors(M, 0.0, 2.0) ==
+              ((0.0, ManifoldDiff.ProjectorOntoVector(M, 0.0, SA[1.0])),)
+        @test ManifoldDiff.jacobi_field(
+            M,
+            0.0,
+            1.0,
+            0.5,
+            2.0,
+            ManifoldDiff.βdifferential_shortest_geodesic_startpoint,
+        ) === 2.0
     end
     TEST_STATIC_SIZED && @testset "Real Circle and static sized arrays" begin
         X = @MArray fill(0.0)

test/manifolds/euclidean.jl

@@ -342,4 +342,27 @@ using Manifolds: induced_basis
         M = Euclidean(4)
         @test_throws DimensionMismatch distance(M, [1, 2, 3, 4], [1 2; 3 4])
     end
+
+    @testset "ManifoldDiff" begin
+        # ManifoldDiff
+        M0 = Euclidean()
+        @test ManifoldDiff.adjoint_Jacobi_field(
+            M0,
+            0.0,
+            1.0,
+            0.5,
+            2.0,
+            ManifoldDiff.βdifferential_shortest_geodesic_startpoint,
+        ) === 2.0
+        @test ManifoldDiff.diagonalizing_projectors(M0, 0.0, 2.0) ==
+              ((0.0, ManifoldDiff.IdentityProjector()),)
+        @test ManifoldDiff.jacobi_field(
+            M0,
+            0.0,
+            1.0,
+            0.5,
+            2.0,
+            ManifoldDiff.βdifferential_shortest_geodesic_startpoint,
+        ) === 2.0
+    end
 end

test/manifolds/hyperbolic.jl

@@ -328,4 +328,15 @@ include("../utils.jl")
         A = change_metric(M, MinkowskiMetric(), p, X)
         @test A == X
     end
+    @testset "ManifoldDiff" begin
+        # ManifoldDiff
+        M = Hyperbolic(2)
+        p = [1.0, 1.0, sqrt(3)]
+        X = [1.0, 2.0, sqrt(3)]
+        dpX = ManifoldDiff.diagonalizing_projectors(M, p, X)
+        @test dpX[1][1] == 0.0
+        @test dpX[1][2].X == 2 * normalize(X)
+        @test dpX[2][1] == -1.0
+        @test dpX[2][2].X == 2 * normalize(X)
+    end
 end