general relativistic HMC #372
Conversation
Signed-off-by: Kai Xu <[email protected]>
Signed-off-by: Kai Xu <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Signed-off-by: Kai Xu <[email protected]>
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Signed-off-by: Kai Xu <[email protected]>
Signed-off-by: Kai Xu <[email protected]>
Signed-off-by: Kai Xu <[email protected]>
![github-actions[bot]](https://avatars.githubusercontent.com/in/15368?s=60&v=4){kind=small}
There was a problem hiding this comment.
Remaining comments which cannot be posted as a review comment to avoid GitHub Rate Limit
JuliaFormatter
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/research/src/sampler_utility.jl
Lines 273 to 279 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/research/test/relativistic.jl
Lines 44 to 49 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/research/test/riemannian.jl
Lines 20 to 23 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/research/test/riemannian.jl
Lines 46 to 49 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/research/test/riemannian.jl
Lines 53 to 56 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/src/riemannian/hamiltonian.jl
Line 311 in 62077db
[JuliaFormatter] reported by reviewdog 🐶
AdvancedHMC.jl/src/riemannian/metric.jl
Line 83 in 62077db
| const axes_grid1 = pyimport("mpl_toolkits.axes_grid1") | ||
| plt.style.use("bmh") | ||
| function subplots(args...; kwargs...) | ||
| retval = plt.subplots(args...; tight_layout=true, kwargs...) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| retval = plt.subplots(args...; tight_layout=true, kwargs...) | |
| retval = plt.subplots(args...; tight_layout = true, kwargs...) |
| return retval | ||
| end | ||
| function savefig(fig, fn; kwargs...) | ||
| fig.savefig(fn; bbox_inches="tight", kwargs...) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| fig.savefig(fn; bbox_inches="tight", kwargs...) | |
| fig.savefig(fn; bbox_inches = "tight", kwargs...) |
| isdistributed() = isdefined(Main, :Distributed) | ||
|
|
||
| "pmap if Distribued is used" | ||
| maybepmap(args...; kwargs...) = |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| maybepmap(args...; kwargs...) = | |
| maybepmap(args...; kwargs...) = |
| for dir_type in ("artifacts", "results", "log") | ||
| function_name = Symbol(dir_type * "dir") | ||
| @eval begin | ||
| $function_name(args...; suffix="") = projectdir($dir_type, args...) * (isempty(suffix) ? "" : "-$suffix") |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| $function_name(args...; suffix="") = projectdir($dir_type, args...) * (isempty(suffix) ? "" : "-$suffix") | |
| $function_name(args...; suffix = "") = | |
| projectdir($dir_type, args...) * (isempty(suffix) ? "" : "-$suffix") |
| return path | ||
| end | ||
|
|
||
| function make_Z(f, is=nothing, js=nothing; xlim=(-6, 2), ylim=(-2, 2)) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| function make_Z(f, is=nothing, js=nothing; xlim=(-6, 2), ylim=(-2, 2)) | |
| function make_Z(f, is = nothing, js = nothing; xlim = (-6, 2), ylim = (-2, 2)) |
|
|
||
| kinetic = get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| kinetic = get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps | |
| kinetic = | |
| get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps |
|
|
||
| kinetic = get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps | ||
| hps.kinetic == :relativistic ? RelativisticKinetic(hps.m, hps.c) : | ||
| hps.kinetic == :dimwise_relativistic ? DimensionwiseRelativisticKinetic(hps.m, hps.c) : |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| hps.kinetic == :dimwise_relativistic ? DimensionwiseRelativisticKinetic(hps.m, hps.c) : | |
| hps.kinetic == :dimwise_relativistic ? | |
| DimensionwiseRelativisticKinetic(hps.m, hps.c) : |
| return (; rng, target, hamiltonian, proposal, θ₀) | ||
| end | ||
|
|
||
| function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool=false) |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool=false) | |
| function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool = false) |
|
|
||
| z₀ = initial_momentum isa AbstractVector ? phasepoint(hamiltonian, θ₀, initial_momentum) : |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| z₀ = initial_momentum isa AbstractVector ? phasepoint(hamiltonian, θ₀, initial_momentum) : | |
| z₀ = | |
| initial_momentum isa AbstractVector ? | |
| phasepoint(hamiltonian, θ₀, initial_momentum) : |
| initial_momentum == :default ? phasepoint(rng, θ₀, hamiltonian) : | ||
| # NOTE The two options below ensures intiial momentum is identical | ||
| initial_momentum == :simple ? phasepoint(hamiltonian, θ₀, randn(rng, size(θ₀)...)) : | ||
| initial_momentum == :gaussian_kinetic ? phasepoint(hamiltonian, θ₀, rand(rng, hamiltonian.metric, GaussianKinetic(), θ₀)) : |
There was a problem hiding this comment.
[JuliaFormatter] reported by reviewdog 🐶
| initial_momentum == :gaussian_kinetic ? phasepoint(hamiltonian, θ₀, rand(rng, hamiltonian.metric, GaussianKinetic(), θ₀)) : | |
| initial_momentum == :gaussian_kinetic ? | |
| phasepoint(hamiltonian, θ₀, rand(rng, hamiltonian.metric, GaussianKinetic(), θ₀)) : |
| ArgCheck = "1, 2" | ||
| CUDA = "3, 4, 5" | ||
| DocStringExtensions = "0.8, 0.9" | ||
| InplaceOps = "0.3" | ||
| LinearAlgebra = "1.6" |
There was a problem hiding this comment.
Duplicate with L53.
| LinearAlgebra = "1.6" |
| @@ -4,6 +4,7 @@ version = "0.6.2" | |||
|
|
|||
| [deps] | |||
| AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" | |||
| AdaptiveRejectionSampling = "c75e803d-635f-53bd-ab7d-544e482d8c75" | |||
There was a problem hiding this comment.
A minor issue: Do we want to depend on ARS? Is it possible to turn this into a weak deps?
| include("relativistic/hamiltonian.jl") | ||
| export RelativisticKinetic, DimensionwiseRelativisticKinetic | ||
|
|
||
| using AdaptiveRejectionSampling: RejectionSampler, run_sampler! |
There was a problem hiding this comment.
Consider turning ARS into an AHMCARSext extension so ARS is a weak dep. See also the above comment about ARS.
| relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) = | ||
| sum(kinetic.c^2 * relativistic_mass(kinetic, r, r′)) | ||
|
|
||
| struct DimensionwiseRelativisticKinetic{T} <: AbstractRelativisticKinetic{T} |
There was a problem hiding this comment.
Consider adding a ref to:
Lu, X., Perrone, V., Hasenclever, L., Teh, Y., & Vollmer, S. (2016). Relativistic Monte Carlo. International Conference on Artificial Intelligence and Statistics, 54, 1236–1245.
|
|
||
| relativistic_mass(kinetic::RelativisticKinetic, r, r′ = r) = | ||
| kinetic.m * sqrt(dot(r, r′) / (kinetic.m^2 * kinetic.c^2) + 1) | ||
| relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) = |
There was a problem hiding this comment.
| relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) = | |
| # Negative kinetic energy | |
| relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) = |
There was a problem hiding this comment.
@xukai92 I added some further comments. Please take a look.
A high-level comment:
- Perhaps we should add a folder
general_relavistic(orriemannian_relativistic), and move all general relavistic code there:general_relavistic/hamiltonian.jlandgeneral_relavistic/metric.jl.
|
|
||
| relativistic_mass(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = | ||
| kinetic.m .* sqrt.(r .* r′ ./ (kinetic.m .^ 2 .* kinetic.c .^ 2) .+ 1) | ||
| relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = |
There was a problem hiding this comment.
| relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = | |
| # Negative kinetic energy | |
| # Eq (5) of Lu et al. (2016) | |
| relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) = |
| return -relativistic_energy(h.kinetic, r, h.metric._temp) - logZ_partial | ||
| end | ||
|
|
||
| # QUES L31 of hamiltonian.jl now reads a bit weird (semantically) |
There was a problem hiding this comment.
I don't follow this comment. Consider removing or clarifying.
| # Gr = G \ r | ||
| # ∂ℓπ∂θ[i] - 1 / 2 * tr(G \ ∂G∂θᵢ) + 1 / 2 * Gr' * ∂G∂θᵢ * Gr | ||
| # 1 / 2 * tr(invG * ∂G∂θᵢ) | ||
| # 1 / 2 * r' * invG * ∂G∂θᵢ * invG * r |
There was a problem hiding this comment.
| # Gr = G \ r | |
| # ∂ℓπ∂θ[i] - 1 / 2 * tr(G \ ∂G∂θᵢ) + 1 / 2 * Gr' * ∂G∂θᵢ * Gr | |
| # 1 / 2 * tr(invG * ∂G∂θᵢ) | |
| # 1 / 2 * r' * invG * ∂G∂θᵢ * invG * r |
| #! (18) of Overleaf note | ||
| mass = relativistic_mass(h.kinetic, r, M \ r) | ||
|
|
||
| # TODO convert the code below to a test case |
There was a problem hiding this comment.
Let's convert this to a test in this PR.
| return return_cache ? (dv, (; ℓπ, ∂ℓπ∂θ, ∂H∂θ, Q, softabsλ, J, term_1_cached, M)) : dv | ||
| end | ||
|
|
||
| # FIXME This implementation for dimensionwise is incorrect |
| return J | ||
| end | ||
|
|
||
| # TODO Add stricter types to block hamiltonian.jl#L37 from working on unknown metric/kinetic |
There was a problem hiding this comment.
This comment is likely fragile due to the concrete reference to line number.
Consider fixing it or use a better reference (e.g. refer to a function name)
No description provided.