Skip to content

Commit

Permalink
Format .jl files (#194)
Browse files Browse the repository at this point in the history 
Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
  • Loading branch information
@github-actions
github-actions[bot] committed Apr 19, 2020
1 parent e17fbee commit ca340b4
Show file tree
Hide file tree
Showing 25 changed files with 837 additions and 587 deletions.
57 changes: 35 additions & 22 deletions src/AdvancedHMC.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,8 @@ module AdvancedHMC
const DEBUG = convert(Bool, parse(Int, get(ENV, "DEBUG_AHMC", "0")))

using Statistics: mean, var, middle
using LinearAlgebra: Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky, UniformScaling
using LinearAlgebra:
Symmetric, UpperTriangular, mul!, ldiv!, dot, I, diag, cholesky, UniformScaling
using StatsFuns: logaddexp, logsumexp
using Random: GLOBAL_RNG, AbstractRNG
using ProgressMeter: ProgressMeter
Expand Down Expand Up @@ -33,48 +34,60 @@ export Leapfrog, JitteredLeapfrog, TemperedLeapfrog

include("trajectory.jl")
@deprecate find_good_eps find_good_stepsize
export EndPointTS, SliceTS, MultinomialTS,
StaticTrajectory, HMCDA, NUTS,
ClassicNoUTurn, GeneralisedNoUTurn,
find_good_stepsize
export EndPointTS,
SliceTS,
MultinomialTS,
StaticTrajectory,
HMCDA,
NUTS,
ClassicNoUTurn,
GeneralisedNoUTurn,
find_good_stepsize

include("adaptation/Adaptation.jl")
using .Adaptation
import .Adaptation: StepSizeAdaptor, MassMatrixAdaptor, StanHMCAdaptor, NesterovDualAveraging
import .Adaptation:
StepSizeAdaptor, MassMatrixAdaptor, StanHMCAdaptor, NesterovDualAveraging

# Helpers for initializing adaptors via AHMC structs

StepSizeAdaptor(δ::AbstractFloat, stepsize::AbstractScalarOrVec{<:AbstractFloat}) =
StepSizeAdaptor(δ::AbstractFloat, stepsize::AbstractScalarOrVec{<:AbstractFloat}) =
NesterovDualAveraging(δ, stepsize)
StepSizeAdaptor(δ::AbstractFloat, i::AbstractIntegrator) = StepSizeAdaptor(δ, nom_step_size(i))
StepSizeAdaptor(δ::AbstractFloat, i::AbstractIntegrator) =
StepSizeAdaptor(δ, nom_step_size(i))

MassMatrixAdaptor(m::UnitEuclideanMetric{T}) where {T} =
UnitMassMatrix{T}()
MassMatrixAdaptor(m::UnitEuclideanMetric{T}) where {T} = UnitMassMatrix{T}()
MassMatrixAdaptor(m::DiagEuclideanMetric{T}) where {T} =
WelfordVar{T}(size(m); var=copy(m.M⁻¹))
WelfordVar{T}(size(m); var = copy(m.M⁻¹))
MassMatrixAdaptor(m::DenseEuclideanMetric{T}) where {T} =
WelfordCov{T}(size(m); cov=copy(m.M⁻¹))
WelfordCov{T}(size(m); cov = copy(m.M⁻¹))

MassMatrixAdaptor(
m::Type{TM},
sz::Tuple{Vararg{Int}}=(2,)
) where {TM<:AbstractMetric} = MassMatrixAdaptor(Float64, m, sz)
MassMatrixAdaptor(m::Type{TM}, sz::Tuple{Vararg{Int}} = (2,)) where {TM<:AbstractMetric} =
MassMatrixAdaptor(Float64, m, sz)

MassMatrixAdaptor(
::Type{T},
::Type{TM},
sz::Tuple{Vararg{Int}}=(2,)
) where {T, TM<:AbstractMetric} = MassMatrixAdaptor(TM(T, sz))
sz::Tuple{Vararg{Int}} = (2,),
) where {T,TM<:AbstractMetric} = MassMatrixAdaptor(TM(T, sz))

# Deprecations

@deprecate StanHMCAdaptor(n_adapts, pc, ssa) initialize!(StanHMCAdaptor(pc, ssa), n_adapts)
@deprecate NesterovDualAveraging(δ::AbstractFloat, i::AbstractIntegrator) StepSizeAdaptor(δ, i)
@deprecate NesterovDualAveraging(δ::AbstractFloat, i::AbstractIntegrator) StepSizeAdaptor(
δ,
i,
)
@deprecate Preconditioner(args...) MassMatrixAdaptor(args...)

export StepSizeAdaptor, NesterovDualAveraging,
MassMatrixAdaptor, UnitMassMatrix, WelfordVar, WelfordCov,
NaiveHMCAdaptor, StanHMCAdaptor
export StepSizeAdaptor,
NesterovDualAveraging,
MassMatrixAdaptor,
UnitMassMatrix,
WelfordVar,
WelfordCov,
NaiveHMCAdaptor,
StanHMCAdaptor

include("diagnosis.jl")

Expand Down
11 changes: 6 additions & 5 deletions src/adaptation/Adaptation.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,11 +28,12 @@ export MassMatrixAdaptor, UnitMassMatrix, WelfordVar, WelfordCov
## TODO: generalise this to a list of adaptors
##

struct NaiveHMCAdaptor{M<:MassMatrixAdaptor, Tssa<:StepSizeAdaptor} <: AbstractAdaptor
pc :: M
ssa :: Tssa
struct NaiveHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAdaptor
pc::M
ssa::Tssa
end
Base.show(io::IO, a::NaiveHMCAdaptor) = print(io, "NaiveHMCAdaptor(pc=$(a.pc), ssa=$(a.ssa))")
Base.show(io::IO, a::NaiveHMCAdaptor) =
print(io, "NaiveHMCAdaptor(pc=$(a.pc), ssa=$(a.ssa))")

getM⁻¹(ca::NaiveHMCAdaptor) = getM⁻¹(ca.pc)
getϵ(ca::NaiveHMCAdaptor) = getϵ(ca.ssa)
Expand All @@ -41,7 +42,7 @@ getϵ(ca::NaiveHMCAdaptor) = getϵ(ca.ssa)
function adapt!(
nca::NaiveHMCAdaptor,
θ::AbstractVecOrMat{<:AbstractFloat},
α::AbstractScalarOrVec{<:AbstractFloat}
α::AbstractScalarOrVec{<:AbstractFloat},
)
adapt!(nca.ssa, θ, α)
adapt!(nca.pc, θ, α)
Expand Down
62 changes: 34 additions & 28 deletions src/adaptation/massmatrix.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -11,7 +11,7 @@ function adapt!(
adaptor::MassMatrixAdaptor,
θ::AbstractVecOrMat{<:AbstractFloat},
α::AbstractScalarOrVec{<:AbstractFloat},
is_update::Bool=true
is_update::Bool = true,
)
resize!(adaptor, θ)
push!(adaptor, θ)
Expand All @@ -38,7 +38,7 @@ adapt!(
::UnitMassMatrix,
::AbstractVecOrMat{<:AbstractFloat},
::AbstractScalarOrVec{<:AbstractFloat},
is_update::Bool=true
is_update::Bool = true,
) = nothing

## Diagonal mass matrix adaptor
Expand All @@ -54,15 +54,16 @@ function update!(ve::DiagMatrixEstimator)
end

# NOTE: this naive variance estimator is used only in testing
struct NaiveVar{T<:AbstractFloat, E<:AbstractVector{<:AbstractVecOrMat{T}}} <: DiagMatrixEstimator{T}
S :: E
NaiveVar(S::E) where {E} = new{eltype(eltype(E)), E}(S)
struct NaiveVar{T<:AbstractFloat,E<:AbstractVector{<:AbstractVecOrMat{T}}} <:
DiagMatrixEstimator{T}
S::E
NaiveVar(S::E) where {E} = new{eltype(eltype(E)),E}(S)
end

NaiveVar{T}(sz::Tuple{Int}) where {T<:AbstractFloat} = NaiveVar(Vector{Vector{T}}())
NaiveVar{T}(sz::Tuple{Int,Int}) where {T<:AbstractFloat} = NaiveVar(Vector{Matrix{T}}())

NaiveVar(sz::Union{Tuple{Int}, Tuple{Int,Int}}) = NaiveVar{Float64}(sz)
NaiveVar(sz::Union{Tuple{Int},Tuple{Int,Int}}) = NaiveVar{Float64}(sz)

Base.push!(nv::NaiveVar, s::AbstractVecOrMat) = push!(nv.S, s)

Expand All @@ -74,28 +75,30 @@ function get_estimation(nv::NaiveVar)
end

# Ref： https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_var_estimator.hpp
mutable struct WelfordVar{T<:AbstractFloat, E<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}
n :: Int
n_min :: Int
μ :: E
M :: E
δ :: E # cache for diff
var :: E # cache for variance
mutable struct WelfordVar{T<:AbstractFloat,E<:AbstractVecOrMat{T}} <: DiagMatrixEstimator{T}
n::Int
n_min::Int
μ::E
M::E
δ::E # cache for diff
var::E # cache for variance
function WelfordVar(n::Int, n_min::Int, μ::E, M::E, δ::E, var::E) where {E}
return new{eltype(E), E}(n, n_min, μ, M, δ, var)
return new{eltype(E),E}(n, n_min, μ, M, δ, var)
end
end

Base.show(io::IO, ::WelfordVar) = print(io, "WelfordVar")

function WelfordVar{T}(
sz::Union{Tuple{Int}, Tuple{Int,Int}};
n_min::Int=10, var=ones(T, sz)
sz::Union{Tuple{Int},Tuple{Int,Int}};
n_min::Int = 10,
var = ones(T, sz),
) where {T<:AbstractFloat}
return WelfordVar(0, n_min, zeros(T, sz), zeros(T, sz), zeros(T, sz), var)
end

WelfordVar(sz::Union{Tuple{Int}, Tuple{Int,Int}}; kwargs...) = WelfordVar{Float64}(sz; kwargs...)
WelfordVar(sz::Union{Tuple{Int},Tuple{Int,Int}}; kwargs...) =
WelfordVar{Float64}(sz; kwargs...)

function Base.resize!(wv::WelfordVar, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat}
if size(θ) != size(wv.var)
Expand Down Expand Up @@ -143,14 +146,15 @@ function update!(ce::DenseMatrixEstimator)
end

# NOTE: This naive covariance estimator is used only in testing.
struct NaiveCov{F<:AbstractFloat, T<:AbstractVector{<:AbstractVector{F}}} <: DenseMatrixEstimator{T}
S :: T
NaiveCov(S::E) where {E} = new{eltype(eltype(E)), E}(S)
struct NaiveCov{F<:AbstractFloat,T<:AbstractVector{<:AbstractVector{F}}} <:
DenseMatrixEstimator{T}
S::T
NaiveCov(S::E) where {E} = new{eltype(eltype(E)),E}(S)
end

NaiveCov{T}(sz::Tuple{Int}) where {T<:AbstractFloat} = NaiveCov(Vector{Vector{T}}())

NaiveCov(sz::Union{Tuple{Int}, Tuple{Int,Int}}; kwargs...) = NaiveCov{Float64}(sz; kwargs...)
NaiveCov(sz::Union{Tuple{Int},Tuple{Int,Int}}; kwargs...) = NaiveCov{Float64}(sz; kwargs...)

Base.push!(nc::NaiveCov, s::AbstractVector) = push!(nc.S, s)

Expand All @@ -163,18 +167,20 @@ end

# Ref: https://github.com/stan-dev/math/blob/develop/stan/math/prim/mat/fun/welford_covar_estimator.hpp
mutable struct WelfordCov{F<:AbstractFloat} <: DenseMatrixEstimator{F}
n :: Int
n_min :: Int
μ :: Vector{F}
M :: Matrix{F}
δ :: Vector{F} # cache for diff
cov :: Matrix{F}
n::Int
n_min::Int
μ::Vector{F}
M::Matrix{F}
δ::Vector{F} # cache for diff
cov::Matrix{F}
end

Base.show(io::IO, ::WelfordCov) = print(io, "WelfordCov")

function WelfordCov{T}(
sz::Tuple{Int}; n_min::Int=10, cov=LinearAlgebra.diagm(0 => ones(T, first(sz)))
sz::Tuple{Int};
n_min::Int = 10,
cov = LinearAlgebra.diagm(0 => ones(T, first(sz))),
) where {T<:AbstractFloat}
d = first(sz)
return WelfordCov(0, n_min, zeros(T, d), zeros(T, d, d), zeros(T, d), cov)
Expand Down
75 changes: 51 additions & 24 deletions src/adaptation/stan_adaptor.jl
View file
Original file line number Diff line number Diff line change
@@ -1,16 +1,22 @@
### Mutable states

mutable struct StanHMCAdaptorState
i :: Int
window_start :: Int # start of mass matrix adaptation
window_end :: Int # end of mass matrix adaptation
window_splits :: Vector{Int} # iterations to update metric using mass matrix
i::Int
window_start::Int # start of mass matrix adaptation
window_end::Int # end of mass matrix adaptation
window_splits::Vector{Int} # iterations to update metric using mass matrix
end

StanHMCAdaptorState() = StanHMCAdaptorState(0, 0, 0, Vector{Int}(undef, 0))

# Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/windowed_adaptation.hpp
function initialize!(state::StanHMCAdaptorState, init_buffer::Int, term_buffer::Int, window_size::Int, n_adapts::Int)
function initialize!(
state::StanHMCAdaptorState,
init_buffer::Int,
term_buffer::Int,
window_size::Int,
n_adapts::Int,
)
window_start = init_buffer + 1
window_end = n_adapts - term_buffer

Expand All @@ -37,55 +43,76 @@ function initialize!(state::StanHMCAdaptorState, init_buffer::Int, term_buffer::
end
end

state.window_start = window_start
state.window_end = window_end
state.window_start = window_start
state.window_end = window_end
state.window_splits = window_splits
end

function Base.show(io::IO, state::StanHMCAdaptorState)
print(io, "window($(state.window_start), $(state.window_end)), window_splits(" * string(join(state.window_splits, ", ")) * ")")
print(
io,
"window($(state.window_start), $(state.window_end)), window_splits(" *
string(join(state.window_splits, ", ")) *
")",
)
end

### Stan's windowed adaptation

# Acknowledgement: this adaption settings is mimicing Stan's 3-phase adaptation.
struct StanHMCAdaptor{M<:MassMatrixAdaptor, Tssa<:StepSizeAdaptor} <: AbstractAdaptor
pc :: M
ssa :: Tssa
init_buffer :: Int
term_buffer :: Int
window_size :: Int
state :: StanHMCAdaptorState
struct StanHMCAdaptor{M<:MassMatrixAdaptor,Tssa<:StepSizeAdaptor} <: AbstractAdaptor
pc::M
ssa::Tssa
init_buffer::Int
term_buffer::Int
window_size::Int
state::StanHMCAdaptorState
end
Base.show(io::IO, a::StanHMCAdaptor) =
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)")
Base.show(io::IO, a::StanHMCAdaptor) = 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)",
)

function StanHMCAdaptor(
pc::MassMatrixAdaptor,
ssa::StepSizeAdaptor;
init_buffer::Int=75,
term_buffer::Int=50,
window_size::Int=25
init_buffer::Int = 75,
term_buffer::Int = 50,
window_size::Int = 25,
)
return StanHMCAdaptor(pc, ssa, init_buffer, term_buffer, window_size, StanHMCAdaptorState())
return StanHMCAdaptor(
pc,
ssa,
init_buffer,
term_buffer,
window_size,
StanHMCAdaptorState(),
)
end

getM⁻¹(ca::StanHMCAdaptor) = getM⁻¹(ca.pc)
getϵ(ca::StanHMCAdaptor) = getϵ(ca.ssa)

function initialize!(adaptor::StanHMCAdaptor, n_adapts::Int)
initialize!(adaptor.state, adaptor.init_buffer, adaptor.term_buffer, adaptor.window_size, n_adapts)
initialize!(
adaptor.state,
adaptor.init_buffer,
adaptor.term_buffer,
adaptor.window_size,
n_adapts,
)
return adaptor
end
finalize!(adaptor::StanHMCAdaptor) = finalize!(adaptor.ssa)

is_in_window(a::StanHMCAdaptor) = a.state.i >= a.state.window_start && a.state.i <= a.state.window_end
is_in_window(a::StanHMCAdaptor) =
a.state.i >= a.state.window_start && a.state.i <= a.state.window_end
is_window_end(a::StanHMCAdaptor) = a.state.i in a.state.window_splits

function adapt!(
tp::StanHMCAdaptor,
θ::AbstractVecOrMat{<:AbstractFloat},
α::AbstractScalarOrVec{<:AbstractFloat}
α::AbstractScalarOrVec{<:AbstractFloat},
)
tp.state.i += 1

Expand Down

0 comments on commit ca340b4

Please sign in to comment.