projected_distribution.jl

"""
    ProjectedPointDistribution(M::AbstractManifold, d, proj!, p)

Generates a random point in ambient space of `M` and projects it to `M`
using function `proj!`. Generated arrays are of type `TResult`, which can be
specified by providing the `p` argument.
"""
struct ProjectedPointDistribution{TResult,TM<:AbstractManifold,TD<:Distribution,TProj} <:
       MPointDistribution{TM}
    manifold::TM
    distribution::TD
    proj!::TProj
end

function ProjectedPointDistribution(
    M::AbstractManifold,
    d::Distribution,
    proj!,
    ::TResult,
) where {TResult}
    return ProjectedPointDistribution{TResult,typeof(M),typeof(d),typeof(proj!)}(
        M,
        d,
        proj!,
    )
end

"""
    projected_distribution(M::AbstractManifold, d, [p=rand(d)])

Wrap the standard distribution `d` into a manifold-valued distribution. Generated
points will be of similar type to `p`. By default, the type is not changed.
"""
function projected_distribution(M::AbstractManifold, d, p=rand(d))
    return ProjectedPointDistribution(M, d, project!, p)
end

function Random.rand(
    rng::AbstractRNG,
    d::ProjectedPointDistribution{TResult},
) where {TResult}
    p = convert(TResult, rand(rng, d.distribution))
    return d.proj!(d.manifold, p, p)
end

function Random.rand(
    rng::AbstractRNG,
    d::ProjectedPointDistribution{TResult},
    n::Int,
) where {TResult}
    ps = [convert(TResult, rand(rng, d.distribution)) for _ in 1:n]
    map(p -> d.proj!(d.manifold, p, p), ps)
    return ps
end

function Distributions._rand!(
    rng::AbstractRNG,
    d::ProjectedPointDistribution,
    p::AbstractArray{<:Number},
)
    Distributions._rand!(rng, d.distribution, p)
    return d.proj!(d.manifold, p, p)
end

Distributions.support(d::ProjectedPointDistribution) = MPointSupport(d.manifold)

"""
    ProjectedFVectorDistribution(type::VectorSpaceFiber, p, d, project!)

Generates a random vector from ambient space of manifold `type.manifold`
at point `p` and projects it to vector space of type `type` using function
`project!`, see [`project`](@ref) for documentation.
Generated arrays are of type `TResult`.
"""
struct ProjectedFVectorDistribution{
    TResult,
    TSpace<:VectorSpaceFiber,
    TD<:Distribution,
    TProj,
} <: FVectorDistribution{TSpace}
    type::TSpace
    distribution::TD
    project!::TProj
end

function ProjectedFVectorDistribution(
    type::VectorSpaceFiber,
    d::Distribution,
    project!,
    ::TResult,
) where {TResult}
    return ProjectedFVectorDistribution{TResult,typeof(type),typeof(d),typeof(project!)}(
        type,
        d,
        project!,
    )
end

function Random.rand(
    rng::AbstractRNG,
    d::ProjectedFVectorDistribution{TResult},
) where {TResult}
    X = convert(TResult, reshape(rand(rng, d.distribution), size(d.type.point)))
    return d.project!(d.type.manifold, X, d.type.point, X)
end

function Distributions._rand!(
    rng::AbstractRNG,
    d::ProjectedFVectorDistribution,
    X::AbstractArray{<:Number},
)
    # calling _rand!(rng, d.d, v) doesn't work for all arrays types
    return copyto!(X, rand(rng, d))
end

"""
    normal_tvector_distribution(M::Euclidean, p, σ)

Normal distribution in ambient space with standard deviation `σ`
projected to tangent space at `p`.
"""
function normal_tvector_distribution(M::AbstractManifold, p, σ)
    d = Distributions.MvNormal(zero(vec(p)), σ * I)
    return ProjectedFVectorDistribution(TangentSpace(M, p), d, project!, p)
end

function Distributions.support(tvd::ProjectedFVectorDistribution)
    return FVectorSupport(tvd.type)
end