diff --git a/Project.toml b/Project.toml index 93315ed..bf9b8a3 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "AdvancedMH" uuid = "5b7e9947-ddc0-4b3f-9b55-0d8042f74170" -version = "0.8.9" +version = "0.8.10" [deps] AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001" diff --git a/src/proposal.jl b/src/proposal.jl index bc445fc..9e7e184 100644 --- a/src/proposal.jl +++ b/src/proposal.jl @@ -110,7 +110,7 @@ end function propose( rng::Random.AbstractRNG, - proposal::Proposal{<:Function}, + proposal::Proposal{<:Function}, model::DensityModelOrLogDensityModel, t ) @@ -129,24 +129,75 @@ end # Multiple proposals #################### +""" + proposal_dim(p::Proposal) + +Return the number of scalar parameters produced by proposal `p`. + +Users with custom proposal types should extend this function to enable +mixed-dimension proposal vectors. +""" +proposal_dim(p::Proposal{<:UnivariateDistribution}) = 1 +proposal_dim(::Proposal{<:Function}) = 1 +proposal_dim(p::Proposal{<:MultivariateDistribution}) = length(p.proposal) +function proposal_dim(p::Proposal{<:AbstractArray}) + return sum(d -> d isa UnivariateDistribution ? 1 : length(d), p.proposal) +end + +function _vcat_proposals(draws) + return vcat(draws...) +end + +""" + _split_params(proposals::AbstractArray{<:Proposal}, params::AbstractVector) + +Split a flat parameter vector `params` into chunks matching each proposal's dimension, +as determined by [`proposal_dim`](@ref). +""" +function _split_params(proposals::AbstractArray{<:Proposal}, params::AbstractVector) + total_dim = sum(proposal_dim, proposals) + if total_dim != length(params) + throw(DimensionMismatch( + "sum of proposal dimensions ($total_dim) does not match parameter length ($(length(params))). " * + "For function-valued proposals, `proposal_dim` defaults to 1; override it for your callable type " * + "to specify a different dimension." + )) + end + result = Vector{Any}(undef, length(proposals)) + offset = 0 + for (i, p) in enumerate(proposals) + dim = proposal_dim(p) + if dim == 1 + result[i] = params[offset+1] + else + result[i] = params[(offset+1):(offset+dim)] + end + offset += dim + end + return result +end + function propose( rng::Random.AbstractRNG, proposals::AbstractArray{<:Proposal}, model::DensityModelOrLogDensityModel, ) - return map(proposals) do proposal + draws = map(proposals) do proposal return propose(rng, proposal, model) end + return _vcat_proposals(draws) end function propose( rng::Random.AbstractRNG, proposals::AbstractArray{<:Proposal}, model::DensityModelOrLogDensityModel, - ts, + ts::AbstractVector, ) - return map(proposals, ts) do proposal, t + split_ts = _split_params(proposals, ts) + draws = map(proposals, split_ts) do proposal, t return propose(rng, proposal, model, t) end + return _vcat_proposals(draws) end @generated function propose( @@ -232,6 +283,16 @@ function logratio_proposal_density(proposals::Tuple, states::Tuple, candidates:: return valfirst + valtail end +function logratio_proposal_density( + proposals::AbstractArray{<:Proposal}, states::AbstractVector, candidates::AbstractVector +) + split_states = _split_params(proposals, states) + split_candidates = _split_params(proposals, candidates) + return sum(zip(proposals, split_states, split_candidates)) do (proposal, state, candidate) + return logratio_proposal_density(proposal, state, candidate) + end +end + # fallback for general iterators (arrays etc.) - possibly not type stable! function logratio_proposal_density(proposals, states, candidates) return sum(zip(proposals, states, candidates)) do (proposal, state, candidate) diff --git a/test/runtests.jl b/test/runtests.jl index b40a852..c7fe344 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -294,6 +294,52 @@ include("util.jl") end end + @testset "Varying dimension proposals" begin + # Issue #120: proposals of varying dimensions can't be combined in a vector + + # Basic reproduce case from the issue: scalar + 2D multivariate + m1 = DensityModel(x -> x[1] + x[2] + x[3]) + p1 = [StaticProposal(Normal()), RandomWalkProposal(MvNormal(zeros(2), I))] + chain1 = sample(m1, MetropolisHastings(p1), 10; chain_type=Any, progress=false) + @test length(chain1) == 10 + @test chain1[1].params isa AbstractVector + @test length(chain1[1].params) == 3 + + # Reversed order: 2D multivariate first, then scalar + p2 = [StaticProposal(MvNormal(zeros(2), I)), RandomWalkProposal(Normal())] + chain2 = sample(m1, MetropolisHastings(p2), 10; chain_type=Any, progress=false) + @test length(chain2) == 10 + @test length(chain2[1].params) == 3 + + # Multiple multivariate proposals: 2D + 3D + m3 = DensityModel(x -> sum(x)) + p3 = [RandomWalkProposal(MvNormal(zeros(2), I)), RandomWalkProposal(MvNormal(zeros(3), I))] + chain3 = sample(m3, MetropolisHastings(p3), 10; chain_type=Any, progress=false) + @test length(chain3) == 10 + @test length(chain3[1].params) == 5 + + # Backward compatibility: all scalar proposals still work + m4 = DensityModel(x -> x[1] + x[2]) + p4 = [StaticProposal(Normal()), RandomWalkProposal(Normal())] + chain4 = sample(m4, MetropolisHastings(p4), 100; chain_type=Any, progress=false) + @test length(chain4) == 100 + @test length(chain4[1].params) == 2 + + # With initial_params + m5 = DensityModel(x -> x[1] + x[2] + x[3]) + p5 = [StaticProposal(Normal()), RandomWalkProposal(MvNormal(zeros(2), I))] + chain5 = sample(m5, MetropolisHastings(p5), 10; chain_type=Any, initial_params=[1.0, 2.0, 3.0], progress=false) + @test chain5[1].params == [1.0, 2.0, 3.0] + + # Single proposal in a vector should still produce a vector of params + m6 = DensityModel(x -> x[1]) + p6 = [StaticProposal(Normal())] + chain6 = sample(m6, MetropolisHastings(p6), 10; chain_type=Any, progress=false) + @test length(chain6) == 10 + @test chain6[1].params isa AbstractVector + @test length(chain6[1].params) == 1 + end + @testset "MALA" begin @testset "basic" begin # Set up the sampler.