keep parameter order

[a1ix2]

Addresses #443

[codecov]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 83.75%. Comparing base (aa6161b) to head (12c5b66).
Report is 1 commits behind head on master.

Additional details and impacted files

@@           Coverage Diff           @@
##           master     #466   +/-   ##
=======================================
  Coverage   83.75%   83.75%           
=======================================
  Files          20       20           
  Lines        1071     1071           
=======================================
  Hits          897      897           
  Misses        174      174           

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[mhauru]

Looks good to me, thanks @a1ix2! Did you run into a case where this caused you trouble? If so, do you have a reproducible example that triggers the issue (even non-minimal)? Would be good to have a test for this too, but the original issue you reference didn't come with an MWE.

[a1ix2]

For me it was mostly an annoyance when trying to interoperate with ComponentArrays when using get/get_params. Here's a more or less useless example but still.

using Turing, Distributions

@model function fakemodel(xs)
    mu1 ~ Normal(10.0)
    mu2 ~ Normal(10.0)
    sigma1 ~ InverseGamma(2.0, 1.0)
    sigma2 ~ InverseGamma(2.0, 1.0)
    for i in eachindex(xs)
        xs[i] ~ MixtureModel([Normal(mu1, sigma1), Normal(mu2, sigma2)], [0.5, 0.5])
    end
end

xs = rand(MixtureModel([Normal(0.1, 1.0), Normal(3.0, 0.5)], [0.5, 0.5]), 1000)

model = fakemodel(xs)

chain = sample(model, NUTS(), 100)
Chains MCMC chain (1000×16×1 Array{Float64, 3}):

Iterations        = 501:1:1500
Number of chains  = 1
Samples per chain = 1000
Wall duration     = 7.47 seconds
Compute duration  = 7.47 seconds
parameters        = mu1, mu2, sigma1, sigma2
internals         = lp, n_steps, is_accept, acceptance_rate, log_density, hamiltonian_energy, hamiltonian_energy_error, max_hamiltonian_energy_error, tree_depth, numerical_error, step_size, nom_step_size

Summary Statistics
  parameters      mean       std      mcse   ess_bulk   ess_tail      rhat   ess_per_sec 
      Symbol   Float64   Float64   Float64    Float64    Float64   Float64       Float64 

         mu1    0.1273    0.0501    0.0021   575.6213   762.3949    0.9998       77.0577
         mu2    2.9832    0.0298    0.0013   549.6575   614.0007    1.0023       73.5820
      sigma1    0.9569    0.0422    0.0016   713.1709   706.1234    1.0021       95.4713
      sigma2    0.5569    0.0230    0.0011   441.0267   636.4294    1.0131       59.0397

Quantiles
  parameters      2.5%     25.0%     50.0%     75.0%     97.5% 
      Symbol   Float64   Float64   Float64   Float64   Float64 

         mu1    0.0251    0.0938    0.1278    0.1611    0.2225
         mu2    2.9281    2.9612    2.9836    3.0043    3.0422
      sigma1    0.8801    0.9258    0.9561    0.9851    1.0417
      sigma2    0.5132    0.5416    0.5561    0.5718    0.6057


all_chains = get_params(chain)

# Returns an NTuple with a jumble of unordered internals and parameters
keys(allchains)
(:hamiltonian_energy, :n_steps, :numerical_error, :max_hamiltonian_energy_error, :mu2, :sigma1, :hamiltonian_energy_error, :mu1, :is_accept, :sigma2, :log_density, :tree_depth, :step_size, :acceptance_rate, :lp, :nom_step_size)

# Also returns a NamedTuple with unordered parameters
all_param_chains = get(chain, section=:parameters)
(mu1 = [0.16170966805695142; 0.1554824980047796; … ; 0.08014927911102815; 0.0922168468995448;;], mu2 = [2.9505553517683385; 3.0352436810243066; … ; 2.9597771408952247; 2.9428157073776444;;], sigma2 = [0.549196145152957; 0.5526604135099159; … ; 0.5646327205926919; 0.6056774396897618;;], sigma1 = [0.8824624114147762; 0.9787977609583128; … ; 0.9162592661201038; 0.931760395213333;;])

keys(all_param_chains)
(:mu1, :mu2, :sigma2, :sigma1)

Ideally I'd want to obtain a chain as a vector of ComponentArrays, one for each iteration, so that I can splat more easily into other blackbox optimizers like CMAEvolutionStrategy or do some more exotic things with the chain and so on. Right now I have to rely on writing a return statement in the model with all parameters manually listed, which for larger models can quickly become cumbersome.

using ComponentArrays: ComponentArray

@model function fakemodel_ret(xs)
    mu1 ~ Normal(10.0)
    mu2 ~ Normal(10.0)
    sigma1 ~ InverseGamma(2.0, 1.0)
    sigma2 ~ InverseGamma(2.0, 1.0)
    for i in eachindex(xs)
        xs[i] ~ MixtureModel([Normal(mu1, sigma1), Normal(mu2, sigma2)], [0.5, 0.5])
    end
    
    return ComponentArray(mu1=mu1, mu2=mu2, sigma1=sigma1, sigma2=sigma2)
end

model_ret = fakemodel_ret(xs)

chain_ret = sample(model_ret, NUTS(), 100)

qs = generated_quantities(model_ret, chain_ret)
100×1 Matrix{ComponentArrays.ComponentVector{Float64, Vector{Float64}, Tuple{ComponentArrays.Axis{(mu1 = 1, mu2 = 2, sigma1 = 3, sigma2 = 4)}}}}:
 ComponentVector{Float64}(mu1 = 0.046532704470605424, mu2 = 2.8039504541617895, sigma1 = 0.9703569337946495, sigma2 = 0.6089110120310475)
 ComponentVector{Float64}(mu1 = 0.040917812049849844, mu2 = 2.841307648729105, sigma1 = 0.9544060429194443, sigma2 = 0.6118699510022713)
 ComponentVector{Float64}(mu1 = 0.09324593833044038, mu2 = 2.946543600442379, sigma1 = 0.9156444311184055, sigma2 = 0.5980596134047582)
 ComponentVector{Float64}(mu1 = 0.077413160845598, mu2 = 2.9853950141108454, sigma1 = 0.9306150268283169, sigma2 = 0.5730219004304182)
 ComponentVector{Float64}(mu1 = 0.077413160845598, mu2 = 2.9853950141108454, sigma1 = 0.9306150268283169, sigma2 = 0.5730219004304182)
 ComponentVector{Float64}(mu1 = 0.12381704501730878, mu2 = 2.943460622188596, sigma1 = 0.9389566709406034, sigma2 = 0.5817716773501175)
 ComponentVector{Float64}(mu1 = 0.14694179738373675, mu2 = 3.014066740591662, sigma1 = 0.9773173434950448, sigma2 = 0.5408948674654698)
 ComponentVector{Float64}(mu1 = 0.16274059104751656, mu2 = 2.961514482764827, sigma1 = 0.9191391634120403, sigma2 = 0.531141409877753)
 ⋮
 ComponentVector{Float64}(mu1 = 0.1336587300237204, mu2 = 3.0101961952383243, sigma1 = 1.006199407718754, sigma2 = 0.5396113955693571)
 ComponentVector{Float64}(mu1 = 0.18123321003279896, mu2 = 2.965394772548506, sigma1 = 0.9432549803459527, sigma2 = 0.566500401681874)
 ComponentVector{Float64}(mu1 = 0.1332449026582564, mu2 = 2.969989165286554, sigma1 = 0.9109307412479161, sigma2 = 0.5640799801027031)
 ComponentVector{Float64}(mu1 = 0.0737339678641216, mu2 = 2.992745081964356, sigma1 = 0.9526279315218458, sigma2 = 0.5889234079236011)
 ComponentVector{Float64}(mu1 = 0.07232030671440184, mu2 = 2.9871692965873744, sigma1 = 0.9576485757330118, sigma2 = 0.5558591474507124)
 ComponentVector{Float64}(mu1 = 0.0725153641558526, mu2 = 2.9491436090825234, sigma1 = 0.9469384994954372, sigma2 = 0.5644147829364776)
 ComponentVector{Float64}(mu1 = 0.2468327616001294, mu2 = 3.0346149145068355, sigma1 = 0.9989299098321219, sigma2 = 0.533038850827515)

or if I want a flat vector of the best chain sample I use something like

_bestidx = findmax(chain.value[var=:log_density])[2][1]
_bestsample = chain.value[_bestidx, 1:end-12, 1]

where the end-12 skips the 12 internals (I think this number is dependent on the particular sampler?), but this gives me an AxisArray I can't easily cast into a ComponentArray (which I find much more intuitive to deal with).

While I haven't yet figured out a way to cleanly do these kinds of things using native MCMCChains.jl functionalities I foresaw all those quirks as potential stumbling blocks.

(sorry for all the edits)

[mhauru]

I added a test that covers the changes here. Thanks @a1ix2, much appreciated!

[radka-j]

LGTM :)

[mhauru]

Live once v6.0.7 is out.