Merge branch 'main' into qqy/NEW_RMHMC

ErikQQY · web-flow · commit 5f9592f5e948 · 2025-06-30T01:20:02.000+08:00
diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
@@ -19,7 +19,7 @@ jobs:
       actions: write
       contents: read
     runs-on: ${{ matrix.os }}
-    continue-on-error: ${{ matrix.version == 'pre' }}
+    continue-on-error: ${{ matrix.version == 'pre' || matrix.version == 'min' }}
     strategy:
       matrix:
         version:
@@ -30,12 +30,6 @@ jobs:
           - ubuntu-latest
           - macOS-latest
           - windows-latest
-        exclude:
-          - version: 'min'
-            os: macOS-latest # Apple Silicon
-        include:
-          - version: 'min'
-            os: macOS-13 # Intel
     steps:
       - uses: actions/checkout@v4
       - uses: julia-actions/setup-julia@v2
@@ -47,3 +41,9 @@ jobs:
         uses: julia-actions/julia-runtest@v1
         env:
           AHMC_TEST_GROUP: AdvancedHMC
+      - uses: julia-actions/julia-processcoverage@v1
+      - uses: codecov/codecov-action@v5
+        with:
+          files: lcov.info
+          token: ${{ secrets.CODECOV_TOKEN }}
+          fail_ci_if_error: true
diff --git a/HISTORY.md b/HISTORY.md
@@ -3,6 +3,8 @@
 ## 0.8.0
 
   - To make an MCMC transtion from phasepoint `z` using trajectory `τ`(or HMCKernel `κ`) under Hamiltonian `h`, use `transition(h, τ, z)` or `transition(rng, h, τ, z)`(if using HMCKernel, use `transition(h, κ, z)` or `transition(rng, h, κ, z)`).
+  - The `initial_step_size` in `find_good_stepsize` for heuristic search of a good initial leap-frog step-size can be manually specified, default as `1//10`.
+  - The printing interface has been upgraded to a more user-friendly design.
 
 ## v0.7.1
 
diff --git a/Project.toml b/Project.toml
@@ -18,29 +18,29 @@ StatsFuns = "4c63d2b9-4356-54db-8cca-17b64c39e42c"
 
 [weakdeps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"
-ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
+ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66"
 MCMCChains = "c7f686f2-ff18-58e9-bc7b-31028e88f75d"
 OrdinaryDiffEq = "1dea7af3-3e70-54e6-95c3-0bf5283fa5ed"
 
 [extensions]
 AdvancedHMCADTypesExt = "ADTypes"
-AdvancedHMCComponentArraysExt = "ComponentArrays"
 AdvancedHMCCUDAExt = "CUDA"
+AdvancedHMCComponentArraysExt = "ComponentArrays"
 AdvancedHMCMCMCChainsExt = "MCMCChains"
 AdvancedHMCOrdinaryDiffEqExt = "OrdinaryDiffEq"
 
 [compat]
 ADTypes = "1"
 AbstractMCMC = "5.6"
 ArgCheck = "1, 2"
-ComponentArrays = "0.15"
 CUDA = "3, 4, 5"
+ComponentArrays = "0.15"
 DocStringExtensions = "0.8, 0.9"
 LinearAlgebra = "<0.1, 1"
 LogDensityProblems = "2"
 LogDensityProblemsAD = "1"
-MCMCChains = "5, 6"
+MCMCChains = "5, 6, 7"
 OrdinaryDiffEq = "6"
 ProgressMeter = "1"
 Random = "<0.1, 1"
diff --git a/README.md b/README.md
@@ -2,7 +2,7 @@
 
 [![CI](https://github.com/TuringLang/AdvancedHMC.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/TuringLang/AdvancedHMC.jl/actions/workflows/CI.yml)
 [![DOI](https://zenodo.org/badge/72657907.svg)](https://zenodo.org/badge/latestdoi/72657907)
-[![Coverage Status](https://coveralls.io/repos/github/TuringLang/AdvancedHMC.jl/badge.svg?branch=main)](https://coveralls.io/github/TuringLang/AdvancedHMC.jl?branch=main)
+[![Coverage Status](https://codecov.io/gh/TuringLang/AdvancedHMC.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/TuringLang/AdvancedHMC.jl)
 [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/stable/)
 [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/dev/)
 [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
diff --git a/docs/src/index.md b/docs/src/index.md
@@ -2,7 +2,7 @@
 
 [![CI](https://github.com/TuringLang/AdvancedHMC.jl/actions/workflows/CI.yml/badge.svg)](https://github.com/TuringLang/AdvancedHMC.jl/actions/workflows/CI.yml)
 [![DOI](https://zenodo.org/badge/72657907.svg)](https://zenodo.org/badge/latestdoi/72657907)
-[![Coverage Status](https://coveralls.io/repos/github/TuringLang/AdvancedHMC.jl/badge.svg?branch=kx%2Fbug-fix)](https://coveralls.io/github/TuringLang/AdvancedHMC.jl?branch=kx%2Fbug-fix)
+[![Coverage Status](https://codecov.io/gh/TuringLang/AdvancedHMC.jl/branch/main/graph/badge.svg)](https://codecov.io/gh/TuringLang/AdvancedHMC.jl)
 [![Stable](https://img.shields.io/badge/docs-stable-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/stable/)
 [![Dev](https://img.shields.io/badge/docs-dev-blue.svg)](https://turinglang.github.io/AdvancedHMC.jl/dev/)
 [![Aqua QA](https://raw.githubusercontent.com/JuliaTesting/Aqua.jl/master/badge.svg)](https://github.com/JuliaTesting/Aqua.jl)
diff --git a/ext/AdvancedHMCADTypesExt.jl b/ext/AdvancedHMCADTypesExt.jl
@@ -1,15 +1,17 @@
 module AdvancedHMCADTypesExt
 
-using AdvancedHMC:
-    AbstractMetric, LogDensityModel, Hamiltonian, LogDensityProblems, LogDensityProblemsAD
+using AdvancedHMC: AdvancedHMC, AbstractMetric, Hamiltonian, LogDensityModel
+using AdvancedHMC: LogDensityProblems, LogDensityProblemsAD
 using ADTypes: AbstractADType
 
-function Hamiltonian(
+function AdvancedHMC.Hamiltonian(
     metric::AbstractMetric, ℓπ::LogDensityModel, kind::AbstractADType; kwargs...
 )
     return Hamiltonian(metric, ℓπ.logdensity, kind; kwargs...)
 end
-function Hamiltonian(metric::AbstractMetric, ℓπ, kind::AbstractADType; kwargs...)
+function AdvancedHMC.Hamiltonian(
+    metric::AbstractMetric, ℓπ, kind::AbstractADType; kwargs...
+)
     if LogDensityProblems.capabilities(ℓπ) === nothing
         throw(
             ArgumentError(
diff --git a/research/tests/Project.toml b/research/tests/Project.toml
@@ -2,7 +2,6 @@
 AdaptiveRejectionSampling = "c75e803d-635f-53bd-ab7d-544e482d8c75"
 AdvancedHMC = "0bf59076-c3b1-5ca4-86bd-e02cd72cde3d"
 BenchmarkTools = "6e4b80f9-dd63-53aa-95a3-0cdb28fa8baf"
-Comonicon = "863f3e99-da2a-4334-8734-de3dacbe5542"
 DiffResults = "163ba53b-c6d8-5494-b064-1a9d43ac40c5"
 Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 DynamicPPL = "366bfd00-2699-11ea-058f-f148b4cae6d8"
diff --git a/research/tests/runtests.jl b/research/tests/runtests.jl
@@ -1,11 +1,9 @@
-using Comonicon, ReTest
+using ReTest
 
 # include the source code for experimental HMC
 include("../src/relativistic_hmc.jl")
 
 # include the tests for experimental HMC
 include("relativistic_hmc.jl")
 
-Comonicon.@main function runtests(patterns...; dry::Bool=false)
-    return retest(patterns...; dry=dry, verbose=Inf)
-end
+retest(; dry=false, verbose=Inf)
diff --git a/src/adaptation/Adaptation.jl b/src/adaptation/Adaptation.jl
@@ -37,8 +37,8 @@ struct NaiveHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAd
     pc::M
     ssa::Tssa
 end
-function Base.show(io::IO, a::NaiveHMCAdaptor)
-    return print(io, "NaiveHMCAdaptor(pc=$(a.pc), ssa=$(a.ssa))")
+function Base.show(io::IO, ::MIME"text/plain", a::NaiveHMCAdaptor)
+    return print(io, "NaiveHMCAdaptor(pc=", a.pc, ", ssa=", a.ssa, ")")
 end
 
 getM⁻¹(ca::NaiveHMCAdaptor) = getM⁻¹(ca.pc)
diff --git a/src/adaptation/massmatrix.jl b/src/adaptation/massmatrix.jl
@@ -23,7 +23,9 @@ end
 
 struct UnitMassMatrix{T<:AbstractFloat} <: MassMatrixAdaptor end
 
-Base.show(io::IO, ::UnitMassMatrix) = print(io, "UnitMassMatrix")
+function Base.show(io::IO, mime::MIME"text/plain", ::UnitMassMatrix{T}) where {T}
+    return print(io, "UnitMassMatrix{", T, "} adaptor")
+end
 
 UnitMassMatrix() = UnitMassMatrix{Float64}()
 
@@ -91,7 +93,9 @@ mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T},V<:AbstractVec
     end
 end
 
-Base.show(io::IO, ::WelfordVar) = print(io, "WelfordVar")
+function Base.show(io::IO, mime::MIME"text/plain", ::WelfordVar{T}) where {T}
+    return print(io, "WelfordVar{", T, "} adaptor")
+end
 
 function WelfordVar{T}(
     sz::Union{Tuple{Int},Tuple{Int,Int}}; n_min::Int=10, var=ones(T, sz)
@@ -190,7 +194,9 @@ mutable struct WelfordCov{F<:AbstractFloat,C<:AbstractMatrix{F}} <: DenseMatrixE
     cov::C
 end
 
-Base.show(io::IO, ::WelfordCov) = print(io, "WelfordCov")
+function Base.show(io::IO, mime::MIME"text/plain", ::WelfordCov{T}) where {T}
+    return print(io, "WelfordCov{", T, "} adaptor")
+end
 
 function WelfordCov{T}(
     sz::Tuple{Int}; n_min::Int=10, cov=LinearAlgebra.diagm(0 => ones(T, first(sz)))
diff --git a/src/adaptation/stan_adaptor.jl b/src/adaptation/stan_adaptor.jl
@@ -49,13 +49,10 @@ function initialize!(
     return nothing
 end
 
-function Base.show(io::IO, state::StanHMCAdaptorState)
-    return print(
-        io,
-        "window($(state.window_start), $(state.window_end)), window_splits(" *
-        string(join(state.window_splits, ", ")) *
-        ")",
-    )
+function Base.show(io::IO, mime::MIME"text/plain", state::StanHMCAdaptorState)
+    print(io, "window(", state.window_start, ", ", state.window_end, "), window_splits(")
+    join(io, state.window_splits, ", ")
+    return print(io, ")")
 end
 
 ### Stan's windowed adaptation
@@ -69,10 +66,22 @@ struct StanHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAda
     window_size::Int
     state::StanHMCAdaptorState
 end
-function Base.show(io::IO, a::StanHMCAdaptor)
+function Base.show(io::IO, mime::MIME"text/plain", a::StanHMCAdaptor)
     return print(
         io,
-        "StanHMCAdaptor(\n    pc=$(a.pc),\n    ssa=$(a.ssa),\n    init_buffer=$(a.init_buffer), term_buffer=$(a.term_buffer), window_size=$(a.window_size),\n    state=$(a.state)\n)",
+        "StanHMCAdaptor(\n    pc=",
+        a.pc,
+        ",\n    ssa=",
+        a.ssa,
+        ",\n    init_buffer=",
+        a.init_buffer,
+        ", term_buffer=",
+        a.term_buffer,
+        ", window_size=",
+        a.window_size,
+        ",\n    state=",
+        a.state,
+        "\n)",
     )
 end
 
diff --git a/src/adaptation/stepsize.jl b/src/adaptation/stepsize.jl
@@ -1,10 +1,24 @@
 ### Mutable states
+"""
+$(TYPEDEF)
+
+Dual Averaging state
+
+Mutable state for storing the current iteration of the dual averaging algorithm.
 
+# Fields
+
+$(TYPEDFIELDS)
+"""
 mutable struct DAState{T<:AbstractScalarOrVec{<:AbstractFloat}}
+    "Adaptation iteration"
     m::Int
     ϵ::T
+    "Asymptotic mean of parameter"
     μ::T
+    "Moving average parameter"
     x_bar::T
+    "Moving average statistic"
     H_bar::T
 end
 
@@ -63,48 +77,68 @@ getϵ(ss::StepSizeAdaptor) = ss.state.ϵ
 struct FixedStepSize{T<:AbstractScalarOrVec{<:AbstractFloat}} <: StepSizeAdaptor
     ϵ::T
 end
-Base.show(io::IO, a::FixedStepSize) = print(io, "FixedStepSize(", a.ϵ, ")")
+function Base.show(io::IO, mime::MIME"text/plain", a::FixedStepSize)
+    return print(io, "FixedStepSize adaptor with step size ", a.ϵ)
+end
 
 getϵ(fss::FixedStepSize) = fss.ϵ
 
 struct ManualSSAdaptor{T<:AbstractScalarOrVec{<:AbstractFloat}} <: StepSizeAdaptor
     state::MSSState{T}
 end
-Base.show(io::IO, a::ManualSSAdaptor) = print(io, "ManualSSAdaptor()")
+function Base.show(io::IO, mime::MIME"text/plain", a::ManualSSAdaptor{T}) where {T}
+    return print(io, "ManualSSAdaptor{", T, "} with step size of ", a.state.ϵ)
+end
 
 function ManualSSAdaptor(initϵ::T) where {T<:AbstractScalarOrVec{<:AbstractFloat}}
     return ManualSSAdaptor{T}(MSSState(initϵ))
 end
 
 """
+$(TYPEDEF)
+
 An implementation of the Nesterov dual averaging algorithm to tune step size.
 
-References
+# Fields
+
+$(TYPEDFIELDS)
+
+# References
 
 Hoffman, M. D., & Gelman, A. (2014). The No-U-Turn Sampler: adaptively setting path lengths in Hamiltonian Monte Carlo. Journal of Machine Learning Research, 15(1), 1593-1623.
 Nesterov, Y. (2009). Primal-dual subgradient methods for convex problems. Mathematical programming, 120(1), 221-259.
 """
 struct NesterovDualAveraging{T<:AbstractFloat,S<:AbstractScalarOrVec{T}} <: StepSizeAdaptor
+    "Adaption scaling"
     γ::T
+    "Effective starting iteration"
     t_0::T
+    "Adaption shrinkage"
     κ::T
+    "Target value of statistic"
     δ::T
     state::DAState{S}
 end
-function Base.show(io::IO, a::NesterovDualAveraging)
+function Base.show(io::IO, mime::MIME"text/plain", a::NesterovDualAveraging{T}) where {T}
     return print(
         io,
-        "NesterovDualAveraging(γ=",
+        "NesterovDualAveraging{",
+        T,
+        "} with\n",
+        "Scaling γ=",
         a.γ,
-        ", t_0=",
+        "\n",
+        "Starting iter t_0=",
         a.t_0,
-        ", κ=",
+        "\n",
+        "Shrinkage κ=",
         a.κ,
-        ", δ=",
+        "\n",
+        "Target statistic δ=",
         a.δ,
-        ", state.ϵ=",
+        "\n",
+        "Curret ϵ=",
         getϵ(a),
-        ")",
     )
 end
 
diff --git a/src/hamiltonian.jl b/src/hamiltonian.jl
@@ -4,8 +4,14 @@ struct Hamiltonian{M<:AbstractMetric,K<:AbstractKinetic,Tlogπ,T∂logπ∂θ}
     ℓπ::Tlogπ
     ∂ℓπ∂θ::T∂logπ∂θ
 end
-function Base.show(io::IO, h::Hamiltonian)
-    return print(io, "Hamiltonian(metric=$(h.metric), kinetic=$(h.kinetic))")
+function Base.show(io::IO, mime::MIME"text/plain", h::Hamiltonian)
+    return print(
+        io,
+        "Hamiltonian with ",
+        nameof(typeof(h.metric)),
+        " and ",
+        nameof(typeof(h.kinetic)),
+    )
 end
 
 # By default we use Gaussian kinetic energy; also to ensure backward compatibility at the time this was introduced
diff --git a/src/integrator.jl b/src/integrator.jl
diff --git a/src/metric.jl b/src/metric.jl
diff --git a/src/riemannian/integrator.jl b/src/riemannian/integrator.jl
diff --git a/src/trajectory.jl b/src/trajectory.jl
diff --git a/test/Project.toml b/test/Project.toml
diff --git a/test/runtests.jl b/test/runtests.jl
