Simplifying `Metric` types by using `AbstractPDMat`

[yebai]

The following 3 types in src/metric.jl:

• AdvancedHMC.jl/src/metric.jl

  Lines 16 to 24 in a0af7cc

  	struct DiagEuclideanMetric{T<:Real,A<:AbstractVector{T}} <: AbstractMetric{T}
  	dim :: Int
  	# Diagnal of the inverse of the mass matrix
  	M⁻¹ :: A
  	# Sqare root of the inverse of the mass matrix
  	sqrtM⁻¹ :: A
  	# Pre-allocation for intermediate variables
  	_temp :: A
  	end

• AdvancedHMC.jl/src/metric.jl

  Lines 41 to 49 in a0af7cc

  	struct DenseEuclideanMetric{T<:Real,AV<:AbstractVector{T},AM<:AbstractMatrix{T}} <: AbstractMetric{T}
  	dim :: Int
  	# Inverse of the mass matrix
  	M⁻¹ :: AM
  	# U of the Cholesky decomposition of the mass matrix
  	cholM⁻¹ :: UpperTriangular{T,AM}
  	# Pre-allocation for intermediate variables
  	_temp :: AV
  	end

• AdvancedHMC.jl/src/metric.jl

  Lines 3 to 5 in a0af7cc

  	struct UnitEuclideanMetric{T<:Real} <: AbstractMetric{T}
  	dim :: Int
  	end

can be merged into one type using AbstractPDMat

struct EuclideanMetric{T<:Real,A<:AbstractPDMat{T}} <: AbstractMetric{T}
    M⁻¹     ::  A # contains dim, M⁻¹ and its Cholesky decomposition. 
    # Pre-allocation for intermediate variables
    _temp   ::  A
end

Reference: https://github.com/JuliaStats/PDMats.jl

[xukai92]

Got it!

[yebai]

Also, the Metric._temp can be removed if we introduce a PhasePoint type (see also #17):

z_new::PhasePoint = rand_momentum(rng::AbstractRNG, h::Hamiltonian, z_old::PhasePoint)

Since rand_momentum can perform in-place update for z::PhasePoint. However, I'm not sure whether this is necessary since momentum variables are not updated as frequent as position variables.

[xukai92]

@yebai There is some werid design in PDMats.jl that is not consistent with AHMC. Basically PDiagMat will take a vector and compute the inv it - not sqrt it, which is what we need in AHMC.

[yebai]

@yebai There is some werid design in PDMats.jl that is not consistent with AHMC. Basically PDiagMat will take a vector and compute the inv it - not sqrt it, which is what we need in AHMC.

Can we create a wrapper type for AbstractPDMat? For example

struct Metric{T<:Real,M<:AbstractPDMat{T},V<:AbstractVector{T}} <: AbstractPDMat{T}
 m::M
 sqrt_diag::V    # the vector of sqrt diagonal elements
end

[xukai92]

This is almost the same is current Metric types right? And the current version we have would be more memory friendly than this because we only store the original M inverse & its sqrt or chol if needed. While warpping things like you proposed, m would store useless invese matrix for diagonal.

[yebai]

Ah, I see. Can you remind me why we need sqrt_diag?

[xukai92]

See: https://github.com/TuringLang/AdvancedHMC.jl/blob/master/src/adaptation/precond.jl#L342

Reference: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/hmc/hamiltonians/diag_e_metric.hpp#L53

[yebai]

sqrt seems like a cheap operation (see wiki), maybe we don't need to cache it? I don't have a strong feeling here - our Metric interface looks good to me actually.

[xukai92]

sqrt seems like a cheap operation (see wiki), maybe we don't need to cache it?

If we remove that field then the interface is a little bit inconsistent?

I don't have a strong feeling here - our Metric interface looks good to me actually.
I agree. I cleaned it up a bit. There was some unused functions that makes the code look worse than it should be.

I prefer to not change it :D

[xukai92]

Shall we close this issue then.