Move AdvancedHMC support to an extension

.github/workflows/IntegrationTests.yml

@@ -44,3 +44,4 @@ jobs:
       - uses: codecov/codecov-action@v5
         with:
           files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}

Project.toml

@@ -1,10 +1,11 @@
 name = "Pathfinder"
 uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 authors = ["Seth Axen <[email protected]> and contributors"]
-version = "0.9.17"
+version = "0.10.0-DEV"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Folds = "41a02a25-b8f0-4f67-bc48-60067656b558"
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
@@ -18,26 +19,28 @@ PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 PSIS = "ce719bf2-d5d0-4fb9-925d-10a81b42ad04"
 ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
-Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 Transducers = "28d57a85-8fef-5791-bfe6-a80928e7c999"
 
 [weakdeps]
 Accessors = "7d9f7c33-5ae7-4f3b-8dc6-eff91059b697"
+AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 Turing = "fce5fe82-541a-59a6-adf8-730c64b5f9a0"
 
 [extensions]
+PathfinderAdvancedHMCExt = "AdvancedHMC"
 PathfinderDynamicHMCExt = "DynamicHMC"
 PathfinderTuringExt = ["Accessors", "DynamicPPL", "MCMCChains", "Turing"]
 
 [compat]
 ADTypes = "0.2.5, 1"
 Accessors = "0.1.12"
+AdvancedHMC = "0.6"
 Distributions = "0.25.87"
 DynamicHMC = "3.4.0"
 DynamicPPL = "0.25.2, 0.27, 0.28, 0.29, 0.30, 0.31, 0.32, 0.33, 0.34, 0.35"
@@ -55,7 +58,6 @@ PDMats = "0.11.35"
 PSIS = "0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9"
 ProgressLogging = "0.1.4"
 Random = "1"
-Requires = "1.1"
 ReverseDiff = "1.15"
 SciMLBase = "2.30"
 Statistics = "1"

docs/make.jl

@@ -1,8 +1,11 @@
+using AdvancedHMC  # triggering loading of PathfinderAdvancedHMCExt
 using Pathfinder
 using Documenter
 using DocumenterCitations
 using DocumenterInterLinks
 
+PathfinderAdvancedHMCExt = Base.get_extension(Pathfinder, :PathfinderAdvancedHMCExt)
+
 DocMeta.setdocmeta!(Pathfinder, :DocTestSetup, :(using Pathfinder); recursive=true)
 
 bib = CitationBibliography(joinpath(@__DIR__, "src", "references.bib"); style=:numeric)
@@ -37,7 +40,7 @@ links = InterLinks(
 )
 
 makedocs(;
-    modules=[Pathfinder],
+    modules=[Pathfinder, PathfinderAdvancedHMCExt],
     authors="Seth Axen <[email protected]> and contributors",
     repo=Remotes.GitHub("mlcolab", "Pathfinder.jl"),
     sitename="Pathfinder.jl",

docs/src/examples/initializing-hmc.md

@@ -209,11 +209,12 @@ samples_ahmc2, stats_ahmc2 = sample(
 
 ### Use Pathfinder's metric estimate for sampling
 
-To use Pathfinder's metric with no metric adaptation, we need to use Pathfinder's own [`Pathfinder.RankUpdateEuclideanMetric`](@ref) type, which just wraps our inverse metric estimate for use with AdvancedHMC:
+To use Pathfinder's metric with no metric adaptation, we need to use Pathfinder's own [`RankUpdateEuclideanMetric`](@ref PathfinderAdvancedHMCExt.RankUpdateEuclideanMetric) type, which just wraps our inverse metric estimate for use with AdvancedHMC.
+We can construct this metric by calling `AdvancedHMC.AbstractMetric`:
 
 ```@example 1
 nadapts = 75
-metric = Pathfinder.RankUpdateEuclideanMetric(inv_metric)
+metric = AdvancedHMC.AbstractMetric(inv_metric)
 hamiltonian = Hamiltonian(metric, ∇P)
 ϵ = find_good_stepsize(hamiltonian, init_params)
 integrator = Leapfrog(ϵ)

docs/src/examples/turing.md

@@ -56,7 +56,7 @@ We can use Pathfinder's estimate of the metric and only perform enough warm-up t
 
 ```@example 1
 inv_metric = result_multi.pathfinder_results[1].fit_distribution.Σ
-metric = Pathfinder.RankUpdateEuclideanMetric(inv_metric)
+metric = AdvancedHMC.AbstractMetric(inv_metric)
 kernel = HMCKernel(Trajectory{MultinomialTS}(Leapfrog(0.0), GeneralisedNoUTurn()))
 adaptor = StepSizeAdaptor(0.8, 1.0)  # adapt only the step size
 nuts = AdvancedHMC.HMCSampler(kernel, metric, adaptor)

docs/src/lib/public.md

@@ -27,3 +27,11 @@ Order = [:function, :type]
 Public = true
 Private = false
 ```
+
+```@meta
+CurrentModule = Base.get_extension(Pathfinder, :PathfinderAdvancedHMCExt)
+```
+
+```@docs
+RankUpdateEuclideanMetric
+```

ext/PathfinderAdvancedHMCExt.jl

@@ -0,0 +1,83 @@
+module PathfinderAdvancedHMCExt
+
+using AdvancedHMC: AdvancedHMC
+using Pathfinder: Pathfinder
+using LinearAlgebra: Diagonal, diag
+using PDMats: PDMats
+using Random: Random
+
+"""
+    RankUpdateEuclideanMetric{T,M} <: AdvancedHMC.AbstractMetric
+
+A Gaussian Euclidean metric (mass matrix) whose inverse is constructed by
+rank-updates.
+
+To construct this metric, call `AdvancedHMC.AbstractMetric` with a
+[`Pathfinder.WoodburyPDMat`](@ref) as the argument.
+
+# Example
+
+```jldoctest
+julia> using LinearAlgebra, Pathfinder, AdvancedHMC
+
+julia> W = Pathfinder.WoodburyPDMat(Diagonal([0.1, 0.2]), [0.7 0.2]', Diagonal([0.3]));
+
+julia> AdvancedHMC.AbstractMetric(W)
+RankUpdateEuclideanMetric(diag=[0.247, 0.21200000000000002])
+```
+
+See also: The AdvancedHMC [metric](@extref AdvancedHMC hamiltonian_mm) documentation.
+"""
+RankUpdateEuclideanMetric
+
+struct RankUpdateEuclideanMetric{T,M<:Pathfinder.WoodburyPDMat{T,<:Diagonal{T}}} <:
+       AdvancedHMC.AbstractMetric
+    M⁻¹::M
+end
+
+Base.@deprecate(
+    Pathfinder.RankUpdateEuclideanMetric(M⁻¹::Pathfinder.WoodburyPDMat),
+    AdvancedHMC.AbstractMetric(M⁻¹),
+    false,
+)
+
+AdvancedHMC.AbstractMetric(M⁻¹::Pathfinder.WoodburyPDMat) = RankUpdateEuclideanMetric(M⁻¹)
+
+AdvancedHMC.renew(::RankUpdateEuclideanMetric, M⁻¹) = RankUpdateEuclideanMetric(M⁻¹)
+
+Base.size(metric::RankUpdateEuclideanMetric, dim...) = size(metric.M⁻¹.A.diag, dim...)
+
+Base.eltype(metric::RankUpdateEuclideanMetric) = eltype(metric.M⁻¹)
+
+function Base.show(io::IO, metric::RankUpdateEuclideanMetric)
+    print(io, "RankUpdateEuclideanMetric(diag=$(diag(metric.M⁻¹)))")
+    return nothing
+end
+
+function Base.rand(
+    rng::Random.AbstractRNG,
+    metric::RankUpdateEuclideanMetric{T},
+    ::AdvancedHMC.GaussianKinetic,
+) where {T}
+    M⁻¹ = metric.M⁻¹
+    r = Random.randn(rng, T, size(metric)...)
+    PDMats.invunwhiten!(r, M⁻¹, r)
+    return r
+end
+
+function AdvancedHMC.neg_energy(
+    h::AdvancedHMC.Hamiltonian{<:RankUpdateEuclideanMetric,<:AdvancedHMC.GaussianKinetic},
+    r::T,
+    θ::T,
+) where {T<:AbstractVecOrMat}
+    return -PDMats.quad(h.metric.M⁻¹, r) / 2
+end
+
+function AdvancedHMC.∂H∂r(
+    h::AdvancedHMC.Hamiltonian{<:RankUpdateEuclideanMetric,<:AdvancedHMC.GaussianKinetic},
+    r::AbstractVecOrMat,
+)
+    return h.metric.M⁻¹ * r
+end
+
+end  # module

src/Pathfinder.jl

@@ -13,7 +13,6 @@ using PDMats: PDMats
 using ProgressLogging: ProgressLogging
 using PSIS: PSIS
 using Random
-using Requires: Requires
 using SciMLBase: SciMLBase
 using Statistics: Statistics
 using StatsBase: StatsBase
@@ -39,6 +38,8 @@ end
 # We depend on Optim, and Optim depends on ForwardDiff, so we can offer it as a default.
 default_ad() = ADTypes.AutoForwardDiff()
 
+function RankUpdateEuclideanMetric end
+
 include("transducers.jl")
 include("woodbury.jl")
 include("optimize.jl")
@@ -49,10 +50,4 @@ include("resample.jl")
 include("singlepath.jl")
 include("multipath.jl")
 
-function __init__()
-    Requires.@require AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d" begin
-        include("integration/advancedhmc.jl")
-    end
-end
-
 end

test/integration/AdvancedHMC/runtests.jl

@@ -67,29 +67,15 @@ end
     ∂ℓπ∂θ(x) = -x
 
     @testset "RankUpdateEuclideanMetric" begin
-        metric = Pathfinder.RankUpdateEuclideanMetric(M⁻¹)
+        metric = @test_deprecated Pathfinder.RankUpdateEuclideanMetric(M⁻¹)
         @test metric.M⁻¹ === M⁻¹
+        @test metric === AdvancedHMC.AbstractMetric(M⁻¹)
         metric_dense = AdvancedHMC.DenseEuclideanMetric(Symmetric(Matrix(M⁻¹)))
         h = AdvancedHMC.Hamiltonian(metric, ℓπ, ∂ℓπ∂θ)
         h_dense = AdvancedHMC.Hamiltonian(metric_dense, ℓπ, ∂ℓπ∂θ)
         r = randn(5)
         θ = randn(5)
 
-        metric2 = Pathfinder.RankUpdateEuclideanMetric(3)
-        @test metric2.M⁻¹ ≈ I
-        @test size(metric2) == (3,)
-        @test size(metric2, 2) == 1
-        @test eltype(metric2) === Float64
-        metric2 = Pathfinder.RankUpdateEuclideanMetric((4,))
-        @test metric2.M⁻¹ ≈ I
-        @test size(metric2) == (4,)
-        @test size(metric2, 2) == 1
-        metric2 = Pathfinder.RankUpdateEuclideanMetric(Float32, (4,))
-        @test metric2.M⁻¹ ≈ I
-        @test size(metric2) == (4,)
-        @test size(metric2, 2) == 1
-        @test eltype(metric2.M⁻¹) === Float32
-
         @test size(metric) == (5,)
         @test size(metric, 2) == 1
         @test AdvancedHMC.renew(metric, M⁻¹2).M⁻¹ === M⁻¹2
@@ -178,7 +164,7 @@ end
         end
 
         @testset "Initial point and final metric" begin
-            metric = Pathfinder.RankUpdateEuclideanMetric(result_pf.fit_distribution.Σ)
+            metric = AdvancedHMC.AbstractMetric(result_pf.fit_distribution.Σ)
             hamiltonian = Hamiltonian(metric, ∇P)
             ϵ = find_good_stepsize(hamiltonian, θ₀)
             integrator = Leapfrog(ϵ)