atlas on the Grassmann manifold

Author: mateuszbaran

src/Manifolds.jl

@@ -736,7 +736,7 @@ export AbstractMetric,
     RiemannianMetric,
     StiefelSubmersionMetric,
     WarpedMetric
-export AbstractAtlas, RetractionAtlas
+export AbstractAtlas, GrassmannAtlas, RetractionAtlas
 # Vector transport types
 export AbstractVectorTransportMethod, ParallelTransport, ProjectionTransport
 # Retraction types

src/manifolds/Grassmann.jl

@@ -82,6 +82,26 @@ function Grassmann(n::Int, k::Int, field::AbstractNumbers = ℝ; parameter::Symb
     return Grassmann{field, typeof(size)}(size)
 end
 
+@doc raw"""
+    GrassmannAtlas()
+
+The standard atlas of the real Grassmann manifold ``\mathrm{Gr}(n,k)``. Its
+``\binom{n}{k}`` charts are indexed by ordered tuples `i` of `k` row indices.
+The chart indexed by `i` contains the subspaces whose corresponding `k`-by-`k`
+minor in the rows `i` is invertible.
+
+For a Stiefel representative `p`, let `pᵢ` denote this minor and let `j` be
+the complementary row indices. The coordinate map is
+
+````math
+\varphi_i([p]) = \operatorname{vec}\left((p p_i^{-1})_j\right).
+````
+
+Its inverse inserts the reshaped coordinates into the rows `j`, inserts the
+identity into the rows `i`, and orthonormalizes the resulting matrix.
+"""
+struct GrassmannAtlas <: AbstractAtlas{ℝ} end
+
 function allocation_promotion_function(::Grassmann{ℂ}, f, args::Tuple)
     return complex
 end
@@ -239,3 +259,100 @@ Return the default vector transport method for the [`Grassmann`](@ref) manifold,
 which is `ParallelTransport``()`.
 """
 default_vector_transport_method(::Grassmann) = ParallelTransport()
+
+function _grassmann_chart_rows(M::Grassmann, i::AbstractVector)
+    n, k = get_parameter(M.size)
+    rows = collect(i)
+    if length(rows) != k || !all(row -> 1 <= row <= n, rows) || length(unique(rows)) != k
+        throw(ArgumentError("A Grassmann chart index must contain $k distinct row indices in 1:$n."))
+    end
+    return rows
+end
+
+function _grassmann_chart_complement(M::Grassmann, i::AbstractVector)
+    n, _ = get_parameter(M.size)
+    return setdiff(collect(1:n), _grassmann_chart_rows(M, i))
+end
+
+@doc raw"""
+    inner(M::Grassmann, A::GrassmannAtlas, i::AbstractVector, a, Xc, Yc)
+
+Compute the Riemannian inner product of coordinate vectors `Xc` and `Yc` at
+coordinates `a` in the chart `i` of the standard [`GrassmannAtlas`](@ref).
+Writing `Z = reshape(a, n-k, k)`, `U = reshape(Xc, n-k, k)`,
+`V = reshape(Yc, n-k, k)`, and `G = I + ZᵀZ`, the formula is
+
+````math
+g_Z(U,V) = \operatorname{tr}\left(G^{-1} U^\mathsf{T}
+\left(I - ZG^{-1}Z^\mathsf{T}\right)V\right).
+````
+"""
+function inner(M::Grassmann{ℝ}, ::GrassmannAtlas, i::AbstractVector, a, Xc, Yc)
+    _grassmann_chart_rows(M, i)
+    n, k = get_parameter(M.size)
+    coordinate_dimension = k * (n - k)
+    length(a) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension chart coordinates."))
+    length(Xc) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension vector coordinates."))
+    length(Yc) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension vector coordinates."))
+    Z = reshape(a, n - k, k)
+    U = reshape(Xc, n - k, k)
+    V = reshape(Yc, n - k, k)
+    G = I + transpose(Z) * Z
+    return dot(U / G, V - Z * (G \ (transpose(Z) * V)))
+end
+
+@doc raw"""
+    affine_connection!(M::Grassmann, Zc, A::GrassmannAtlas, i, a, Xc, Yc)
+
+Store the Levi-Civita covariant derivative of `Yc` in direction `Xc` in
+`Zc`, using the standard [`GrassmannAtlas`](@ref). Writing
+`Z = reshape(a, n-k, k)`, `U = reshape(Xc, n-k, k)`,
+`V = reshape(Yc, n-k, k)`, and `G = I + ZᵀZ`, the coordinate expression is
+
+````math
+\nabla_U V = -UG^{-1}Z^\mathsf{T}V - VG^{-1}Z^\mathsf{T}U.
+````
+"""
+function affine_connection!(
+        M::Grassmann{ℝ},
+        Zc,
+        ::GrassmannAtlas,
+        i::AbstractVector,
+        a,
+        Xc,
+        Yc,
+    )
+    _grassmann_chart_rows(M, i)
+    n, k = get_parameter(M.size)
+    coordinate_dimension = k * (n - k)
+    length(a) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension chart coordinates."))
+    length(Xc) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension vector coordinates."))
+    length(Yc) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension vector coordinates."))
+    length(Zc) == coordinate_dimension || throw(DimensionMismatch("Expected $coordinate_dimension output coordinates."))
+    Z = reshape(a, n - k, k)
+    U = reshape(Xc, n - k, k)
+    V = reshape(Yc, n - k, k)
+    G = I + transpose(Z) * Z
+    Zc .= vec(-(U * (G \ (transpose(Z) * V)) + V * (G \ (transpose(Z) * U))))
+    return Zc
+end
+
+@doc raw"""
+    det_local_metric(M::Grassmann, A::GrassmannAtlas, i::AbstractVector, a)
+
+Return the determinant of the local metric in the standard
+[`GrassmannAtlas`](@ref) at coordinates `a` in chart `i`.
+"""
+function det_local_metric(M::Grassmann{ℝ}, A::GrassmannAtlas, i::AbstractVector, a)
+    return det(local_metric(M, A, i, a))
+end
+
+@doc raw"""
+    inverse_chart_injectivity_radius(M::Grassmann, A::GrassmannAtlas, i)
+
+Return the injectivity radius of an affine chart in the standard
+[`GrassmannAtlas`](@ref), which is infinite.
+"""
+function inverse_chart_injectivity_radius(M::Grassmann{ℝ}, ::GrassmannAtlas, i::AbstractVector)
+    return Inf
+end

src/manifolds/GrassmannStiefel.jl

@@ -518,3 +518,126 @@ which is given by a zero matrix the same size as `p`.
 zero_vector(::Grassmann, ::Any...)
 
 zero_vector!(::Grassmann, X, p) = fill!(X, 0)
+
+# GrassmannAtlas
+
+function _grassmann_largest_minor_rows(p, k)
+    n = size(p, 1)
+    rows = collect(1:k)
+    best_rows = copy(rows)
+    best_minor = zero(real(eltype(p)))
+    while true
+        minor = abs(det(p[rows, :]))
+        if minor > best_minor
+            best_minor = minor
+            best_rows .= rows
+        end
+
+        position = k
+        while position >= 1 && rows[position] == n - k + position
+            position -= 1
+        end
+        position == 0 && break
+        rows[position] += 1
+        for next_position in (position + 1):k
+            rows[next_position] = rows[next_position - 1] + 1
+        end
+    end
+    iszero(best_minor) && throw(DomainError(p, "The point has no invertible $k-by-$k row minor."))
+    return best_rows
+end
+
+@doc raw"""
+    get_chart_index(M::Grassmann, A::GrassmannAtlas, p)
+    get_chart_index(M::Grassmann, A::GrassmannAtlas, i, a)
+
+Return a chart index suitable for a point `p` or coordinates `a` in chart `i`.
+For a point, the index of the row minor with the largest absolute determinant is
+returned, breaking ties lexicographically. For coordinates, the represented
+point is reconstructed and the same selection rule is applied.
+"""
+function get_chart_index(M::Grassmann{ℝ}, ::GrassmannAtlas, p)
+    _, k = get_parameter(M.size)
+    return _grassmann_largest_minor_rows(p, k)
+end
+function get_chart_index(M::Grassmann{ℝ}, A::GrassmannAtlas, i::AbstractVector, a)
+    return get_chart_index(M, A, get_point(M, A, i, a))
+end
+
+@doc raw"""
+    get_parameters!(M::Grassmann, a, A::GrassmannAtlas, i, p)
+
+Store the standard affine coordinates of `p` in `a` for the chart indexed by
+the row tuple `i`. The selected row minor of `p` must be invertible.
+"""
+function get_parameters!(M::Grassmann{ℝ}, a, ::GrassmannAtlas, i::AbstractVector, p)
+    rows = _grassmann_chart_rows(M, i)
+    complement = _grassmann_chart_complement(M, i)
+    p_rows = p[rows, :]
+    iszero(det(p_rows)) && throw(DomainError(p, "The point does not belong to chart $i."))
+    a .= vec((p / p_rows)[complement, :])
+    return a
+end
+
+@doc raw"""
+    get_point!(M::Grassmann, p, A::GrassmannAtlas, i::AbstractVector, a)
+
+Store in `p` the point represented by affine coordinates `a` in the chart of
+the standard [`GrassmannAtlas`](@ref) indexed by `i`.
+"""
+function get_point!(M::Grassmann{ℝ}, p, ::GrassmannAtlas, i::AbstractVector, a)
+    rows = _grassmann_chart_rows(M, i)
+    complement = _grassmann_chart_complement(M, i)
+    n, k = get_parameter(M.size)
+    length(a) == k * (n - k) || throw(DimensionMismatch("Expected $(k * (n - k)) chart coordinates."))
+    fill!(p, zero(eltype(p)))
+    p[rows, :] .= Matrix{eltype(p)}(I, k, k)
+    p[complement, :] .= reshape(a, n - k, k)
+    return project!(M, p, p)
+end
+
+@doc raw"""
+    get_coordinates_induced_basis!(M::Grassmann, c, p, X, B::InducedBasis{<:Any, <:Any, <:GrassmannAtlas})
+
+Store in `c` the coordinates of a tangent vector `X` at `p` with respect to
+the basis induced by the standard [`GrassmannAtlas`](@ref).
+"""
+function get_coordinates_induced_basis!(
+        M::Grassmann{ℝ},
+        c,
+        p,
+        X,
+        B::InducedBasis{ℝ, TangentSpaceType, <:GrassmannAtlas},
+    )
+    rows = _grassmann_chart_rows(M, B.i)
+    complement = _grassmann_chart_complement(M, B.i)
+    p_rows = p[rows, :]
+    normalized_p = p / p_rows
+    normalized_X = X / p_rows
+    c .= vec(normalized_X[complement, :] .- normalized_p[complement, :] * normalized_X[rows, :])
+    return c
+end
+
+@doc raw"""
+    get_vector_induced_basis!(M::Grassmann, X, p, c, B::InducedBasis{ℝ, TangentSpaceType, <:GrassmannAtlas})
+
+Store in `X` the tangent vector at `p` represented by coordinates `c` with
+respect to the basis induced by the standard [`GrassmannAtlas`](@ref).
+"""
+function get_vector_induced_basis!(
+        M::Grassmann{ℝ},
+        X,
+        p,
+        c,
+        B::InducedBasis{ℝ, TangentSpaceType, <:GrassmannAtlas},
+    )
+    rows = _grassmann_chart_rows(M, B.i)
+    complement = _grassmann_chart_complement(M, B.i)
+    n, k = get_parameter(M.size)
+    length(c) == k * (n - k) || throw(DimensionMismatch("Expected $(k * (n - k)) basis coordinates."))
+    fill!(X, zero(eltype(X)))
+    X[complement, :] .= reshape(c, n - k, k)
+    X .= X * p[rows, :]
+    project!(M, X, p, X)
+    return X
+end

test/manifolds-old/grassmann.jl

@@ -1,4 +1,6 @@
 include("../header.jl")
+using DiffEqCallbacks, OrdinaryDiffEq
+using ForwardDiff
 
 @testset "Grassmann" begin
     @testset "Real" begin
@@ -409,4 +411,65 @@ include("../header.jl")
         @test get_embedding(M, typeof(p2)) == Euclidean(3, 3; parameter = :field)
         @test get_total_space(M) == Stiefel(3, 2; parameter = :field)
     end
+
+    @testset "GrassmannAtlas" begin
+        M = Grassmann(4, 2)
+        A = GrassmannAtlas()
+        charts = ([1, 2], [1, 3], [1, 4], [2, 3], [2, 4], [3, 4])
+        a = [0.1, -0.2, 0.3, -0.4]
+        c = [0.2, -0.1, 0.4, 0.3]
+        d = [-0.3, 0.5, 0.1, -0.2]
+
+        @test Manifolds.get_chart_index(M, A, [1, 2], 3 * a) == [1, 4]
+        for i in charts
+            @testset "chart $i" begin
+                p = get_point(M, A, i, a)
+                @test is_point(M, p; error = :error)
+                @test get_parameters(M, A, i, p) ≈ a
+                @test Manifolds.get_chart_index(M, A, i, a) ==
+                    Manifolds.get_chart_index(M, A, p)
+                @test Manifolds.inverse_chart_injectivity_radius(M, A, i) == Inf
+
+                selected_i = Manifolds.get_chart_index(M, A, p)
+                @test selected_i in charts
+                selected_a = get_parameters(M, A, selected_i, p)
+                @test isapprox(M, p, get_point(M, A, selected_i, selected_a))
+
+                B = induced_basis(M, A, i)
+                X = get_vector(M, p, c, B)
+                Y = get_vector(M, p, d, B)
+                @test is_vector(M, p, X; error = :error, atol = 1.0e-14)
+                @test get_coordinates(M, p, X, B) ≈ c
+                @test Manifolds.inner(M, A, i, a, c, c) ≈ Manifolds.inner(M, p, X, X)
+                @test Manifolds.inner(M, A, i, a, c, d) ≈ Manifolds.inner(M, p, X, Y)
+                @test Manifolds.det_local_metric(M, A, i, a) > 0
+
+                # TODO: check against the embedding-based implementation of the Levi-Civita connection
+                Zc = affine_connection(M, A, i, a, c, d)
+                @test Zc ≈ Manifolds.levi_civita_affine_connection(M, A, i, a, c, d)
+                affine_connection!(M, Zc, A, i, a, c, d)
+                @test Zc ≈ Manifolds.levi_civita_affine_connection(M, A, i, a, c, d)
+            end
+        end
+
+        @test_throws DomainError get_parameters(M, A, [1, 2], [0.0 0.0; 0.0 0.0; 1.0 0.0; 0.0 1.0])
+        @test_throws ArgumentError get_point(M, A, [1, 1], a)
+
+        @testset "chart integration" begin
+            i = [1, 2]
+            Xc = [0.04, -0.03, 0.02, 0.01]
+            Yc = [-0.02, 0.03, 0.01, -0.04]
+            p = get_point(M, A, i, a)
+            B = induced_basis(M, A, i)
+            X = get_vector(M, p, Xc, B)
+            Y = get_vector(M, p, Yc, B)
+            q = exp(M, p, X)
+
+            exp_solution = Manifolds.solve_chart_exp_ode(
+                M, a, Xc, A, i; final_time = 1.0, abstol = 1.0e-10, reltol = 1.0e-10,
+            )
+            q_chart, X_chart = exp_solution(1.0)
+            @test isapprox(M, q_chart, q; atol = 1.0e-8)
+        end
+    end
 end