CholeskySpace.jl

@doc raw"""
    CholeskySpace{T} <: AbstractManifold{ℝ}

The manifold of lower triangular matrices with positive diagonal and
a metric based on the cholesky decomposition. The formulae for this manifold
are for example summarized in Table 1 of [Lin:2019](@cite).

# Constructor

    CholeskySpace(n; parameter::Symbol=:type)

Generate the manifold of $n× n$ lower triangular matrices with positive diagonal.
"""
struct CholeskySpace{T} <: AbstractManifold{ℝ}
    size::T
end

function CholeskySpace(n::Int; parameter::Symbol=:type)
    size = wrap_type_parameter(parameter, (n,))
    return CholeskySpace{typeof(size)}(size)
end

@doc raw"""
    check_point(M::CholeskySpace, p; kwargs...)

Check whether the matrix `p` lies on the [`CholeskySpace`](@ref) `M`, i.e.
it's size fits the manifold, it is a lower triangular matrix and has positive
entries on the diagonal.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_point(M::CholeskySpace, p; kwargs...)
    cks = check_size(M, p)
    cks === nothing || return cks
    if !isapprox(norm(strictlyUpperTriangular(p)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(p) - Diagonal(p)),
            "The point $(p) does not lie on $(M), since it strictly upper triangular nonzero entries",
        )
    end
    if any(diag(p) .<= 0)
        return DomainError(
            min(diag(p)...),
            "The point $(p) does not lie on $(M), since it hast nonpositive entries on the diagonal",
        )
    end
    return nothing
end

"""
    check_vector(M::CholeskySpace, p, X; kwargs... )

Check whether `v` is a tangent vector to `p` on the [`CholeskySpace`](@ref) `M`, i.e.
after [`check_point`](@ref)`(M,p)`, `X` has to have the same dimension as `p`
and a symmetric matrix.
The tolerance for the tests can be set using the `kwargs...`.
"""
function check_vector(M::CholeskySpace, p, X; kwargs...)
    if !isapprox(norm(strictlyUpperTriangular(X)), 0.0; kwargs...)
        return DomainError(
            norm(UpperTriangular(X) - Diagonal(X)),
            "The matrix $(X) is not a tangent vector at $(p) (represented as an element of the Lie algebra) since it is not lower triangular.",
        )
    end
    return nothing
end

@doc raw"""
    distance(M::CholeskySpace, p, q)

Compute the Riemannian distance on the [`CholeskySpace`](@ref) `M` between two
matrices `p`, `q` that are lower triangular with positive diagonal. The formula
reads

````math
d_{\mathcal M}(p,q) = \sqrt{\sum_{i>j} (p_{ij}-q_{ij})^2 +
\sum_{j=1}^m (\log p_{jj} - \log q_{jj})^2
}
````
"""
function distance(::CholeskySpace, p, q)
    return sqrt(
        sum((strictlyLowerTriangular(p) - strictlyLowerTriangular(q)) .^ 2) +
        sum((log.(diag(p)) - log.(diag(q))) .^ 2),
    )
end

@doc raw"""
    exp(M::CholeskySpace, p, X)

Compute the exponential map on the [`CholeskySpace`](@ref) `M` emanating from the lower
triangular matrix with positive diagonal `p` towards the lower triangular matrix `X`
The formula reads

````math
\exp_p X = ⌊ p ⌋ + ⌊ X ⌋ + \operatorname{diag}(p)
\operatorname{diag}(p)\exp\bigl( \operatorname{diag}(X)\operatorname{diag}(p)^{-1}\bigr),
````

where $⌊\cdot⌋$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
exp(::CholeskySpace, ::Any...)

function exp!(::CholeskySpace, q, p, X)
    q .= (
        strictlyLowerTriangular(p) +
        strictlyLowerTriangular(X) +
        Diagonal(diag(p)) * Diagonal(exp.(diag(X) ./ diag(p)))
    )
    return q
end

@doc raw"""
    inner(M::CholeskySpace, p, X, Y)

Compute the inner product on the [`CholeskySpace`](@ref) `M` at the
lower triangular matric with positive diagonal `p` and the two tangent vectors
`X`,`Y`, i.e they are both lower triangular matrices with arbitrary diagonal.
The formula reads

````math
g_p(X,Y) = \sum_{i>j} X_{ij}Y_{ij} + \sum_{j=1}^m X_{ii}Y_{ii}p_{ii}^{-2}
````
"""
function inner(::CholeskySpace, p, X, Y)
    return (
        sum(strictlyLowerTriangular(X) .* strictlyLowerTriangular(Y)) +
        sum(diag(X) .* diag(Y) ./ (diag(p) .^ 2))
    )
end

"""
    is_flat(::CholeskySpace)

Return false. [`CholeskySpace`](@ref) is not a flat manifold.
"""
is_flat(M::CholeskySpace) = false

@doc raw"""
    log(M::CholeskySpace, X, p, q)

Compute the logarithmic map on the [`CholeskySpace`](@ref) `M` for the geodesic emanating
from the lower triangular matrix with positive diagonal `p` towards `q`.
The formula reads

````math
\log_p q = ⌊ p ⌋ - ⌊ q ⌋ + \operatorname{diag}(p)\log\bigl(\operatorname{diag}(q)\operatorname{diag}(p)^{-1}\bigr),
````

where $⌊\cdot⌋$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
log(::Cholesky, ::Any...)

function log!(::CholeskySpace, X, p, q)
    return copyto!(
        X,
        strictlyLowerTriangular(q) - strictlyLowerTriangular(p) +
        Diagonal(diag(p) .* log.(diag(q) ./ diag(p))),
    )
end

@doc raw"""
    manifold_dimension(M::CholeskySpace)

Return the manifold dimension for the [`CholeskySpace`](@ref) `M`, i.e.

````math
    \dim(\mathcal M) = \frac{N(N+1)}{2}.
````
"""
function manifold_dimension(M::CholeskySpace)
    N = get_parameter(M.size)[1]
    return div(N * (N + 1), 2)
end

@doc raw"""
    representation_size(M::CholeskySpace)

Return the representation size for the [`CholeskySpace`](@ref)`{N}` `M`, i.e. `(N,N)`.
"""
function representation_size(M::CholeskySpace)
    N = get_parameter(M.size)[1]
    return (N, N)
end

function Base.show(io::IO, ::CholeskySpace{TypeParameter{Tuple{n}}}) where {n}
    return print(io, "CholeskySpace($(n))")
end
function Base.show(io::IO, M::CholeskySpace{Tuple{Int}})
    n = get_parameter(M.size)[1]
    return print(io, "CholeskySpace($(n); parameter=:field)")
end

# two small helpers for strictly lower and upper triangulars
strictlyLowerTriangular(p) = LowerTriangular(p) - Diagonal(diag(p))

strictlyUpperTriangular(p) = UpperTriangular(p) - Diagonal(diag(p))

@doc raw"""
    parallel_transport_to(M::CholeskySpace, p, X, q)

Parallely transport the tangent vector `X` at `p` along the geodesic to `q`
on the [`CholeskySpace`](@ref) manifold `M`. The formula reads

````math
\mathcal P_{q←p}(X) = ⌊ X ⌋
+ \operatorname{diag}(q)\operatorname{diag}(p)^{-1}\operatorname{diag}(X),
````

where $⌊\cdot⌋$ denotes the strictly lower triangular matrix,
and $\operatorname{diag}$ extracts the diagonal matrix.
"""
parallel_transport_to(::CholeskySpace, ::Any, ::Any, ::Any)

function parallel_transport_to!(::CholeskySpace, Y, p, X, q)
    return copyto!(Y, strictlyLowerTriangular(p) + Diagonal(diag(q) .* diag(X) ./ diag(p)))
end

@doc raw"""
    zero_vector(M::CholeskySpace, p)

Return the zero tangent vector on the [`CholeskySpace`](@ref) `M` at `p`.
"""
zero_vector(::CholeskySpace, ::Any...)

zero_vector!(M::CholeskySpace, X, p) = fill!(X, 0)