Add ComponentArrays extension (#407)

* Add ComponentArrays extension

* Change to more general arrays

* Fix stale compat

* More general similar

* Add proper errors and correct test cases

* Fix errors in mass matrix adaption

* Remove PhasePoint fix

* Correct use ComponentMatrix

* Dont use rand matrix in dense metric

* Apply reviews

* Relax type strict

* Ensure axes matching of general arrays

* Only focus on axes

* Clean up

* Clean extension

* More informative error message

---------

Co-authored-by: Hong Ge <3279477+yebai@users.noreply.github.com>

Author: ErikQQY

Project.toml

@@ -18,12 +18,14 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [weakdeps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 
 [extensions]
 AdvancedHMCADTypesExt = "ADTypes"
+AdvancedHMCComponentArraysExt = "ComponentArrays"
 AdvancedHMCCUDAExt = "CUDA"
 AdvancedHMCMCMCChainsExt = "MCMCChains"
 AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
@@ -32,6 +34,7 @@ AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
 ADTypes = "1"
 AbstractMCMC = "5.6"
 ArgCheck = "1, 2"
+ComponentArrays = "0.15"
 CUDA = "3, 4, 5"
 DocStringExtensions = "0.8, 0.9"
 LinearAlgebra = "<0.1, 1"

ext/AdvancedHMCComponentArraysExt.jl

@@ -0,0 +1,8 @@
+module AdvancedHMCComponentArraysExt
+
+using AdvancedHMC: AdvancedHMC, __axes
+using ComponentArrays: ComponentVecOrMat, getaxes
+
+AdvancedHMC.__axes(r::ComponentVecOrMat) = getaxes(r)
+
+end # module

src/adaptation/massmatrix.jl

@@ -78,15 +78,16 @@ function get_estimation(nv::NaiveVar)
 end
 
 # Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
-mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}
+mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVecOrMat{T}} <:
+               DiagMatrixEstimator{T}
     n::Int
     n_min::Int
     μ::E
     M::E
     δ::E    # cache for diff
-    var::E    # cache for variance
-    function WelfordVar(n::Int, n_min::Int, μ::E, M::E, δ::E, var::E) where {E}
-        return new{eltype(E),E}(n, n_min, μ, M, δ, var)
+    var::V    # cache for variance
+    function WelfordVar(n::Int, n_min::Int, μ::E, M::E, δ::E, var::V) where {E,V}
+        return new{eltype(E),E,V}(n, n_min, μ, M, δ, var)
     end
 end
 
@@ -182,13 +183,13 @@ function get_estimation(nc::NaiveCov)
 end
 
 # Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
-mutable struct WelfordCov{F<:AbstractFloat} <: DenseMatrixEstimator{F}
+mutable struct WelfordCov{F<:AbstractFloat,C<:AbstractMatrix{F}} <: DenseMatrixEstimator{F}
     n::Int
     n_min::Int
     μ::Vector{F}
     M::Matrix{F}
     δ::Vector{F}  # cache for diff
-    cov::Matrix{F}
+    cov::C
 end
 
 Base.show(io::IO, ::WelfordCov) = print(io, "WelfordCov")

src/hamiltonian.jl

@@ -43,10 +43,22 @@ end
 
 ∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat) = copy(r)
 function ∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat)
-    return h.metric.M⁻¹ .* r
+    (; M⁻¹) = h.metric
+    axes_M⁻¹ = __axes(M⁻¹)
+    axes_r = __axes(r)
+    (first(axes_M⁻¹) !== first(axes_r)) && throw(
+        ArgumentError("AxesMismatch: M⁻¹ has axes $(axes_M⁻¹) but r has axes $(axes_r)")
+    )
+    return M⁻¹ .* r
 end
 function ∂H∂r(h::Hamiltonian{<:DenseEuclideanMetric,<:GaussianKinetic}, r::AbstractVecOrMat)
-    return h.metric.M⁻¹ * r
+    (; M⁻¹) = h.metric
+    axes_M⁻¹ = __axes(M⁻¹)
+    axes_r = __axes(r)
+    (last(axes_M⁻¹) !== first(axes_r)) && throw(
+        ArgumentError("AxesMismatch: M⁻¹ has axes $(axes_M⁻¹) but r has axes $(axes_r)")
+    )
+    return M⁻¹ * r
 end
 
 # TODO (kai) make the order of θ and r consistent with neg_energy

src/utilities.jl

@@ -22,6 +22,13 @@ function _randn(
     return out
 end
 
+"""
+    __axes(r::AbstractVecOrMat)
+
+Return the axes of input `r`, where `r` can be `AbstractArrays`, `ComponentArrays` or other custom arrays.
+"""
+@inline __axes(r::AbstractVecOrMat) = axes(r)
+
 """ 
 `rand_coupled` produces coupled randomness given a vector of RNGs. For example, 
 when a vector of RNGs is provided, `rand_coupled` peforms a single `rand` call 

test/demo.jl

@@ -71,39 +71,84 @@ end
 
     ℓπ = DemoProblemComponentArrays()
 
-    # Define a Hamiltonian system
-    D = length(p1)          # number of parameters
-    metric = DiagEuclideanMetric(D)
+    @testset "Test Diagonal ComponentArray metric" begin
 
-    # choose AD framework or provide a function manually
-    hamiltonian = Hamiltonian(metric, ℓπ, Val(:ForwardDiff); x=p1)
+        # Define a Hamiltonian system
+        M⁻¹ = ComponentArray(; μ=1.0, σ=1.0)
+        metric = DiagEuclideanMetric(M⁻¹)
 
-    # Define a leapfrog solver, with initial step size chosen heuristically
-    initial_ϵ = find_good_stepsize(hamiltonian, p1)
-    integrator = Leapfrog(initial_ϵ)
+        # choose AD framework or provide a function manually
+        hamiltonian = Hamiltonian(metric, ℓπ, Val(:ForwardDiff))
 
-    # Define an HMC sampler, with the following components
-    proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
-    adaptor = StanHMCAdaptor(MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator))
+        # Define a leapfrog solver, with initial step size chosen heuristically
+        initial_ϵ = find_good_stepsize(hamiltonian, p1)
+        integrator = Leapfrog(initial_ϵ)
 
-    # -- run sampler
-    n_samples, n_adapts = 100, 50
-    samples, stats = sample(
-        hamiltonian,
-        proposal,
-        p1,
-        n_samples,
-        adaptor,
-        n_adapts;
-        progress=false,
-        verbose=false,
-    )
+        # Define an HMC sampler, with the following components
+        proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+        adaptor = StanHMCAdaptor(
+            MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)
+        )
 
-    @test length(samples) == n_samples
-    @test length(stats) == n_samples
-    labels = ComponentArrays.labels(samples[1])
-    @test "μ" ∈ labels
-    @test "σ" ∈ labels
+        # -- run sampler
+        n_samples, n_adapts = 100, 50
+        samples, stats = sample(
+            hamiltonian,
+            proposal,
+            p1,
+            n_samples,
+            adaptor,
+            n_adapts;
+            progress=false,
+            verbose=false,
+        )
+
+        @test length(samples) == n_samples
+        @test length(stats) == n_samples
+        lab = ComponentArrays.labels(samples[1])
+        @test "μ" ∈ lab
+        @test "σ" ∈ lab
+    end
+
+    @testset "Test Dense ComponentArray metric" begin
+
+        # Define a Hamiltonian system
+        ax = getaxes(p1)[1]
+        M⁻¹ = ComponentArray([2.0 1.0; 1.0 2.0], ax, ax)
+        metric = DenseEuclideanMetric(M⁻¹)
+
+        # choose AD framework or provide a function manually
+        hamiltonian = Hamiltonian(metric, ℓπ, Val(:ForwardDiff))
+
+        # Define a leapfrog solver, with initial step size chosen heuristically
+        initial_ϵ = find_good_stepsize(hamiltonian, p1)
+        integrator = Leapfrog(initial_ϵ)
+
+        # Define an HMC sampler, with the following components
+        proposal = HMCKernel(Trajectory{MultinomialTS}(integrator, GeneralisedNoUTurn()))
+        adaptor = StanHMCAdaptor(
+            MassMatrixAdaptor(metric), StepSizeAdaptor(0.8, integrator)
+        )
+
+        # -- run sampler
+        n_samples, n_adapts = 100, 50
+        samples, stats = sample(
+            hamiltonian,
+            proposal,
+            p1,
+            n_samples,
+            adaptor,
+            n_adapts;
+            progress=false,
+            verbose=false,
+        )
+
+        @test length(samples) == n_samples
+        @test length(stats) == n_samples
+        lab = ComponentArrays.labels(samples[1])
+        @test "μ" ∈ lab
+        @test "σ" ∈ lab
+    end
 end
 
 @testset "ADTypes" begin

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)
@@ -76,3 +77,35 @@ end
         end
     end
 end
+
+@testset "Energy with ComponentArrays" begin
+    n_tests = 10
+    for T in [Float32, Float64]
+        for _ in 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⁻¹ = ComponentArray(;
+                a=ones(T, D) + abs.(randn(T, D)), b=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, 2 * D, 2 * D)
+            ax = getaxes(r_init)[1]
+            M⁻¹ = ComponentArray(m' * m, ax, ax)
+            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