Skip to content

Commit

Permalink
Update MCMCChain ordering
Browse files Browse the repository at this point in the history
  • Loading branch information
paschermayr committed Dec 28, 2022
1 parent 5ec10fe commit ece8366
Show file tree
Hide file tree
Showing 4 changed files with 28 additions and 21 deletions.
4 changes: 2 additions & 2 deletions Project.toml
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
name = "Baytes"
uuid = "72ddfcfc-6e9d-43df-829b-7aed7c549d4f"
authors = ["Patrick Aschermayr <p.aschermayr@gmail.com>"]
version = "0.2.1"
version = "0.2.2"

[deps]
ArgCheck = "dce04be8-c92d-5529-be00-80e4d2c0e197"
Expand Down Expand Up @@ -30,7 +30,7 @@ BaytesPMCMC = "0.2"
BaytesSMC = "0.2"
DocStringExtensions = "0.8, 0.9"
JLD2 = "0.4"
MCMCDiagnosticTools = "0.1"
MCMCDiagnosticTools = "0.2"
ModelWrappers = "0.3"
PrettyTables = "2"
ProgressMeter = "1.7"
Expand Down
14 changes: 7 additions & 7 deletions src/sampling/chain.jl
Original file line number Diff line number Diff line change
Expand Up @@ -54,7 +54,7 @@ function paramdiagnostics(vec::AbstractVector{T}) where {T<:Real}
end
function paramdiagnostics(arr::Array{T,3}) where {T<:Real}
return map(
iter -> paramdiagnostics(vec(view(arr, :, iter, :))), Base.OneTo(size(arr, 2))
iter -> paramdiagnostics(vec(view(arr, :, :, iter))), Base.OneTo(size(arr, 3))
)
end

Expand All @@ -75,8 +75,8 @@ function paramquantiles(vec::AbstractVector{T}, printdefault::PrintDefault) wher
end
function paramquantiles(arr::Array{T,3}, printdefault::PrintDefault) where {T<:Real}
return map(
iter -> paramquantiles(vec(view(arr, :, iter, :)), printdefault),
Base.OneTo(size(arr, 2)),
iter -> paramquantiles(vec(view(arr, :, :, iter)), printdefault),
Base.OneTo(size(arr, 3)),
)
end

Expand Down Expand Up @@ -113,10 +113,10 @@ Check if any parameter has been stuck at each iteration in any chain, in which c
"""
function is_stuck(arr3D::AbstractArray)
# Loop through chains and parameter to check if first parameter is equal to all samples == chain stuck
for Nchains in Base.OneTo( size(arr3D,3) )
for Nparams in Base.OneTo( size(arr3D,2) )
_benchmark = arr3D[begin,Nparams,Nchains]
stuck = all(val -> val == _benchmark, @view( arr3D[:,Nparams,Nchains] ))
for Nparams in Base.OneTo( size(arr3D, 3) )
for Nchains in Base.OneTo( size(arr3D, 2) )
_benchmark = arr3D[begin, Nchains, Nparams]
stuck = all(val -> val == _benchmark, @view( arr3D[:, Nchains, Nparams] ))
if stuck
return true, (Nparams, Nchains)
end
Expand Down
24 changes: 14 additions & 10 deletions src/sampling/inference.jl
Original file line number Diff line number Diff line change
Expand Up @@ -127,7 +127,8 @@ end
################################################################################
"""
$(SIGNATURES)
Change trace.val to 3d Array that is consistent with MCMCCHains dimensons. First dimension is iterations, second number of parameter, third number of chains.
Change trace.val to 3d Array that is consistent with MCMCCHains dimensons.
First dimension is iterations, second number of parameter, third number of chains.
# Examples
```julia
Expand All @@ -141,13 +142,15 @@ function trace_to_3DArray(
## Get trace information
@unpack tagged, chains, effective_iterations = transform
## Preallocate array
mcmcchain = zeros(length(effective_iterations), length(tagged), length(chains))
mcmcchain = zeros(length(effective_iterations), length(chains), length(tagged))
## Flatten corresponding parameter
#!NOTE: This is threadsave, but chain is not flattened in correct ordered, which might be troublesome for MCMC chain analysis.
#!NOTE: Commented-out loop below is threadsave, but chain is not flattened in correct order, which might be troublesome for MCMC chain analysis.
# Threads.@threads for (idx, chain) in collect(enumerate(chains))
for (idx, chain) in collect(enumerate(chains))
#!NOTE: idx and chain are separated in case not all chains are used for summary
for (iter0, iterburnin) in enumerate(effective_iterations)
mcmcchain[iter0, :, idx] .= flatten(tagged.info.reconstruct, subset(trace.val[chain][iterburnin], tagged.parameter))
#!NOTE: New way of structuring parameter: (ndraws, nchains, nparams)
mcmcchain[iter0, idx, :] .= flatten(tagged.info.reconstruct, subset(trace.val[chain][iterburnin], tagged.parameter))
end
end
## Return MCMCChain
Expand All @@ -170,13 +173,14 @@ function trace_to_3DArrayᵤ(
## Get trace information
@unpack tagged, chains, effective_iterations = transform
## Preallocate array
mcmcchain = zeros(length(effective_iterations), length(tagged), length(chains))
mcmcchain = zeros(length(effective_iterations), length(chains), length(tagged))
## Flatten corresponding parameter
#!NOTE: This is threadsave, but chain is not flattened in correct ordered, which might be troublesome for MCMC chain analysis.
#!NOTE: This is threadsave, but chain is not flattened in correct order, which might be troublesome for MCMC chain analysis.
# Threads.@threads for (idx, chain) in collect(enumerate(chains))
for (idx, chain) in collect(enumerate(chains))
for (iter0, iterburnin) in enumerate(effective_iterations)
mcmcchain[iter0, :, idx] .=
# mcmcchain[iter0, :, idx] .=
mcmcchain[iter0, idx, :] .=
flatten(tagged.info.reconstruct,
unconstrain(tagged.info.transform, subset(trace.val[chain][iterburnin], tagged.parameter) )
)
Expand Down Expand Up @@ -205,7 +209,7 @@ function trace_to_2DArray(
## Preallocate array
mcmcchain = zeros(length(effective_iterations) * length(chains), length(tagged))
## Flatten corresponding parameter
#!NOTE: This is threadsave, but chain is not flattened in correct ordered, which might be troublesome for MCMC chain analysis.
#!NOTE: This is threadsave, but chain is not flattened in correct order, which might be troublesome for MCMC chain analysis.
# Threads.@threads for (idx, chain) in collect(enumerate(chains))
iter = 0
for (idx, chain) in collect(enumerate(chains))
Expand Down Expand Up @@ -265,7 +269,7 @@ function trace_to_posteriormean(
transform::TraceTransform
)
@unpack tagged = transform
mod_array_mean = map(iter -> mean(view(mod_array, :, iter, :)), Base.OneTo(size(mod_array, 2)))
mod_array_mean = map(iter -> mean(view(mod_array, :, :, iter)), Base.OneTo(size(mod_array, 3)))
mod_nt_mean = ModelWrappers.unflatten(tagged.info.reconstruct, mod_array_mean)
return mod_array_mean, mod_nt_mean
end
Expand Down Expand Up @@ -329,7 +333,7 @@ function flatten_chainvals(
## Preallocate array
mcmcchain = [ [ zeros(tagged.info.reconstruct.default.output, length(tagged)) for _ in eachindex(effective_iterations) ] for _ in eachindex(chains) ]
## Flatten corresponding parameter
#!NOTE: This is threadsave, but chain is not flattened in correct ordered, which might be troublesome for MCMC chain analysis, hence we opt out of it.
#!NOTE: This is threadsave, but chain is not flattened in correct order, which might be troublesome for MCMC chain analysis, hence we opt out of it.
# Threads.@threads for (idx, chain) in collect(enumerate(chains))
for (idx, chain) in collect(enumerate(chains))
for (iter0, iterburnin) in enumerate(effective_iterations)
Expand Down
7 changes: 5 additions & 2 deletions test/test-construction.jl
Original file line number Diff line number Diff line change
Expand Up @@ -277,11 +277,14 @@ end
end

@testset "Utility, check if chain stuck" begin
_chain = randn(_rng, 10, 15, 4)
_Ndraws = 1000
_Nparams = 10
_Nchains = 4
_chain = randn(_rng, _Ndraws, _Nchains, _Nparams)
_chain2 = deepcopy(_chain)
param_stuck = 7
chain_stuck = 3
_chain2[:, param_stuck, chain_stuck] .= 1.0
_chain2[:, chain_stuck, param_stuck] .= 1.0

stuck, paramchain = Baytes.is_stuck(_chain)
stuck2, paramchain2 = Baytes.is_stuck(_chain2)
Expand Down

2 comments on commit ece8366

@paschermayr
Copy link
Owner Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

@JuliaRegistrator
Copy link

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Registration pull request created: JuliaRegistries/General/74722

After the above pull request is merged, it is recommended that a tag is created on this repository for the registered package version.

This will be done automatically if the Julia TagBot GitHub Action is installed, or can be done manually through the github interface, or via:

git tag -a v0.2.2 -m "<description of version>" ece83665597764fa1dadf8296a4bc98d79ded529
git push origin v0.2.2

Please sign in to comment.