Extracting sampled parameters is broken

[luiarthur]

I know in Turing v0.12.0 (not sure about v0.13.0), after sampling from joint posterior distribution of some parameters, you can extract parameter vectors with chain[:sym], where chain is the output to sample(...), and :sym is a Symbol for some vector (possibly higher-order Array?) of parameters. It seems that's no longer possible in the master branch of Turing?

Here's an example:

using Turing  # from master branch

# Simulate data.
N = 1000
mu_true = [5.0 10.0 2.0]
sigma_true = 0.6
y = randn(N, length(mu_true)) * sigma_true .+ mu_true  # size(y) = (1000, 3)

# Model definition.
@model function demo_model(y)
    K = size(y, 2)
    mu ~ filldist(Normal(0, 10), K)
    sigma ~ LogNormal(0, 2)
    for k in 1:K
        y[:, k] .~ Normal(mu[k], sigma)
    end
end

# Sample from posterior via NUTS.
nadapt, target_accept_ratio = (1000, 0.8)
iterations = 1000 + nadapt
chain = sample(demo_model(y), NUTS(nadapt, target_accept_ratio), iterations)

Everything works fine up to this point.

This also works as expected, as :sigma is a univariate parameter.

# Get sigma
chain[:sigma]  # works for univariate parameters

# 2-dimensional AxisArray{Float64,2,...} with axes:
#     :iter, 1:1:1000
#     :chain, 1:1
# And data, a 1000×1 Array{Float64,2}:
#  0.5891087703064164
#  0.5982078191819813
#  0.6014525731322462
#  0.6049135394153814
#  0.6135096911950264
#  ⋮
#  0.6045248767595608
#  0.6043856761562517
#  0.6197436168542185
#  0.6036546321663407

This doesn't work as :mu is a vector parameter.

# FIXME: Get mu
chain[:mu]  # broken

Here's the error:

ERROR: ArgumentError: index mu not found
Stacktrace:
 [1] axisindexes at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:317 [inlined]
 [2] axisindexes at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:199 [inlined]
 [3] macro expansion at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:394 [inlined]
 [4] _to_index(::AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},Ax
isArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}}, ::Tuple{AxisArrays.Unsupported,AxisArra
ys.Unsupported,AxisArrays.Unsupported}, ::Colon, ::Symbol, ::Colon) at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/sr
c/indexing.jl:350
 [5] to_index at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:347 [inlined]
 [6] getindex at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:123 [inlined]
 [7] getindex(::Chains{Float64,AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{In
t64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:param
eters, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}, ::Function, ::Symbol, ::Function) a
t /home/ubuntu/.julia/packages/MCMCChains/EaAHc/src/chains.jl:113
 [8] getindex(::Chains{Float64,AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{In
t64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:param
eters, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}, ::Symbol) at /home/ubuntu/.julia/pa
ckages/MCMCChains/EaAHc/src/chains.jl:110
 [9] top-level scope at REPL[87]:1

Here's my slightly awkward workaround. But is there another way to achieve this right now?

# Workaround for getting array parameters
function extract_array_param(chain, sym)
    rgx = Regex("$(sym)\\[\\d+(,\\d+)*\\]")
    syms = filter(x -> x != nothing, match.(rgx, String.(chain.name_map[1])))
    syms = map(x -> Symbol(x.match), syms)
    return chain[syms]
end

extract_array_param(chain, :mu)  # This works.

I think this is related, but I can no longer use these macro strings to extract loglikelihood information.

# FIXME: Get log-likelihood.
loglike = logprob"y=y | model=demo_model, chain=chain"  # broken

Here is the error:

ERROR: ArgumentError: index mu not found
Stacktrace:
 [1] axisindexes at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:317 [inlined]
 [2] axisindexes at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:199 [inlined]
 [3] macro expansion at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:394 [inlined]
 [4] _to_index(::AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int
64}}}}, ::Tuple{AxisArrays.Unsupported,AxisArrays.Unsupported,AxisArrays.Unsupported}, ::Colon, ::Symbol, ::Colon) at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:350
 [5] to_index at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:347 [inlined]
 [6] getindex at /home/ubuntu/.julia/packages/AxisArrays/IFpjG/src/indexing.jl:123 [inlined]
 [7] getindex(::Chains{Float64,AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain
,UnitRange{Int64}}}},Missing,NamedTuple{(:parameters, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}, ::Function, ::Symbol, ::Function) at /home/ubuntu/.julia/pac
kages/MCMCChains/EaAHc/src/chains.jl:113
 [8] getindex at /home/ubuntu/.julia/packages/MCMCChains/EaAHc/src/chains.jl:110 [inlined]
 [9] macro expansion at /home/ubuntu/.julia/packages/DynamicPPL/MRwtL/src/prob_macro.jl:236 [inlined]
 [10] _setval!(::NamedTuple{(:mu, :sigma),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:mu,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{Dyn
amicPPL.VarName{:mu,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:sigma,Tuple{}},Int64},Array{LogNormal{Float64},1},Array{Dynami
cPPL.VarName{:sigma,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}}, ::DynamicPPL.VarInfo{NamedTuple{(:mu, :sigma),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:mu,Tupl
e{}},Int64},Array{DistributionsAD.TuringScalMvNormal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:mu,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metad
ata{Dict{DynamicPPL.VarName{:sigma,Tuple{}},Int64},Array{LogNormal{Float64},1},Array{DynamicPPL.VarName{:sigma,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64}, ::Ch
ains{Float64,AxisArrays.AxisArray{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}
}},Missing,NamedTuple{(:parameters, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}) at /home/ubuntu/.julia/packages/DynamicPPL/MRwtL/src/prob_macro.jl:233
 [11] _setval! at /home/ubuntu/.julia/packages/DynamicPPL/MRwtL/src/prob_macro.jl:231 [inlined]
 [12] #105 at /home/ubuntu/.julia/packages/DynamicPPL/MRwtL/src/prob_macro.jl:192 [inlined]
 [13] iterate at ./generator.jl:47 [inlined]
 [14] _collect(::UnitRange{Int64}, ::Base.Generator{UnitRange{Int64},DynamicPPL.var"#105#106"{DynamicPPL.Model{var"###evaluator#450",(:y,),Tuple{Array{Float64,2}},(),DynamicPPL.ModelGen{var"#
##generator#451",(:y,),(),Tuple{}}},DynamicPPL.VarInfo{NamedTuple{(:mu, :sigma),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:mu,Tuple{}},Int64},Array{DistributionsAD.TuringScalMvNormal{
Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:mu,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:sigma,Tuple{}},Int64},Ar
ray{LogNormal{Float64},1},Array{DynamicPPL.VarName{:sigma,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},Chains{Float64,AxisArrays.AxisArray{Float64,3,Array{Float6
4,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:parameters, :internals),Tuple{Ar
ray{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}}}, ::Base.EltypeUnknown, ::Base.HasShape{1}) at ./array.jl:678
 [15] collect_similar(::UnitRange{Int64}, ::Base.Generator{UnitRange{Int64},DynamicPPL.var"#105#106"{DynamicPPL.Model{var"###evaluator#450",(:y,),Tuple{Array{Float64,2}},(),DynamicPPL.ModelGe
n{var"###generator#451",(:y,),(),Tuple{}}},DynamicPPL.VarInfo{NamedTuple{(:mu, :sigma),Tuple{DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:mu,Tuple{}},Int64},Array{DistributionsAD.TuringScalMv
Normal{Array{Float64,1},Float64},1},Array{DynamicPPL.VarName{:mu,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}},DynamicPPL.Metadata{Dict{DynamicPPL.VarName{:sigma,Tuple{}},In
t64},Array{LogNormal{Float64},1},Array{DynamicPPL.VarName{:sigma,Tuple{}},1},Array{Float64,1},Array{Set{DynamicPPL.Selector},1}}}},Float64},Chains{Float64,AxisArrays.AxisArray{Float64,3,Array
{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:parameters, :internals),T
uple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}}}) at ./array.jl:607
 [16] map(::Function, ::UnitRange{Int64}) at ./abstractarray.jl:2072
 [17] loglikelihood(::NamedTuple{(:y,),Tuple{Array{Float64,2}}}, ::NamedTuple{(:model, :chain),Tuple{DynamicPPL.ModelGen{var"###generator#451",(:y,),(),Tuple{}},Chains{Float64,AxisArrays.Axis
Array{Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:para
meters, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}}}, ::DynamicPPL.ModelGen{var"###generator#451",(:y,),(),Tuple{}}, ::Nothing) at /home/ubuntu/.julia/package
s/DynamicPPL/MRwtL/src/prob_macro.jl:190
 [18] logprob(::NamedTuple{(:y,),Tuple{Array{Float64,2}}}, ::NamedTuple{(:model, :chain),Tuple{DynamicPPL.ModelGen{var"###generator#451",(:y,),(),Tuple{}},Chains{Float64,AxisArrays.AxisArray{
Float64,3,Array{Float64,3},Tuple{AxisArrays.Axis{:iter,StepRange{Int64,Int64}},AxisArrays.Axis{:var,Array{Symbol,1}},AxisArrays.Axis{:chain,UnitRange{Int64}}}},Missing,NamedTuple{(:parameters
, :internals),Tuple{Array{Symbol,1},Array{Symbol,1}}},NamedTuple{(),Tuple{}}}}}) at /home/ubuntu/.julia/packages/DynamicPPL/MRwtL/src/prob_macro.jl:29
 [19] top-level scope at REPL[91]:1

[devmotion]

The first error is due to an intentional change in MCMCChains 4.0 (see TuringLang/MCMCChains.jl#197 (comment)). In each Chains object the single entries of vector-valued parameters are named and handled separately (by default they are named mu[1], mu[2] etc.), and only these can be extracted by chain["mu[1]"] (symbols work as well) or chain[:, "mu[1]", :]. To extract a group of parameters you can use group(chain, :mu) or group(chain, "mu") which implements the old indexing that always looked for groups of parameters. However, this is prone to errors if you rearrange or rename parameters (which is one reason for why indexing works differently in MCMCChains 4.0, apart from type stability and other inconsistencies).

The second error should be fixed by TuringLang/DynamicPPL.jl#147.

[luiarthur]

Great, thanks!

I'll close this issue then, since there's already a fix to the first problem, and there's a fix in the works for the other.

[ajinkya-k]

I am facing this exact same issue:

@model rirc(x,y, a₁, b₁, b₂) = begin
    σ² ~ InverseGamma(1/2, 3/2)
    β̄ ~ Normal(0, b₁.*sqrt(σ²))

    M = length(x)
    α = zeros(Real, M)
    β = zeros(Real, M)
    for i in 1:length(M)
        # println("iter: ", α[i])
        α[i] ~ Normal(0, a₁.*sqrt(σ²))
        β[i] ~ Normal(β̄, b₂.*sqrt(σ²))
        for j in 1:length(x[i])
            y[i][j] ~ Normal(α[i] .+ β[i].*x[i][j], sqrt(σ²))
        end
    end
    return σ², α, β
end

The vector parameters alpha and beta dont work properly. For example here the output from names:

julia> names(chain)
13-element Vector{Symbol}:
 :acceptance_rate
 :hamiltonian_energy
 :hamiltonian_energy_error
 :is_accept
 :log_density
 :lp
 :n_steps
 :nom_step_size
 :step_size
 Symbol("α[1]")
 Symbol("β[1]")
 :β̄

Now chain[:,"β[1]",:] works but chain[:,"β[2]",:] doesn't , also group(chain, β) only gives the first component.

Anyone know what i am doing wrong?