From dbb5d16b257a7022bdf62da72d48b8823f9b459b Mon Sep 17 00:00:00 2001 From: erathorn Date: Thu, 3 Aug 2023 11:55:26 +0200 Subject: [PATCH 1/4] correct length check for phase point --- src/hamiltonian.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index f2bbd230..1daa5605 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -51,7 +51,7 @@ struct PhasePoint{T<:AbstractVecOrMat{<:AbstractFloat},V<:DualValue} ℓπ::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) + @argcheck length(θ) == length(r) == length(ℓπ.gradient) == length(ℓk.gradient) if any(isfinite.((θ, r, ℓπ, ℓκ)) .== false) # @warn "The current proposal will be rejected due to numerical error(s)." isfinite.((θ, r, ℓπ, ℓκ)) # NOTE eltype has to be inlined to avoid type stability issue; see #267 From 18dc7f0efb754dbf9b869ed23e7565e93f2c33d1 Mon Sep 17 00:00:00 2001 From: erathorn Date: Thu, 3 Aug 2023 12:49:42 +0200 Subject: [PATCH 2/4] fix typo --- src/hamiltonian.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 1daa5605..7a39bf9f 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -51,7 +51,7 @@ struct PhasePoint{T<:AbstractVecOrMat{<:AbstractFloat},V<:DualValue} ℓπ::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(ℓk.gradient) + @argcheck length(θ) == length(r) == length(ℓπ.gradient) == length(ℓκ.gradient) if any(isfinite.((θ, r, ℓπ, ℓκ)) .== false) # @warn "The current proposal will be rejected due to numerical error(s)." isfinite.((θ, r, ℓπ, ℓκ)) # NOTE eltype has to be inlined to avoid type stability issue; see #267 From 449a5893f9bd5622b58004b0a2fce9125bf26e49 Mon Sep 17 00:00:00 2001 From: erathorn Date: Thu, 3 Aug 2023 13:57:03 +0200 Subject: [PATCH 3/4] DenseMetric and Comoponent Arrays --- src/hamiltonian.jl | 7 +++++-- test/hamiltonian.jl | 35 ++++++++++++++++++++++++++++++++++- 2 files changed, 39 insertions(+), 3 deletions(-) diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl index 7a39bf9f..4047f30e 100644 --- a/src/hamiltonian.jl +++ b/src/hamiltonian.jl @@ -42,8 +42,11 @@ end ∂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 +function ∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat) + out = similar(r) # Make sure the output of this function is of the same type as r + mul!(out, h.metric.M⁻¹, r) + out +end struct PhasePoint{T<:AbstractVecOrMat{<:AbstractFloat},V<:DualValue} θ::T # Position variables / model parameters. diff --git a/test/hamiltonian.jl b/test/hamiltonian.jl index 47948325..f24330a9 100644 --- a/test/hamiltonian.jl +++ b/test/hamiltonian.jl @@ -1,6 +1,7 @@ using ReTest, AdvancedHMC using AdvancedHMC: GaussianKinetic, DualValue, PhasePoint using LinearAlgebra: dot, diagm +using ComponentArrays @testset "Hamiltonian" begin f = x -> dot(x, x) @@ -38,7 +39,7 @@ end end end -@testset "Metric" begin +@testset "Metric Base Array" begin n_tests = 10 for T in [Float32, Float64] @@ -49,18 +50,50 @@ end h = Hamiltonian(UnitEuclideanMetric(T, D), ℓπ, ∂ℓπ∂θ) @test -AdvancedHMC.neg_energy(h, r_init, θ_init) == sum(abs2, r_init) / 2 @test AdvancedHMC.∂H∂r(h, r_init) == r_init + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) M⁻¹ = ones(T, D) + abs.(randn(T, D)) h = Hamiltonian(DiagEuclideanMetric(M⁻¹), ℓπ, ∂ℓπ∂θ) @test -AdvancedHMC.neg_energy(h, r_init, θ_init) ≈ r_init' * diagm(0 => M⁻¹) * r_init / 2 @test AdvancedHMC.∂H∂r(h, r_init) == M⁻¹ .* r_init + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) m = randn(T, D, D) M⁻¹ = m' * m h = Hamiltonian(DenseEuclideanMetric(M⁻¹), ℓπ, ∂ℓπ∂θ) @test -AdvancedHMC.neg_energy(h, r_init, θ_init) ≈ r_init' * M⁻¹ * r_init / 2 @test AdvancedHMC.∂H∂r(h, r_init) == M⁻¹ * r_init + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) + end + end +end + +@testset "Metric ComponentArrays" begin + n_tests = 10 + for T in [Float32, Float64] + for _ = 1:n_tests + θ_init = ComponentArray(a = randn(T, D), b= randn(T, D)) + r_init = ComponentArray(a = randn(T, D), b= randn(T, D)) + + h = Hamiltonian(UnitEuclideanMetric(T, 2*D), ℓπ, ∂ℓπ∂θ) + @test -AdvancedHMC.neg_energy(h, r_init, θ_init) == sum(abs2, r_init) / 2 + @test AdvancedHMC.∂H∂r(h, r_init) == r_init + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) + + M⁻¹ = ones(T, 2*D) + abs.(randn(T, 2*D)) + h = Hamiltonian(DiagEuclideanMetric(M⁻¹), ℓπ, ∂ℓπ∂θ) + @test -AdvancedHMC.neg_energy(h, r_init, θ_init) ≈ + r_init' * diagm(0 => M⁻¹) * r_init / 2 + @test AdvancedHMC.∂H∂r(h, r_init) == M⁻¹ .* r_init + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) + + m = randn(T, 2*D, 2*D) + M⁻¹ = m' * m + h = Hamiltonian(DenseEuclideanMetric(M⁻¹), ℓπ, ∂ℓπ∂θ) + @test -AdvancedHMC.neg_energy(h, r_init, θ_init) ≈ r_init' * M⁻¹ * r_init / 2 + @test all(AdvancedHMC.∂H∂r(h, r_init) .== M⁻¹ * r_init) + @test typeof(AdvancedHMC.∂H∂r(h, r_init)) == typeof(r_init) end end end From 7c585240b26c5019caaefc20d6e1912160d0bc6d Mon Sep 17 00:00:00 2001 From: erathorn Date: Thu, 3 Aug 2023 15:42:41 +0200 Subject: [PATCH 4/4] Update test/hamiltonian.jl Co-authored-by: Tor Erlend Fjelde --- test/hamiltonian.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/test/hamiltonian.jl b/test/hamiltonian.jl index f24330a9..fd23e784 100644 --- a/test/hamiltonian.jl +++ b/test/hamiltonian.jl @@ -73,8 +73,8 @@ end n_tests = 10 for T in [Float32, Float64] for _ = 1:n_tests - θ_init = ComponentArray(a = randn(T, D), b= randn(T, D)) - r_init = ComponentArray(a = randn(T, D), b= randn(T, D)) + θ_init = ComponentArray(a = randn(T, D), b = randn(T, D)) + r_init = ComponentArray(a = randn(T, D), b = randn(T, D)) h = Hamiltonian(UnitEuclideanMetric(T, 2*D), ℓπ, ∂ℓπ∂θ) @test -AdvancedHMC.neg_energy(h, r_init, θ_init) == sum(abs2, r_init) / 2