Extending RNG API started in symplectic manifold (#467)

* Extending RNG API started in symplectic manifold

* more work on random point generation

* more distributions

* some tangent space rand

* random tvectors for product, rotations and Stiefel

* mostly complete RNG

* fixing and missing tests

* small adjustment

* minor change in SPD RNG

* bump version

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.7.7"
+version = "0.7.8"
 
 [deps]
 Colors = "5ae59095-9a9b-59fe-a467-6f913c188581"
@@ -31,7 +31,7 @@ FiniteDiff = "2"
 Graphs = "1.4"
 HybridArrays = "0.4"
 Kronecker = "0.4, 0.5"
-ManifoldsBase = "0.12.12"
+ManifoldsBase = "0.12.13"
 Plots = "1"
 RecipesBase = "1.1"
 RecursiveArrayTools = "2"

src/distributions.jl

@@ -66,3 +66,41 @@ end
 function uniform_distribution(M::AbstractManifold)
     return uniform_distribution(M, allocate_result(M, uniform_distribution))
 end
+
+"""
+    Random.rand(M::AbstractManifold, [d::Integer]; vector_at=nothing)
+    Random.rand(rng::AbstractRNG, M::AbstractManifold, [d::Integer]; vector_at=nothing)
+
+
+Generate a random point on manifold `M` (when `vector_at` is `nothing`) or a tangent
+vector at point `vector_at` (when it is not `nothing`).
+
+Optionally a random number generator `rng` to be used can be specified. An optional integer
+`d` indicates that a vector of `d` points or tangent vectors is to be generated.
+
+!!! note
+
+    Usually a uniform distribution should be expected for compact manifolds and a
+    Gaussian-like distribution for non-compact manifolds and tangent vectors, although it is
+    not guaranteed. The distribution may change between releases. 
+
+    `rand` methods for specific manifolds may take additional keyword arguments.
+
+"""
+Random.rand(M::AbstractManifold)
+function Random.rand(M::AbstractManifold, d::Integer; kwargs...)
+    return [rand(M; kwargs...) for _ in 1:d]
+end
+function Random.rand(rng::AbstractRNG, M::AbstractManifold, d::Integer; kwargs...)
+    return [rand(rng, M; kwargs...) for _ in 1:d]
+end
+function Random.rand(M::AbstractManifold; kwargs...)
+    p = allocate_result(M, rand)
+    rand!(M, p; kwargs...)
+    return p
+end
+function Random.rand(rng::AbstractRNG, M::AbstractManifold; kwargs...)
+    p = allocate_result(M, rand)
+    rand!(rng, M, p; kwargs...)
+    return p
+end

src/manifolds/Circle.jl

@@ -425,6 +425,46 @@ project(::Circle{ℂ}, p::Number, X::Number) = X - complex_dot(p, X) * p
 project!(::Circle{ℝ}, Y, p, X) = (Y .= X)
 project!(::Circle{ℂ}, Y, p, X) = (Y .= X - complex_dot(p, X) * p)
 
+@doc raw"""
+    Random.rand(M::Circle{ℝ}; vector_at = nothing, σ::Real=1.0)
+
+If `vector_at` is `nothing`, return a random point on the [`Circle`](@ref) ``\mathbb S^1``
+by picking a random element from ``[-\pi,\pi)`` uniformly.
+
+If `vector_at` is not `nothing`, return a random tangent vector from the tangent space of
+the point `vector_at` on the [`Circle``](@ref) by using a normal distribution with
+mean 0 and standard deviation `σ`.
+"""
+function Random.rand(::Circle{ℝ}; vector_at=nothing, σ::Real=1.0)
+    if vector_at === nothing
+        return sym_rem(rand() * 2 * π)
+    else
+        return σ * randn()
+    end
+end
+function Random.rand(rng::AbstractRNG, ::Circle{ℝ}; vector_at=nothing, σ::Real=1.0)
+    if vector_at === nothing
+        return sym_rem(rand(rng) * 2 * π)
+    else
+        return σ * randn(rng)
+    end
+end
+
+function Random.rand!(M::Circle{ℝ}, pX; vector_at=nothing, σ::Real=one(eltype(pX)))
+    pX .= rand(M; vector_at, σ)
+    return pX
+end
+function Random.rand!(
+    rng::AbstractRNG,
+    M::Circle{ℝ},
+    pX;
+    vector_at=nothing,
+    σ::Real=one(eltype(pX)),
+)
+    pX .= rand(rng, M; vector_at, σ)
+    return pX
+end
+
 retract(M::Circle, p, q) = retract(M, p, q, ExponentialRetraction())
 retract(M::Circle, p, q, m::ExponentialRetraction) = exp(M, p, q)
 

src/manifolds/Euclidean.jl

@@ -495,6 +495,21 @@ project(::Euclidean{Tuple{}}, ::Number, X::Number) = X
 
 project!(::Euclidean, Y, p, X) = copyto!(Y, X)
 
+function Random.rand!(::Euclidean, pX; σ=one(eltype(pX)), vector_at=nothing)
+    pX .= randn(size(pX)) .* σ
+    return pX
+end
+function Random.rand!(
+    rng::AbstractRNG,
+    ::Euclidean,
+    pX;
+    σ=one(eltype(pX)),
+    vector_at=nothing,
+)
+    pX .= randn(rng, size(pX)) .* σ
+    return pX
+end
+
 """
     representation_size(M::Euclidean)
 

src/manifolds/FixedRankMatrices.jl

@@ -391,6 +391,114 @@ function project!(::FixedRankMatrices, Y::UMVTVector, p::SVDMPoint, A::AbstractM
     return Y
 end
 
+@doc raw"""
+    Random.rand(M::FixedRankMatrices; vector_at=nothing, kwargs...)
+
+If `vector_at` is `nothing`, return a random point on the [`FixedRankMatrices`](@ref)
+manifold. The orthogonal matrices are sampled from the [`Stiefel`(@ref) manifold
+and the singular values are sampled uniformly at random.
+
+If `vector_at` is not `nothing`, generate a random tangent vector in the tangent space of
+the point `vector_at` on the `FixedRankMatrices` manifold `M`.
+"""
+function Random.rand(
+    M::FixedRankMatrices{m,n,k};
+    vector_at=nothing,
+    kwargs...,
+) where {m,n,k}
+    if vector_at === nothing
+        p = SVDMPoint(
+            Matrix{Float64}(undef, m, k),
+            Vector{Float64}(undef, k),
+            Matrix{Float64}(undef, k, n),
+        )
+        return rand!(M, p; kwargs...)
+    else
+        X = UMVTVector(
+            Matrix{Float64}(undef, m, k),
+            Matrix{Float64}(undef, k, k),
+            Matrix{Float64}(undef, k, n),
+        )
+        return rand!(M, X; vector_at, kwargs...)
+    end
+end
+function Random.rand(
+    rng::AbstractRNG,
+    M::FixedRankMatrices{m,n,k};
+    vector_at=nothing,
+    kwargs...,
+) where {m,n,k}
+    if vector_at === nothing
+        p = SVDMPoint(
+            Matrix{Float64}(undef, m, k),
+            Vector{Float64}(undef, k),
+            Matrix{Float64}(undef, k, n),
+        )
+        return rand!(rng, M, p; kwargs...)
+    else
+        X = UMVTVector(
+            Matrix{Float64}(undef, m, k),
+            Matrix{Float64}(undef, k, k),
+            Matrix{Float64}(undef, k, n),
+        )
+        return rand!(rng, M, X; vector_at, kwargs...)
+    end
+end
+
+function Random.rand!(
+    M::FixedRankMatrices{m,n,k},
+    pX;
+    vector_at=nothing,
+    kwargs...,
+) where {m,n,k}
+    if vector_at === nothing
+        U = rand(Stiefel(m, k); kwargs...)
+        S = sort(rand(k); rev=true)
+        V = rand(Stiefel(n, k); kwargs...)
+        copyto!(M, pX, SVDMPoint(U, S, V'))
+    else
+        Up = randn(m, k)
+        Vp = randn(n, k)
+        A = randn(k, k)
+        copyto!(
+            pX,
+            UMVTVector(
+                Up - vector_at.U * vector_at.U' * Up,
+                A,
+                Vp' - Vp' * vector_at.Vt' * vector_at.Vt,
+            ),
+        )
+    end
+    return pX
+end
+function Random.rand!(
+    rng::AbstractRNG,
+    ::FixedRankMatrices{m,n,k},
+    pX;
+    vector_at=nothing,
+    kwargs...,
+) where {m,n,k}
+    if vector_at === nothing
+        U = rand(rng, Stiefel(m, k); kwargs...)
+        S = sort(rand(rng, k); rev=true)
+        V = rand(rng, Stiefel(n, k); kwargs...)
+        copyto!(pX, SVDMPoint(U, S, V'))
+    else
+        Up = randn(rng, m, k)
+        Vp = randn(rng, n, k)
+        A = randn(rng, k, k)
+        copyto!(
+            pX,
+            UMVTVector(
+                Up - vector_at.U * vector_at.U' * Up,
+                A,
+                Vp' - Vp' * vector_at.Vt' * vector_at.Vt,
+            ),
+        )
+    end
+    return pX
+end
+
 @doc raw"""
     representation_size(M::FixedRankMatrices{m,n,k})
 

src/manifolds/Grassmann.jl

@@ -321,6 +321,53 @@ project(::Grassmann, ::Any...)
 
 project!(::Grassmann, Y, p, X) = copyto!(Y, X - p * p' * X)
 
+@doc raw"""
+    rand(M::Grassmann; σ::Real=1.0, vector_at=nothing)
+
+When `vector_at` is `nothing`, return a random point `p` on [`Grassmann`](@ref) manifold `M`
+by generating a random (Gaussian) matrix with standard deviation `σ` in matching
+size, which is orthonormal.
+
+When `vector_at` is not `nothing`, return a (Gaussian) random vector from the tangent space
+``T_p\mathrm{Gr}(n,k)`` with mean zero and standard deviation `σ` by projecting a random
+Matrix onto the tangent space at `vector_at`.
+"""
+rand(M::Grassmann; σ::Real=1.0)
+
+function Random.rand!(
+    M::Grassmann{n,k,𝔽},
+    pX;
+    σ::Real=one(real(eltype(pX))),
+    vector_at=nothing,
+) where {n,k,𝔽}
+    if vector_at === nothing
+        V = σ * randn(𝔽 === ℝ ? Float64 : ComplexF64, (n, k))
+        pX .= qr(V).Q[:, 1:k]
+    else
+        Z = σ * randn(eltype(pX), size(pX))
+        project!(M, pX, vector_at, Z)
+        pX .= pX ./ norm(pX)
+    end
+    return pX
+end
+function Random.rand!(
+    rng::AbstractRNG,
+    M::Grassmann{n,k,𝔽},
+    pX;
+    σ::Real=one(real(eltype(pX))),
+    vector_at=nothing,
+) where {n,k,𝔽}
+    if vector_at === nothing
+        V = σ * randn(rng, 𝔽 === ℝ ? Float64 : ComplexF64, (n, k))
+        pX .= qr(V).Q[:, 1:k]
+    else
+        Z = σ * randn(rng, eltype(pX), size(pX))
+        project!(M, pX, vector_at, Z)
+        pX .= pX ./ norm(pX)
+    end
+    return pX
+end
+
 @doc raw"""
     representation_size(M::Grassmann{n,k})
 

src/manifolds/HyperbolicHyperboloid.jl

@@ -406,6 +406,35 @@ function project!(
     return (Y.value .= X.value .+ minkowski_metric(p.value, X.value) .* p.value)
 end
 
+function Random.rand!(M::Hyperbolic{N}, pX; vector_at=nothing, σ=one(eltype(pX))) where {N}
+    if vector_at === nothing
+        a = σ .* randn(N)
+        pX[1:(end - 1)] .= a
+        pX[end] = sqrt(1 + dot(a, a))
+    else
+        Y = σ * randn(eltype(vector_at), size(vector_at))
+        project!(M, pX, vector_at, Y)
+    end
+    return pX
+end
+function Random.rand!(
+    rng::AbstractRNG,
+    M::Hyperbolic{N},
+    pX;
+    vector_at=nothing,
+    σ=one(eltype(pX)),
+) where {N}
+    if vector_at === nothing
+        a = σ .* randn(rng, N)
+        pX[1:(end - 1)] .= a
+        pX[end] = sqrt(1 + dot(a, a))
+    else
+        Y = σ * randn(rng, eltype(vector_at), size(vector_at))
+        project!(M, pX, vector_at, Y)
+    end
+    return pX
+end
+
 function vector_transport_to!(M::Hyperbolic, Y, p, X, q, ::ParallelTransport)
     w = log(M, p, q)
     wn = norm(M, p, w)