Reduce allocations in stepsize.jl (#390)

* Reduce allocations in stepsize.jl

* Fix format

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Fix finalize and add type parameter

* Fix format

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Apply suggestions from code review

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>

* Update src/adaptation/stepsize.jl

---------

Co-authored-by: github-actions[bot] <41898282+github-actions[bot]@users.noreply.github.com>
Co-authored-by: Hong Ge <[email protected]>

Author: 3 people

src/adaptation/stepsize.jl

@@ -8,7 +8,7 @@ mutable struct DAState{T<:AbstractScalarOrVec{<:AbstractFloat}}
     H_bar::T
 end
 
-computeμ(ϵ::AbstractScalarOrVec{<:AbstractFloat}) = log.(10 * ϵ)
+computeμ(ϵ::AbstractFloat) = log(10 * ϵ)
 
 function DAState(ϵ::T) where {T}
     μ = computeμ(ϵ)
@@ -17,22 +17,34 @@ end
 
 function DAState(ϵ::AbstractVector{T}) where {T}
     n = length(ϵ)
-    μ = computeμ(ϵ)
+    μ = map(computeμ, ϵ)
     return DAState(0, ϵ, μ, zeros(T, n), zeros(T, n))
 end
 
 function reset!(das::DAState{T}) where {T<:AbstractFloat}
     das.m = 0
     das.μ = computeμ(das.ϵ)
     das.x_bar = zero(T)
-    return das.H_bar = zero(T)
+    das.H_bar = zero(T)
+    return das
 end
 
 function reset!(das::DAState{<:AbstractVector{T}}) where {T<:AbstractFloat}
     das.m = 0
-    das.μ .= computeμ(das.ϵ)
-    das.x_bar .= zero(T)
-    return das.H_bar .= zero(T)
+    map!(computeμ, das.μ, das.ϵ)
+    fill!(das.x_bar, zero(T))
+    fill!(das.H_bar, zero(T))
+    return das
+end
+
+function finalize!(das::DAState{<:AbstractFloat})
+    das.ϵ = exp(das.x_bar)
+    return das
+end
+
+function finalize!(das::DAState{<:AbstractVector{<:AbstractFloat}})
+    map!(exp, das.ϵ, das.x_bar)
+    return das
 end
 
 mutable struct MSSState{T<:AbstractScalarOrVec{<:AbstractFloat}}
@@ -51,7 +63,7 @@ getϵ(ss::StepSizeAdaptor) = ss.state.ϵ
 struct FixedStepSize{T<:AbstractScalarOrVec{<:AbstractFloat}} <: StepSizeAdaptor
     ϵ::T
 end
-Base.show(io::IO, a::FixedStepSize) = print(io, "FixedStepSize($(a.ϵ))")
+Base.show(io::IO, a::FixedStepSize) = print(io, "FixedStepSize(", a.ϵ, ")")
 
 getϵ(fss::FixedStepSize) = fss.ϵ
 
@@ -82,7 +94,17 @@ end
 function Base.show(io::IO, a::NesterovDualAveraging)
     return print(
         io,
-        "NesterovDualAveraging(γ=$(a.γ), t_0=$(a.t_0), κ=$(a.κ), δ=$(a.δ), state.ϵ=$(getϵ(a)))",
+        "NesterovDualAveraging(γ=",
+        a.γ,
+        ", t_0=",
+        a.t_0,
+        ", κ=",
+        a.κ,
+        ", δ=",
+        a.δ,
+        ", state.ϵ=",
+        getϵ(a),
+        ")",
     )
 end
 
@@ -95,35 +117,29 @@ end
 function NesterovDualAveraging(
     δ::T, ϵ::VT
 ) where {T<:AbstractFloat,VT<:AbstractScalarOrVec{T}}
-    return NesterovDualAveraging(T(0.05), T(10.0), T(0.75), δ, ϵ)
+    return NesterovDualAveraging(T(1//20), T(10), T(3//4), δ, ϵ)
 end
 
 # Ref: https://github.com/stan-dev/stan/blob/develop/src/stan/mcmc/stepsize_adaptation.hpp
 # Note: This function is not merged with `adapt!` to empahsize the fact that
 #       step size adaptation is not dependent on `θ`.
+# Note 2: `da.state` and `α` support vectorised HMC but should do so together. 
 function adapt_stepsize!(
-    da::NesterovDualAveraging{T}, α::AbstractScalarOrVec{<:T}
+    da::NesterovDualAveraging{T}, α::AbstractScalarOrVec{T}
 ) where {T<:AbstractFloat}
     @debug "Adapting step size..." α
 
-    # Clip average MH acceptance probability
-    if α isa AbstractVector
-        α[α .> 1] .= one(T)
-    else
-        α = α > 1 ? one(T) : α
-    end
-
     (; state, γ, t_0, κ, δ) = da
     (; μ, m, x_bar, H_bar) = state
 
     m = m + 1
 
     η_H = one(T) / (m + t_0)
-    H_bar = (one(T) - η_H) * H_bar .+ η_H * (δ .- α)
+    H_bar = (one(T) - η_H) .* H_bar .+ η_H .* (δ .- min.(one(T), α))
 
-    x = μ .- H_bar * sqrt(m) / γ     # x ≡ logϵ
+    x = μ .- H_bar .* (sqrt(m) / γ)     # x ≡ logϵ
     η_x = m^(-κ)
-    x_bar = (one(T) - η_x) * x_bar .+ η_x * x
+    x_bar = (one(T) - η_x) .* x_bar .+ η_x .* x
 
     ϵ = exp.(x)
     @debug "Adapting step size..." new_ϵ = ϵ old_ϵ = da.state.ϵ
@@ -151,9 +167,12 @@ function adapt!(
     return nothing
 end
 
-reset!(da::NesterovDualAveraging) = reset!(da.state)
+function reset!(da::NesterovDualAveraging)
+    reset!(da.state)
+    return da
+end
 
 function finalize!(da::NesterovDualAveraging)
-    da.state.ϵ = exp.(da.state.x_bar)
-    return nothing
+    finalize!(da.state)
+    return da
 end