Add DynamicHMC compatibility

.github/workflows/CI.yml

@@ -1,7 +1,17 @@
 name: CI
 on:
-  - push
-  - pull_request
+  push:
+    branches:
+      - main
+    tags: '*'
+  pull_request:
+  schedule:
+    - cron: "0 0 * * *"
+concurrency:
+  # Skip intermediate builds: always.
+  # Cancel intermediate builds: only if it is a pull request build.
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: ${{ startsWith(github.ref, 'refs/pull/') }}
 jobs:
   test:
     name: Julia ${{ matrix.version }} - ${{ matrix.os }} - ${{ matrix.arch }} - ${{ github.event_name }}
@@ -22,22 +32,13 @@ jobs:
         with:
           version: ${{ matrix.version }}
           arch: ${{ matrix.arch }}
-      - uses: actions/cache@v1
-        env:
-          cache-name: cache-artifacts
-        with:
-          path: ~/.julia/artifacts
-          key: ${{ runner.os }}-test-${{ env.cache-name }}-${{ hashFiles('**/Project.toml') }}
-          restore-keys: |
-            ${{ runner.os }}-test-${{ env.cache-name }}-
-            ${{ runner.os }}-test-
-            ${{ runner.os }}-
+      - uses: julia-actions/cache@v1
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
       - uses: julia-actions/julia-processcoverage@v1
-      - uses: codecov/codecov-action@v1
+      - uses: codecov/codecov-action@v2
         with:
-          file: lcov.info
+          files: lcov.info
   docs:
     name: Documentation
     runs-on: ubuntu-latest
@@ -46,20 +47,8 @@ jobs:
       - uses: julia-actions/setup-julia@v1
         with:
           version: '1'
-      - run: |
-          julia --project=docs -e '
-            using Pkg
-            Pkg.develop(PackageSpec(path=pwd()))
-            Pkg.instantiate()'
-      - run: |
-          julia --project=docs -e '
-            using Documenter: DocMeta, doctest
-            using Pathfinder
-            DocMeta.setdocmeta!(Pathfinder, :DocTestSetup, :(using Pathfinder); recursive=true)
-            doctest(Pathfinder)'
-      - run: julia --project=docs docs/make.jl
+      - uses: julia-actions/julia-buildpkg@v1
+      - uses: julia-actions/julia-docdeploy@v1
         env:
-          JULIA_PKG_SERVER: ""
           GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }}
           DOCUMENTER_KEY: ${{ secrets.DOCUMENTER_KEY }}
-          GKSwstype: "100" # https://discourse.julialang.org/t/generation-of-documentation-fails-qt-qpa-xcb-could-not-connect-to-display/60988

.github/workflows/IntegrationTests.yml

@@ -0,0 +1,37 @@
+name: IntegrationTest
+on:
+  push:
+    branches: [master]
+    tags: [v*]
+  pull_request:
+jobs:
+  test:
+    name: ${{ matrix.package }}
+    runs-on: ${{ matrix.os }}
+    strategy:
+      fail-fast: false
+      matrix:
+        version: [1]
+        os: [ubuntu-latest]
+        arch: [x64]
+        package:
+          - DynamicHMC
+    steps:
+      - uses: actions/checkout@v2
+      - uses: julia-actions/setup-julia@v1
+        with:
+          version: ${{ matrix.version }}
+          arch: ${{ matrix.arch }}
+      - uses: julia-actions/julia-buildpkg@v1
+      - run: |
+          julia --code-coverage=user -e '
+            using Pkg
+            test_path = joinpath(pwd(), "test", "integration", "${{ matrix.package }}")
+            Pkg.activate(test_path)
+            Pkg.develop(PackageSpec(path=pwd()))
+            Pkg.instantiate()
+            include(joinpath(test_path, "runtests.jl"))'
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v2
+        with:
+          files: lcov.info

Project.toml

@@ -1,7 +1,7 @@
 name = "Pathfinder"
 uuid = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
 authors = ["Seth Axen <[email protected]> and contributors"]
-version = "0.3.2"
+version = "0.3.3"
 
 [deps]
 AbstractDifferentiation = "c29ec348-61ec-40c8-8164-b8c60e9d9f3d"
@@ -13,6 +13,7 @@ Optim = "429524aa-4258-5aef-a3af-852621145aeb"
 PDMats = "90014a1f-27ba-587c-ab20-58faa44d9150"
 PSIS = "ce719bf2-d5d0-4fb9-925d-10a81b42ad04"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
 StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
@@ -25,6 +26,7 @@ GalacticOptim = "2"
 Optim = "1.4"
 PDMats = "0.11"
 PSIS = "0.2, 0.3"
+Requires = "1"
 StatsBase = "0.33"
 StatsFuns = "0.9"
 julia = "1.6"

src/Pathfinder.jl

@@ -10,6 +10,7 @@ using Optim: Optim, LineSearches
 using PDMats: PDMats
 using PSIS: PSIS
 using Random
+using Requires: Requires
 using Statistics: Statistics
 using StatsBase: StatsBase
 using StatsFuns: log2π
@@ -33,4 +34,10 @@ include("resample.jl")
 include("singlepath.jl")
 include("multipath.jl")
 
+function __init__()
+    Requires.@require DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb" begin
+        include("integration/dynamichmc.jl")
+    end
+end
+
 end

src/integration/dynamichmc.jl

@@ -0,0 +1,42 @@
+using .DynamicHMC: DynamicHMC
+
+function DynamicHMC.GaussianKineticEnergy(M⁻¹::WoodburyPDMat)
+    return DynamicHMC.GaussianKineticEnergy(M⁻¹, WoodburyLeftInvFactor(M⁻¹))
+end
+
+function DynamicHMC.kinetic_energy(
+    κ::DynamicHMC.GaussianKineticEnergy{<:WoodburyPDMat}, p, q=nothing
+)
+    return PDMats.quad(κ.M⁻¹, p) / 2
+end
+
+# utility object so we can use DynamicHMC.GaussianKineticEnergy
+struct WoodburyLeftInvFactor{T<:Real,TW<:PDMats.AbstractPDMat{T}} <: AbstractMatrix{T}
+    A::TW
+end
+
+Base.size(L::WoodburyLeftInvFactor, dims::Int...) = size(L.A, dims...)
+
+function LinearAlgebra.mul!(
+    r::StridedVecOrMat{T}, L::WoodburyLeftInvFactor{T}, x::StridedVecOrMat{T}
+) where {T}
+    return invunwhiten!(r, L.A, x)
+end
+
+function Base.Matrix(L::WoodburyLeftInvFactor)
+    W = L.A
+    n, m = size(W.B)
+    k = min(m, n)
+    Lmat = zeros(n, n)
+    Lmat[diagind(Lmat)] .= true
+    @views ldiv!(W.UC, Lmat[1:k, 1:k])
+    lmul!(W.Q, Lmat)
+    ldiv!(W.UA, Lmat)
+    return Lmat
+end
+
+function Base.AbstractMatrix{T}(L::WoodburyLeftInvFactor) where {T}
+    return WoodburyLeftInvFactor(AbstractMatrix{T}(L.A))
+end
+
+Base.getindex(L::WoodburyLeftInvFactor, inds...) = getindex(Matrix(L), inds...)

src/woodbury.jl

@@ -151,6 +151,21 @@ function PDMats.invquad!(r::AbstractArray, W::WoodburyPDMat, x::StridedMatrix{T}
     return r
 end
 
+function PDMats.quad!(r::AbstractArray, W::WoodburyPDMat, x::StridedMatrix{T}) where {T}
+    v = lmul!(W.Q', W.UA * x)
+    k = minimum(size(W.B))
+    @views lmul!(W.UC, v[1:k, :])
+    colwise_sumsq!(r, v)
+    return r
+end
+
+function PDMats.quad(W::WoodburyPDMat, x::AbstractVector{T}) where {T}
+    v = W.Q' * (W.UA * x)
+    n, m = size(W.B)
+    k = min(m, n)
+    return @views sum(abs2, W.UC * v[1:k]) + sum(abs2, v[(k + 1):n])
+end
+
 function PDMats.unwhiten!(
     r::StridedVecOrMat{T}, W::WoodburyPDMat, x::StridedVecOrMat{T}
 ) where {T}
@@ -162,6 +177,17 @@ function PDMats.unwhiten!(
     return r
 end
 
+function invunwhiten!(
+    r::StridedVecOrMat{T}, W::WoodburyPDMat, x::StridedVecOrMat{T}
+) where {T}
+    k = minimum(size(W.B))
+    copyto!(r, x)
+    @views ldiv!(W.UC, x isa AbstractVector ? r[1:k] : r[1:k, :])
+    lmul!(W.Q, r)
+    ldiv!(W.UA, r)
+    return r
+end
+
 # adapted from https://github.com/JuliaStats/PDMats.jl/blob/master/src/utils.jl
 function colwise_sumsq!(r::AbstractArray, a::AbstractMatrix)
     n = length(r)

test/integration/DynamicHMC/Project.toml

@@ -0,0 +1,21 @@
+[deps]
+Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
+DynamicHMC = "bbc10e6e-7c05-544b-b16e-64fede858acb"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
+LogDensityProblems = "6fdf6af0-433a-55f7-b3ed-c6c6e0b8df7c"
+Optim = "429524aa-4258-5aef-a3af-852621145aeb"
+Pathfinder = "b1d3bc72-d0e7-4279-b92f-7fa5d6d2d454"
+Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
+StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+TransformVariables = "84d833dd-6860-57f9-a1a7-6da5db126cff"
+
+[compat]
+Distributions = "0.25"
+DynamicHMC = "3"
+LogDensityProblems = "0.11"
+Optim = "1.4"
+Pathfinder = "0.3"
+StatsBase = "0.33"
+TransformVariables = "0.6"
+julia = "1.6"

test/integration/DynamicHMC/runtests.jl

@@ -0,0 +1,109 @@
+using Distributions,
+    DynamicHMC,
+    LinearAlgebra,
+    LogDensityProblems,
+    Optim,
+    Pathfinder,
+    Random,
+    StatsBase,
+    Test,
+    TransformVariables
+
+struct RegressionProblem{X,Y}
+    x::X
+    y::Y
+end
+
+function (prob::RegressionProblem)(θ)
+    σ = θ.σ
+    α = θ.α
+    β = θ.β
+    x = prob.x
+    y = prob.y
+    lp = logpdf(truncated(Normal(); lower=0), σ)
+    lp += logpdf(Normal(), α)
+    lp += sum(b -> logpdf(Normal(), b), β)
+    y_hat = muladd(x, β, α)
+    lp += sum(eachindex(y_hat, y)) do i
+        return loglikelihood(Normal(y_hat[i], σ), y[i])
+    end
+    return lp
+end
+
+@testset "DynamicHMC integration" begin
+    A = Diagonal(rand(5))
+    B = randn(5, 2)
+    D = exp(Symmetric(randn(2, 2)))
+    M⁻¹ = Pathfinder.WoodburyPDMat(A, B, D)
+
+    @testset "Pathfinder.WoodburyLeftInvFactor" begin
+        L = Pathfinder.WoodburyLeftInvFactor(M⁻¹)
+        v = randn(5)
+        V = randn(5, 2)
+        @test L.A === M⁻¹
+        @test size(L) === size(M⁻¹)
+        Lmat = @inferred Matrix(L)
+        @test Lmat isa Matrix
+        @test Lmat * Lmat' ≈ inv(M⁻¹)
+        @test L * v ≈ Lmat * v
+        @test L * V ≈ Lmat * V
+        @test AbstractMatrix{Float32}(L) isa Pathfinder.WoodburyLeftInvFactor{Float32}
+        for i in axes(L, 1), j in axes(L, 2)
+            @test L[i, j] ≈ Lmat[i, j]
+        end
+    end
+
+    @testset "DynamicHMC.GaussianKineticEnergy" begin
+        κ = DynamicHMC.GaussianKineticEnergy(M⁻¹)
+        @test κ.M⁻¹ === M⁻¹
+        @test κ.W === Pathfinder.WoodburyLeftInvFactor(M⁻¹)
+        κdense = DynamicHMC.GaussianKineticEnergy(Symmetric(Matrix(M⁻¹)))
+        κdense2 = DynamicHMC.GaussianKineticEnergy(Symmetric(Matrix(M⁻¹)), Matrix(κ.W))
+        p = randn(5)
+        q = randn(5)
+
+        @test size(κ) == size(κdense)
+        @test DynamicHMC.kinetic_energy(κ, p, q) ≈ DynamicHMC.kinetic_energy(κdense, p, q)
+        @test DynamicHMC.calculate_p♯(κ, p, q) ≈ DynamicHMC.calculate_p♯(κdense, p, q)
+        @test DynamicHMC.∇kinetic_energy(κ, p, q) ≈ DynamicHMC.∇kinetic_energy(κdense, p, q)
+        @test DynamicHMC.rand_p(MersenneTwister(42), κ, q) ≈
+            DynamicHMC.rand_p(MersenneTwister(42), κdense2, q)
+    end
+
+    @testset "DynamicHMC.mcmc_with_warmup" begin
+        ndraws = 1_000
+        rng = MersenneTwister(42)
+        x = 0:0.01:1
+        y = sin.(x) .+ randn.(rng) .* 0.2 .+ x
+        X = [x x .^ 2 x .^ 3]
+        prob = RegressionProblem(X, y)
+        trans = as((σ=asℝ₊, α=asℝ, β=as(Array, size(X, 2))))
+        P = TransformedLogDensity(trans, prob)
+        ∇P = ADgradient(:ForwardDiff, P)
+
+        result_hmc1 = mcmc_with_warmup(rng, ∇P, ndraws; reporter=NoProgressReport())
+
+        logp(x) = LogDensityProblems.logdensity(P, x)
+        ∇logp(x) = LogDensityProblems.logdensity_and_gradient(∇P, x)[2]
+        θ₀ = rand(rng, Uniform(-2, 2), LogDensityProblems.dimension(P))
+        result_pf = pathfinder(logp, ∇logp, θ₀, 1; rng)
+        result_hmc2 = mcmc_with_warmup(
+            rng,
+            ∇P,
+            ndraws;
+            initialization=(;
+                q=result_pf[2][:, 1], κ=GaussianKineticEnergy(result_pf[1].Σ)
+            ),
+            warmup_stages=default_warmup_stages(; middle_steps=0, doubling_stages=0),
+            reporter=NoProgressReport(),
+        )
+        # check that the posterior means are approximately equal
+        m1, v1 = mean_and_var(result_hmc1.chain)
+        m2, v2 = mean_and_var(result_hmc2.chain)
+        # NOTE: a more strict test would compute MCSE here
+        s = sqrt.(v1 .+ v2)
+        m = m1 - m2
+        bounds = quantile.(Normal.(0, s), 0.05)
+        @test all(bounds .< m .< -bounds)
+    end
+end