Update SSMProblems interface (#119)

* Integrate new SSMProblems interface

* Fix unit tests

* Removed type piracy

* Fix examples and streamline Levy SSM

* Formatter

* Fix Literate Errors

* Remove compat entry for Documenter

* Increase HTML size threshold

* Suppress output in Levy model

* Fix update rate calculation

* Update Project.toml

---------

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

Author: charlesknipp

Project.toml

@@ -1,7 +1,7 @@
 name = "AdvancedPS"
 uuid = "576499cb-2369-40b2-a588-c64705576edc"
 authors = ["TuringLang"]
-version = "0.7"
+version = "0.7.1"
 
 [deps]
 AbstractMCMC = "80f14c24-f653-4e6a-9b94-39d6b0f70001"
@@ -26,7 +26,7 @@ Random = "<0.0.1, 1"
 Random123 = "1.3"
 Requires = "1.0"
 StatsFuns = "0.9, 1"
-SSMProblems = "0.5"
+SSMProblems = "0.6"
 julia = "1.10.8"
 
 [extras]

docs/Project.toml

@@ -2,6 +2,3 @@
 AdvancedPS = "576499cb-2369-40b2-a588-c64705576edc"
 Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4"
 Literate = "98b081ad-f1c9-55d3-8b20-4c87d4299306"
-
-[compat]
-Documenter = "0.27"

docs/make.jl

@@ -48,7 +48,7 @@ DocMeta.setdocmeta!(AdvancedPS, :DocTestSetup, :(using AdvancedPS); recursive=tr
 
 makedocs(;
     sitename="AdvancedPS",
-    format=Documenter.HTML(),
+    format=Documenter.HTML(; size_threshold=1000 * 2^11), # 1Mb per page
     modules=[AdvancedPS],
     pages=[
         "Home" => "index.md",

examples/gaussian-process/script.jl

@@ -8,59 +8,59 @@ using Distributions
 using Libtask
 using SSMProblems
 
-struct GaussianProcessDynamics{T<:Real,KT<:Kernel} <: LatentDynamics{T,T}
+struct GaussianProcessDynamics{T<:Real,KT<:Kernel} <: SSMProblems.LatentDynamics
     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
 
-struct LinearGaussianDynamics{T<:Real} <: LatentDynamics{T,T}
-    a::T
-    b::T
-    q::T
+struct GaussianPrior{ΣT<:Real} <: SSMProblems.StatePrior
+    σ::ΣT
 end
 
-function SSMProblems.distribution(proc::LinearGaussianDynamics{T}) where {T<:Real}
-    return Normal(zero(T), proc.q)
+SSMProblems.distribution(proc::GaussianPrior) = Normal(0, proc.σ)
+
+struct LinearGaussianDynamics{AT<:Real,BT<:Real,QT<:Real} <: SSMProblems.LatentDynamics
+    a::AT
+    b::BT
+    q::QT
 end
 
 function SSMProblems.distribution(proc::LinearGaussianDynamics, ::Int, state)
     return Normal(proc.a * state + proc.b, proc.q)
 end
 
-struct StochasticVolatility{T<:Real} <: ObservationProcess{T,T} end
+struct StochasticVolatility <: SSMProblems.ObservationProcess end
 
-function SSMProblems.distribution(::StochasticVolatility{T}, ::Int, state) where {T<:Real}
-    return Normal(zero(T), exp((1 / 2) * state))
+function SSMProblems.distribution(::StochasticVolatility, ::Int, state)
+    return Normal(0, exp(state / 2))
 end
 
-function LinearGaussianStochasticVolatilityModel(a::T, q::T) where {T<:Real}
-    dyn = LinearGaussianDynamics(a, zero(T), q)
-    obs = StochasticVolatility{T}()
-    return SSMProblems.StateSpaceModel(dyn, obs)
+function LinearGaussianStochasticVolatilityModel(a, q)
+    prior = GaussianPrior(q)
+    dyn = LinearGaussianDynamics(a, 0, q)
+    obs = StochasticVolatility()
+    return SSMProblems.StateSpaceModel(prior, dyn, obs)
 end
 
 function GaussianProcessStateSpaceModel(::Type{T}, kernel::KT) where {T<:Real,KT<:Kernel}
+    prior = GaussianPrior(one(T))
     dyn = GaussianProcessDynamics(T, kernel)
-    obs = StochasticVolatility{T}()
-    return SSMProblems.StateSpaceModel(dyn, obs)
+    obs = StochasticVolatility()
+    return SSMProblems.StateSpaceModel(prior, dyn, obs)
 end
 
 const GPSSM{T,KT<:Kernel} = SSMProblems.StateSpaceModel{
-    T,
-    GaussianProcessDynamics{T,KT},
-    StochasticVolatility{T}
+    <:GaussianPrior,<:GaussianProcessDynamics{T,KT},StochasticVolatility
 };
 
 # 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}
+function AdvancedPS.dynamics(ssm::AdvancedPS.TracedSSM{<:GPSSM}, step::Int)
     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]))
+    post = posterior(prior, ssm.X[1:(step - 1)])
+    μ, σ = mean_and_cov(post, [step])
+    return LinearGaussianDynamics(0, μ[1], sqrt(σ[1]))
 end
 
 # Everything is now ready to simulate some data. 
@@ -70,9 +70,9 @@ _, x, y = sample(rng, true_model, 100);
 
 # Create the model and run the sampler
 gpssm = GaussianProcessStateSpaceModel(Float64, SqExponentialKernel());
-model = gpssm(y);
+model = AdvancedPS.TracedSSM(gpssm, y);
 pg = AdvancedPS.PGAS(20);
-chains = sample(rng, model, pg, 250; progress=false);
+chains = sample(rng, model, pg, 250);
 #md nothing #hide
 
 particles = hcat([chain.trajectory.model.X for chain in chains]...);

examples/gaussian-ssm/script.jl

@@ -28,27 +28,31 @@ 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`.
-mutable struct Parameters{T<:Real}
-    a::T
-    q::T
-    r::T
+mutable struct Parameters{AT<:Real,QT<:Real,RT<:Real}
+    a::AT
+    q::QT
+    r::RT
 end
 
-struct LinearGaussianDynamics{T<:Real} <: SSMProblems.LatentDynamics{T,T}
-    a::T
-    q::T
+struct GaussianPrior{ΣT<:Real} <: SSMProblems.StatePrior
+    σ::ΣT
 end
 
-function SSMProblems.distribution(dyn::LinearGaussianDynamics{T}; kwargs...) where {T<:Real}
-    return Normal(zero(T), sqrt(dyn.q^2 / (1 - dyn.a^2)))
+struct LinearGaussianDynamics{AT<:Real,QT<:Real} <: SSMProblems.LatentDynamics
+    a::AT
+    q::QT
+end
+
+function SSMProblems.distribution(prior::GaussianPrior; kwargs...)
+    return Normal(0, prior.σ)
 end
 
 function SSMProblems.distribution(dyn::LinearGaussianDynamics, step::Int, state; kwargs...)
     return Normal(dyn.a * state, dyn.q)
 end
 
-struct LinearGaussianObservation{T<:Real} <: SSMProblems.ObservationProcess{T,T}
-    r::T
+struct LinearGaussianObservation{RT<:Real} <: SSMProblems.ObservationProcess
+    r::RT
 end
 
 function SSMProblems.distribution(
@@ -58,9 +62,10 @@ function SSMProblems.distribution(
 end
 
 function LinearGaussianStateSpaceModel(θ::Parameters)
+    prior = GaussianPrior(sqrt(θ.q^2 / (1 - θ.a^2)))
     dyn = LinearGaussianDynamics(θ.a, θ.q)
     obs = LinearGaussianObservation(θ.r)
-    return SSMProblems.StateSpaceModel(dyn, obs)
+    return SSMProblems.StateSpaceModel(prior, dyn, obs)
 end
 
 # Everything is now ready to simulate some data. 
@@ -75,8 +80,9 @@ 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. 
-pgas = AdvancedPS.PGAS(20)
-chains = sample(rng, true_model(y), pgas, 500; progress=false);
+N = 20
+pgas = AdvancedPS.PGAS(N)
+chains = sample(rng, AdvancedPS.TracedSSM(true_model, y), pgas, 500; progress=false);
 #md nothing #hide
 
 # 
@@ -104,4 +110,4 @@ plot(
     xlabel="Iteration",
     ylabel="Update rate",
 )
-hline!([1 - 1 / length(chains)]; label="N: $(length(chains))")
+hline!([1 - 1 / N]; label="N: $(N)")

examples/levy-ssm/script.jl

@@ -27,17 +27,16 @@ function simulate(
         t = t0
         truncated = last_jump < tolerance
         while !truncated
-            t += rand(rng, Exponential(one(T) / rate))
-            xi = one(T) / (β * (exp(t / C) - one(T)))
-            prob = (one(T) + β * xi) * exp(-β * xi)
+            t += rand(rng, Exponential(1 / rate))
+            xi = 1 / (β * (exp(t / C) - 1))
+            prob = (1 + β * xi) * exp(-β * xi)
             if rand(rng) < prob
                 push!(jumps, xi)
                 last_jump = xi
             end
             truncated = last_jump < tolerance
         end
-        times = rand(rng, Uniform(start, finish), length(jumps))
-        return GammaPath(jumps, times)
+        return GammaPath(jumps, rand(rng, Uniform(start, finish), length(jumps)))
     end
 end
 
@@ -47,85 +46,66 @@ function integral(times::Array{<:Real}, path::GammaPath)
     end
 end
 
-struct LangevinDynamics{T}
-    A::Matrix{T}
-    L::Vector{T}
-    θ::T
-    H::Vector{T}
-    σe::T
+struct LangevinDynamics{AT<:AbstractMatrix,LT<:AbstractVector,θT<:Real}
+    A::AT
+    L::LT
+    θ::θT
 end
 
-struct NormalMeanVariance{T}
-    μ::T
-    σ::T
+function Base.exp(dyn::LangevinDynamics, dt)
+    f_val = exp(dyn.θ * dt)
+    return [1 (f_val - 1)/dyn.θ; 0 f_val]
 end
 
-f(dt, θ) = exp(θ * dt)
-function Base.exp(dyn::LangevinDynamics{T}, dt::T) where {T<:Real}
-    let θ = dyn.θ
-        f_val = f(dt, θ)
-        return [one(T) (f_val - 1)/θ; zero(T) f_val]
-    end
+function meancov(t, dyn::LangevinDynamics, path::GammaPath, dist::Normal)
+    fts = exp.(Ref(dyn), (t .- path.times)) .* Ref(dyn.L)
+    μ = sum(@. fts * mean(dist) * path.jumps)
+    Σ = sum(@. fts * transpose(fts) * var(dist) * path.jumps)
+    return μ, Σ + eltype(Σ)(1e-6) * I
 end
 
-function meancov(
-    t::T, dyn::LangevinDynamics, path::GammaPath, nvm::NormalMeanVariance
-) where {T<:Real}
-    μ = zeros(T, 2)
-    Σ = zeros(T, (2, 2))
-    let times = path.times, jumps = path.jumps, μw = nvm.μ, σw = nvm.σ
-        for (v, z) in zip(times, jumps)
-            ft = exp(dyn, (t - v)) * dyn.L
-            μ += ft .* μw .* z
-            Σ += ft * transpose(ft) .* σw^2 .* z
-        end
-
-        # Guarantees positive semi-definiteness
-        return μ, Σ + T(1e-6) * I
-    end
+struct LevyPrior{XT<:AbstractVector,ΣT<:AbstractMatrix} <: StatePrior
+    μ::XT
+    Σ::ΣT
 end
 
-struct LevyLangevin{T} <: LatentDynamics{T,Vector{T}}
-    dt::T
-    dyn::LangevinDynamics{T}
-    process::GammaProcess{T}
-    nvm::NormalMeanVariance{T}
-end
+SSMProblems.distribution(proc::LevyPrior) = MvNormal(proc.μ, proc.Σ)
 
-function SSMProblems.distribution(proc::LevyLangevin{T}) where {T<:Real}
-    return MultivariateNormal(zeros(T, 2), I)
+struct LevyLangevin{T<:Real,LT<:LangevinDynamics,ΓT<:GammaProcess,DT<:Normal} <:
+       SSMProblems.LatentDynamics
+    dt::T
+    dyn::LT
+    process::ΓT
+    dist::DT
 end
 
-function SSMProblems.distribution(proc::LevyLangevin{T}, step::Int, state) where {T<:Real}
+function SSMProblems.distribution(proc::LevyLangevin, step::Int, state)
     dt = proc.dt
     path = simulate(rng, proc.process, dt, (step - 1) * dt, step * dt)
-    μ, Σ = meancov(step * dt, proc.dyn, path, proc.nvm)
-    return MultivariateNormal(exp(proc.dyn, dt) * state + μ, Σ)
+    μ, Σ = meancov(step * dt, proc.dyn, path, proc.dist)
+    return MvNormal(exp(proc.dyn, dt) * state + μ, Σ)
 end
 
-struct LinearGaussianObservation{T<:Real} <: ObservationProcess{T,T}
-    H::Vector{T}
-    R::T
+struct LinearGaussianObservation{HT<:AbstractVector,RT<:Real} <:
+       SSMProblems.ObservationProcess
+    H::HT
+    R::RT
 end
 
-function SSMProblems.distribution(proc::LinearGaussianObservation, step::Int, state)
+function SSMProblems.distribution(proc::LinearGaussianObservation, ::Int, state)
     return Normal(transpose(proc.H) * state, proc.R)
 end
 
-function LevyModel(dt, θ, σe, C, β, μw, σw; ϵ=1e-10)
-    A = [0.0 1.0; 0.0 θ]
-    L = [0.0; 1.0]
-    H = [1.0, 0]
-
+function LevyModel(dt, θ, σe, C, β, μw, σw; kwargs...)
     dyn = LevyLangevin(
         dt,
-        LangevinDynamics(A, L, θ, H, σe),
-        GammaProcess(C, β; ϵ),
-        NormalMeanVariance(μw, σw),
+        LangevinDynamics([0 1; 0 θ], [0; 1], θ),
+        GammaProcess(C, β; kwargs...),
+        Normal(μw, σw),
     )
 
-    obs = LinearGaussianObservation(H, σe)
-    return StateSpaceModel(dyn, obs)
+    obs = LinearGaussianObservation([1; 0], σe)
+    return SSMProblems.StateSpaceModel(LevyPrior(zeros(Bool, 2), I(2)), dyn, obs)
 end
 
 # Levy SSM with Langevin dynamics
@@ -139,18 +119,18 @@ end
 # Simulation parameters
 N = 200
 ts = range(0, 100; length=N)
-levyssm = LevyModel(step(ts), θ, 1.0, 1.0, 1.0, 0.0, 1.0);
+levyssm = LevyModel(step(ts), -0.5, 1, 1.0, 1.0, 0, 1);
 
 # Simulate data
 rng = Random.MersenneTwister(1234);
 _, X, Y = sample(rng, levyssm, N);
 
 # Run sampler
 pg = AdvancedPS.PGAS(50);
-chains = sample(rng, levyssm(Y), pg, 100);
+chains = sample(rng, AdvancedPS.TracedSSM(levyssm, Y), pg, 100; progress=false);
 
 # Concat all sampled states
-marginal_states = hcat([chain.trajectory.model.X for chain in chains]...)
+marginal_states = hcat([chain.trajectory.model.X for chain in chains]...);
 
 # Plot marginal state and jump intensities for one trajectory
 p1 = plot(
@@ -166,7 +146,6 @@ plot!(
     label="Marginal State (x2)",
 )
 
-# TODO: collect jumps from the model
 p2 = scatter([], []; color=:darkorange, label="Jumps")
 
 plot(