chore: unpack kinetic.jl (#303)

* chore: unpack kinetic.jl * chore: cherry pick more changes from #300 Co-authored-by: Kai Xu <[email protected]>
TuringLang · Nov 29, 2022 · 403c7e5 · 403c7e5
1 parent 11380fc
commit 403c7e5
Show file tree

Hide file tree

Showing 4 changed files with 20 additions and 15 deletions.
diff --git a/src/AdvancedHMC.jl b/src/AdvancedHMC.jl
@@ -30,7 +30,13 @@ include("utilities.jl")
 # r: momentum variables
 # z: phase point / a pair of θ and r
 
-include("kinetic.jl")
+# TODO Move it back to hamiltonian.jl after the rand interface is updated
+abstract type AbstractKinetic end
+
+struct GaussianKinetic <: AbstractKinetic end
+
+export GaussianKinetic
+
 include("metric.jl")
 export UnitEuclideanMetric, DiagEuclideanMetric, DenseEuclideanMetric
 

diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
@@ -34,9 +34,9 @@ function ∂H∂θ(h::Hamiltonian, θ::AbstractVecOrMat)
     return DualValue(res[1], -res[2])
 end
 
-∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric}, r::AbstractVecOrMat) = copy(r)
-∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric}, r::AbstractVecOrMat) = h.metric.M⁻¹ .* r
-∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric}, r::AbstractVecOrMat) = h.metric.M⁻¹ * r
+∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric, <:GaussianKinetic}, r::AbstractVecOrMat) = copy(r)
+∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric, <:GaussianKinetic}, r::AbstractVecOrMat) = h.metric.M⁻¹ .* r
+∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric, <:GaussianKinetic}, r::AbstractVecOrMat) = h.metric.M⁻¹ * r
 
 struct PhasePoint{T<:AbstractVecOrMat{<:AbstractFloat}, V<:DualValue}
     θ::T  # Position variables / model parameters.
@@ -109,32 +109,34 @@ neg_energy(z::PhasePoint) = z.ℓπ.value + z.ℓκ.value
 
 neg_energy(h::Hamiltonian, θ::AbstractVecOrMat) = h.ℓπ(θ)
 
+# GaussianKinetic
+
 neg_energy(
-    h::Hamiltonian{<:UnitEuclideanMetric},
+    h::Hamiltonian{<:UnitEuclideanMetric, <:GaussianKinetic},
     r::T,
     θ::T
 ) where {T<:AbstractVector} = -sum(abs2, r) / 2
 
 neg_energy(
-    h::Hamiltonian{<:UnitEuclideanMetric},
+    h::Hamiltonian{<:UnitEuclideanMetric, <:GaussianKinetic},
     r::T,
     θ::T
 ) where {T<:AbstractMatrix} = -vec(sum(abs2, r; dims=1)) / 2
 
 neg_energy(
-    h::Hamiltonian{<:DiagEuclideanMetric},
+    h::Hamiltonian{<:DiagEuclideanMetric, <:GaussianKinetic},
     r::T,
     θ::T
 ) where {T<:AbstractVector} = -sum(abs2.(r) .* h.metric.M⁻¹) / 2
 
 neg_energy(
-    h::Hamiltonian{<:DiagEuclideanMetric},
+    h::Hamiltonian{<:DiagEuclideanMetric, <:GaussianKinetic},
     r::T,
     θ::T
 ) where {T<:AbstractMatrix} = -vec(sum(abs2.(r) .* h.metric.M⁻¹; dims=1) ) / 2
 
 function neg_energy(
-    h::Hamiltonian{<:DenseEuclideanMetric},
+    h::Hamiltonian{<:DenseEuclideanMetric, <:GaussianKinetic},
     r::T,
     θ::T
 ) where {T<:AbstractVecOrMat}

diff --git a/src/kinetic.jl b/src/kinetic.jl
diff --git a/src/metric.jl b/src/metric.jl
@@ -87,7 +87,7 @@ getname(m::T) where {T<:AbstractMetric} = getname(T)
 function _rand(
     rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}},
     metric::UnitEuclideanMetric{T},
-    kinetic,
+    kinetic::GaussianKinetic,
 ) where {T}
     r = randn(rng, T, size(metric)...)
     return r
@@ -96,7 +96,7 @@ end
 function _rand(
     rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}},
     metric::DiagEuclideanMetric{T},
-    kinetic,
+    kinetic::GaussianKinetic,
 ) where {T}
     r = randn(rng, T, size(metric)...)
     r ./= metric.sqrtM⁻¹
@@ -106,7 +106,7 @@ end
 function _rand(
     rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}},
     metric::DenseEuclideanMetric{T},
-    kinetic,
+    kinetic::GaussianKinetic,
 ) where {T}
     r = randn(rng, T, size(metric)...)
     ldiv!(metric.cholM⁻¹, r)