Merge branch 'main' into qqy/sghmc

Author: sunxd3

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"

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, a::NaiveHMCAdaptor)
+function Base.show(io::IO, ::MIME"text/plain", a::NaiveHMCAdaptor)
     return print(io, "NaiveHMCAdaptor(pc=$(a.pc), ssa=$(a.ssa))")
 end
 

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

src/adaptation/stan_adaptor.jl

@@ -49,7 +49,7 @@ function initialize!(
     return nothing
 end
 
-function Base.show(io::IO, state::StanHMCAdaptorState)
+function Base.show(io::IO, mime::MIME"text/plain", state::StanHMCAdaptorState)
     return print(
         io,
         "window($(state.window_start), $(state.window_end)), window_splits(" *
@@ -69,7 +69,7 @@ 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)",

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,66 @@ 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
 

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

src/integrator.jl

@@ -72,7 +72,9 @@ struct Leapfrog{T<:AbstractScalarOrVec{<:AbstractFloat}} <: AbstractLeapfrog{T}
     "Step size."
     ϵ::T
 end
-Base.show(io::IO, l::Leapfrog) = print(io, "Leapfrog(ϵ=$(round.(l.ϵ; sigdigits=3)))")
+function Base.show(io::IO, mime::MIME"text/plain", l::Leapfrog)
+    return print(io, "Leapfrog with step size ϵ=$(round.(l.ϵ; sigdigits=3))")
+end
 integrator_eltype(i::AbstractLeapfrog{T}) where {T<:AbstractFloat} = T
 
 ### Jittering
@@ -118,10 +120,10 @@ end
 
 JitteredLeapfrog(ϵ0, jitter) = JitteredLeapfrog(ϵ0, jitter, ϵ0)
 
-function Base.show(io::IO, l::JitteredLeapfrog)
+function Base.show(io::IO, mime::MIME"text/plain", l::JitteredLeapfrog)
     return print(
         io,
-        "JitteredLeapfrog(ϵ0=$(round.(l.ϵ0; sigdigits=3)), jitter=$(round.(l.jitter; sigdigits=3)), ϵ=$(round.(l.ϵ; sigdigits=3)))",
+        "JitteredLeapfrog with step size $(round.(l.ϵ0; sigdigits=3)), jitter $(round.(l.jitter; sigdigits=3)), jittered step size $(round.(l.ϵ; sigdigits=3))",
     )
 end
 
@@ -171,9 +173,10 @@ struct TemperedLeapfrog{FT<:AbstractFloat,T<:AbstractScalarOrVec{FT}} <: Abstrac
     α::FT
 end
 
-function Base.show(io::IO, l::TemperedLeapfrog)
+function Base.show(io::IO, mime::MIME"text/plain", l::TemperedLeapfrog)
     return print(
-        io, "TemperedLeapfrog(ϵ=$(round.(l.ϵ; sigdigits=3)), α=$(round.(l.α; sigdigits=3)))"
+        io,
+        "TemperedLeapfrog with step size ϵ=$(round.(l.ϵ; sigdigits=3)) and temperature parameter α=$(round.(l.α; sigdigits=3))",
     )
 end
 

src/metric.jl

@@ -33,8 +33,12 @@ 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)
-    return print(io, "UnitEuclideanMetric($(_string_M⁻¹(ones(uem.size))))")
+function Base.show(io::IO, ::MIME"text/plain", uem::UnitEuclideanMetric{T}) where {T}
+    return print(
+        io,
+        "UnitEuclideanMetric{$T} with size $(size(uem)) mass matrix:\n",
+        _string_M⁻¹(ones(uem.size)),
+    )
 end
 
 struct DiagEuclideanMetric{T,A<:AbstractVecOrMat{T}} <: AbstractMetric
@@ -58,8 +62,12 @@ 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)
-    return print(io, "DiagEuclideanMetric($(_string_M⁻¹(dem.M⁻¹)))")
+function Base.show(io::IO, ::MIME"text/plain", dem::DiagEuclideanMetric{T}) where {T}
+    return print(
+        io,
+        "DiagEuclideanMetric{$T} with size $(size(dem)) mass matrix:\n",
+        _string_M⁻¹(dem.M⁻¹),
+    )
 end
 
 struct DenseEuclideanMetric{
@@ -94,8 +102,12 @@ 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)
-    return print(io, "DenseEuclideanMetric(diag=$(_string_M⁻¹(dem.M⁻¹)))")
+function Base.show(io::IO, ::MIME"text/plain", dem::DenseEuclideanMetric{T}) where {T}
+    return print(
+        io,
+        "DenseEuclideanMetric{$T} with size $(size(dem)) mass matrix:\n",
+        _string_M⁻¹(dem.M⁻¹),
+    )
 end
 
 # `rand` functions for `metric` types.

src/trajectory.jl

@@ -108,7 +108,12 @@ struct SliceTS{F<:AbstractFloat,P<:PhasePoint} <: AbstractTrajectorySampler
     n::Int
 end
 
-Base.show(io::IO, s::SliceTS) = print(io, "SliceTS(ℓu=$(s.ℓu), n=$(s.n))")
+function Base.show(io::IO, mime::MIME"text/plain", s::SliceTS)
+    return print(
+        io,
+        "SliceTS with slice variable ℓu=$(s.ℓu) and number of acceptable candiadtes n=$(s.n)",
+    )
+end
 
 """
 $(TYPEDEF)
@@ -217,9 +222,10 @@ end
 
 ConstructionBase.constructorof(::Type{<:Trajectory{TS}}) where {TS} = Trajectory{TS}
 
-function Base.show(io::IO, τ::Trajectory{TS}) where {TS}
+function Base.show(io::IO, mime::MIME"text/plain", τ::Trajectory{TS}) where {TS}
     return print(
-        io, "Trajectory{$TS}(integrator=$(τ.integrator), tc=$(τ.termination_criterion))"
+        io,
+        "Trajectory{$TS} with $(τ.integrator) and termination criterion $(τ.termination_criterion)",
     )
 end
 
@@ -468,8 +474,10 @@ struct Termination
     numerical::Bool
 end
 
-function Base.show(io::IO, d::Termination)
-    return print(io, "Termination(dynamic=$(d.dynamic), numerical=$(d.numerical))")
+function Base.show(io::IO, mime::MIME"text/plain", d::Termination)
+    return print(
+        io, "Termination reasons of (dynamic=$(d.dynamic), numerical=$(d.numerical))"
+    )
 end
 function Base.:*(d1::Termination, d2::Termination)
     return Termination(d1.dynamic || d2.dynamic, d1.numerical || d2.numerical)
@@ -484,12 +492,6 @@ Check termination of a Hamiltonian trajectory.
 function Termination(s::SliceTS, nt::Trajectory, H0::F, H′::F) where {F<:AbstractFloat}
     return Termination(false, !(s.ℓu < nt.termination_criterion.Δ_max + -H′))
 end
-
-"""
-$(SIGNATURES)
-
-Check termination of a Hamiltonian trajectory.
-"""
 function Termination(
     s::MultinomialTS, nt::Trajectory, H0::F, H′::F
 ) where {F<:AbstractFloat}