Skip to content

Commit 6a55a3f

Browse files
xukai92Kai XuKai Xuyebai
authored
Relativistic HMC (#302)
* feat: relativistic HMC * chore: move codes to experimental folder * chore: commit unsaved change * chore: warp experimental codes into a seperate module * chore: remove new the module * chore: create experimental tests and set up CI * Moved all experimental code including tests into research folder. (#304) * Moved all experimental code into research folder. * Minor changes in deps. * Improve tests for experimental code. * Change CI name. * Fix path for research includes. * More fixes to includes * Minor fixes. * Update research/src/relativistic_hmc.jl * Apply suggestions from code review * Create README.md * Update research/src/relativistic_hmc.jl * Update relativistic_hmc.jl Co-authored-by: Kai Xu <[email protected]> Co-authored-by: Kai Xu <[email protected]> Co-authored-by: Hong Ge <[email protected]>
1 parent 403c7e5 commit 6a55a3f

File tree

7 files changed

+177
-0
lines changed

7 files changed

+177
-0
lines changed

‎.github/workflows/ExperimentalTests.yml

+40
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,40 @@
1+
name: TestExperimentalFeatures
2+
3+
on:
4+
push:
5+
branches:
6+
- master
7+
pull_request:
8+
9+
jobs:
10+
test:
11+
runs-on: ${{ matrix.os }}
12+
strategy:
13+
matrix:
14+
version:
15+
- '1'
16+
os:
17+
- ubuntu-latest
18+
- macOS-latest
19+
- windows-latest
20+
arch:
21+
- x86
22+
- x64
23+
exclude:
24+
- os: ubuntu-latest
25+
arch: x86
26+
- os: macOS-latest
27+
arch: x86
28+
- os: windows-latest
29+
arch: x86
30+
steps:
31+
- uses: actions/checkout@v2
32+
- uses: julia-actions/setup-julia@v1
33+
with:
34+
version: ${{ matrix.version }}
35+
arch: ${{ matrix.arch }}
36+
- uses: julia-actions/julia-buildpkg@latest
37+
- name: Run integration tests
38+
uses: julia-actions/julia-runtest@latest
39+
env:
40+
AHMC_TEST_GROUP: Experimental

‎research/Project.toml

+28
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,28 @@
1+
[deps]
2+
AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
3+
AdaptiveRejectionSampling = "c75e803d-635f-53bd-ab7d-544e482d8c75"
4+
ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
5+
DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
6+
InplaceOps = "505f98c9-085e-5b2c-8e89-488be7bf1f34"
7+
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
8+
ProgressMeter = "92933f4c-e287-5a05-a399-4b506db050ca"
9+
Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
10+
Requires = "ae029012-a4dd-5104-9daa-d747884805df"
11+
Setfield = "efcf1570-3423-57d1-acb7-fd33fddbac46"
12+
Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
13+
StatsBase = "2913bbd2-ae8a-5f71-8c99-4fb6c76f3a91"
14+
StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
15+
UnPack = "3a884ed6-31ef-47d7-9d2a-63182c4928ed"
16+
17+
[compat]
18+
AbstractMCMC = "3.2, 4"
19+
ArgCheck = "1, 2"
20+
DocStringExtensions = "0.8, 0.9"
21+
InplaceOps = "0.3"
22+
ProgressMeter = "1"
23+
Requires = "0.5, 1"
24+
Setfield = "0.7, 0.8, 1"
25+
StatsBase = "0.31, 0.32, 0.33"
26+
StatsFuns = "0.8, 0.9, 1"
27+
UnPack = "1"
28+
julia = "1"

‎research/README.md

+8
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
In order to use algorithms in this folder, please navigate to the AdvancedHMC folder and run
2+
3+
4+
```
5+
] activate research/
6+
] develop src/
7+
] instantiate
8+
```

‎research/src/relativistic_hmc.jl

+70
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,70 @@
1+
using AdvancedHMC
2+
import AdvancedHMC: ∂H∂r, neg_energy, AbstractKinetic
3+
import Random: AbstractRNG
4+
5+
struct RelativisticKinetic{T} <: AbstractKinetic
6+
"Mass"
7+
m::T
8+
"Speed of light"
9+
c::T
10+
end
11+
12+
13+
function ∂H∂r(h::Hamiltonian{<:UnitEuclideanMetric,<:RelativisticKinetic}, r::AbstractVecOrMat)
14+
mass = h.kinetic.m .* sqrt.(r.^2 ./ (h.kinetic.m.^2 * h.kinetic.c.^2) .+ 1)
15+
return r ./ mass
16+
end
17+
function ∂H∂r(h::Hamiltonian{<:DiagEuclideanMetric,<:RelativisticKinetic}, r::AbstractVecOrMat)
18+
r = h.metric.sqrtM⁻¹ .* r
19+
mass = h.kinetic.m .* sqrt.(r.^2 ./ (h.kinetic.m.^2 * h.kinetic.c.^2) .+ 1)
20+
retval = r ./ mass # red part of (15)
21+
return h.metric.sqrtM⁻¹ .* retval # (15)
22+
end
23+
24+
25+
function neg_energy(
26+
h::Hamiltonian{<:UnitEuclideanMetric,<:RelativisticKinetic},
27+
r::T,
28+
θ::T
29+
) where {T<:AbstractVector}
30+
return -sum(h.kinetic.m .* h.kinetic.c.^2 .* sqrt.(r.^2 ./ (h.kinetic.m.^2 .* h.kinetic.c.^2) .+ 1))
31+
end
32+
33+
function neg_energy(
34+
h::Hamiltonian{<:DiagEuclideanMetric,<:RelativisticKinetic},
35+
r::T,
36+
θ::T
37+
) where {T<:AbstractVector}
38+
r = h.metric.sqrtM⁻¹ .* r
39+
return -sum(h.kinetic.m .* h.kinetic.c.^2 .* sqrt.(r.^2 ./ (h.kinetic.m.^2 .* h.kinetic.c.^2) .+ 1))
40+
end
41+
42+
43+
using AdaptiveRejectionSampling: RejectionSampler, run_sampler!
44+
45+
# TODO Support AbstractVector{<:AbstractRNG}
46+
function _rand(
47+
rng::AbstractRNG,
48+
metric::UnitEuclideanMetric{T},
49+
kinetic::RelativisticKinetic{T},
50+
) where {T}
51+
h_temp = Hamiltonian(metric, kinetic, identity, identity)
52+
densityfunc = x -> exp(neg_energy(h_temp, [x], [x]))
53+
sampler = RejectionSampler(densityfunc, (-Inf, Inf); max_segments=5)
54+
sz = size(metric)
55+
r = run_sampler!(rng, sampler, prod(sz))
56+
r = reshape(r, sz)
57+
return r
58+
end
59+
60+
# TODO Support AbstractVector{<:AbstractRNG}
61+
function _rand(
62+
rng::AbstractRNG,
63+
metric::DiagEuclideanMetric{T},
64+
kinetic::RelativisticKinetic{T},
65+
) where {T}
66+
r = _rand(rng, UnitEuclideanMetric(size(metric)), kinetic)
67+
# p' = A p where A = sqrtM
68+
r ./= metric.sqrtM⁻¹
69+
return r
70+
end

‎research/tests/relativistic_hmc.jl

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
using ReTest, AdvancedHMC
2+
using AdvancedHMC
3+
using LinearAlgebra: dot
4+
5+
@testset "Hamiltonian" begin
6+
f = x -> dot(x, x)
7+
g = x -> 2x
8+
metric = UnitEuclideanMetric(10)
9+
h = Hamiltonian(metric, RelativisticKinetic(1.0, 1.0), f, g)
10+
@test h.kinetic isa RelativisticKinetic
11+
end

‎research/tests/runtests.jl

+11
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,11 @@
1+
using ReTest
2+
3+
# include the source code for relativistic HMC
4+
include("../src/relativistic_hmc.jl")
5+
6+
# include the tests for relativistic HMC
7+
include("relativistic_hmc.jl")
8+
9+
@main function runtests(patterns...; dry::Bool=false)
10+
retest(patterns...; dry=dry, verbose=Inf)
11+
end

‎test/runtests.jl

+9
Original file line numberDiff line numberDiff line change
@@ -33,6 +33,15 @@ if GROUP == "All" || GROUP == "AdvancedHMC"
3333
end
3434
end
3535

36+
if GROUP == "All" || GROUP == "Experimental"
37+
using Pkg
38+
# activate separate test environment
39+
Pkg.activate(joinpath(DIRECTORY_AdvancedHMC, "research"))
40+
Pkg.develop(PackageSpec(; path=DIRECTORY_AdvancedHMC))
41+
Pkg.instantiate()
42+
include(joinpath(DIRECTORY_AdvancedHMC, "research/tests", "runtests.jl"))
43+
end
44+
3645
if GROUP == "All" || GROUP == "Downstream"
3746
using Pkg
3847
try

0 commit comments

Comments
 (0)