Type design for PhasePoint and DualValue.

Project.toml

@@ -1,6 +1,5 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-authors = ["The Turing Team"]
 version = "0.1.5"
 
 [deps]
@@ -10,6 +9,8 @@ LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
+Parameters = "d96e819e-fc66-5662-9728-84c9c7592b0a"
 
 [extras]
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"

src/AdvancedHMC.jl

@@ -6,6 +6,8 @@ using Statistics: mean, var, middle
 using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky
 using LazyArrays: BroadcastArray
 using Random: GLOBAL_RNG, AbstractRNG
+using Parameters: @unpack
+using ArgCheck: @argcheck
 
 import StatsBase: sample
 

src/adaptation/Adaptation.jl

@@ -1,6 +1,7 @@
 module Adaptation
 
-import Base: string
+import Base: string, rand
+using Random: GLOBAL_RNG, AbstractRNG
 using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky
 import LinearAlgebra, Statistics
 using ..AdvancedHMC: DEBUG

src/adaptation/precond.jl

@@ -282,6 +282,35 @@ function Base.getproperty(dem::DenseEuclideanMetric, d::Symbol)
     return d === :dim ? size(getfield(dem, :M⁻¹), 1) : getfield(dem, d)
 end
 
+# `rand` functions for `metric` types.
+function Base.rand(
+    rng::AbstractRNG,
+    metric::UnitEuclideanMetric
+)
+    r = randn(rng, metric.dim)
+    return r
+end
+
+function Base.rand(
+    rng::AbstractRNG,
+    metric::DiagEuclideanMetric
+)
+    r = randn(rng, metric.dim)
+    r ./= metric.sqrtM⁻¹
+    return r
+end
+
+function Base.rand(
+    rng::AbstractRNG,
+    metric::DenseEuclideanMetric
+)
+    r = randn(rng, metric.dim)
+    ldiv!(metric.cholM⁻¹, r)
+    return r
+end
+
+Base.rand(metric::AbstractMetric) = rand(GLOBAL_RNG, metric)
+
 ####
 #### Preconditioner constructors
 ####

src/hamiltonian.jl

@@ -1,60 +1,100 @@
 # TODO: add a type for kinetic energy
+# TODO: add cache for gradients by letting ∂logπ∂θ return both log density and gradient
 
 struct Hamiltonian{M<:AbstractMetric, Tlogπ, T∂logπ∂θ}
     metric::M
-    logπ::Tlogπ
-    ∂logπ∂θ::T∂logπ∂θ
+    ℓπ::Tlogπ
+    ∂ℓπ∂θ::T∂logπ∂θ
 end
 
 # Create a `Hamiltonian` with a new `M⁻¹`
-(h::Hamiltonian)(M⁻¹) = Hamiltonian(h.metric(M⁻¹), h.logπ, h.∂logπ∂θ)
+(h::Hamiltonian)(M⁻¹) = Hamiltonian(h.metric(M⁻¹), h.ℓπ, h.∂ℓπ∂θ)
 
-∂H∂θ(h::Hamiltonian, θ::AbstractVector) = -h.∂logπ∂θ(θ)
+∂H∂θ(h::Hamiltonian, θ::AbstractVector) = -h.∂ℓπ∂θ(θ)
 
 ∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric}, r::AbstractVector) = copy(r)
 ∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric}, r::AbstractVector) = h.metric.M⁻¹ .* r
 ∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric}, r::AbstractVector) = h.metric.M⁻¹ * r
 
-function hamiltonian_energy(h::Hamiltonian, θ::AbstractVector, r::AbstractVector)
-    K = kinetic_energy(h, r, θ)
-    if isnan(K)
-        K = Inf
-        @warn "Kinetic energy is `NaN` and is set to `Inf`."
-    end
-    V = potential_energy(h, θ)
-    if isnan(V)
-        V = Inf
-        @warn "Potential energy is `NaN` and is set to `Inf`."
+struct DualValue{Tv<:AbstractFloat, Tg<:AbstractVector{Tv}}
+    value::Tv    # Cached value, e.g. logπ(θ).
+    gradient::Tg # Cached gradient, e.g. ∇logπ(θ).
+end
+
+struct PhasePoint{T<:AbstractVector, V<:DualValue}
+    θ::T  # Position variables / model parameters.
+    r::T  # Momentum variables
+    ℓπ::V # Cached neg potential energy for the current θ.
+    ℓκ::V # Cached neg kinect energy for the current r.
+    function PhasePoint(θ::T, r::T, ℓπ::V, ℓκ::V) where {T,V}
+        @argcheck length(θ) == length(r) == length(ℓπ.gradient) == length(ℓπ.gradient)
+        # if !(all(isfinite, θ) && all(isfinite, r) && all(isfinite, ℓπ) && all(isfinite, ℓκ))
+        if !(isfinite(θ) && isfinite(r) && isfinite(ℓπ) && isfinite(ℓκ))
+            @warn "The current proposal will be rejected (due to numerical error(s))..."
+            ℓκ = DualValue(-Inf, ℓκ.gradient)
+            ℓπ = DualValue(-Inf, ℓπ.gradient)
+        end
+        new{T,V}(θ, r, ℓπ, ℓκ)
     end
-    return K + V
 end
 
-potential_energy(h::Hamiltonian, θ::AbstractVector) = -h.logπ(θ)
+phasepoint(
+    h::Hamiltonian,
+    θ::T,
+    r::T;
+    ℓπ = DualValue(neg_energy(h, r, θ), ∂H∂θ(h, θ)),
+    ℓκ = DualValue(neg_energy(h, θ), ∂H∂r(h, r))
+) where {T<:AbstractVector} = PhasePoint(θ, r, ℓπ, ℓκ)
+
+
+Base.isfinite(v::DualValue) = all(isfinite, v.value) && all(isfinite, v.gradient)
+Base.isfinite(v::AbstractVector) = all(isfinite, v)
+Base.isfinite(z::PhasePoint) = isfinite(z.ℓπ) && isfinite(z.ℓκ)
+
+###
+### Negative energy (or log probability) functions.
+### NOTE: the general form (i.e. non-Euclidean) of K depends on both θ and r.
+###
+
+neg_energy(z::PhasePoint) = z.ℓπ.value + z.ℓκ.value
 
-# Kinetic energy
-# NOTE: the general form of K depends on both θ and r
-kinetic_energy(h::Hamiltonian{<:UnitEuclideanMetric}, r, θ) = sum(abs2, r) / 2
-function kinetic_energy(h::Hamiltonian{<:DiagEuclideanMetric}, r, θ)
-    return sum(abs2(r[i]) * h.metric.M⁻¹[i] for i in 1:length(r)) / 2
+neg_energy(h::Hamiltonian, θ::AbstractVector) = h.ℓπ(θ)
+
+neg_energy(
+    h::Hamiltonian{<:UnitEuclideanMetric},
+    r::T,
+    θ::T
+) where {T<:AbstractVector} = -sum(abs2, r) / 2
+
+function neg_energy(
+    h::Hamiltonian{<:DiagEuclideanMetric},
+    r::T,
+    θ::T
+) where {T<:AbstractVector}
+    _r = [abs2(r[i]) * h.metric.M⁻¹[i] for i in 1:length(r)]
+    return -sum(_r) / 2
 end
-function kinetic_energy(h::Hamiltonian{<:DenseEuclideanMetric}, r, θ)
+
+function neg_energy(
+    h::Hamiltonian{<:DenseEuclideanMetric},
+    r::T,
+    θ::T
+) where {T<:AbstractVector}
     mul!(h.metric._temp, h.metric.M⁻¹, r)
-    return dot(r, h.metric._temp) / 2
+    return -dot(r, h.metric._temp) / 2
 end
 
-# Momentum sampler
-function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:UnitEuclideanMetric})
-    return randn(rng, h.metric.dim)
-end
-function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:DiagEuclideanMetric})
-    r = randn(rng, h.metric.dim)
-    r ./= h.metric.sqrtM⁻¹
-    return r
-end
-function rand_momentum(rng::AbstractRNG, h::Hamiltonian{<:DenseEuclideanMetric})
-    r = randn(rng, h.metric.dim)
-    ldiv!(h.metric.cholM⁻¹, r)
-    return r
-end
+####
+#### Momentum sampler
+####
+
+rand_momentum(
+    rng::AbstractRNG,
+    z::PhasePoint,
+    h::Hamiltonian
+) = phasepoint(h, z.θ, rand(rng, h.metric))
 
-rand_momentum(h::Hamiltonian) = rand_momentum(GLOBAL_RNG, h)
+rand_momentum(
+    z::PhasePoint,
+    h::Hamiltonian
+) = phasepoint(h, z.θ, rand(GLOBAL_RNG, h.metric))