general relativistic HMC#372

Open

xukai92 wants to merge 10 commits intomainfrom

kx/general-relativistic

Member

xukai92 commented Jul 25, 2024

No description provided.

xukai92 and others added 3 commits

July 25, 2024 12:45


          feat: port relativistic hmc from research repo

7fc04d2

Signed-off-by: Kai Xu <xuk@ibm.com>


          feat(wip): port reemannian hmc from research repo

3813ba0

Signed-off-by: Kai Xu <xuk@ibm.com>


          Update src/AdvancedHMC.jl

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

yebai reviewed

View reviewed changes

src/AdvancedHMC.jl Show resolved Hide resolved

yebai and others added 7 commits

July 26, 2024 21:00


          Update AdvancedHMC.jl

6b023f7


          feat(wip): port reemannian hmc from research repo

3bbb17f

Signed-off-by: Kai Xu <xuk@ibm.com>


          Apply suggestions from code review

cdef753

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>


          chore: remove experimental notebook

4ff42cf

Signed-off-by: Kai Xu <xuk@ibm.com>


          feat: port rest of the code from research repo

1acbf8f

Signed-off-by: Kai Xu <xuk@ibm.com>


          feat: port rest of the code from research repo

f446f0e

Signed-off-by: Kai Xu <xuk@ibm.com>


          Merge branch 'master' into kx/general-relativistic

62077db

xukai92 marked this pull request as ready for review

October 26, 2024 16:07

xukai92 requested a review from yebai

October 26, 2024 16:07

github-actions bot reviewed

View reviewed changes

Contributor

github-actions bot left a 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

    
           z_lst = step( 
        
               proposal.τ.integrator, 
        
               hamiltonian, 
        
               z₀, 
        
               hps.L; 
        
               full_trajectory = Val(true), 
        
           )

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/relativistic.jl

Lines 44 to 49 in 62077db

    
               UnitEuclideanMetric(D), 
        
               DiagEuclideanMetric(D), 
        
               DiagEuclideanMetric([0.5, 4.0]), 
        
               DenseEuclideanMetric(D), 
        
               DenseEuclideanMetric([1.32 0.79; 0.79 0.48]), 
        
           ]

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/relativistic.jl

Line 51 in 62077db

RelativisticKinetic(hps.m, hps.c),

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/relativistic.jl

Line 69 in 62077db

finite_difference_gradient(Hamifuncr, r),

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Lines 20 to 23 in 62077db

    
           @testset "$(nameof(typeof(target)))" for target in [ 
        
               HighDimGaussian(2),  
        
               Funnel(), 
        
           ]

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 43 in 62077db

    
           @test δ(reshape_∂G∂θ(finite_difference_jacobian(Gfunc, x)), ∂G∂θfunc(x)) < TARGET_TOL

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Lines 46 to 49 in 62077db

    
           @testset "$(nameof(typeof(hessmap)))" for hessmap in [ 
        
                   IdentityMap(),  
        
                   SoftAbsMap(hps.α), 
        
               ]

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Lines 53 to 56 in 62077db

    
               GaussianKinetic(),  
        
               RelativisticKinetic(hps.m, hps.c),  
        
               DimensionwiseRelativisticKinetic(hps.m, hps.c), 
        
           ]

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 60 in 62077db

    
           (hessmap isa SoftAbsMap || all(iszero.(x))) # that of IdentityMap can be non-PD for x==0 that I know it's PD

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 63 in 62077db

@test neg_energy(hamiltonian, r, x) ≈ logpdf(MvNormal(zeros(D), Σ), r)

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 70 in 62077db

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 80 in 62077db

) < HAMILTONIAN_TOL broken=broken

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/research/test/riemannian.jl

Line 86 in 62077db

finite_difference_gradient(Hamifuncr, r),

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/src/relativistic/hamiltonian.jl

Line 75 in 62077db

end

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/src/riemannian/hamiltonian.jl

Line 311 in 62077db

end

[JuliaFormatter] _{reported by reviewdog 🐶}

AdvancedHMC.jl/src/riemannian/metric.jl

Line 83 in 62077db

) = _rand(rng, metric, kinetic, θ)

research/src/experiment_utility.jl

+              const axes_grid1 = pyimport("mpl_toolkits.axes_grid1")
+              plt.style.use("bmh")
+              function subplots(args...; kwargs...)
+                  retval = plt.subplots(args...; tight_layout=true, kwargs...)

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                retval = plt.subplots(args...; tight_layout=true, kwargs...)
          
                retval = plt.subplots(args...; tight_layout = true, kwargs...)

research/src/experiment_utility.jl

+                  return retval
+              end
+              function savefig(fig, fn; kwargs...)
+                  fig.savefig(fn; bbox_inches="tight", kwargs...)

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                fig.savefig(fn; bbox_inches="tight", kwargs...)
          
                fig.savefig(fn; bbox_inches = "tight", kwargs...)

research/src/experiment_utility.jl

+              isdistributed() = isdefined(Main, :Distributed)
+              "pmap if Distribued is used"
+              maybepmap(args...; kwargs...) =

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
            maybepmap(args...; kwargs...) = 
          
            maybepmap(args...; kwargs...) =

research/src/experiment_utility.jl

+              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")

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                    $function_name(args...; suffix="") = projectdir($dir_type, args...) * (isempty(suffix) ? "" : "-$suffix")
          
                    $function_name(args...; suffix = "") =
          
                        projectdir($dir_type, args...) * (isempty(suffix) ? "" : "-$suffix")

research/src/experiment_utility.jl

+                  return path
+              end
+              function make_Z(f, is=nothing, js=nothing; xlim=(-6, 2), ylim=(-2, 2))

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
            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))

research/src/sampler_utility.jl

Comment on lines +216 to +217


		kinetic = get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                kinetic = get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps
          
                kinetic =
          
                    get(hps, :kinetic, :gaussian) == :gaussian ? GaussianKinetic() : # support missing kinetic in hps

research/src/sampler_utility.jl

+                  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) :

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                    hps.kinetic == :dimwise_relativistic ? DimensionwiseRelativisticKinetic(hps.m, hps.c) :
          
                    hps.kinetic == :dimwise_relativistic ?
          
                    DimensionwiseRelativisticKinetic(hps.m, hps.c) :

research/src/sampler_utility.jl

+                  return (; rng, target, hamiltonian, proposal, θ₀)
+              end
+              function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool=false)

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
            function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool=false)
          
            function sample_target(hps; rng = MersenneTwister(1110), output_only::Bool = false)

research/src/sampler_utility.jl

Comment on lines +266 to +267


		z₀ = initial_momentum isa AbstractVector ? phasepoint(hamiltonian, θ₀, initial_momentum) :

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                z₀ = initial_momentum isa AbstractVector ? phasepoint(hamiltonian, θ₀, initial_momentum) :
          
                z₀ =
          
                    initial_momentum isa AbstractVector ?
          
                    phasepoint(hamiltonian, θ₀, initial_momentum) :

research/src/sampler_utility.jl

+                      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(), θ₀)) :

Contributor

github-actions bot Oct 26, 2024

[JuliaFormatter] _{reported by reviewdog 🐶}

Suggested change

      
                    initial_momentum == :gaussian_kinetic ? phasepoint(hamiltonian, θ₀, rand(rng, hamiltonian.metric, GaussianKinetic(), θ₀)) :
          
                    initial_momentum == :gaussian_kinetic ?
          
                    phasepoint(hamiltonian, θ₀, rand(rng, hamiltonian.metric, GaussianKinetic(), θ₀)) :

yebai reviewed

View reviewed changes

Member

yebai left a comment

I left some comments above. Will continue later.

Project.toml

               CUDA = "3, 4, 5"
               DocStringExtensions = "0.8, 0.9"
               InplaceOps = "0.3"
+              LinearAlgebra = "1.6"

Member

yebai Oct 30, 2024

Duplicate with L53.

Suggested change

LinearAlgebra = "1.6"

Project.toml

               [deps]
               AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
+              AdaptiveRejectionSampling = "c75e803d-635f-53bd-ab7d-544e482d8c75"

Member

yebai Oct 30, 2024

A minor issue: Do we want to depend on ARS? Is it possible to turn this into a weak deps?

src/AdvancedHMC.jl

+              include("relativistic/hamiltonian.jl")
+              export RelativisticKinetic, DimensionwiseRelativisticKinetic
+              using AdaptiveRejectionSampling: RejectionSampler, run_sampler!

Member

yebai Oct 30, 2024

Consider turning ARS into an AHMCARSext extension so ARS is a weak dep. See also the above comment about ARS.

src/relativistic/hamiltonian.jl

+              relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) =
+                  sum(kinetic.c^2 * relativistic_mass(kinetic, r, r′))
+              struct DimensionwiseRelativisticKinetic{T} <: AbstractRelativisticKinetic{T}

Member

yebai Oct 30, 2024

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.

yebai reviewed

View reviewed changes

src/relativistic/hamiltonian.jl

+              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) =

Member

yebai Nov 1, 2024

Suggested change

      
            relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) =
          
            # Negative kinetic energy
          
            relativistic_energy(kinetic::RelativisticKinetic, r, r′ = r) =

yebai requested changes

View reviewed changes

Member

yebai left a comment

@xukai92 I added some further comments. Please take a look.

A high-level comment:

Perhaps we should add a folder general_relavistic (or riemannian_relativistic), and move all general relavistic code there: general_relavistic/hamiltonian.jl and general_relavistic/metric.jl.

src/relativistic/hamiltonian.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) =

Member

yebai Nov 1, 2024

Suggested change

      
            relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) =
          
            # Negative kinetic energy
          
            # Eq (5) of Lu et al. (2016)
          
            relativistic_energy(kinetic::DimensionwiseRelativisticKinetic, r, r′ = r) =

src/riemannian/hamiltonian.jl

+                  return -relativistic_energy(h.kinetic, r, h.metric._temp) - logZ_partial
+              end
+              # QUES L31 of hamiltonian.jl now reads a bit weird (semantically)

Member

yebai Nov 1, 2024

I don't follow this comment. Consider removing or clarifying.

src/riemannian/hamiltonian.jl

Comment on lines +61 to +64

+                          # 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

Member

yebai Nov 1, 2024

Suggested change

      
                        # 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

src/riemannian/hamiltonian.jl

+                  #! (18) of Overleaf note
+                  mass = relativistic_mass(h.kinetic, r, M \ r)
+                  # TODO convert the code below to a test case

Member

yebai Nov 1, 2024

Let's convert this to a test in this PR.

src/riemannian/hamiltonian.jl

+                  return return_cache ? (dv, (; ℓπ, ∂ℓπ∂θ, ∂H∂θ, Q, softabsλ, J, term_1_cached, M)) : dv
+              end
+              # FIXME This implementation for dimensionwise is incorrect

Member

yebai Nov 1, 2024

@torfjelde is this fixed now?

src/riemannian/hamiltonian.jl

+                  return J
+              end
+              # TODO Add stricter types to block hamiltonian.jl#L37 from working on unknown metric/kinetic

Member

yebai Nov 1, 2024

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)

yebai assigned yebai and ErikQQY

Contributor

nsiccha commented Jan 16, 2026

Leaving a comment to remember to benchmark my upcoming implementation against this one, @yebai.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet