Merge branch 'main' into qqy/NEW_RMHMC

ErikQQY · web-flow · commit 6df40b52ff6b · 2025-11-21T23:09:33.000+08:00
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -31,7 +31,7 @@ jobs:
           - macOS-latest
           - windows-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
diff --git a/.github/workflows/DocsNav.yml b/.github/workflows/DocsNav.yml
@@ -17,7 +17,7 @@ jobs:
 
     steps:
       - name: Checkout gh-pages branch
-        uses: actions/checkout@v4
+        uses: actions/checkout@v5
         with:
           ref: gh-pages
 
diff --git a/.github/workflows/ExperimentalTests.yml b/.github/workflows/ExperimentalTests.yml
@@ -28,7 +28,7 @@ jobs:
           - macOS-latest
           - windows-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
diff --git a/.github/workflows/Format.yml b/.github/workflows/Format.yml
@@ -16,7 +16,7 @@ jobs:
   format:
     runs-on: ubuntu-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: 1
diff --git a/.github/workflows/IntegrationTests.yml b/.github/workflows/IntegrationTests.yml
@@ -28,7 +28,7 @@ jobs:
           - macOS-latest
           - windows-latest
     steps:
-      - uses: actions/checkout@v4
+      - uses: actions/checkout@v5
       - uses: julia-actions/setup-julia@v2
         with:
           version: ${{ matrix.version }}
diff --git a/Project.toml b/Project.toml
@@ -1,6 +1,6 @@
 name = "AdvancedHMC"
 uuid = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
-version = "0.8.0"
+version = "0.8.3"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -32,7 +32,7 @@ AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
 
 [compat]
 ADTypes = "1"
-AbstractMCMC = "5.6"
+AbstractMCMC = "5.9"
 ArgCheck = "1, 2"
 CUDA = "3, 4, 5"
 ComponentArrays = "0.15"
diff --git a/src/abstractmcmc.jl b/src/abstractmcmc.jl
@@ -33,6 +33,7 @@ getintegrator(state::HMCState) = state.κ.τ.integrator
 function AbstractMCMC.getparams(state::HMCState)
     return state.transition.z.θ
 end
+AbstractMCMC.getstats(state::AdvancedHMC.HMCState) = state.transition.stat
 
 function AbstractMCMC.setparams!!(
     model::AbstractMCMC.LogDensityModel, state::HMCState, params
@@ -55,8 +56,8 @@ function AbstractMCMC.sample(
     sampler::AbstractHMCSampler,
     N::Integer;
     n_adapts::Int=min(div(N, 10), 1_000),
-    progress=true,
-    verbose=false,
+    progress::Bool=true,
+    verbose::Bool=false,
     callback=nothing,
     kwargs...,
 )
@@ -94,8 +95,8 @@ function AbstractMCMC.sample(
     N::Integer,
     nchains::Integer;
     n_adapts::Int=min(div(N, 10), 1_000),
-    progress=true,
-    verbose=false,
+    progress::Bool=true,
+    verbose::Bool=false,
     callback=nothing,
     kwargs...,
 )
@@ -217,7 +218,7 @@ logging behavior of the non-AbstractMCMC [`sample`](@ref).
 # Fields
 $(FIELDS)
 """
-struct HMCProgressCallback{P}
+struct HMCProgressCallback{P<:Union{ProgressMeter.Progress,Nothing}}
     "`Progress` meter from ProgressMeters.jl, or `nothing`."
     pm::P
     "If `pm === nothing` and this is `true` some information will be logged upon completion of adaptation."
@@ -227,12 +228,21 @@ struct HMCProgressCallback{P}
     num_divergent_transitions_during_adaption::Base.RefValue{Int}
 end
 
-function HMCProgressCallback(n_samples; progress=true, verbose=false)
+function HMCProgressCallback(n_samples::Integer; progress::Bool=true, verbose::Bool=false)
     pm = progress ? ProgressMeter.Progress(n_samples; desc="Sampling", barlen=31) : nothing
     return HMCProgressCallback(pm, verbose, Ref(0), Ref(0))
 end
 
-function (cb::HMCProgressCallback)(rng, model, spl, t, state, i; n_adapts::Int=0, kwargs...)
+function (cb::HMCProgressCallback)(
+    rng::AbstractRNG,
+    model::AbstractMCMC.LogDensityModel,
+    spl::AbstractHMCSampler,
+    t::Transition,
+    state::HMCState,
+    i::Int;
+    n_adapts::Int=0,
+    kwargs...,
+)
     verbose = cb.verbose
     pm = cb.pm
 
@@ -260,16 +270,18 @@ function (cb::HMCProgressCallback)(rng, model, spl, t, state, i; n_adapts::Int=0
         # Do include current iteration and mass matrix
         pm_next!(
             pm,
-            (
-                iterations=i,
-                ratio_divergent_transitions=round(
-                    percentage_divergent_transitions; digits=2
-                ),
-                ratio_divergent_transitions_during_adaption=round(
-                    percentage_divergent_transitions_during_adaption; digits=2
+            merge(
+                (;
+                    iterations=i,
+                    ratio_divergent_transitions=round(
+                        percentage_divergent_transitions; digits=2
+                    ),
+                    ratio_divergent_transitions_during_adaption=round(
+                        percentage_divergent_transitions_during_adaption; digits=2
+                    ),
+                    mass_matrix=metric,
                 ),
-                tstat...,
-                mass_matrix=metric,
+                tstat,
             ),
         )
         # Report finish of adapation
diff --git a/src/adaptation/Adaptation.jl b/src/adaptation/Adaptation.jl
@@ -37,7 +37,7 @@ struct NaiveHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAd
     pc::M
     ssa::Tssa
 end
-function Base.show(io::IO, ::MIME"text/plain", a::NaiveHMCAdaptor)
+function Base.show(io::IO, a::NaiveHMCAdaptor)
     return print(io, "NaiveHMCAdaptor(pc=", a.pc, ", ssa=", a.ssa, ")")
 end
 
diff --git a/src/adaptation/massmatrix.jl b/src/adaptation/massmatrix.jl
@@ -23,7 +23,7 @@ end
 
 struct UnitMassMatrix{T<:AbstractFloat} <: MassMatrixAdaptor end
 
-function Base.show(io::IO, mime::MIME"text/plain", ::UnitMassMatrix{T}) where {T}
+function Base.show(io::IO, ::UnitMassMatrix{T}) where {T}
     return print(io, "UnitMassMatrix{", T, "} adaptor")
 end
 
@@ -93,7 +93,7 @@ mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVec
     end
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", ::WelfordVar{T}) where {T}
+function Base.show(io::IO, ::WelfordVar{T}) where {T}
     return print(io, "WelfordVar{", T, "} adaptor")
 end
 
@@ -194,7 +194,7 @@ mutable struct WelfordCov{F<:AbstractFloat,C<:AbstractMatrix{F}} <: DenseMatrixE
     cov::C
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", ::WelfordCov{T}) where {T}
+function Base.show(io::IO, ::WelfordCov{T}) where {T}
     return print(io, "WelfordCov{", T, "} adaptor")
 end
 
diff --git a/src/adaptation/stan_adaptor.jl b/src/adaptation/stan_adaptor.jl
@@ -49,7 +49,7 @@ function initialize!(
     return nothing
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", state::StanHMCAdaptorState)
+function Base.show(io::IO, state::StanHMCAdaptorState)
     print(io, "window(", state.window_start, ", ", state.window_end, "), window_splits(")
     join(io, state.window_splits, ", ")
     return print(io, ")")
@@ -66,6 +66,23 @@ struct StanHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAda
     window_size::Int
     state::StanHMCAdaptorState
 end
+
+function Base.show(io::IO, a::StanHMCAdaptor)
+    return print(
+        io,
+        "StanHMCAdaptor(",
+        a.pc,
+        ", ",
+        a.ssa,
+        "; init_buffer=",
+        a.init_buffer,
+        ", term_buffer=",
+        a.term_buffer,
+        ", window_size=",
+        a.window_size,
+        ")",
+    )
+end
 function Base.show(io::IO, mime::MIME"text/plain", a::StanHMCAdaptor)
     return print(
         io,
diff --git a/src/adaptation/stepsize.jl b/src/adaptation/stepsize.jl
@@ -77,7 +77,7 @@ getϵ(ss::StepSizeAdaptor) = ss.state.ϵ
 struct FixedStepSize{T<:AbstractScalarOrVec{<:AbstractFloat}} <: StepSizeAdaptor
     ϵ::T
 end
-function Base.show(io::IO, mime::MIME"text/plain", a::FixedStepSize)
+function Base.show(io::IO, a::FixedStepSize)
     return print(io, "FixedStepSize adaptor with step size ", a.ϵ)
 end
 
@@ -86,7 +86,7 @@ getϵ(fss::FixedStepSize) = fss.ϵ
 struct ManualSSAdaptor{T<:AbstractScalarOrVec{<:AbstractFloat}} <: StepSizeAdaptor
     state::MSSState{T}
 end
-function Base.show(io::IO, mime::MIME"text/plain", a::ManualSSAdaptor{T}) where {T}
+function Base.show(io::IO, a::ManualSSAdaptor{T}) where {T}
     return print(io, "ManualSSAdaptor{", T, "} with step size of ", a.state.ϵ)
 end
 
@@ -119,6 +119,23 @@ struct NesterovDualAveraging{T<:AbstractFloat,S<:AbstractScalarOrVec{T}} <: Step
     δ::T
     state::DAState{S}
 end
+
+function Base.show(io::IO, a::NesterovDualAveraging)
+    print(
+        io,
+        "NesterovDualAveraging(",
+        a.γ,
+        ", ",
+        a.t_0,
+        ", ",
+        a.κ,
+        ", ",
+        a.δ,
+        ", ",
+        a.state.ϵ,
+        ")",
+    )
+end
 function Base.show(io::IO, mime::MIME"text/plain", a::NesterovDualAveraging{T}) where {T}
     return print(
         io,
diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
@@ -4,7 +4,7 @@ struct Hamiltonian{M<:AbstractMetric,K<:AbstractKinetic,Tlogπ,T∂logπ∂θ}
     ℓπ::Tlogπ
     ∂ℓπ∂θ::T∂logπ∂θ
 end
-function Base.show(io::IO, mime::MIME"text/plain", h::Hamiltonian)
+function Base.show(io::IO, h::Hamiltonian)
     return print(
         io,
         "Hamiltonian with ",
diff --git a/src/integrator.jl b/src/integrator.jl
@@ -72,8 +72,8 @@ struct Leapfrog{T<:AbstractScalarOrVec{<:AbstractFloat}} <: AbstractLeapfrog{T}
     "Step size."
     ϵ::T
 end
-function Base.show(io::IO, mime::MIME"text/plain", l::Leapfrog)
-    return print(io, "Leapfrog with step size ϵ=", round.(l.ϵ; sigdigits=3), ")")
+function Base.show(io::IO, l::Leapfrog)
+    return print(io, "Leapfrog with step size ϵ=", round.(l.ϵ; sigdigits=3))
 end
 integrator_eltype(i::AbstractLeapfrog{T}) where {T<:AbstractFloat} = T
 
@@ -120,7 +120,7 @@ end
 
 JitteredLeapfrog(ϵ0, jitter) = JitteredLeapfrog(ϵ0, jitter, ϵ0)
 
-function Base.show(io::IO, mime::MIME"text/plain", l::JitteredLeapfrog)
+function Base.show(io::IO, l::JitteredLeapfrog)
     return print(
         io,
         "JitteredLeapfrog with step size ",
@@ -178,7 +178,7 @@ struct TemperedLeapfrog{FT<:AbstractFloat,T<:AbstractScalarOrVec{FT}} <: Abstrac
     α::FT
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", l::TemperedLeapfrog)
+function Base.show(io::IO, l::TemperedLeapfrog)
     return print(
         io,
         "TemperedLeapfrog with step size ϵ=",
diff --git a/src/metric.jl b/src/metric.jl
@@ -7,7 +7,7 @@ abstract type AbstractMetric end
 
 _string_M⁻¹(mat::AbstractMatrix, n_chars::Int=32) = _string_M⁻¹(diag(mat), n_chars)
 function _string_M⁻¹(vec::AbstractVector, n_chars::Int=32)
-    s_vec = string(vec)
+    s_vec = repr(vec; context=(:compact => true))
     l = length(s_vec)
     s_dots = " ...]"
     n_diag_chars = n_chars - length(s_dots)
@@ -33,6 +33,10 @@ renew(ue::UnitEuclideanMetric, M⁻¹) = UnitEuclideanMetric(M⁻¹, ue.size)
 Base.eltype(::UnitEuclideanMetric{T}) where {T} = T
 Base.size(e::UnitEuclideanMetric) = e.size
 Base.size(e::UnitEuclideanMetric, dim::Int) = e.size[dim]
+
+function Base.show(io::IO, uem::UnitEuclideanMetric{T}) where {T}
+    print(io, "UnitEuclideanMetric(", T, ", ", uem.size, ")")
+end
 function Base.show(io::IO, ::MIME"text/plain", uem::UnitEuclideanMetric{T}) where {T}
     return print(
         io,
@@ -66,6 +70,10 @@ renew(ue::DiagEuclideanMetric, M⁻¹) = DiagEuclideanMetric(M⁻¹)
 
 Base.eltype(::DiagEuclideanMetric{T}) where {T} = T
 Base.size(e::DiagEuclideanMetric, dim...) = size(e.M⁻¹, dim...)
+
+function Base.show(io::IO, dem::DiagEuclideanMetric)
+    print(io, "DiagEuclideanMetric(", _string_M⁻¹(dem.M⁻¹), ")")
+end
 function Base.show(io::IO, ::MIME"text/plain", dem::DiagEuclideanMetric{T}) where {T}
     return print(
         io,
@@ -110,6 +118,10 @@ renew(ue::DenseEuclideanMetric, M⁻¹) = DenseEuclideanMetric(M⁻¹)
 
 Base.eltype(::DenseEuclideanMetric{T}) where {T} = T
 Base.size(e::DenseEuclideanMetric, dim...) = size(e._temp, dim...)
+
+function Base.show(io::IO, dem::DenseEuclideanMetric)
+    print(io, "DenseEuclideanMetric(", _string_M⁻¹(dem.M⁻¹), ")")
+end
 function Base.show(io::IO, ::MIME"text/plain", dem::DenseEuclideanMetric{T}) where {T}
     return print(
         io,
diff --git a/src/sampler.jl b/src/sampler.jl
@@ -93,15 +93,15 @@ end
 """
 Progress meter update with all trajectory stats, iteration number and metric shown.
 """
-function pm_next!(pm, stat::NamedTuple)
-    ProgressMeter.next!(pm; showvalues=map(tuple, values(stat), keys(stat)))
+function pm_next!(pm::ProgressMeter.Progress, stat::NamedTuple)
+    ProgressMeter.next!(pm; showvalues=map(tuple, keys(stat), values(stat)))
     return nothing
 end
 
 """
 Simple progress meter update without any show values.
 """
-simple_pm_next!(pm, stat::NamedTuple) = ProgressMeter.next!(pm)
+simple_pm_next!(pm::ProgressMeter.Progress, ::NamedTuple) = ProgressMeter.next!(pm)
 
 ##
 ## Sampling functions
diff --git a/src/trajectory.jl b/src/trajectory.jl
@@ -108,7 +108,7 @@ struct SliceTS{F<:AbstractFloat,P<:PhasePoint} <: AbstractTrajectorySampler
     n::Int
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", s::SliceTS)
+function Base.show(io::IO, s::SliceTS)
     return print(
         io,
         "SliceTS with slice variable ℓu=",
@@ -225,7 +225,7 @@ end
 
 ConstructionBase.constructorof(::Type{<:Trajectory{TS}}) where {TS} = Trajectory{TS}
 
-function Base.show(io::IO, mime::MIME"text/plain", τ::Trajectory{TS}) where {TS}
+function Base.show(io::IO, τ::Trajectory{TS}) where {TS}
     return print(
         io,
         "Trajectory{",
@@ -482,7 +482,7 @@ struct Termination
     numerical::Bool
 end
 
-function Base.show(io::IO, mime::MIME"text/plain", d::Termination)
+function Base.show(io::IO, d::Termination)
     return print(
         io, "Termination reasons of (dynamic=", d.dynamic, ", numerical=", d.numerical, ")"
     )
diff --git a/test/Project.toml b/test/Project.toml
diff --git a/test/common.jl b/test/common.jl
diff --git a/test/sampler-vec.jl b/test/sampler-vec.jl
diff --git a/test/sampler.jl b/test/sampler.jl
