Setting up multiple chains in Turing.jl

General Usage
21

Thank you. This is indeed much simpler. It might be helpful to other users to add a pmap example to the documentation. Thanks again!

22

Hi Cameron,

It’s not only simpler, but it actually works better! I used something similar to a few functions I am trying to add to TuringModels to redo the sections in a Chains object created by Turing. It seems to work, but at the end doing chains[:a] shows it is not correct. A MWE (well ‘minimum’?):

using Turing

function flatten_name_map(chn::MCMCChains.AbstractChains)
  parms = values(chn.name_map)
  [String(parms[i][j]) for i in 1:length(parms) for j in 1:length(parms[i])]
end

function construct_a3d(chn::MCMCChains.AbstractChains)
  d, p, c = size(chn.value)
  a3d = fill(0.0, d, p, c);
  for (i, par) in enumerate(flatten_name_map(chn))
    a3d[:, i, 1] = reshape(chn[par], d)
  end
  a3d
end

function move_parameters_to_new_section(chn::MCMCChains.AbstractChains,
    new_section::Symbol, parameters_to_move::Vector{String})

  a3d = construct_a3d(chn)
  parms = values(chn.name_map.parameters)
  moved_parameters = String[]
  for par in parameters_to_move
    (par in parms) ? append!(moved_parameters, [par]) : @warn "$par not in $parms, ignored"
  end
  
  return MCMCChains.Chains(a3d,
    Symbol.(flatten_name_map(chn)),
    Dict(
      :parameters => Symbol.(filter(x -> !(x in parameters_to_move), parms)),
      new_section => Symbol.(moved_parameters),
      :internals => Symbol.(values(chn.name_map.internals))))
end

@model model() = begin
    a ~ Normal(-10,.1)
    b ~ Normal(-8,.1)
    c ~ Normal(-6,.1)
    d ~ Normal(-4,.1)
    return nothing
end

Nsamples = 2000
Nadapt = 1000
draws = Nadapt+1 : Nsamples
δ = .85
sampler = NUTS(Nsamples, Nadapt, δ)

sampler = Turing.NUTS(2000, 1000, 0.65)
chn = sample(model(), sampler)

chn1 = move_parameters_to_new_section(chn, :pooled, ["b", "d", "zeta"])

describe(chn1)
describe(chn1, section=:pooled)

mean(chn1[:a])

You mind taking a look to spot my error? Or if this functionality is better achieved in some other way?

Thanks!

23

I think I can serve this better on the back end by providing the functions set_name_map(c::Chains, new_name_map::Dict) and set_name_map(c::Chains, new_name_map::NamedTuple). I can also call them set_sections if you’d prefer. People probably shouldn’t have to write all this stuff outside of the library to do basic stuff like this.

24

Fully agreed. My intention was to test these and suggest it via a PR. Something like set_sections or update_sections all work fine. As an end user I tend to think in sections, as a developer probably more in terms of name_map.

It is interesting that the longer I try to build packages the more I think that newer (good!) stuff like NamedTuples is on the developer side (the developer would probably use Parameters.jl and stuff like @unpack …), Dicts and DataFrames are end user. Just a side note.

In line with that, one other thought/suggestion. In StatisticalRethinking I have a function to_df(chain::MCMCChains.Chains) which converts a chain to a DataFrame. As an example, you did that great example in TuringExamples using training and test observations (that I hijacked for 04/clip-38s.jl), it would be nice to continue using DataFrames containing the samples.

Right now I use:

function to_df(chain::MCMCChains.AbstractChains)
  d, p, c = size(chain.value)
  df = DataFrame()
  a3d = zeros(d*c, p);
  indx = 1
  for name in Symbol.(keys(chain))
    a3d[:,  indx] = reshape(convert(Array{Float64}, chain[name]), d*c)
    indx += 1
  end
  cnames = keys(chain)
  snames = [Symbol(cnames[i]) for i in 1: length(cnames)]
  DataFrame(a3d[:,:], snames)
end