diff --git a/docs/make.jl b/docs/make.jl index 9dbde1302d..5b1285d036 100755 --- a/docs/make.jl +++ b/docs/make.jl @@ -88,7 +88,9 @@ if run_quarto || run_on_CI # For a breaking release -> also set the tutorials folder to the most recent version Pkg.instantiate() Pkg.activate(@__DIR__) # but return to the docs one before - run(`quarto render $(tutorials_folder)`) + run(`quarto render $(tutorials_folder) --log-level debug`) + # Info to know in the event of stalling if quarto is the culprit + @info "Finished rendering Quarto" end # (d) load necessary packages for the docs diff --git a/docs/src/features/atlases.md b/docs/src/features/atlases.md index e00c3fde48..72ca43fa38 100644 --- a/docs/src/features/atlases.md +++ b/docs/src/features/atlases.md @@ -55,7 +55,18 @@ Order = [:type, :function] ```@docs Manifolds.IntegratorTerminatorNearChartBoundary Manifolds.estimate_distance_from_bvp +Manifolds.solve_chart_adjoint_differential_exp_argument +Manifolds.solve_chart_adjoint_differential_exp_basepoint +Manifolds.solve_chart_adjoint_differential_log_argument +Manifolds.solve_chart_adjoint_differential_log_basepoint +Manifolds.solve_chart_differential_exp_argument +Manifolds.solve_chart_differential_exp_basepoint +Manifolds.solve_chart_differential_log_argument +Manifolds.solve_chart_differential_log_basepoint Manifolds.solve_chart_exp_ode +Manifolds.solve_chart_jacobi_field Manifolds.solve_chart_log_bvp Manifolds.solve_chart_parallel_transport_ode +Manifolds.solve_chart_volume_density +Manifolds._jacobi_exp_matrix ``` diff --git a/ext/ManifoldsOrdinaryDiffEqDiffEqCallbacksExt.jl b/ext/ManifoldsOrdinaryDiffEqDiffEqCallbacksExt.jl index a7909884bc..7495785807 100644 --- a/ext/ManifoldsOrdinaryDiffEqDiffEqCallbacksExt.jl +++ b/ext/ManifoldsOrdinaryDiffEqDiffEqCallbacksExt.jl @@ -5,9 +5,24 @@ using Manifolds: IntegratorTerminatorNearChartBoundary, affine_connection, get_chart_index, + riemann_tensor, transition_map!, transition_map_diff! -import Manifolds: solve_chart_exp_ode, solve_chart_parallel_transport_ode +import Manifolds: + solve_chart_adjoint_differential_exp_argument, + solve_chart_adjoint_differential_exp_basepoint, + solve_chart_adjoint_differential_log_argument, + solve_chart_adjoint_differential_log_basepoint, + solve_chart_differential_exp_argument, + solve_chart_differential_exp_basepoint, + solve_chart_differential_log_argument, + solve_chart_differential_log_basepoint, + solve_chart_exp_ode, + solve_chart_jacobi_field, + solve_chart_parallel_transport_ode, + solve_chart_volume_density, + _adjoint_coordinate_map, + _jacobi_exp_matrix using ManifoldsBase using DiffEqCallbacks @@ -15,6 +30,7 @@ using OrdinaryDiffEqRosenbrock: Rodas5P using OrdinaryDiffEqVerner: AutoVern9 using SciMLBase: SciMLBase, ODEProblem, solve +using LinearAlgebra using RecursiveArrayTools: ArrayPartition """ @@ -105,17 +121,40 @@ function (scs::StitchedChartSolution{:PT})(t::Real) ), ) end +function (scs::StitchedChartSolution{:Jacobi})(t::Real) + if t < scs.sols[1][1].t[1] + throw(DomainError("Time $t is outside of the solution.")) + end + for (sol, i) in scs.sols + if t <= sol.t[end] + B = induced_basis(scs.M, scs.A, i) + solt = sol(t) + p = get_point(scs.M, scs.A, i, solt.x[1]) + X = get_vector(scs.M, p, solt.x[2], B) + Y = get_vector(scs.M, p, solt.x[3], B) + dY = get_vector(scs.M, p, solt.x[4], B) + return (p, X, Y, dY) + end + end + throw( + DomainError( + "Time $t is outside of the solution (solution time range is [$(scs.sols[1][1].t[1]), $(scs.sols[end][1].t[end])]).", + ), + ) +end function (scs::StitchedChartSolution)(t::AbstractArray) return map(scs, t) end -function chart_exp_problem(u, params, t) +function chart_exp_problem!(du, u, params, t) M, A, i = params a = u.x[1] dx = u.x[2] - ddx = -affine_connection(M, A, i, a, dx, dx) - return ArrayPartition(dx, ddx) + copyto!(du.x[1], dx) + affine_connection!(M, du.x[2], A, i, a, dx, dx) + du.x[2] .*= -1 + return nothing end """ @@ -160,8 +199,9 @@ function solve_chart_exp_ode( sols = StitchedChartSolution(M, A, :Exp, typeof(i0)) while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time params = (M, A, cur_i) - prob = - ODEProblem(chart_exp_problem, u0, (init_time, final_time), params; callback = cb) + prob = ODEProblem{true}( + chart_exp_problem!, u0, (init_time, final_time), params; callback = cb + ) sol = solve(prob, solver; kwargs...) retcode = sol.retcode init_time = sol.t[end]::typeof(final_time) @@ -180,21 +220,24 @@ function solve_chart_exp_ode( return sols end -function chart_pt_problem(u, params, t) +function chart_pt_problem!(du, u, params, t) M, A, i = params a = u.x[1] dx = u.x[2] dY = u.x[3] - ddx = -affine_connection(M, A, i, a, dx, dx) - ddY = -affine_connection(M, A, i, a, dx, dY) - return ArrayPartition(dx, ddx, ddY) + copyto!(du.x[1], dx) + affine_connection!(M, du.x[2], A, i, a, dx, dx) + du.x[2] .*= -1 + affine_connection!(M, du.x[3], A, i, a, dx, dY) + du.x[3] .*= -1 + return nothing end """ solve_chart_parallel_transport_ode( M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; - solver=AutoVern9(Rodas5P()), check_chart_switch_kwargs=NamedTuple(), final_time=1.0, + solver=AutoVern9(Rodas5P()), check_chart_switch_kwargs=NamedTuple(), final_time::Real=1.0, kwargs... ) @@ -204,7 +247,7 @@ coordinates `Xc` in the induced basis. """ function solve_chart_parallel_transport_ode( M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; - solver = AutoVern9(Rodas5P()), final_time = 1.0, check_chart_switch_kwargs = NamedTuple(), + solver = AutoVern9(Rodas5P()), final_time::Real = 1.0, check_chart_switch_kwargs = NamedTuple(), kwargs... ) u0 = ArrayPartition(copy(a), copy(Xc), copy(Yc)) @@ -218,8 +261,9 @@ function solve_chart_parallel_transport_ode( sols = StitchedChartSolution(M, A, :PT, typeof(i0)) while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time params = (M, A, cur_i) - prob = - ODEProblem(chart_pt_problem, u0, (init_time, final_time), params; callback = cb) + prob = ODEProblem{true}( + chart_pt_problem!, u0, (init_time, final_time), params; callback = cb + ) sol = solve(prob, solver; kwargs...) retcode = sol.retcode init_time = sol.t[end]::typeof(final_time) @@ -239,4 +283,395 @@ function solve_chart_parallel_transport_ode( return sols end +function chart_jacobi_field_problem!(du, u, params, t) + M, A, i = params + a = u.x[1] + dx = u.x[2] + Y = u.x[3] + dY = u.x[4] + + copyto!(du.x[1], dx) + affine_connection!(M, du.x[2], A, i, a, dx, dx) + du.x[2] .*= -1 + affine_connection!(M, du.x[4], A, i, a, dx, dY) + du.x[4] .*= -1 + # temporarily save Riemann tensor value in du.x[3], then overwrite it with the final value later + riemann_tensor!(M, du.x[3], A, i, a, Y, dx, dx) + du.x[4] .-= du.x[3] + affine_connection!(M, du.x[3], A, i, a, dx, Y) + du.x[3] .*= -1 + du.x[3] .+= dY + return nothing +end + +@doc raw""" + solve_chart_jacobi_field( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc, dYc; + solver=AutoVern9(Rodas5P()), final_time::Real = 1.0, + check_chart_switch_kwargs = NamedTuple(), kwargs... + ) + +Solve the Jacobi equation along the geodesic starting at parameters `a` in chart `i0` with +initial velocity coordinates `Xc`. `Yc` and `dYc` are, respectively, the coordinates of the +initial Jacobi field and its initial covariant derivative in the induced basis of the chart. + +In chart coordinates, this solves the system +```math +\begin{aligned} +\dot a^k &= X^k, & +\dot X^k &= -\Gamma^k_{ij}(a)X^iX^j, \\ +\dot Y^k &= dY^k - \Gamma^k_{ij}(a)X^iY^j, & +\dot{dY}^k &= -\Gamma^k_{ij}(a)X^i dY^j - R^k_{\ell ij}(a)Y^\ell X^iX^j, +\end{aligned} +``` +with initial conditions ``a(0) = a``, ``X(0)`` is equal to `Xc`, ``Y(0)`` is equal to `Yc`, and +``dY(0) = dYc``. Here, `dY` represents the coordinates of +``\nabla_{\dot\gamma}Y``. ``\Gamma^k_{ij}`` are the Christoffel symbols of the affine +connection calculated using the mutating variant of +[`affine_connection`](@ref affine_connection(::AbstractManifold, ::AbstractAtlas, ::Any, ::Any, ::Any, ::Any)) +and ``R^k_{\ell ij}`` are the components of the Riemann curvature tensor calculated using +the mutating variant of [`riemann_tensor`](@ref riemann_tensor(::AbstractManifold, ::AbstractAtlas, ::Any, ::Any, ::Any, ::Any, ::Any)). + +The returned `StitchedChartSolution{:Jacobi}` returns `(p, X, Y, dY)` at time `t`, where `p` +is the point on the geodesic, `X` its velocity, `Y` the Jacobi field, and `dY` its covariant +derivative. +""" +function solve_chart_jacobi_field( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc, dYc; + solver = AutoVern9(Rodas5P()), final_time::Real = 1.0, + check_chart_switch_kwargs = NamedTuple(), kwargs... + ) + u0 = ArrayPartition(copy(a), copy(Xc), copy(Yc), copy(dYc)) + cur_i = i0 + cb = FunctionCallingCallback( + IntegratorTerminatorNearChartBoundary(check_chart_switch_kwargs); func_start = false + ) + retcode = SciMLBase.ReturnCode.Terminated + init_time = zero(final_time) + sols = StitchedChartSolution(M, A, :Jacobi, typeof(i0)) + while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time + params = (M, A, cur_i) + prob = ODEProblem{true}( + chart_jacobi_field_problem!, u0, (init_time, final_time), params; callback = cb + ) + sol = solve(prob, solver; kwargs...) + retcode = sol.retcode + init_time = sol.t[end]::typeof(final_time) + push!(sols.sols, (sol, cur_i)) + a_final = sol.u[end].x[1]::typeof(a) + new_i = get_chart_index(M, A, cur_i, a_final) + if new_i !== cur_i + transition_map!(M, u0.x[1], A, cur_i, new_i, a_final) + transition_map_diff!( + M, u0.x[2], A, cur_i, a_final, sol.u[end].x[2]::typeof(Xc), new_i + ) + transition_map_diff!( + M, u0.x[3], A, cur_i, a_final, sol.u[end].x[3]::typeof(Yc), new_i + ) + transition_map_diff!( + M, u0.x[4], A, cur_i, a_final, sol.u[end].x[4]::typeof(dYc), new_i + ) + cur_i = new_i + end + end + return sols +end + +function _jacobi_endpoint_coordinates(solution, final_time) + sol, _ = solution.sols[end] + return sol(final_time).x[3] +end + +function _chart_jacobi_field_matrix_problem!(du, u, params, t) + M, A, i = params + a = u.x[1] + dx = u.x[2] + Y = u.x[3] + dY = u.x[4] + + copyto!(du.x[1], dx) + affine_connection!(M, du.x[2], A, i, a, dx, dx) + du.x[2] .*= -1 + for j in axes(Y, 2) + affine_connection!(M, view(du.x[4], :, j), A, i, a, dx, view(dY, :, j)) + view(du.x[4], :, j) .*= -1 + riemann_tensor!(M, view(du.x[3], :, j), A, i, a, view(Y, :, j), dx, dx) + view(du.x[4], :, j) .-= view(du.x[3], :, j) + affine_connection!(M, view(du.x[3], :, j), A, i, a, dx, view(Y, :, j)) + view(du.x[3], :, j) .*= -1 + view(du.x[3], :, j) .+= view(dY, :, j) + end + return nothing +end + +function _transition_map_diff_matrix!(M::AbstractManifold, C_out, A::AbstractAtlas, i_from, a, C_in, i_to) + for j in axes(C_in, 2) + transition_map_diff!(M, view(C_out, :, j), A, i_from, a, view(C_in, :, j), i_to) + end + return C_out +end + +@doc raw""" + _jacobi_exp_matrix(M::AbstractManifold, a, Xc, A::AbstractAtlas, i0; kwargs...) + +Solve the chart-coordinate geodesic and a matrix-valued Jacobi equation to compute the +coordinate matrix of the differential of the exponential map with respect to either its +argument (if `wrt` is set to `:argument`) or its basepoint (if `wrt` is set to `:basepoint`). + +The geodesic coordinates satisfy +```math +\dot a^k = X^k, +\qquad +\dot X^k = -\Gamma^k_{ij}(a)X^iX^j. +``` +For the matrices ``Y`` and ``dY``, whose columns are Jacobi fields and their covariant +derivatives, respectively, the system is +```math +\begin{aligned} +\dot Y^k{}_r &= dY^k{}_r - \Gamma^k_{ij}(a)X^iY^j{}_r, \\ +\dot{dY}^k{}_r &= -\Gamma^k_{ij}(a)X^i dY^j{}_r + - R^k_{\ell ij}(a)Y^\ell{}_rX^iX^j. +\end{aligned} +``` +The initial conditions are ``a(0) = a``, ``X(0)`` is set to `Xc`. +If the keyword argument `wrt` is set to `:basepoint`, then ``Y(0)`` is set to `0`, and +``dY(0) = I``. If the keyword argument `wrt` is set to `:argument`, then ``Y(0) = I``, and ``dY(0) = 0``. +Thus, the returned matrix ``Y(1)`` represents +``D_X\exp_p(X)`` in the chart-induced bases if `wrt` is set to `:argument` and +``D_p\exp_p(X)`` if `wrt` is set to `:basepoint`. The function also returns the final chart +index and the final point coordinates. +""" +function _jacobi_exp_matrix( + M::AbstractManifold, + a, + Xc, + A::AbstractAtlas, + i0; + solver = AutoVern9(Rodas5P()), + final_time::Real = 1.0, + check_chart_switch_kwargs = NamedTuple(), + wrt::Symbol, + kwargs..., + ) + n = length(Xc) + if wrt === :argument + u0 = ArrayPartition(copy(a), copy(Xc), zeros(eltype(Xc), n, n), Matrix{eltype(Xc)}(I, n, n)) + elseif wrt === :basepoint + u0 = ArrayPartition(copy(a), copy(Xc), Matrix{eltype(Xc)}(I, n, n), zeros(eltype(Xc), n, n)) + else + throw(ArgumentError("`wrt` must be either `:basepoint` or `:argument`.")) + end + cur_i = i0 + cb = FunctionCallingCallback( + IntegratorTerminatorNearChartBoundary(check_chart_switch_kwargs); + func_start = false, + ) + retcode = SciMLBase.ReturnCode.Terminated + init_time = zero(final_time) + while retcode === SciMLBase.ReturnCode.Terminated && init_time < final_time + params = (M, A, cur_i) + prob = ODEProblem{true}( + _chart_jacobi_field_matrix_problem!, u0, (init_time, final_time), params; callback = cb + ) + sol = solve(prob, solver; kwargs...) + retcode = sol.retcode + init_time = sol.t[end]::typeof(final_time) + a_final = sol.u[end].x[1]::typeof(a) + new_i = get_chart_index(M, A, cur_i, a_final) + if new_i !== cur_i + transition_map!(M, u0.x[1], A, cur_i, new_i, a_final) + transition_map_diff!(M, u0.x[2], A, cur_i, a_final, sol.u[end].x[2]::typeof(Xc), new_i) + _transition_map_diff_matrix!(M, u0.x[3], A, cur_i, a_final, sol.u[end].x[3], new_i) + _transition_map_diff_matrix!(M, u0.x[4], A, cur_i, a_final, sol.u[end].x[4], new_i) + cur_i = new_i + elseif retcode !== SciMLBase.ReturnCode.Terminated + return sol.u[end].x[3], cur_i, a_final + end + end + return u0.x[3], cur_i, u0.x[1] +end + +@doc raw""" + _adjoint_coordinate_map(M::AbstractManifold, A::AbstractAtlas, + i_from, a_from, L, i_to, a_to, Yc) + +Apply the Riemannian adjoint of a linear map represented in chart-induced bases. +If `L` represents a map from the tangent space at the point with coordinates `a_from` in +chart `i_from` to the tangent space at the point with coordinates `a_to` in chart `i_to`, +this function returns the coordinates of its adjoint applied to `Yc`. The adjoint is computed +using the local metric matrices as + +```math +L^* = G_{\mathrm{from}}^{-1}L^\mathsf{T}G_{\mathrm{to}}. +``` +""" +function _adjoint_coordinate_map(M::AbstractManifold, A::AbstractAtlas, i_from, a_from, L, i_to, a_to, Yc) + G_from = local_metric(M, A, i_from, a_from) + G_to = local_metric(M, A, i_to, a_to) + return G_from \ (transpose(L) * G_to * Yc) +end + +@doc raw""" + solve_chart_volume_density( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0; kwargs... + ) + +Compute the volume density of the exponential map in chart coordinates. The coordinates `a` +and `Xc` are represented in the induced basis of chart `i0` from atlas `A`. +""" +function solve_chart_volume_density( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0; kwargs... + ) + E, final_i, a_final = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + return abs(det(E)) * sqrt( + det_local_metric(M, A, final_i, a_final) / det_local_metric(M, A, i0, a) + ) +end + +@doc raw""" + solve_chart_differential_exp_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Solve the Jacobi equation for ``D_p\exp_p(X)[Y]``. The coordinate vectors +`Xc` and `Yc` are represented in the chart-induced basis at `p`. + +The ODE is solved by `solve_chart_jacobi_field` with `dYc` set to `0`. +""" +function solve_chart_differential_exp_basepoint( + M::AbstractManifold, + a, + Xc, + A::AbstractAtlas, + i0, + Yc; + kwargs..., + ) + return solve_chart_jacobi_field(M, a, Xc, A, i0, Yc, zero(Yc); kwargs...) +end + +@doc raw""" + solve_chart_differential_exp_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Solve the Jacobi equation for ``D_X\exp_p(X)[Y]``. The coordinate vectors +`Xc` and `Yc` are represented in the chart-induced basis at `p`. +""" +function solve_chart_differential_exp_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + return solve_chart_jacobi_field(M, a, Xc, A, i0, zero(Yc), Yc; kwargs...) +end + +@doc raw""" + solve_chart_differential_log_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Solve the Jacobi equation for ``D_p\log_p(q)[Y]``, where +``q = \exp_p(X)``. The coordinate vector `Yc` is represented in the +chart-induced basis at `p`; the differential is the covariant derivative in +`solution(0)[4]`. +""" +function solve_chart_differential_log_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + baseline = solve_chart_jacobi_field(M, a, Xc, A, i0, Yc, zero(Yc); kwargs...) + E, _, _ = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + final_time = get(kwargs, :final_time, 1.0) + dYc = -E \ _jacobi_endpoint_coordinates(baseline, final_time) + return solve_chart_jacobi_field(M, a, Xc, A, i0, Yc, dYc; kwargs...) +end + +@doc raw""" + solve_chart_differential_log_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Solve the Jacobi equation for ``D_q\log_p(q)[Y]``, where +``q = \exp_p(X)``. The coordinate vector `Yc` is represented in the +chart-induced basis at `q`; the differential is the covariant derivative in +`solution(0)[4]`. +""" +function solve_chart_differential_log_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + E, _, _ = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + dYc = E \ Yc + return solve_chart_jacobi_field(M, a, Xc, A, i0, zero(Yc), dYc; kwargs...) +end + +@doc raw""" + solve_chart_adjoint_differential_exp_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Compute the chart coordinates at the base point of the adjoint of +``D_p\exp_p(X)`` applied to `Yc`. Here `Yc` contains coordinates in the induced basis +of the final chart reached by the geodesic. `p` is the point with coordinates `a` in chart +`i0`. `X` is the tangent vector with coordinates `Xc` in the induced basis of chart `i0` +at `p`. +""" +function solve_chart_adjoint_differential_exp_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + B, final_i, a_final = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :basepoint, kwargs...) + return _adjoint_coordinate_map(M, A, i0, a, B, final_i, a_final, Yc) +end + +@doc raw""" + solve_chart_adjoint_differential_exp_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Compute the chart coordinates at the base point of the adjoint of +``D_X\exp_p(X)`` applied to `Yc`. Here `Yc` contains coordinates in the induced basis +of the final chart reached by the geodesic. `p` is the point with coordinates `a` in chart +`i0`. `X` is the tangent vector with coordinates `Xc` in the induced basis of chart `i0` +at `p`. +""" +function solve_chart_adjoint_differential_exp_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + E, final_i, a_final = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + return _adjoint_coordinate_map(M, A, i0, a, E, final_i, a_final, Yc) +end + +@doc raw""" + solve_chart_adjoint_differential_log_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Compute the chart coordinates at the base point of the adjoint of +``D_p\log_p(q)``, where ``q = \exp_p(X)``. Both the input and output use the induced +basis of the initial chart. `p` is the point with coordinates `a` in chart `i0`. +`X` is the tangent vector with coordinates `Xc` in the induced basis of chart `i0` at `p`. +""" +function solve_chart_adjoint_differential_log_basepoint( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + E, _, _ = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + B, _, _ = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :basepoint, kwargs...) + return _adjoint_coordinate_map(M, A, i0, a, -(E \ B), i0, a, Yc) +end + +@doc raw""" + solve_chart_adjoint_differential_log_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + +Compute the chart coordinates at `q = exp_p(X)` of the adjoint of +``D_q\log_p(q)`` applied to `Yc`. The input uses the induced basis of the initial chart; +the output uses that of the final chart reached by the geodesic. +`p` is the point with coordinates `a` in chart `i0`. `X` is the tangent vector with +coordinates `Xc` in the induced basis of chart `i0` at `p`. +""" +function solve_chart_adjoint_differential_log_argument( + M::AbstractManifold, a, Xc, A::AbstractAtlas, i0, Yc; kwargs... + ) + E, final_i, a_final = _jacobi_exp_matrix(M, a, Xc, A, i0; wrt = :argument, kwargs...) + return _adjoint_coordinate_map(M, A, final_i, a_final, E \ I, i0, a, Yc) +end + end diff --git a/src/Manifolds.jl b/src/Manifolds.jl index 85258986d5..3164185f2f 100644 --- a/src/Manifolds.jl +++ b/src/Manifolds.jl @@ -572,6 +572,18 @@ function estimate_distance_from_bvp end function solve_chart_exp_ode end function solve_chart_parallel_transport_ode end +function solve_chart_jacobi_field end +function solve_chart_volume_density end +function solve_chart_differential_exp_basepoint end +function solve_chart_differential_exp_argument end +function solve_chart_differential_log_basepoint end +function solve_chart_differential_log_argument end +function solve_chart_adjoint_differential_exp_basepoint end +function solve_chart_adjoint_differential_exp_argument end +function solve_chart_adjoint_differential_log_basepoint end +function solve_chart_adjoint_differential_log_argument end +function _jacobi_exp_matrix end +function _adjoint_coordinate_map end # TODO: Remove once the new interface is done function find_eps end @@ -723,7 +735,7 @@ export AbstractMetric, RiemannianMetric, StiefelSubmersionMetric, WarpedMetric -export AbstractAtlas, RetractionAtlas +export AbstractAtlas, GrassmannAtlas, RetractionAtlas # Vector transport types export AbstractVectorTransportMethod, ParallelTransport, ProjectionTransport # Retraction types diff --git a/src/atlases.jl b/src/atlases.jl index eb130ccd7a..b3154d4694 100644 --- a/src/atlases.jl +++ b/src/atlases.jl @@ -60,7 +60,7 @@ RetractionAtlas() = RetractionAtlas(ExponentialRetraction(), LogarithmicInverseR Calculate the affine connection on manifold `M` at point with parameters `a` in chart `i` of [`AbstractAtlas`](@ref) `A` of vectors with coefficients `Xc` and `Yc` in induced basis. """ -function affine_connection(M::AbstractManifold, A, i, a, Xc, Yc) +function affine_connection(M::AbstractManifold, A::AbstractAtlas, i, a, Xc, Yc) Zc = similar(Xc, Base.promote_type(eltype(Xc), eltype(Yc), eltype(a))) return affine_connection!(M, Zc, A, i, a, Xc, Yc) end @@ -1127,7 +1127,7 @@ This function returns the vector `W (in induced-chart coordinates) given by # See also - `riemann_tensor(M, A, i, a)` which returns the full 4-tensor -- []`affine_connection`](@ref) used to obtain connection coefficients (see [`levi_civita_affine_connection!`](@ref) for a generic implementation) +- [`affine_connection`](@ref) used to obtain connection coefficients (see [`levi_civita_affine_connection!`](@ref) for a generic implementation) """ function riemann_tensor( M::AbstractManifold, A::AbstractAtlas, i, a, Xc, Yc, Zc; diff --git a/src/manifolds/Grassmann.jl b/src/manifolds/Grassmann.jl index 347465bddc..d29325fc23 100644 --- a/src/manifolds/Grassmann.jl +++ b/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 diff --git a/src/manifolds/GrassmannStiefel.jl b/src/manifolds/GrassmannStiefel.jl index f58da6bb9a..6faea39279 100644 --- a/src/manifolds/GrassmannStiefel.jl +++ b/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 diff --git a/src/manifolds/Sphere.jl b/src/manifolds/Sphere.jl index 314171fee6..f605615dbf 100644 --- a/src/manifolds/Sphere.jl +++ b/src/manifolds/Sphere.jl @@ -656,11 +656,19 @@ end StereographicAtlas() The stereographic atlas of ``S^n`` with two charts: one with the singular -point (-1, 0, ..., 0) (called `:north`) and one with the singular -point (1, 0, ..., 0) (called `:south`). +point (-1, 0, ..., 0) (called `:south`) and one with the singular +point (1, 0, ..., 0) (called `:north`). """ struct StereographicAtlas <: AbstractAtlas{ℝ} end +@doc raw""" + get_chart_index(M::Sphere, A::StereographicAtlas, p) + get_chart_index(M::Sphere, A::StereographicAtlas, i, a) + +Return the preferred chart index for a point `p` on `M`, or for coordinates `a` in +chart `i`. The `:south` chart is selected for points with a negative first coordinate; +the equator belongs to the `:north` chart. +""" function get_chart_index(::Sphere{ℝ}, ::StereographicAtlas, p) if p[1] < 0 return :south @@ -668,7 +676,27 @@ function get_chart_index(::Sphere{ℝ}, ::StereographicAtlas, p) return :north end end +function get_chart_index(::Sphere{ℝ}, ::StereographicAtlas, i::Symbol, a) + anorm2 = dot(a, a) + if i === :north + return anorm2 > 1 ? :south : :north + else + return anorm2 < 1 ? :south : :north + end +end + +@doc raw""" + get_parameters!(M::Sphere, x, A::StereographicAtlas, i, p) +Store in `x` the stereographic coordinates of the point `p` in chart `i` of `A`. +For ``p = (p_1,p_{2:n+1})``, the coordinate maps are + +````math +\varphi_{\mathrm{north}}(p) = \frac{p_{2:n+1}}{1+p_1}, +\qquad +\varphi_{\mathrm{south}}(p) = \frac{p_{2:n+1}}{1-p_1}. +```` +""" function get_parameters!(::Sphere{ℝ}, x, ::StereographicAtlas, i::Symbol, p) if i === :north return x .= p[2:end] ./ (1 + p[1]) @@ -677,6 +705,20 @@ function get_parameters!(::Sphere{ℝ}, x, ::StereographicAtlas, i::Symbol, p) end end +@doc raw""" + get_point!(M::Sphere, p, A::StereographicAtlas, i, x) + +Store in `p` the point on `M` represented by stereographic coordinates `x` in chart +`i` of `A`. For ``r^2 = \lVert x \rVert^2``, the inverse coordinate maps are + +````math +\varphi_{\mathrm{north}}^{-1}(x) = +\left(\frac{1-r^2}{1+r^2}, \frac{2x}{1+r^2}\right), +\qquad +\varphi_{\mathrm{south}}^{-1}(x) = +\left(\frac{r^2-1}{1+r^2}, \frac{2x}{1+r^2}\right). +```` +""" function get_point!(::Sphere{ℝ}, p, ::StereographicAtlas, i::Symbol, x) xnorm2 = dot(x, x) if i === :north @@ -688,6 +730,12 @@ function get_point!(::Sphere{ℝ}, p, ::StereographicAtlas, i::Symbol, x) return p end +""" + get_coordinates_induced_basis!(M::Sphere, Y, p, X, B::InducedBasis{<:Any, <:Any, <:StereographicAtlas}) + +Store in `Y` the stereographic coordinate representation of the tangent vector `X` +at `p` with respect to the induced basis `B`. +""" function get_coordinates_induced_basis!( M::Sphere{ℝ}, Y, @@ -708,6 +756,12 @@ function get_coordinates_induced_basis!( return Y end +""" + get_vector_induced_basis!(M::Sphere, Y, p, X, B::InducedBasis{<:Any, <:Any, <:StereographicAtlas}) + +Store in `Y` the tangent vector at `p` represented by stereographic coordinates `X` +with respect to the induced basis `B`. +""" function get_vector_induced_basis!( M::Sphere{ℝ}, Y, @@ -739,11 +793,50 @@ function get_vector_induced_basis!( return Y end + +@doc raw""" + local_metric(M::Sphere{ℝ}, A::StereographicAtlas, i, a) + +Return the local representation of the spherical metric in the stereographic atlas +at coordinates ``a``. The formula reads + +````math +g_a = \frac{4}{(1 + \lVert a \rVert^2)^2} I. +```` +""" +function local_metric(M::Sphere{ℝ}, A::StereographicAtlas, i, a) + return (4 / (1 + dot(a, a))^2) * I +end +function det_local_metric(M::Sphere{ℝ}, ::StereographicAtlas, i, a) + return (4 / (1 + dot(a, a))^2)^manifold_dimension(M) +end +# The InducedBasis variant of `local_metric` is deprecated, kept for compatibility with older versions of Manifolds.jl function local_metric( M::Sphere{ℝ}, p, B::InducedBasis{ℝ, TangentSpaceType, StereographicAtlas, Symbol}, ) a = get_parameters(M, B.A, B.i, p) - return (4 / (1 + dot(a, a))^2) * I + return local_metric(M, B.A, B.i, a) +end + +@doc raw""" + affine_connection!(M::Sphere, Zc, A::StereographicAtlas, i, a, Xc, Yc) + +Store in `Zc` the covariant derivative of the coordinate vector field `Yc` in the +coordinate direction `Xc` at stereographic coordinates `a` in chart `i` of `A`. Its +coordinate expression is + +````math +\nabla_{X_c}Y_c = -\frac{2}{1 + \lVert a \rVert^2} +\left(X_c\langle a,Y_c\rangle + Y_c\langle a,X_c\rangle +- a\langle X_c,Y_c\rangle\right). +```` +""" +function affine_connection!( + M::Sphere{ℝ}, Zc, A::StereographicAtlas, i, a, Xc, Yc + ) + factor = -2 / (1 + dot(a, a)) + Zc .= factor .* (Xc .* dot(a, Yc) .+ Yc .* dot(a, Xc) .- a .* dot(Xc, Yc)) + return Zc end diff --git a/test/manifolds-old/embedded_torus.jl b/test/manifolds-old/embedded_torus.jl index b10e16ad0a..b34c0451f1 100644 --- a/test/manifolds-old/embedded_torus.jl +++ b/test/manifolds-old/embedded_torus.jl @@ -100,12 +100,12 @@ using LinearAlgebra @test isapprox( pexp_3[1], [2.701765894057119, 2.668437820810143, -1.8341712552932237]; - atol = 1.0e-5, + atol = 2.0e-5, ) @test isapprox( pexp_3[2], [-0.41778834843865575, 2.935021992911625, 0.7673987137187901]; - atol = 1.0e-5, + atol = 2.0e-5, ) @test isapprox( pexp_3[3], diff --git a/test/manifolds-old/grassmann.jl b/test/manifolds-old/grassmann.jl index 48e1cde1f7..5443a0e93b 100644 --- a/test/manifolds-old/grassmann.jl +++ b/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 diff --git a/test/manifolds/test_sphere.jl b/test/manifolds/test_sphere.jl index 227751fb4d..fd3b597fcb 100644 --- a/test/manifolds/test_sphere.jl +++ b/test/manifolds/test_sphere.jl @@ -221,6 +221,18 @@ using ManifoldDiff @testset "StereographicAtlas" begin M = Sphere(2) A = Manifolds.StereographicAtlas() + @testset "chart index from coordinates" begin + for (i, a, expected) in [ + (:north, [0.5, 0.0], :north), + (:north, [1.0, 0.0], :north), + (:north, [2.0, 0.0], :south), + (:south, [0.5, 0.0], :south), + (:south, [1.0, 0.0], :north), + (:south, [2.0, 0.0], :north), + ] + @test Manifolds.get_chart_index(M, A, i, a) === expected + end + end p = [1 / sqrt(3), 1 / sqrt(3), 1 / sqrt(3)] for k in [1, -1] p *= k diff --git a/test/runtests.jl b/test/runtests.jl index ca6867810d..792796c2de 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -24,6 +24,7 @@ end include_test("test_statistics.jl") end if TEST_SET ∈ ["all", "integration"] + include_test("test_atlases.jl") include_test("approx_inverse_retraction.jl") # manifolds requiring ODE solvers include_test("manifolds-old/embedded_torus.jl") diff --git a/test/test_atlases.jl b/test/test_atlases.jl new file mode 100644 index 0000000000..e9e0c53f39 --- /dev/null +++ b/test/test_atlases.jl @@ -0,0 +1,134 @@ +using Manifolds, Test +using ManifoldDiff +using DiffEqCallbacks, OrdinaryDiffEq, RecursiveArrayTools + +using LinearAlgebra + +@testset "Atlases" begin + M = Sphere(2) + A = Manifolds.StereographicAtlas() + a = [0.2, -0.3] + i = :north + p = get_point(M, A, i, a) + B = induced_basis(M, A, i) + for (Xc, Yc) in [([0.8, 0.9], [-0.8, 0.25]), ([0.3, 0.2], [-0.2, 0.25])] + # tests with & without chart switching + X = get_vector(M, p, Xc, B) + Y = get_vector(M, p, Yc, B) + + @test Manifolds.solve_chart_volume_density(M, a, Xc, A, i) ≈ + volume_density(M, p, X) atol = 1.0e-8 + + solution = Manifolds.solve_chart_differential_exp_basepoint( + M, a, Xc, A, i, Yc; final_time = 1.0 + ) + p_final, _, Y_final, _ = solution(1.0) + + expected = zero_vector(M, p_final) + ManifoldDiff.jacobi_field!( + M, + expected, + p, + exp(M, p, X), + 1.0, + Y, + ManifoldDiff.βdifferential_exp_basepoint, + ) + @test isapprox(p_final, exp(M, p, X); atol = 1.0e-6) + @test isapprox(Y_final, expected; atol = 1.0e-5) + + solution = Manifolds.solve_chart_differential_exp_argument( + M, a, Xc, A, i, Yc; final_time = 1.0 + ) + p_final, _, Y_final, _ = solution(1.0) + + expected = zero_vector(M, p_final) + ManifoldDiff.jacobi_field!( + M, + expected, + p, + exp(M, p, X), + 1.0, + Y, + ManifoldDiff.βdifferential_exp_argument, + ) + @test isapprox(M, p_final, exp(M, p, X); atol = 1.0e-7) + @test isapprox(M, p_final, Y_final, expected; atol = 1.0e-6) + end + # tests without chart switching + Xc, Yc = [0.3, 0.2], [-0.2, 0.25] + X = get_vector(M, p, Xc, B) + Y = get_vector(M, p, Yc, B) + + solution = Manifolds.solve_chart_differential_log_basepoint( + M, a, Xc, A, i, Yc; final_time = 1.0 + ) + _, _, _, dY_initial = solution(0.0) + + expected = zero_vector(M, p) + ManifoldDiff.jacobi_field!( + M, + expected, + p, + exp(M, p, X), + 0.0, + Y, + ManifoldDiff.βdifferential_log_basepoint, + ) + @test isapprox(M, p, dY_initial, expected; atol = 1.0e-7) + + q = exp(M, p, X) + Bq = induced_basis(M, A, i) + Yq = get_vector(M, q, Yc, Bq) + solution = Manifolds.solve_chart_differential_log_argument( + M, a, Xc, A, i, Yc; final_time = 1.0 + ) + _, _, _, dY_initial = solution(0.0) + + @test_throws DomainError solution(-1.0) + @test_throws DomainError solution(2.0) + + expected = zero_vector(M, p) + ManifoldDiff.jacobi_field!( + M, + expected, + q, + p, + 1.0, + Yq, + ManifoldDiff.βdifferential_log_argument, + ) + @test isapprox(M, p, dY_initial, expected; atol = 1.0e-8) + + + Yq = get_vector(M, q, Yc, Bq) + Zc = Manifolds.solve_chart_adjoint_differential_exp_basepoint(M, a, Xc, A, i, Yc) + Z = get_vector(M, p, Zc, B) + @test isapprox(M, p, Z, ManifoldDiff.adjoint_differential_exp_basepoint(M, p, X, Yq); atol = 1.0e-8) + + + Zc = Manifolds.solve_chart_adjoint_differential_exp_argument(M, a, Xc, A, i, Yc) + Z = get_vector(M, p, Zc, B) + @test isapprox(M, p, Z, ManifoldDiff.adjoint_differential_exp_argument(M, p, X, Yq); atol = 1.0e-8) + + + Zc = Manifolds.solve_chart_adjoint_differential_log_basepoint(M, a, Xc, A, i, Yc) + Z = get_vector(M, p, Zc, B) + @test isapprox(M, p, Z, ManifoldDiff.adjoint_differential_log_basepoint(M, p, q, Y); atol = 1.0e-7) + + Zc = Manifolds.solve_chart_adjoint_differential_log_argument(M, a, Xc, A, i, Yc) + Z = get_vector(M, q, Zc, Bq) + @test isapprox(M, q, Z, ManifoldDiff.adjoint_differential_log_argument(M, p, q, Y); atol = 1.0e-8) + + (m_jebm, i_jebm, a_jebm) = Manifolds._jacobi_exp_matrix(M, a, Xc, A, i; wrt = :basepoint, final_time = 0.0) + @test m_jebm ≈ I + @test i_jebm == i + @test a_jebm ≈ a + + (m_jeam, i_jeam, a_jeam) = Manifolds._jacobi_exp_matrix(M, a, Xc, A, i; wrt = :argument, final_time = 0.0) + @test norm(m_jeam) < 1.0e-16 + @test i_jeam == i + @test a_jeam ≈ a + + @test_throws ArgumentError Manifolds._jacobi_exp_matrix(M, a, Xc, A, i; wrt = :teapot, final_time = -1.0) +end diff --git a/tutorials/Project.toml b/tutorials/Project.toml index 4c955709b6..41c2801711 100644 --- a/tutorials/Project.toml +++ b/tutorials/Project.toml @@ -1,5 +1,5 @@ [deps] -BoundaryValueDiffEq = "764a87c0-6b3e-53db-9096-fe964310641d" +BoundaryValueDiffEqMIRK = "1a22d4ce-7765-49ea-b6f2-13c8438986a6" CSV = "336ed68f-0bac-5ca0-87d4-7b16caf5d00b" CairoMakie = "13f3f980-e62b-5c42-98c6-ff1f3baf88f0" DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0" @@ -20,7 +20,7 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182" path = ".." [compat] -BoundaryValueDiffEq = "4, 5" +BoundaryValueDiffEqMIRK = "1.17" CSV = "0.10" CairoMakie = "0.15" DataFrames = "1" diff --git a/tutorials/_quarto.yml b/tutorials/_quarto.yml index ddbbbe835b..d06accff6c 100644 --- a/tutorials/_quarto.yml +++ b/tutorials/_quarto.yml @@ -3,6 +3,7 @@ project: output-dir: ../docs/src/tutorials render: - "*.qmd" + - "!working-in-charts.qmd" crossref: fig-prefix: Figure diff --git a/tutorials/working-in-charts.qmd b/tutorials/working-in-charts.qmd index 46aaa080b4..8223d31cd4 100644 --- a/tutorials/working-in-charts.qmd +++ b/tutorials/working-in-charts.qmd @@ -25,7 +25,7 @@ In this tutorial we focus on chart-based computation. ```{julia} #| output: false -using Manifolds, RecursiveArrayTools, OrdinaryDiffEq, DiffEqCallbacks, BoundaryValueDiffEq +using Manifolds, RecursiveArrayTools, OrdinaryDiffEq, DiffEqCallbacks, BoundaryValueDiffEqMIRK ``` The manifold we consider is the `M` is the torus in form of the [`EmbeddedTorus`](https://juliamanifolds.github.io/Manifolds.jl/latest/manifolds/torus.html#Manifolds.EmbeddedTorus), that is the representation defined as a surface of revolution of a circle of radius 2 around a circle of radius 3. @@ -47,8 +47,8 @@ The torus will be colored according to its Gaussian curvature stored in `gcs`. W In the documentation this tutorial represents a static situation (without interactivity). `Makie.jl` rendering is turned off. ```{julia} -# using GLMakie, Makie -# GLMakie.activate!() +using GLMakie, Makie +GLMakie.activate!() """ torus_figure() @@ -69,7 +69,6 @@ function torus_figure() Y1, Z1; shading=true, - ambient=Vec3f(0.65, 0.65, 0.65), backlight=1.0f0, color=gcs, colormap=Reverse(:RdBu), @@ -81,6 +80,27 @@ function torus_figure() Colorbar(fig[1, 2], pltobj, height=Relative(0.5), label="Gaussian curvature") return ax, fig end + +function jacobi_figure(geodesic, vector_fields; colors) + ax, fig = torus_figure() + times = 0.0:0.05:1.0 + curve = geodesic.(times) + points = Point3f.(first.(curve)) + lines!(ax, points; linewidth=4.0, color=:green, label="geodesic") + indices = 1:4:length(times) + for (vectors, color) in zip(vector_fields, colors) + arrows3d!( + ax, + Point3f.(first.(vectors[indices])), + Point3f.(last.(vectors[indices])); + shaftradius=0.04, + tiplength=0.1, + tipradius=0.1, + color, + ) + end + return fig +end ``` @@ -143,15 +163,15 @@ We also parametrise the start point and direction. φy = -0.1 geo = solve_for([θₚ, φₚ], [θₓ, φₓ], [θy, φy], t_end)(0.0:dt:t_end); -# geo_ps = [Point3f(s[1]) for s in geo] -# pt_indices = 1:div(length(geo), 10):length(geo) -# geo_ps_pt = [Point3f(s[1]) for s in geo[pt_indices]] -# geo_Ys = [Point3f(s[3]) for s in geo[pt_indices]] - -# ax1, fig1 = torus_figure() -# arrows!(ax1, geo_ps_pt, geo_Ys, linewidth=0.05, color=:blue) -# lines!(geo_ps; linewidth=4.0, color=:green) -# fig1 +geo_ps = [Point3f(s[1]) for s in geo] +pt_indices = 1:div(length(geo), 10):length(geo) +geo_ps_pt = [Point3f(s[1]) for s in geo[pt_indices]] +geo_Ys = [Point3f(s[3]) for s in geo[pt_indices]] + +ax1, fig1 = torus_figure() +arrows3d!(ax1, geo_ps_pt, geo_Ys, linewidth=0.05, color=:blue) +lines!(geo_ps; linewidth=4.0, color=:green) +fig1 ``` ![fig-pt](working-in-charts/working-in-charts-transport.png) @@ -167,8 +187,10 @@ geo = solve_for([θₚ, φₚ], [θₓ, φₓ], [θy, φy], t_end)(0.0:dt:t_end) bvp_i = (0, 0) bvp_a1 = [θ₁, φ₁] bvp_a2 = [θ₂, φ₂] -bvp_sol = Manifolds.solve_chart_log_bvp(M, bvp_a1, bvp_a2, A, bvp_i); -# geo_r = [Point3f(get_point(M, A, bvp_i, p[1:2])) for p in bvp_sol(0.0:0.05:1.0)] +# bvp_sol = Manifolds.solve_chart_log_bvp(M, bvp_a1, bvp_a2, A, bvp_i); +# pts_interp = collect(bvp_sol(0.0:0.05:1.0; idxs=1:2)) +# geo_r = [Point3f(get_point(M, A, bvp_i, p[1:2])) for p in eachcol(pts_interp)] + # ax2, fig2 = torus_figure() # lines!(geo_r; linewidth=4.0, color=:green) @@ -177,4 +199,214 @@ bvp_sol = Manifolds.solve_chart_log_bvp(M, bvp_a1, bvp_a2, A, bvp_i); ![fig-geodesic](working-in-charts/working-in-charts-geodesic.png) +## Jacobi fields and differentials in charts + +Jacobi fields describe how a geodesic changes when its initial point or initial velocity is +perturbed. The chart-based solvers below work with the coordinate vector `Xc` of the geodesic +velocity and coordinate vectors `Yc` in the induced basis of the current chart. They switch +charts automatically when necessary, just as the parallel-transport solver does. + +For this example, choose a short geodesic which remains in its initial chart. This makes the +coordinates at its initial and final points directly comparable. The point `p0` determines a +chart index, and `[0.0, 0.0]` are the coordinates of `p0` in that chart. + +```{julia} +p0 = [Manifolds._torus_param(M, θₚ, φₚ)...] +jacobi_i = Manifolds.get_chart_index(M, A, p0) +jacobi_a = [0.0, 0.0] +jacobi_Xc = [2.3, -1.4] +jacobi_Yc = [-0.25, 0.6] + +``` + +### Solving a Jacobi field + +[`solve_chart_jacobi_field`](@ref) solves the geodesic and a Jacobi field simultaneously. +Its last two arguments are the initial coordinates of the field, `Yc`, and of its covariant +derivative, `dYc`. Evaluating the returned solution yields `(p, X, Y, dY)`: the geodesic point +and velocity, followed by the Jacobi field and its covariant derivative, all in the embedding. + +```{julia} +jacobi_solution = Manifolds.solve_chart_jacobi_field( + M, + jacobi_a, + jacobi_Xc, + A, + jacobi_i, + jacobi_Yc, + [0.05, -0.1]; + final_time=1.0, +) +p1, X1, J1, ∇J1 = jacobi_solution(1.0) +``` + +Plot the geodesic and the Jacobi field. The blue arrows show the field $J$ and the orange +arrows show its covariant derivative $\nabla_{\dot\gamma}J$. + +```{julia} +jacobi_times = 0.0:0.05:1.0 +jacobi_values = jacobi_solution(jacobi_times) +jacobi_vectors = [(value[1], value[3]) for value in jacobi_values] +jacobi_derivatives = [(value[1], value[4]) for value in jacobi_values] +jacobi_figure( + t -> jacobi_solution(t)[1:2], + [jacobi_vectors, jacobi_derivatives]; + colors=[:dodgerblue, :darkorange], +) +``` + +The two exponential-map differential helpers select the appropriate initial conditions for +common variations. `solve_chart_differential_exp_basepoint` uses $Y(0)=Y_c$ and +$\nabla_{\dot\gamma}Y(0)=0$, while `solve_chart_differential_exp_argument` uses +$Y(0)=0$ and $\nabla_{\dot\gamma}Y(0)=Y_c$. Consequently, the Jacobi field at time `1.0` +is respectively $D_p\exp_p(X)[Y]$ or $D_X\exp_p(X)[Y]$. + +```{julia} +dexp_basepoint_solution = Manifolds.solve_chart_differential_exp_basepoint( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc; final_time=1.0 +) +dexp_argument_solution = Manifolds.solve_chart_differential_exp_argument( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc; final_time=1.0 +) + +_, _, differential_exp_basepoint, _ = dexp_basepoint_solution(1.0) +_, _, differential_exp_argument, _ = dexp_argument_solution(1.0) +``` + +The following plot compares the two endpoint fields. The blue field corresponds to moving the +base point and the orange field corresponds to changing the initial velocity. + +```{julia} +dexp_basepoint_values = dexp_basepoint_solution(jacobi_times) +dexp_argument_values = dexp_argument_solution(jacobi_times) +dexp_basepoint_vectors = [(value[1], value[3]) for value in dexp_basepoint_values] +dexp_argument_vectors = [(value[1], value[3]) for value in dexp_argument_values] +jacobi_figure( + t -> dexp_basepoint_solution(t)[1:2], + [dexp_basepoint_vectors, dexp_argument_vectors]; + colors=[:dodgerblue, :darkorange], +) +``` + +The determinant of $D_X\exp_p(X)$, corrected by the local metric determinants at the start +and end of the geodesic, is the volume density of the exponential map. It is available without +having to construct all Jacobi fields separately. + +```{julia} +chart_volume_density = Manifolds.solve_chart_volume_density( + M, jacobi_a, jacobi_Xc, A, jacobi_i +) +``` + +The volume density is a scalar, so represent it by coloring the geodesic according to the +volume density computed for scaled initial velocities $tX$. + +```{julia} +volume_densities = [ + Manifolds.solve_chart_volume_density(M, jacobi_a, t .* jacobi_Xc, A, jacobi_i) + for t in jacobi_times +] +volume_geodesic = Manifolds.solve_chart_exp_ode(M, jacobi_a, jacobi_Xc, A, jacobi_i) +volume_points = Point3f.(first.(volume_geodesic(jacobi_times))) +ax_volume, fig_volume = torus_figure() +lines!( + ax_volume, + volume_points; + color=volume_densities, + colormap=:viridis, + colorrange=extrema(volume_densities), + linewidth=6.0, +) +Colorbar(fig_volume[1, 3], limits=extrema(volume_densities), colormap=:viridis, label="volume density") +fig_volume +``` + +### Differentials and adjoints of the logarithmic map + +Let $q=\exp_p(X)$. The logarithmic-map helpers return Jacobi-field solutions whose covariant +derivative at time `0.0` is the requested differential. For the base-point differential, +`jacobi_Yc` represents a vector at $p$; for the argument differential it represents a vector +at $q$ in the final chart. Since the short geodesic above does not switch charts, both use +the same induced basis here. + +```{julia} +dlog_basepoint_solution = Manifolds.solve_chart_differential_log_basepoint( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc; final_time=1.0 +) +dlog_argument_solution = Manifolds.solve_chart_differential_log_argument( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc; final_time=1.0 +) + +_, _, _, differential_log_basepoint = dlog_basepoint_solution(0.0) +_, _, _, differential_log_argument = dlog_argument_solution(0.0) +``` + +The logarithmic-map fields are fixed at the endpoint and evaluated backwards along the same +geodesic. The arrows at the base point show their values, which are the two requested +differentials of `log`. + +```{julia} +dlog_basepoint_values = dlog_basepoint_solution(jacobi_times) +dlog_argument_values = dlog_argument_solution(jacobi_times) +dlog_basepoint_vectors = [(value[1], value[3]) for value in dlog_basepoint_values] +dlog_argument_vectors = [(value[1], value[3]) for value in dlog_argument_values] +jacobi_figure( + t -> dlog_basepoint_solution(t)[1:2], + [dlog_basepoint_vectors, dlog_argument_vectors]; + colors=[:dodgerblue, :darkorange], +) +``` + +The adjoint helpers return coordinates directly, rather than a time-dependent Jacobi-field +solution. The adjoints of the exponential-map differentials map a vector at $q$ back to $p$. +The adjoints of the logarithmic-map differentials have the converse domain and codomain. + +```{julia} +adjoint_differential_exp_basepoint = Manifolds.solve_chart_adjoint_differential_exp_basepoint( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc +) +adjoint_differential_exp_argument = Manifolds.solve_chart_adjoint_differential_exp_argument( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc +) +adjoint_differential_log_basepoint = Manifolds.solve_chart_adjoint_differential_log_basepoint( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc +) +adjoint_differential_log_argument = Manifolds.solve_chart_adjoint_differential_log_argument( + M, jacobi_a, jacobi_Xc, A, jacobi_i, jacobi_Yc +) +``` + +Visualize the adjoints at their respective source and target points. Blue arrows are the input +coordinates interpreted at their source point; orange arrows are the returned coordinates at +the target point. This makes the direction reversal of each adjoint explicit. + +```{julia} +q = dexp_basepoint_solution(1.0)[1] +Bp = induced_basis(M, A, jacobi_i) +Bq = induced_basis(M, A, jacobi_i) +adjoint_pairs = [ + (q, get_vector(M, q, jacobi_Yc, Bq), p0, get_vector(M, p0, adjoint_differential_exp_basepoint, Bp)), + (q, get_vector(M, q, jacobi_Yc, Bq), p0, get_vector(M, p0, adjoint_differential_exp_argument, Bp)), + (p0, get_vector(M, p0, jacobi_Yc, Bp), p0, get_vector(M, p0, adjoint_differential_log_basepoint, Bp)), + (p0, get_vector(M, p0, jacobi_Yc, Bp), q, get_vector(M, q, adjoint_differential_log_argument, Bq)), +] +adjoint_labels = ["adjoint d exp base point", "adjoint d exp argument", "adjoint d log base point", "adjoint d log argument"] + +fig_adjoint = Figure(size=(1400, 1000), fontsize=16) +for (k, (source, input, target, output)) in enumerate(adjoint_pairs) + row, column = div(k - 1, 2) + 1, mod(k - 1, 2) + 1 + grid = GridLayout(fig_adjoint[row, column]) + Label(grid[1, 1], adjoint_labels[k]) + ax = LScene(grid[2, 1], show_axis=true) + arrows3d!(ax, [Point3f(source)], [Point3f(input)]; color=:dodgerblue, linewidth=0.05) + arrows3d!(ax, [Point3f(target)], [Point3f(output)]; color=:darkorange, linewidth=0.05) +end +fig_adjoint +``` + +For a geodesic that crosses a chart boundary, the returned `StitchedChartSolution` still +evaluates to embedding-space points and tangent vectors. When supplying or interpreting raw +coordinates to the logarithmic or adjoint routines, use the induced basis of the initial chart +at $p$ or of the final chart at $q$, as specified by the corresponding function. + An interactive Pluto version of this tutorial is available in file [`tutorials/working-in-charts.jl`](https://github.com/JuliaManifolds/Manifolds.jl/blob/616855447996fb1ee7dfb2a779341b962a1323f8/tutorials/working-in-charts.jl).