diff --git a/Project.toml b/Project.toml index d422312..acf4ccf 100644 --- a/Project.toml +++ b/Project.toml @@ -10,6 +10,7 @@ CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba" DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" FlexiChains = "4a37a8b9-6e57-4b92-8664-298d46e639f7" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" +OrderedCollections = "bac558e1-5e72-5ebc-8fee-abe8a469f55d" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" @@ -39,6 +40,7 @@ ForwardDiff = "1" LinearAlgebra = "1" LogDensityProblems = "2" Mooncake = "0.5.26" +OrderedCollections = "1" Random = "1" Statistics = "1" Zygote = "0.7.10" diff --git a/ext/DynamicPPLExt.jl b/ext/DynamicPPLExt.jl index 7800d81..cc33ad4 100644 --- a/ext/DynamicPPLExt.jl +++ b/ext/DynamicPPLExt.jl @@ -136,13 +136,12 @@ end function ParallelMCMC._construct_flexichain( ::Type{VarName}, vals::AbstractMatrix{<:Real}, - internals::AbstractMatrix{<:Real}, + internals::NamedTuple, ::Any, - internal_names::Vector{Symbol}, model::DensityModelLDF, ) - pwss = map(zip(eachrow(vals), eachrow(internals))) do (val, internal) - stats = NamedTuple{Tuple(internal_names)}(internal) + pwss = map(enumerate(eachrow(vals))) do (i, val) + stats = map(v -> v[i], internals) DynamicPPL.ParamsWithStats(val, model.logdensity.ld, stats) end return AbstractMCMC.from_samples(VNChain, hcat(pwss)) diff --git a/src/ParallelMCMC.jl b/src/ParallelMCMC.jl index 7a9a557..4ad005e 100644 --- a/src/ParallelMCMC.jl +++ b/src/ParallelMCMC.jl @@ -4,6 +4,7 @@ using AbstractMCMC using CUDA using FlexiChains using LinearAlgebra +using OrderedCollections: OrderedDict using Random using Statistics diff --git a/src/interface.jl b/src/interface.jl index 7523297..41d1de7 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -205,20 +205,23 @@ for TKey in (Symbol, VarName) ) N = length(samples) D = model.dim + # Follow the sampler's working precision instead of pinning `Float64`, so a + # Float32 (e.g. GPU) run yields a Float32 chain rather than a silently widened one. + FP = typeof(sampler.epsilon) - internal_names = [:logp, :accepted] - - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 2) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + accepted = Vector{Bool}(undef, N) for i in 1:N s = samples[i] vals[i, :] .= s.x - internals[i, 1] = s.logp - internals[i, 2] = s.accepted ? 1.0 : 0.0 + logp[i] = s.logp + accepted[i] = s.accepted end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp, accepted=accepted) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end @@ -543,7 +546,7 @@ function _parallel_mala_update_progress!( end function _copy_trajectory_rows!( - vals::AbstractMatrix{Float64}, first_row::Int, S::AbstractMatrix, ncols::Int + vals::AbstractMatrix{<:Real}, first_row::Int, S::AbstractMatrix, ncols::Int ) rows = first_row:(first_row + ncols - 1) S_host = Array(view(S, :, 1:ncols)) @@ -563,10 +566,10 @@ function _sample_parallel_mala_chain( progressname="Sampling", ) where {TKey} D = model.dim - internal_names = [:logp] + FP = typeof(sampler.epsilon) - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 1) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) progress = _parallel_mala_progress(progress, progressname) x0 = _parallel_mala_initial_x(rng, model, sampler, initial_params) @@ -586,7 +589,7 @@ function _sample_parallel_mala_chain( rows = first_row:(first_row + nkeep - 1) _copy_trajectory_rows!(vals, first_row, S, nkeep) - internals[rows, 1] .= view(logps, 1:nkeep) + logp[rows] .= view(logps, 1:nkeep) x0 = copy(view(S, :, nkeep)) nsteps += nkeep @@ -599,15 +602,15 @@ function _sample_parallel_mala_chain( end end - return _construct_flexichain(TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp,) + return _construct_flexichain(TKey, vals, internals, param_names, model) end function _construct_flexichain( ::Type{TKey}, vals::AbstractMatrix{<:Real}, - internals::AbstractMatrix{<:Real}, + internals::NamedTuple, param_names::Any, - internal_names::Vector{Symbol}, model::DensityModel, ) where {TKey} # Wrap user-supplied names in `Parameter`. This allows people to specify, e.g., @@ -616,19 +619,14 @@ function _construct_flexichain( to_parameter(vn::VarName) = FlexiChains.Parameter(vn) to_parameter(s::Symbol) = FlexiChains.Parameter(TKey <: VarName ? VarName{s}() : s) - # vals and internals are both `iters x params` - # FlexiChains expects `iters x chains x params` - arr = hcat(vals, internals) - arr = reshape(arr, size(arr, 1), 1, size(arr, 2)) + N, D = size(vals) - # Wrap parameter names in the format that FlexiChains expects. Ref: - # https://pysm.dev/FlexiChains.jl/stable/arrays/#api-fromarray if param_names === nothing param_names = model.param_names end wrapped_param_names = if param_names === nothing # Wasn't defined either as `param_names` or `model.param_names`, so make some up - (to_parameter(:x) => (size(vals, 2),),) + (to_parameter(:x) => (D,),) else map(param_names) do n if n isa Pair @@ -641,11 +639,23 @@ function _construct_flexichain( end end - all_names = ( - wrapped_param_names..., - FlexiChains.Extra.(internal_names)... + arr = reshape(vals, N, 1, D) + param_chain = FlexiChains.FlexiChain{TKey}(arr, Tuple(wrapped_param_names)) + + isempty(internals) && return param_chain + + #= + The internals have mixed element types (e.g. `Bool` for `accepted`/`is_warmup`, + `FP` for `logp`), so they cannot share a single stacked array without widening. + Store each in its own column via the dict-of-arrays constructor, which preserves + the natural element type (see issue #49), then merge with the parameters. + =# + extras = OrderedDict{FlexiChains.ParameterOrExtra{<:TKey},AbstractArray}( + FlexiChains.Extra(name) => internals[name] for name in keys(internals) ) - return FlexiChains.FlexiChain{TKey}(arr, all_names) + extra_chain = FlexiChains.FlexiChain{TKey}(N, 1, extras) + + return merge(param_chain, extra_chain) end function _sample_parallel_mala_blocks( @@ -870,18 +880,18 @@ for TKey in (Symbol, VarName) ) N = length(samples) D = model.dim + FP = typeof(sampler.epsilon) - internal_names = [:logp] - - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 1) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) for i in 1:N vals[i, :] .= samples[i].x - internals[i, 1] = samples[i].logp + logp[i] = samples[i].logp end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp,) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end @@ -1075,21 +1085,24 @@ for TKey in (Symbol, VarName) filtered = discard_warmup ? filter(s -> !s.is_warmup, samples) : samples N = length(filtered) D = model.dim + FP = typeof(sampler.epsilon_init) - internal_names = [:logp, :accepted, :step_size, :is_warmup] - - vals = Matrix{Float64}(undef, N, D) - internals = Matrix{Float64}(undef, N, 4) + vals = Matrix{FP}(undef, N, D) + logp = Vector{FP}(undef, N) + accepted = Vector{Bool}(undef, N) + step_size = Vector{FP}(undef, N) + is_warmup = Vector{Bool}(undef, N) for i in 1:N s = filtered[i] vals[i, :] .= s.x - internals[i, 1] = s.logp - internals[i, 2] = s.accepted ? 1.0 : 0.0 - internals[i, 3] = s.step_size - internals[i, 4] = s.is_warmup ? 1.0 : 0.0 + logp[i] = s.logp + accepted[i] = s.accepted + step_size[i] = s.step_size + is_warmup[i] = s.is_warmup end - return _construct_flexichain($TKey, vals, internals, param_names, internal_names, model) + internals = (logp=logp, accepted=accepted, step_size=step_size, is_warmup=is_warmup) + return _construct_flexichain($TKey, vals, internals, param_names, model) end end diff --git a/test/test-AbstractMCMC-Interface.jl b/test/test-AbstractMCMC-Interface.jl index c033993..033dc27 100644 --- a/test/test-AbstractMCMC-Interface.jl +++ b/test/test-AbstractMCMC-Interface.jl @@ -152,9 +152,49 @@ gradlogp_iface(x) = -x # logp values should be finite @test all(isfinite, chain[:logp]) - # accepted values should be 0 or 1 + # accepted is stored with its natural `Bool` type (issue #49) acc = chain[:accepted] - @test all(a -> a == 0.0 || a == 1.0, acc) + @test eltype(acc) === Bool + @test all(a -> a == false || a == true, acc) + end + + @testset "chain element type follows sampler precision" begin + # Scalar param names so `chain[:a]` is a scalar-valued column whose eltype is the + # numeric type directly (a vector-valued param would nest inside a `Vector`). + model = DensityModel(logp_iface, gradlogp_iface, 2) + + @testset "Float32 sampler → Float32 chain" begin + sampler = MALASampler(0.1f0) + @test sampler.epsilon isa Float32 + chain = sample( + MersenneTwister(1), + model, + sampler, + 50; + chain_type=SymChain, + progress=false, + param_names=[:a, :b], + ) + @test eltype(chain[:a]) === Float32 + @test eltype(chain[:b]) === Float32 + @test eltype(chain[:logp]) === Float32 + @test eltype(chain[:accepted]) === Bool + end + + @testset "Float64 sampler → Float64 chain" begin + sampler = MALASampler(0.1) + chain = sample( + MersenneTwister(1), + model, + sampler, + 50; + chain_type=SymChain, + progress=false, + param_names=[:a, :b], + ) + @test eltype(chain[:a]) === Float64 + @test eltype(chain[:logp]) === Float64 + end end @testset "sample() with custom param_names" begin diff --git a/test/test-Adaptive-MALA.jl b/test/test-Adaptive-MALA.jl index d4c5a61..5bac2b0 100644 --- a/test/test-Adaptive-MALA.jl +++ b/test/test-Adaptive-MALA.jl @@ -201,13 +201,16 @@ end @test FlexiChains.Extra(:is_warmup) in extras_names @test all(isfinite, chain[:logp]) - @test all(x -> x == 0.0 || x == 1.0, chain[:accepted]) + # accepted / is_warmup keep their natural `Bool` type (issue #49) + @test eltype(chain[:accepted]) === Bool + @test eltype(chain[:is_warmup]) === Bool + @test all(x -> x == false || x == true, chain[:accepted]) @test all(s -> s > 0, chain[:step_size]) # Warmup samples appear at the start - @test chain[:is_warmup][1] == 1.0 - # Post-warmup samples have is_warmup == 0 - @test chain[:is_warmup][end] == 0.0 + @test chain[:is_warmup][1] == true + # Post-warmup samples have is_warmup == false + @test chain[:is_warmup][end] == false end @testset "step_size is constant after warmup" begin