feat(wip): port reemannian hmc from research repo

Signed-off-by: Kai Xu <xuk@ibm.com>

Author: xukai92

src/relativistic/hamiltonian.jl

@@ -0,0 +1,74 @@
+abstract type AbstractRelativisticKinetic{T} <: AbstractKinetic end
+
+struct RelativisticKinetic{T} <: AbstractRelativisticKinetic{T}
+    "Mass"
+    m::T
+    "Speed of light"
+    c::T
+end
+
+relativistic_mass(kinetic::RelativisticKinetic, r, r′ = r) = 
+    kinetic.m * sqrt(dot(r, r′) / (kinetic.m ^ 2 * kinetic.c ^ 2) + 1)
+relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) = sum(
+    kinetic.c ^ 2 * relativistic_mass(kinetic, r, r′)
+)
+
+struct DimensionwiseRelativisticKinetic{T} <: AbstractRelativisticKinetic{T}
+    "Mass"
+    m::T
+    "Speed of light"
+    c::T
+end
+
+relativistic_mass(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = 
+    kinetic.m .* sqrt.(r .* r′ ./ (kinetic.m .^ 2 .* kinetic.c .^ 2) .+ 1)
+relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = sum(
+    kinetic.c .^ 2 .* relativistic_mass(kinetic, r, r′)
+)
+
+function ∂H∂r(
+    h::Hamiltonian{<:UnitEuclideanMetric,<:AbstractRelativisticKinetic},
+    r::AbstractVecOrMat,
+)
+    mass = relativistic_mass(h.kinetic, r)
+    return r ./ mass
+end
+function ∂H∂r(
+    h::Hamiltonian{<:DiagEuclideanMetric,<:AbstractRelativisticKinetic},
+    r::AbstractVecOrMat,
+)
+    r = h.metric.sqrtM⁻¹ .* r
+    mass = relativistic_mass(h.kinetic, r)
+    red_term = r ./ mass # red part of (15)
+    return h.metric.sqrtM⁻¹ .* red_term # (15)
+end
+function ∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric, <:AbstractRelativisticKinetic}, r::AbstractVecOrMat)
+    r = h.metric.cholM⁻¹ * r
+    mass = relativistic_mass(h.kinetic, r)
+    red_term = r ./ mass
+    return h.metric.cholM⁻¹' * red_term
+end
+
+function neg_energy(
+    h::Hamiltonian{<:UnitEuclideanMetric,<:AbstractRelativisticKinetic},
+    r::T,
+    θ::T,
+) where {T<:AbstractVector}
+    return -relativistic_energy(h.kinetic, r)
+end
+function neg_energy(
+    h::Hamiltonian{<:DiagEuclideanMetric,<:AbstractRelativisticKinetic},
+    r::T,
+    θ::T,
+) where {T<:AbstractVector}
+    r = h.metric.sqrtM⁻¹ .* r
+    return -relativistic_energy(h.kinetic, r)
+end
+function neg_energy(
+    h::Hamiltonian{<:DenseEuclideanMetric,<:AbstractRelativisticKinetic},
+    r::T,
+    θ::T
+) where {T<:AbstractVector}
+    r = h.metric.cholM⁻¹ * r
+    return -relativistic_energy(h.kinetic, r)
+end

src/relativistic/metric.jl

@@ -0,0 +1,71 @@
+function rand_angles(rng::AbstractRNG, dim)
+    rand(rng, dim - 1) .* vcat(fill(π, dim - 2), 2 * π)
+end
+
+"Special case of `polar2spherical` with dimension equal 2"
+polar2cartesian(θ, d) = d * [cos(θ), sin(θ)]
+
+# ref: https://en.wikipedia.org/wiki/N-sphere#Spherical_coordinates
+function polar2spherical(θs, d)
+    cos_lst, sin_lst = cos.(θs), sin.(θs)
+    suffixed_cos_lst = vcat(cos_lst, 1) # [cos(θ[1]), cos(θ[2]), ..., cos(θ[d-1]), 1]
+    prefixed_cumprod_sin_lst = vcat(1, cumprod(sin_lst)) # [1, sin(θ[1]), sin(θ[1]) * sin(θ[2]), ..., sin(θ[1]) * ... * sin(θ[d-1])]
+    return d * prefixed_cumprod_sin_lst .* suffixed_cos_lst
+end
+
+momentum_mode(m, c) = sqrt((1 / c^2 + sqrt(1 / c^2 + 4 * m^2)) / 2) # mode of the momentum distribution
+
+function _rand(
+    rng::AbstractRNG,
+    metric::UnitEuclideanMetric{T},
+    kinetic::RelativisticKinetic{T},
+) where {T}
+    densityfunc = x -> exp(-relativistic_energy(kinetic, [x])) * x
+    mm = momentum_mode(kinetic.m, kinetic.c)
+    sampler = RejectionSampler(densityfunc, (0.0, Inf), (mm / 2, mm * 2); max_segments = 5)
+    sz = size(metric)
+    θs = rand_angles(rng, prod(sz))
+    d = only(run_sampler!(rng, sampler, 1))
+    r = polar2spherical(θs, d * rand(rng, [-1, +1])) # TODO Double check if +/- is needed
+    r = reshape(r, sz)
+    return r
+end
+
+# TODO Support AbstractVector{<:AbstractRNG}
+# FIXME Unit-test this using slice sampler or HMC sampler
+function _rand(
+    rng::AbstractRNG,
+    metric::UnitEuclideanMetric{T},
+    kinetic::DimensionwiseRelativisticKinetic{T},
+) where {T}
+    h_temp = Hamiltonian(metric, kinetic, identity, identity)
+    densityfunc = x -> exp(neg_energy(h_temp, [x], [x]))
+    sampler = RejectionSampler(densityfunc, (-Inf, Inf); max_segments = 5)
+    sz = size(metric)
+    r = run_sampler!(rng, sampler, prod(sz))
+    r = reshape(r, sz)
+    return r
+end
+
+# TODO Support AbstractVector{<:AbstractRNG}
+function _rand(
+    rng::AbstractRNG,
+    metric::DiagEuclideanMetric{T},
+    kinetic::AbstractRelativisticKinetic{T},
+) where {T}
+    r = _rand(rng, UnitEuclideanMetric(size(metric)), kinetic)
+    # p' = A p where A = sqrtM
+    r ./= metric.sqrtM⁻¹
+    return r
+end
+# TODO Support AbstractVector{<:AbstractRNG}
+function _rand(
+    rng::AbstractRNG,
+    metric::DenseEuclideanMetric{T},
+    kinetic::AbstractRelativisticKinetic{T},
+) where {T}
+    r = _rand(rng, UnitEuclideanMetric(size(metric)), kinetic)
+    # p' = A p where A = cholM
+    ldiv!(metric.cholM⁻¹, r)
+    return r
+end