format

Author: ErikQQY

src/AdvancedHMC.jl

@@ -2,7 +2,18 @@ module AdvancedHMC
 
 using Statistics: mean, var, middle
 using LinearAlgebra:
-    Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, diagm, cholesky, UniformScaling, logdet, tr
+    Symmetric,
+    UpperTriangular,
+    mul!,
+    ldiv!,
+    dot,
+    I,
+    diag,
+    diagm,
+    cholesky,
+    UniformScaling,
+    logdet,
+    tr
 using StatsFuns: logaddexp, logsumexp, loghalf
 using Random: Random, AbstractRNG
 using ProgressMeter: ProgressMeter

src/riemannian/hamiltonian.jl

@@ -1,6 +1,8 @@
 #! Eq (14) of Girolami & Calderhead (2011)
 function ∂H∂r(
-    h::Hamiltonian{<:DenseRiemannianMetric,<:GaussianKinetic}, θ::AbstractVecOrMat, r::AbstractVecOrMat
+    h::Hamiltonian{<:DenseRiemannianMetric,<:GaussianKinetic},
+    θ::AbstractVecOrMat,
+    r::AbstractVecOrMat,
 )
     H = h.metric.G(θ)
     G = h.metric.map(H)
@@ -87,15 +89,16 @@ function ∂H∂θ_cache(
     d = length(∂ℓπ∂θ)
     D = diagm((Q' * r) ./ softabsλ)
     term_2_cached = Q * D * J * D * Q'
-    g =
-        isdiag ?
-        -(∂ℓπ∂θ - 1 / 2 * diag(term_1_cached * ∂H∂θ) + 1 / 2 * diag(term_2 * ∂H∂θ)) :
+    g = if isdiag
+        -(∂ℓπ∂θ - 1 / 2 * diag(term_1_cached * ∂H∂θ) + 1 / 2 * diag(term_2 * ∂H∂θ))
+    else
         -mapreduce(vcat, 1:d) do i
-            ∂H∂θᵢ = ∂H∂θ[:, :, i]
-            # ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * M' * (J .* (Q' * ∂H∂θᵢ * Q)) * M # (v1)
-            # NOTE Some further optimization can be done here: cache the 1st product all together
-            ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * tr(term_2_cached * ∂H∂θᵢ) # (v2) cache friendly
-        end
+        ∂H∂θᵢ = ∂H∂θ[:, :, i]
+        # ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * M' * (J .* (Q' * ∂H∂θᵢ * Q)) * M # (v1)
+        # NOTE Some further optimization can be done here: cache the 1st product all together
+        ∂ℓπ∂θ[i] - 1 / 2 * tr(term_1_cached * ∂H∂θᵢ) + 1 / 2 * tr(term_2_cached * ∂H∂θᵢ) # (v2) cache friendly
+    end
+    end
 
     dv = DualValue(ℓπ, g)
     return return_cache ? (dv, (; ℓπ, ∂ℓπ∂θ, ∂H∂θ, Q, softabsλ, J, term_1_cached)) : dv

src/riemannian/metric.jl

@@ -44,7 +44,7 @@ function DenseRiemannianMetric(size, G, ∂G∂θ, map=IdentityMap())
 end
 
 # Convenient constructor
-function DenseRiemannianMetric(size, ℓπ, initial_θ, λ, map = IdentityMap())
+function DenseRiemannianMetric(size, ℓπ, initial_θ, λ, map=IdentityMap())
     _Hfunc = VecTargets.gen_hess(x -> -ℓπ(x), initial_θ) # x -> (value, gradient, hessian)
     Hfunc = x -> copy.(_Hfunc(x)) # _Hfunc do in-place computation, copy to avoid bug
 
@@ -96,7 +96,9 @@ end
 
 Base.size(e::DenseRiemannianMetric) = e.size
 Base.size(e::DenseRiemannianMetric, dim::Int) = e.size[dim]
-Base.show(io::IO, drm::DenseRiemannianMetric) = print(io, "DenseRiemannianMetric$(drm.size) with $(drm.map) metric")
+function Base.show(io::IO, drm::DenseRiemannianMetric)
+    return print(io, "DenseRiemannianMetric$(drm.size) with $(drm.map) metric")
+end
 
 function rand_momentum(
     rng::Union{AbstractRNG,AbstractVector{<:AbstractRNG}},

test/riemannian.jl

@@ -25,9 +25,7 @@ using LinearAlgebra
     # Run the sampler to draw samples from the specified Gaussian, where
     #   - `samples` will store the samples
     #   - `stats` will store diagnostic statistics for each sample
-    samples, stats = sample(
-        rng, hamiltonian, kernel, initial_θ, n_samples; progress=true
-    )
+    samples, stats = sample(rng, hamiltonian, kernel, initial_θ, n_samples; progress=true)
     @test length(samples) == n_samples
     @test length(stats) == n_samples
 end
@@ -55,9 +53,7 @@ end
     # Run the sampler to draw samples from the specified Gaussian, where
     #   - `samples` will store the samples
     #   - `stats` will store diagnostic statistics for each sample
-    samples, stats = sample(
-        rng, hamiltonian, kernel, initial_θ, n_samples; progress=true
-    )
+    samples, stats = sample(rng, hamiltonian, kernel, initial_θ, n_samples; progress=true)
     @test length(samples) == n_samples
     @test length(stats) == n_samples
-end
+end