improved GP fix

Author: charlesknipp

examples/gaussian-process/script.jl

@@ -11,81 +11,74 @@ using Distributions
 using Libtask
 using SSMProblems
 
-# Gaussian process encoded transition dynamics
-mutable struct GaussianProcessDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T}
-    proc::AbstractGPs.AbstractGP
-    q::T
-    function GaussianProcessDynamics(q::T, kernel::KT) where {T<:Real,KT<:Kernel}
-        return new{T}(GP(ZeroMean{T}(), kernel), q)
+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
 
-function SSMProblems.distribution(dyn::GaussianProcessDynamics{T}) where {T<:Real}
-    return Normal(zero(T), dyn.q)
-end
-
-# TODO: broken...
-function SSMProblems.simulate(
-    rng::AbstractRNG, dyn::GaussianProcessDynamics, step::Int, state
-)
-    dyn.proc = posterior(dyn.proc(step:step, 0.1), [state])
-    μ, σ = mean_and_cov(dyn.proc, [step])
-    return rand(rng, Normal(μ[1], sqrt(σ[1])))
-end
-
-function SSMProblems.logdensity(dyn::GaussianProcessDynamics, step::Int, state, prev_state)
-    μ, σ = mean_and_cov(dyn.proc, [step])
-    return logpdf(Normal(μ[1], sqrt(σ[1])), state)
-end
-
-# Linear Gaussian dynamics used for simulation
-struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T}
+struct LinearGaussianDynamics{T<:Real} <: LatentDynamics{T,T}
     a::T
+    b::T
     q::T
 end
 
-function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}) where {T<:Real}
-    return Normal(zero(T), dyn.q)
+function SSMProblems.distribution(proc::LinearGaussianDynamics{T}) where {T<:Real}
+    return Normal(zero(T), proc.q)
 end
 
-function SSMProblems.distribution(dyn::LinearGaussianDynamics, ::Int, state)
-    return Normal(dyn.a * state, dyn.q)
+function SSMProblems.distribution(proc::LinearGaussianDynamics, ::Int, state)
+    return Normal(proc.a * state + proc.b, proc.q)
 end
 
-# Observation process used in both variants of the model
-struct StochasticVolatility{T<:Real} <: SSMProblems.ObservationProcess{T,T} end
+struct StochasticVolatility{T<:Real} <: ObservationProcess{T,T} end
 
 function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real}
     return Normal(zero(T), exp((1 / 2) * state))
 end
 
-# Baseline model (for simulation)
 function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real}
-    dyn = LinearGaussianDynamics(a, q)
+    dyn = LinearGaussianDynamics(a, zero(T), q)
     obs = StochasticVolatility{T}()
     return SSMProblems.StateSpaceModel(dyn, obs)
 end
 
-# Gaussian process model (for sampling)
-function GaussianProcessStateSpaceModel(q::T, kernel::KT) where {T<:Real,KT<:Kernel}
-    dyn = GaussianProcessDynamics(q, kernel)
+function GaussianProcessStateSpaceModel(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel}
+    dyn = GaussianProcessDynamics(T, kernel)
     obs = StochasticVolatility{T}()
     return SSMProblems.StateSpaceModel(dyn, obs)
 end
 
+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
+
 # Everything is now ready to simulate some data. 
-rng = Random.MersenneTwister(1234)
-true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5)
+rng = MersenneTwister(1234);
+true_model = LinearGaussianStochasticVolatilityModel(0.9, 0.5);
 _, x, y = sample(rng, true_model, 100);
 
 # Create the model and run the sampler
-gpssm = GaussianProcessStateSpaceModel(0.5, SqExponentialKernel())
-model = gpssm(y)
-pg = AdvancedPS.PGAS(20)
-chains = sample(rng, model, pg, 50)
+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
 

src/model.jl

@@ -31,8 +31,8 @@ mutable struct TracedSSM{SSM,XT,YT} <: SSMProblems.AbstractStateSpaceModel
 end
 
 (model::SSMProblems.StateSpaceModel)(Y::AbstractVector) = TracedSSM(model, Y)
-dynamics(ssm::TracedSSM) = ssm.model.dyn
-observation(ssm::TracedSSM) = ssm.model.obs
+dynamics(ssm::TracedSSM, step::Int) = ssm.model.dyn
+observation(ssm::TracedSSM, step::Int) = ssm.model.obs
 
 isdone(ssm::TracedSSM, step::Int) = step > length(ssm.Y)
 

src/pgas.jl

@@ -26,7 +26,7 @@ Get the log weight of the transition from previous state of `model` to `x`
 function transition_logweight(particle::SSMTrace, x; kwargs...)
     iter = current_step(particle) - 1
     score = SSMProblems.logdensity(
-        dynamics(particle.model), iter, particle.model.X[iter - 1], x, kwargs...
+        dynamics(particle.model, iter), iter, particle.model.X[iter - 1], x, kwargs...
     )
     return score
 end
@@ -59,11 +59,11 @@ function advance!(particle::SSMTrace, isref::Bool=false)
 
     if !isref
         if running_step == 1
-            new_state = SSMProblems.simulate(particle.rng, dynamics(model))
+            new_state = SSMProblems.simulate(particle.rng, dynamics(model, running_step))
         else
             current_state = model.X[running_step - 1]
             new_state = SSMProblems.simulate(
-                particle.rng, dynamics(model), running_step, current_state
+                particle.rng, dynamics(model, running_step), running_step, current_state
             )
         end
     else
@@ -72,7 +72,7 @@ function advance!(particle::SSMTrace, isref::Bool=false)
     end
 
     score = SSMProblems.logdensity(
-        observation(model), running_step, new_state, model.Y[running_step]
+        observation(model, running_step), running_step, new_state, model.Y[running_step]
     )
 
     # Accept transition and move the time index/rng counter