Update to most recent SSMProblems interface (#116)

* update to most recent SSMProblems interface

* fixed linear Gaussian uinit test

* fixed Levy SSM

* temporary fix, that yields inefficient sampler

* improved GP fix

* remove environment path

* Update Project.toml

* remove unused constructor

* update Levy SSM

---------

Co-authored-by: Hong Ge <[email protected]>

Author: charlesknipp

.gitignore

@@ -1,5 +1,5 @@
 *.jl.cov
 *.jl.*.cov
 *.jl.mem
-/Manifest.toml
+Manifest.toml
 /test/Manifest.toml

Project.toml

@@ -26,7 +26,7 @@ Random = "<0.0.1, 1"
 Random123 = "1.3"
 Requires = "1.0"
 StatsFuns = "0.9, 1"
-SSMProblems = "0.1"
+SSMProblems = "0.5"
 julia = "1.7"
 
 [extras]

examples/gaussian-process/script.jl

@@ -8,74 +8,76 @@ using Distributions
 using Libtask
 using SSMProblems
 
-Parameters = @NamedTuple begin
-    a::Float64
-    q::Float64
-    kernel
+struct GaussianProcessDynamics{T<:Real,KT<:Kernel} <: LatentDynamics{T,T}
+    proc::GP{ZeroMean{T},KT}
+    function GaussianProcessDynamics(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel}
+        return new{T,KT}(GP(ZeroMean{T}(), kernel))
+    end
 end
 
-mutable struct GPSSM <: SSMProblems.AbstractStateSpaceModel
-    X::Vector{Float64}
-    observations::Vector{Float64}
-    θ::Parameters
-
-    GPSSM(params::Parameters) = new(Vector{Float64}(), params)
-    GPSSM(y::Vector{Float64}, params::Parameters) = new(Vector{Float64}(), y, params)
+struct LinearGaussianDynamics{T<:Real} <: LatentDynamics{T,T}
+    a::T
+    b::T
+    q::T
 end
 
-seed = 1
-T = 100
-Nₚ = 20
-Nₛ = 250
-a = 0.9
-q = 0.5
-
-params = Parameters((a, q, SqExponentialKernel()))
+function SSMProblems.distribution(proc::LinearGaussianDynamics{T}) where {T<:Real}
+    return Normal(zero(T), proc.q)
+end
 
-f(θ::Parameters, x, t) = Normal(θ.a * x, θ.q)
-h(θ::Parameters) = Normal(0, θ.q)
-g(θ::Parameters, x, t) = Normal(0, exp(0.5 * x)^2)
+function SSMProblems.distribution(proc::LinearGaussianDynamics, ::Int, state)
+    return Normal(proc.a * state + proc.b, proc.q)
+end
 
-rng = Random.MersenneTwister(seed)
+struct StochasticVolatility{T<:Real} <: ObservationProcess{T,T} end
 
-x = zeros(T)
-y = similar(x)
-x[1] = rand(rng, h(params))
-for t in 1:T
-    if t < T
-        x[t + 1] = rand(rng, f(params, x[t], t))
-    end
-    y[t] = rand(rng, g(params, x[t], t))
+function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real}
+    return Normal(zero(T), exp((1 / 2) * state))
 end
 
-function gp_update(model::GPSSM, state, step)
-    gp = GP(model.θ.kernel)
-    prior = gp(1:(step - 1))
-    post = posterior(prior, model.X[1:(step - 1)])
-    μ, σ = mean_and_cov(post, [step])
-    return Normal(μ[1], σ[1])
+function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real}
+    dyn = LinearGaussianDynamics(a, zero(T), q)
+    obs = StochasticVolatility{T}()
+    return SSMProblems.StateSpaceModel(dyn, obs)
 end
 
-SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM) = rand(rng, h(model.θ))
-function SSMProblems.transition!!(rng::AbstractRNG, model::GPSSM, state, step)
-    return rand(rng, gp_update(model, state, step))
+function GaussianProcessStateSpaceModel(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel}
+    dyn = GaussianProcessDynamics(T, kernel)
+    obs = StochasticVolatility{T}()
+    return SSMProblems.StateSpaceModel(dyn, obs)
 end
 
-function SSMProblems.emission_logdensity(model::GPSSM, state, step)
-    return logpdf(g(model.θ, state, step), model.observations[step])
-end
-function SSMProblems.transition_logdensity(model::GPSSM, prev_state, current_state, step)
-    return logpdf(gp_update(model, prev_state, step), current_state)
+const GPSSM{T,KT<:Kernel} = SSMProblems.StateSpaceModel{
+    T,
+    GaussianProcessDynamics{T,KT},
+    StochasticVolatility{T}
+};
+
+# for non-markovian models, we can redefine dynamics to reference the trajectory
+function AdvancedPS.dynamics(
+    ssm::AdvancedPS.TracedSSM{<:GPSSM{T},T,T}, step::Int
+) where {T<:Real}
+    prior = ssm.model.dyn.proc(1:(step - 1))
+    post  = posterior(prior, ssm.X[1:(step - 1)])
+    μ, σ  = mean_and_cov(post, [step])
+    return LinearGaussianDynamics(zero(T), μ[1], sqrt(σ[1]))
 end
 
-AdvancedPS.isdone(::GPSSM, step) = step > T
+# Everything is now ready to simulate some data. 
+rng = MersenneTwister(1234);
+true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5);
+_, x, y = sample(rng, true_model, 100);
 
-model = GPSSM(y, params)
-pg = AdvancedPS.PGAS(Nₚ)
-chains = sample(rng, model, pg, Nₛ)
+# Create the model and run the sampler
+gpssm = GaussianProcessStateSpaceModel(Float64, SqExponentialKernel());
+model = gpssm(y);
+pg = AdvancedPS.PGAS(20);
+chains = sample(rng, model, pg, 250; progress=false);
+#md nothing #hide
 
-particles = hcat([chain.trajectory.model.X for chain in chains]...)
+particles = hcat([chain.trajectory.model.X for chain in chains]...);
 mean_trajectory = mean(particles; dims=2);
+#md nothing #hide
 
 scatter(particles; label=false, opacity=0.01, color=:black, xlabel="t", ylabel="state")
 plot!(x; color=:darkorange, label="Original Trajectory")

examples/gaussian-ssm/script.jl

@@ -28,81 +28,55 @@ using SSMProblems
 # as well as the initial distribution $f_0(x) = \mathcal{N}(0, q^2/(1-a^2))$.
 
 # To use `AdvancedPS` we first need to define a model type that subtypes `AdvancedPS.AbstractStateSpaceModel`.
-Parameters = @NamedTuple begin
-    a::Float64
-    q::Float64
-    r::Float64
+mutable struct Parameters{T<:Real}
+    a::T
+    q::T
+    r::T
 end
 
-mutable struct LinearSSM <: SSMProblems.AbstractStateSpaceModel
-    X::Vector{Float64}
-    observations::Vector{Float64}
-    θ::Parameters
-    LinearSSM(θ::Parameters) = new(Vector{Float64}(), θ)
-    LinearSSM(y::Vector, θ::Parameters) = new(Vector{Float64}(), y, θ)
+struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T}
+    a::T
+    q::T
 end
 
-# and the densities defined above.
-f(θ::Parameters, state, t) = Normal(θ.a * state, θ.q) # Transition density
-g(θ::Parameters, state, t) = Normal(state, θ.r)         # Observation density
-f₀(θ::Parameters) = Normal(0, θ.q^2 / (1 - θ.a^2))    # Initial state density
-#md nothing #hide
+function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}; kwargs...) where {T<:Real}
+    return Normal(zero(T), sqrt(dyn.q^2 / (1 - dyn.a^2)))
+end
 
-# We also need to specify the dynamics of the system through the transition equations:
-# - `AdvancedPS.initialization`: the initial state density
-# - `AdvancedPS.transition`: the state transition density
-# - `AdvancedPS.observation`: the observation score given the observed data
-# - `AdvancedPS.isdone`: signals the end of the execution for the model
-SSMProblems.transition!!(rng::AbstractRNG, model::LinearSSM) = rand(rng, f₀(model.θ))
-function SSMProblems.transition!!(
-    rng::AbstractRNG, model::LinearSSM, state::Float64, step::Int
-)
-    return rand(rng, f(model.θ, state, step))
+function SSMProblems.distribution(dyn::LinearGaussianDynamics, step::Int, state; kwargs...)
+    return Normal(dyn.a * state, dyn.q)
 end
 
-function SSMProblems.emission_logdensity(modeL::LinearSSM, state::Float64, step::Int)
-    return logpdf(g(model.θ, state, step), model.observations[step])
+struct LinearGaussianObservation{T<:Real} <: SSMProblems.ObservationProcess{T,T}
+    r::T
 end
-function SSMProblems.transition_logdensity(
-    model::LinearSSM, prev_state, current_state, step
+
+function SSMProblems.distribution(
+    obs::LinearGaussianObservation, step::Int, state; kwargs...
 )
-    return logpdf(f(model.θ, prev_state, step), current_state)
+    return Normal(state, obs.r)
 end
 
-# We need to think seriously about how the data is handled
-AdvancedPS.isdone(::LinearSSM, step) = step > Tₘ
+function LinearGaussianStateSpaceModel(θ::Parameters)
+    dyn = LinearGaussianDynamics(θ.a, θ.q)
+    obs = LinearGaussianObservation(θ.r)
+    return SSMProblems.StateSpaceModel(dyn, obs)
+end
 
 # Everything is now ready to simulate some data. 
-a = 0.9   # Scale
-q = 0.32  # State variance
-r = 1     # Observation variance
-Tₘ = 200  # Number of observation
-Nₚ = 20   # Number of particles
-Nₛ = 500  # Number of samples
-seed = 1  # Reproduce everything
-
-θ₀ = Parameters((a, q, r))
-rng = Random.MersenneTwister(seed)
-
-x = zeros(Tₘ)
-y = zeros(Tₘ)
-x[1] = rand(rng, f₀(θ₀))
-for t in 1:Tₘ
-    if t < Tₘ
-        x[t + 1] = rand(rng, f(θ₀, x[t], t))
-    end
-    y[t] = rand(rng, g(θ₀, x[t], t))
-end
+rng = Random.MersenneTwister(1234)
+θ = Parameters(0.9, 0.32, 1.0)
+true_model = LinearGaussianStateSpaceModel(θ)
+_, x, y = sample(rng, true_model, 200);
 
 # Here are the latent and obseravation timeseries
 plot(x; label="x", xlabel="t")
 plot!(y; seriestype=:scatter, label="y", xlabel="t", mc=:red, ms=2, ma=0.5)
 
 # `AdvancedPS` subscribes to the `AbstractMCMC` API. To sample we just need to define a Particle Gibbs kernel
 # and a model interface. 
-model = LinearSSM(y, θ₀)
-pgas = AdvancedPS.PGAS(Nₚ)
-chains = sample(rng, model, pgas, Nₛ; progress=false);
+pgas = AdvancedPS.PGAS(20)
+chains = sample(rng, true_model(y), pgas, 500; progress=false);
 #md nothing #hide
 
 # 
@@ -118,7 +92,7 @@ plot!(mean_trajectory; color=:dodgerblue, label="Mean trajectory", opacity=0.9)
 # We used a particle gibbs kernel with the ancestor updating step which should help with the particle 
 # degeneracy problem and improve the mixing. 
 # We can compute the update rate of $x_t$ vs $t$ defined as the proportion of times $t$ where $x_t$ gets updated:
-update_rate = sum(abs.(diff(particles; dims=2)) .> 0; dims=2) / Nₛ
+update_rate = sum(abs.(diff(particles; dims=2)) .> 0; dims=2) / length(chains)
 #md nothing #hide
 
 # and compare it to the theoretical value of $1 - 1/Nₚ$. 
@@ -130,4 +104,4 @@ plot(
     xlabel="Iteration",
     ylabel="Update rate",
 )
-hline!([1 - 1 / Nₚ]; label="N: $(Nₚ)")
+hline!([1 - 1 / length(chains)]; label="N: $(length(chains))")