`sample_array` converts discrete variables to Float64

[penelopeysm]

In here:

Pigeons.jl/src/pt/process_sample.jl

Lines 19 to 32 in 24ee2e1

	function sample_array(pt::PT)
	chains = chains_with_samples(pt)
	dim, size = sample_dim_size(pt, chains)
	n_chains_with_samples = count(!isnothing, chains) # iterators have no `length` method
	result = zeros(size, dim, n_chains_with_samples)
	for (chain_index, t) in enumerate(chains)
	sample = get_sample(pt, t)
	for i in 1:size
	vector = sample[i]
	result[i, :, chain_index] .= vector
	end
	end
	return result
	end

The call to zeros(...) without a type parameter generates an array of Float64, and then when you fill it in in the loop below, this causes discrete variables' values to be promoted to Float64.

For example, consider

using Turing, Pigeons
@model function g()
    x ~ Normal()
    y ~ Poisson(1.0)
    1.0 ~ Normal(x + y)
end
pt = pigeons(; target = TuringLogPotential(g()), record=[traces])

On the present state of #442, Pigeons.get_sample correctly returns an integer for y:

julia> get_sample(pt)
1024-element Pigeons.SampleArray{Vector{Real}, PT}:
 [-0.230514080693772, 1, -2.8910138078074405]
 [-0.6134264488692323, 0, -4.327595523450894]
 [-0.5564237683839224, 1, -3.1474844764319103]
 [-0.507119467528689, 0, -4.1021666882846155]
 [0.7790781993671052, 0, -3.165761707771331]
 [-0.3371460237015391, 1, -2.951544507707104]
 [0.2582402415169788, 2, -4.355952510825017]
 [-1.0033213313226617, 1, -3.844530760296424]
 [-0.011186415041692932, 1, -2.8380022022908307]
 [0.034753058710123685, 0, -3.3043317827889314]
 ⋮
 [0.20754579295651432, 2, -4.281645296099754]
 [-0.3853018051540893, 1, -2.9863345474643452]
 [0.3799820628230435, 1, -2.982263434476601]
 [1.2652043247982583, 1, -4.438619049897563]
 [-0.7392357956892281, 2, -3.838258012908349]
 [1.2823434905416835, 0, -3.6999384036022906]
 [0.4691154815669054, 1, -3.057946401455095]
 [1.0968098077859958, 0, -3.4440590130789026]
 [-1.3111886134136284, 2, -4.439051213501216]

But Pigeons.sample_array returns Array{Float64}:

julia> sample_array(pt)
1024×3×1 Array{Float64, 3}:
[:, :, 1] =
 -0.230514   1.0  -2.89101
 -0.613426   0.0  -4.3276
 -0.556424   1.0  -3.14748
 -0.507119   0.0  -4.10217
  0.779078   0.0  -3.16576
 -0.337146   1.0  -2.95154
  0.25824    2.0  -4.35595
 -1.00332    1.0  -3.84453
 -0.0111864  1.0  -2.838
  0.0347531  0.0  -3.30433
  ⋮
  0.207546   2.0  -4.28165
 -0.385302   1.0  -2.98633
  0.379982   1.0  -2.98226
  1.2652     1.0  -4.43862
 -0.739236   2.0  -3.83826
  1.28234    0.0  -3.69994
  0.469115   1.0  -3.05795
  1.09681    0.0  -3.44406
 -1.31119    2.0  -4.43905

I'm not sure if this is intentional. There are certainly downsides to returning Array{Real}, but personally I think it's best to just inherit the datatype that get_samples returns, and then if the user wants to convert to Float they can do so.

For example, the following implementation is simpler to read and retains the element type faithfully. It does come with some intermediate allocations, but since sample_array probably isn't on a hot path (if I'm not mistaken, it's a postprocessing step), maybe it's fine to be slightly unperformant?

function sample_array(pt::PT)
    chains = chains_with_samples(pt)
    # Vector of Vector of samples, where each sample is itself a Vector.
    vec_vec_samples = [get_sample(pt, chn) for chn in chains]
    # This is nparams x niters x nchains
    arr = stack(stack(vec_vec_samples))
    return permutedims(arr, (2, 1, 3))
end