Unifying trajectories #214
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,4 +1,4 @@ | ||
| .vscode | ||
| .history | ||
| .DS_Store | ||
| Manifest.toml | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -12,8 +12,7 @@ os: | |
| - osx | ||
|
|
||
| julia: | ||
| - 1.3 | ||
| - nightly | ||
|
|
||
| matrix: | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
|
|
@@ -8,13 +8,15 @@ function reconstruct( | |
| return reconstruct(h, metric=metric) | ||
| end | ||
|
|
||
| reconstruct(κ::AbstractKernel, ::AbstractAdaptor) = κ | ||
| function reconstruct( | ||
| κ::AbstractKernel, adaptor::Union{StepSizeAdaptor, NaiveHMCAdaptor, StanHMCAdaptor} | ||
| ) | ||
| # FIXME: this does not support change type of `ϵ` (e.g. Float to Vector) | ||
| τ = reconstruct( | ||
| κ.τ, integrator=update_nom_step_size(κ.τ.integrator, getϵ(adaptor)) | ||
| ) | ||
| return reconstruct(κ, τ=τ) | ||
| end | ||
|
|
||
| function resize(h::Hamiltonian, θ::AbstractVecOrMat{T}) where {T<:AbstractFloat} | ||
|
|
@@ -46,28 +48,28 @@ end | |
| function step( | ||
xukai92 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}}, | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| z::PhasePoint | ||
| ) | ||
| # Refresh momentum | ||
| z = refresh(rng, z, h) | ||
| # Make transition | ||
| return transition(rng, κ, h, z) | ||
| end | ||
|
|
||
| Adaptation.adapt!( | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| adaptor::Adaptation.NoAdaptation, | ||
| i::Int, | ||
| n_adapts::Int, | ||
| θ::AbstractVecOrMat{<:AbstractFloat}, | ||
| α::AbstractScalarOrVec{<:AbstractFloat} | ||
| ) = h, κ, false | ||
|
|
||
| function Adaptation.adapt!( | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| adaptor::AbstractAdaptor, | ||
| i::Int, | ||
| n_adapts::Int, | ||
|
|
@@ -80,10 +82,10 @@ function Adaptation.adapt!( | |
| adapt!(adaptor, θ, α) | ||
| i == n_adapts && finalize!(adaptor) | ||
| h = reconstruct(h, adaptor) | ||
| κ = reconstruct(κ, adaptor) | ||
| isadapted = true | ||
| end | ||
| return h, κ, isadapted | ||
| end | ||
|
|
||
| """ | ||
|
|
@@ -104,7 +106,7 @@ simple_pm_next!(pm, stat::NamedTuple) = ProgressMeter.next!(pm) | |
|
|
||
| sample( | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| θ::AbstractVecOrMat{<:AbstractFloat}, | ||
| n_samples::Int, | ||
| adaptor::AbstractAdaptor=NoAdaptation(), | ||
|
|
@@ -116,7 +118,7 @@ sample( | |
| ) = sample( | ||
| GLOBAL_RNG, | ||
| h, | ||
| κ, | ||
| θ, | ||
| n_samples, | ||
| adaptor, | ||
|
|
@@ -131,7 +133,7 @@ sample( | |
| sample( | ||
| rng::AbstractRNG, | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| θ::AbstractVecOrMat{T}, | ||
| n_samples::Int, | ||
| adaptor::AbstractAdaptor=NoAdaptation(), | ||
|
|
@@ -141,7 +143,7 @@ sample( | |
| progress::Bool=false | ||
| ) | ||
|
|
||
| Sample `n_samples` samples using the proposal `κ` under Hamiltonian `h`. | ||
| - The randomness is controlled by `rng`. | ||
| - If `rng` is not provided, `GLOBAL_RNG` will be used. | ||
| - The initial point is given by `θ`. | ||
|
|
@@ -154,7 +156,7 @@ Sample `n_samples` samples using the proposal `τ` under Hamiltonian `h`. | |
| function sample( | ||
| rng::Union{AbstractRNG, AbstractVector{<:AbstractRNG}}, | ||
| h::Hamiltonian, | ||
| κ::AbstractKernel, | ||
| θ::T, | ||
| n_samples::Int, | ||
| adaptor::AbstractAdaptor=NoAdaptation(), | ||
|
|
@@ -174,18 +176,18 @@ function sample( | |
| pm = progress ? ProgressMeter.Progress(n_samples, desc="Sampling", barlen=31) : nothing | ||
| time = @elapsed for i = 1:n_samples | ||
| # Make a step | ||
| t = transition(rng, h, κ, t.z) | ||
xukai92 marked this conversation as resolved.
Show resolved
Hide resolved
|
||
| # Adapt h and κ; what mutable is the adaptor | ||
| tstat = stat(t) | ||
| h, κ, isadapted = adapt!(h, κ, adaptor, i, n_adapts, t.z.θ, tstat.acceptance_rate) | ||
| tstat = merge(tstat, (is_adapt=isadapted,)) | ||
| # Update progress meter | ||
| if progress | ||
| # Do include current iteration and mass matrix | ||
| pm_next!(pm, (iterations=i, tstat..., mass_matrix=h.metric)) | ||
| # Report finish of adapation | ||
| elseif verbose && isadapted && i == n_adapts | ||
| @info "Finished $n_adapts adapation steps" adaptor κ.τ.integrator h.metric | ||
| end | ||
| # Store sample | ||
| if !drop_warmup || i > n_adapts | ||
|
|
@@ -204,7 +206,7 @@ function sample( | |
| EBFMI_est = "[" * join(EBFMI_est, ", ") * "]" | ||
| average_acceptance_rate = "[" * join(average_acceptance_rate, ", ") * "]" | ||
| end | ||
| @info "Finished $n_samples sampling steps for $n_chains chains in $time (s)" h κ EBFMI_est average_acceptance_rate | ||
| end | ||
| return θs, stats | ||
| end | ||
Oops, something went wrong.
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Maybe consider the following for clarity and performance?