Remove rand overloads for sampling of momentum (#400)

* Remove `rand` overloads for sampling of momentum

* Fix tests

* Fix format

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

Author: devmotion

research/src/riemannian_hmc.jl

@@ -117,15 +117,15 @@ phasepoint(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     θ::AbstractVecOrMat{T},
     h::Hamiltonian,
-) where {T<:Real} = phasepoint(h, θ, rand(rng, h.metric, h.kinetic, θ))
+) where {T<:Real} = phasepoint(h, θ, rand_momentum(rng, h.metric, h.kinetic, θ))
 
 # To change L191 of hamiltonian.jl
 refresh(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     ::FullMomentumRefreshment,
     h::Hamiltonian,
     z::PhasePoint,
-) = phasepoint(h, z.θ, rand(rng, h.metric, h.kinetic, z.θ))
+) = phasepoint(h, z.θ, rand_momentum(rng, h.metric, h.kinetic, z.θ))
 
 # To change L215 of hamiltonian.jl
 refresh(
@@ -136,17 +136,9 @@ refresh(
 ) = phasepoint(
     h,
     z.θ,
-    ref.α * z.r + sqrt(1 - ref.α^2) * rand(rng, h.metric, h.kinetic, z.θ),
+    ref.α * z.r + sqrt(1 - ref.α^2) * rand_momentum(rng, h.metric, h.kinetic, z.θ),
 )
 
-# To change L146 of metric.jl
-# Ignore θ by default (i.e. not position-dependent)
-Base.rand(rng::AbstractRNG, metric::AbstractMetric, kinetic, θ) =
-    rand_momentum(rng, metric, kinetic)    # this disambiguity is required by Random.rand
-Base.rand(rng::AbstractVector{<:AbstractRNG}, metric::AbstractMetric, kinetic, θ) =
-    rand_momentum(rng, metric, kinetic)
-Base.rand(metric::AbstractMetric, kinetic, θ) = rand(Random.default_rng(), metric, kinetic)
-
 ### metric.jl
 
 import AdvancedHMC: _rand
@@ -212,7 +204,7 @@ function rand_momentum(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     metric::DenseRiemannianMetric{T},
     kinetic,
-    θ,
+    θ::AbstractVecOrMat,
 ) where {T}
     r = _randn(rng, T, size(metric)...)
     G⁻¹ = inv(metric.map(metric.G(θ)))
@@ -221,15 +213,6 @@ function rand_momentum(
     return r
 end
 
-Base.rand(rng::AbstractRNG, metric::AbstractRiemannianMetric, kinetic, θ) =
-    rand_momentum(rng, metric, kinetic, θ)
-Base.rand(
-    rng::AbstractVector{<:AbstractRNG},
-    metric::AbstractRiemannianMetric,
-    kinetic,
-    θ,
-) = rand_momentum(rng, metric, kinetic, θ)
-
 ### hamiltonian.jl
 
 import AdvancedHMC: phasepoint, neg_energy, ∂H∂θ, ∂H∂r

src/hamiltonian.jl

@@ -161,7 +161,7 @@ phasepoint(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     θ::AbstractVecOrMat{T},
     h::Hamiltonian,
-) where {T<:Real} = phasepoint(h, θ, rand(rng, h.metric, h.kinetic, θ))
+) where {T<:Real} = phasepoint(h, θ, rand_momentum(rng, h.metric, h.kinetic, θ))
 
 abstract type AbstractMomentumRefreshment end
 
@@ -173,7 +173,7 @@ refresh(
     ::FullMomentumRefreshment,
     h::Hamiltonian,
     z::PhasePoint,
-) = phasepoint(h, z.θ, rand(rng, h.metric, h.kinetic, z.θ))
+) = phasepoint(h, z.θ, rand_momentum(rng, h.metric, h.kinetic, z.θ))
 
 """
 $(TYPEDEF)
@@ -204,5 +204,5 @@ refresh(
 ) = phasepoint(
     h,
     z.θ,
-    ref.α * z.r + sqrt(1 - ref.α^2) * rand(rng, h.metric, h.kinetic, z.θ),
+    ref.α * z.r + sqrt(1 - ref.α^2) * rand_momentum(rng, h.metric, h.kinetic, z.θ),
 )

src/metric.jl

@@ -98,6 +98,7 @@ function rand_momentum(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     metric::UnitEuclideanMetric{T},
     kinetic::GaussianKinetic,
+    ::AbstractVecOrMat,
 ) where {T}
     r = _randn(rng, T, size(metric)...)
     return r
@@ -107,6 +108,7 @@ function rand_momentum(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     metric::DiagEuclideanMetric{T},
     kinetic::GaussianKinetic,
+    ::AbstractVecOrMat,
 ) where {T}
     r = _randn(rng, T, size(metric)...)
     r ./= metric.sqrtM⁻¹
@@ -117,35 +119,9 @@ function rand_momentum(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},
     metric::DenseEuclideanMetric{T},
     kinetic::GaussianKinetic,
+    ::AbstractVecOrMat,
 ) where {T}
     r = _randn(rng, T, size(metric)...)
     ldiv!(metric.cholM⁻¹, r)
     return r
 end
-
-# TODO (kai) The rand interface should be updated as "rand from momentum distribution + optional affine transformation by metric"
-Base.rand(rng::AbstractRNG, metric::AbstractMetric, kinetic::AbstractKinetic) =
-    rand_momentum(rng, metric, kinetic)    # this disambiguity is required by Random.rand
-Base.rand(
-    rng::AbstractVector{<:AbstractRNG},
-    metric::AbstractMetric,
-    kinetic::AbstractKinetic,
-) = rand_momentum(rng, metric, kinetic)
-Base.rand(metric::AbstractMetric, kinetic::AbstractKinetic) =
-    rand(Random.default_rng(), metric, kinetic)
-
-# ignore θ by default unless defined by the specific kinetic (i.e. not position-dependent)
-Base.rand(
-    rng::AbstractRNG,
-    metric::AbstractMetric,
-    kinetic::AbstractKinetic,
-    θ::AbstractVecOrMat,
-) = rand(rng, metric, kinetic)    # this disambiguity is required by Random.rand
-Base.rand(
-    rng::AbstractVector{<:AbstractRNG},
-    metric::AbstractMetric,
-    kinetic::AbstractKinetic,
-    θ::AbstractVecOrMat,
-) = rand(rng, metric, kinetic)
-Base.rand(metric::AbstractMetric, kinetic::AbstractKinetic, θ::AbstractVecOrMat) =
-    rand(metric, kinetic)

src/trajectory.jl

@@ -773,7 +773,7 @@ function find_good_stepsize(
     a_min, a_cross, a_max = T(0.25), T(0.5), T(0.75) # minimal, crossing, maximal accept ratio
     d = T(2.0)
     # Create starting phase point
-    r = rand(rng, h.metric, h.kinetic) # sample momentum variable
+    r = rand_momentum(rng, h.metric, h.kinetic, θ) # sample momentum variable
     z = phasepoint(h, θ, r)
     H = energy(z)
 

test/integrator.jl

@@ -10,7 +10,7 @@ using Statistics: mean
 
     θ_init = randn(D)
     h = Hamiltonian(UnitEuclideanMetric(D), ℓπ, ∂ℓπ∂θ)
-    r_init = AdvancedHMC.rand(h.metric, h.kinetic)
+    r_init = AdvancedHMC.rand_momentum(Random.default_rng(), h.metric, h.kinetic, θ_init)
 
     n_steps = 10
 
@@ -122,7 +122,8 @@ using Statistics: mean
         for lf in [Leapfrog(ϵ), DiffEqIntegrator(ϵ, VerletLeapfrog())]
             q_init = randn(1)
             h = Hamiltonian(UnitEuclideanMetric(1), negU, ForwardDiff)
-            p_init = AdvancedHMC.rand(h.metric, h.kinetic)
+            p_init =
+                AdvancedHMC.rand_momentum(Random.default_rng(), h.metric, h.kinetic, q_init)
 
             q, p = copy(q_init), copy(p_init)
             z = AdvancedHMC.phasepoint(h, q, p)

test/metric.jl

@@ -4,13 +4,14 @@ using ReTest, Random, AdvancedHMC
     @testset "Sample momentum variables from metric via vector of RNGs" begin
         D = 10
         n_chains = 5
+        θ = randn(D, n_chains)
         rng = [MersenneTwister(1) for _ = 1:n_chains]
         for metric in [
             UnitEuclideanMetric((D, n_chains)),
             DiagEuclideanMetric((D, n_chains)),
             # DenseEuclideanMetric((D, n_chains)) # not supported ATM
         ]
-            r = rand(rng, metric, GaussianKinetic())
+            r = AdvancedHMC.rand_momentum(rng, metric, GaussianKinetic(), θ)
             all_same = true
             for i = 2:n_chains
                 all_same = all_same && r[:, i] == r[:, 1]

test/trajectory.jl

@@ -50,7 +50,7 @@ function gettraj(rng, h, ϵ = 0.1, n_steps = 50)
     lf = Leapfrog(ϵ)
 
     q_init = randn(rng, D)
-    p_init = AdvancedHMC.rand(rng, h.metric, h.kinetic)
+    p_init = AdvancedHMC.rand_momentum(rng, h.metric, h.kinetic, q_init)
     z = AdvancedHMC.phasepoint(h, q_init, p_init)
 
     traj_z = Vector(undef, n_steps)
@@ -127,7 +127,7 @@ end
         Leapfrog(find_good_stepsize(h, θ_init)),
         GeneralisedNoUTurn(),
     )
-    r_init = AdvancedHMC.rand(h.metric, h.kinetic)
+    r_init = AdvancedHMC.rand_momentum(Random.default_rng(), h.metric, h.kinetic, θ_init)
 
     @testset "Passing RNG" begin
         τ_with_jittered_lf = Trajectory{MultinomialTS}(