TuringLang / AdvancedPS.jl

Implementation of advanced Sequential Monte Carlo and particle MCMC algorithms
https://turinglang.org/AdvancedPS.jl/
MIT License
56 stars 9 forks source link

Multi-dimensional latent space? #73

Closed davidJasperForman closed 1 year ago

davidJasperForman commented 1 year ago

Does this package support multidimensional state-spaces?

Reading the code it doesn't seem so, and trying to run the following block with a 2D model it doesn't seem so, but I'm new to Julia programming and unsure whether I have bugs.

model = Log2DRatePoissonSSM(θ₀)
pgas = AdvancedPS.PGAS(Nₚ)
chains = sample(rng, model, pgas, Nₛ; progress=false);
FredericWantiez commented 1 year ago

Hi David, it should be able to handle n-d models. Do you have the definition of your model ? Or something similar that's causing you issues, I can have a look.

davidJasperForman commented 1 year ago

Thanks, @FredericWantiez!

The following code works until the final line. It's modeled off of the Gaussian state-space example; it worked in an initial 1D version.

image
### Initialize Model Structure
Parameters = @NamedTuple begin
    A::Matrix{Float64}
    Q::Matrix{Float64}
end

mutable struct LogRatePoissonSSM <: AdvancedPS.AbstractStateSpaceModel
    X::Vector{Float64}
    θ::Parameters
    LogRatePoissonSSM(θ::Parameters) = new(Vector{Float64}(), θ)
end

### Define generative model
f(m::LogRatePoissonSSM, state, t) = MvNormal(m.θ.A * state, m.θ.Q) # Transition density
g(m::LogRatePoissonSSM, stateVal, t)  = Poisson(exp(stateVal))
f₀(m::LogRatePoissonSSM) = MvNormal(zeros(dim), m.θ.Q)    # Initial state density

### Set parameters and generate data
dim = 2
A = [[0.45 0.45]; 0.45 0.45]
Q = [1 0; 0 1]
Tₘ = 24*5  # Number of observation #*4*12
Nₚ = 20   # Number of particles
Nₛ = 500  # Number of samples
seed = 1  # Reproduce everything

θ₀ = Parameters((A, Q))
rng = Random.MersenneTwister(seed)

x = zeros(Float64, Tₘ, dim)
y = zeros(Float64, Tₘ, dim)

reference = LogRatePoissonSSM(θ₀)

x[1,:] = rand(rng, f₀(reference))
for t in 1:Tₘ
    if t < Tₘ
        x[t + 1,:] = rand(rng, f(reference, x[t,:], t))
    end
    for (index, rate) in enumerate(x[t,:])
        y[t,index] = rand(rng, g(reference, rate, t))  # Observation density
    end   
end

### Plot
plot(exp.(x); label="explogr", xlabel="t")
plot!(y[:,1]; seriestype=:scatter,label="y", xlabel="t",mc=:blue, ms=2, ma=0.5)
plot!(y[:,2]; seriestype=:scatter,label="y", xlabel="t",mc=:red, ms=2, ma=0.5)
ylims!((0,40))
plot!(size=(800,600))

### Inform particle Gibbs of the model
AdvancedPS.initialization(model::LogRatePoissonSSM) = f₀(model)
AdvancedPS.transition(model::LogRatePoissonSSM, state, step) = f(model, state, step)
function AdvancedPS.observation(model::LogRatePoissonSSM, state, step)
    rr = 0
    for (index, observation) in enumerate(y[step,:])
        rr += logpdf(g(model, state[index], step), observation)
    end   
    print("AdvancedPS.observation run once. ")
    return rr
end
AdvancedPS.isdone(::LogRatePoissonSSM, step) = step > Tₘ

###
model = LogRatePoissonSSM(θ₀)
pgas = AdvancedPS.PGAS(Nₚ)

### Run and fail!
chains = sample(rng, model, pgas, Nₛ; progress=false);

Output:

AdvancedPS.observation run once. MethodError: Cannot convert an object of type Vector{Float64} to an object of type Float64 Closest candidates are: convert(::Type{T}, ::Gray24) where T<:Real at ~/.julia/packages/ColorTypes/1dGw6/src/conversions.jl:114 convert(::Type{T}, ::Gray) where T<:Real at ~/.julia/packages/ColorTypes/1dGw6/src/conversions.jl:113 convert(::Type{T}, ::Base.TwicePrecision) where T<:Number at twiceprecision.jl:273 ...

Stacktrace: [1] push!(a::Vector{Float64}, item::Vector{Float64}) @ Base ./array.jl:1057 [2] advance!(particle::AdvancedPS.Trace{LogRatePoissonSSM, AdvancedPS.TracedRNG{UInt64, 1, Random123.Philox2x{UInt64, 10}}}, isref::Bool) @ AdvancedPS ~/.julia/packages/AdvancedPS/Vox9w/src/pgas.jl:108 [3] reweight!(pc::AdvancedPS.ParticleContainer{AdvancedPS.Trace{LogRatePoissonSSM, AdvancedPS.TracedRNG{UInt64, 1, Random123.Philox2x{UInt64, 10}}}, AdvancedPS.TracedRNG{UInt64, 1, Random123.Philox2x{UInt64, 10}}}, ref::Nothing) @ AdvancedPS ~/.julia/packages/AdvancedPS/Vox9w/src/container.jl:276 [4] sweep!(rng::MersenneTwister, pc::AdvancedPS.ParticleContainer{AdvancedPS.Trace{LogRatePoissonSSM, AdvancedPS.TracedRNG{UInt64, 1, Random123.Philox2x{UInt64, 10}}}, AdvancedPS.TracedRNG{UInt64, 1, Random123.Philox2x{UInt64, 10}}}, resampler::AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Float64}, ref::Nothing) @ AdvancedPS ~/.julia/packages/AdvancedPS/Vox9w/src/container.jl:338 [5] step(rng::MersenneTwister, model::LogRatePoissonSSM, sampler::AdvancedPS.PGAS{AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Float64}}, state::Nothing; kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ AdvancedPS ~/.julia/packages/AdvancedPS/Vox9w/src/smc.jl:153 [6] step (repeats 2 times) @ ~/.julia/packages/AdvancedPS/Vox9w/src/smc.jl:130 [inlined] [7] macro expansion @ ~/.julia/packages/AbstractMCMC/F9Hbk/src/sample.jl:125 [inlined] [8] macro expansion @ ~/.julia/packages/AbstractMCMC/F9Hbk/src/logging.jl:16 [inlined] [9] mcmcsample(rng::MersenneTwister, model::LogRatePoissonSSM, sampler::AdvancedPS.PGAS{AdvancedPS.ResampleWithESSThreshold{typeof(AdvancedPS.resample_systematic), Float64}}, N::Int64; progress::Bool, progressname::String, callback::Nothing, discard_initial::Int64, thinning::Int64, chain_type::Type, kwargs::Base.Pairs{Symbol, Union{}, Tuple{}, NamedTuple{(), Tuple{}}}) @ AbstractMCMC ~/.julia/packages/AbstractMCMC/F9Hbk/src/sample.jl:116 [10] #sample#16 @ ~/.julia/packages/AbstractMCMC/F9Hbk/src/sample.jl:51 [inlined] [11] top-level scope @ In[10]:1 [12] eval @ ./boot.jl:368 [inlined] [13] include_string(mapexpr::typeof(REPL.softscope), mod::Module, code::String, filename::String) @ Base ./loading.jl:1428

FredericWantiez commented 1 year ago

@davidJasperForman The main issue is that the state dimension of your model does not match your data:

mutable struct LogRatePoissonSSM <: AdvancedPS.AbstractStateSpaceModel
    X::Vector{Float64}
    θ::Parameters
    LogRatePoissonSSM(θ::Parameters) = new(Vector{Float64}(), θ)
end

should be something like:

mutable struct LogRatePoissonSSM <: AdvancedPS.AbstractStateSpaceModel
    X::Vector{Array{Float64, 2}} # Since your latent space is 2-d
    θ::Parameters
    LogRatePoissonSSM(θ::Parameters) = new(Vector{Float64}(), θ)
end
image

The API is a bit weird here I agree, the underlying Vector we use to store the trajectory shouldn't be exposed.

I'll try to add a tutorial for n-d models to highlight the issue.

davidJasperForman commented 1 year ago

Thanks, @FredericWantiez ! This direction of bug-fixing solved my problem.

The variant below is what worked for me for the program above:

mutable struct LogRatePoissonSSM <: AdvancedPS.AbstractStateSpaceModel
    X::Vector{Vector{Float64}} # Since your latent space is 2-d
    θ::Parameters
    LogRatePoissonSSM(θ::Parameters) = new(Vector{Float64}(), θ)
end
FredericWantiez commented 1 year ago

Nice ! Closing this, but feel free to reach out if you find sharp edges, always happy to help.