Tests

TuringLang · Mar 24, 2024 · 2249933 · 2249933
1 parent 3643802
commit 2249933
Show file tree

Hide file tree

Showing 4 changed files with 14 additions and 12 deletions.
diff --git a/test/Project.toml b/test/Project.toml
@@ -4,6 +4,7 @@ Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f"
 Libtask = "6f1fad26-d15e-5dc8-ae53-837a1d7b8c9f"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Random123 = "74087812-796a-5b5d-8853-05524746bad3"
+SSMProblems = "26aad666-b158-4e64-9d35-0e672562fa48"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]

diff --git a/test/container.jl b/test/container.jl
@@ -1,14 +1,17 @@
 @testset "container.jl" begin
 
     # Since the extension would hide the low level function call API
-    mutable struct LogPModel{T} <: AdvancedPS.AbstractStateSpaceModel
+    mutable struct LogPModel{T} <: SSMProblems.AbstractStateSpaceModel
         logp::T
         X::Array{T}
     end
 
-    AdvancedPS.initialization(model::LogPModel) = Uniform()
-    AdvancedPS.transition(model::LogPModel, state, step) = Uniform()
-    AdvancedPS.observation(model::LogPModel, state, step) = model.logp
+    SSMProblems.transition!!(rng::AbstractRNG, model::LogPModel) = rand(rng, Uniform())
+    function SSMProblems.transition!!(rng::AbstractRNG, model::LogPModel, state, step)
+        return rand(rng, Uniform())
+    end
+    SSMProblems.emission_logdensity(model::LogPModel, state, step) = model.logp
+
     AdvancedPS.isdone(model::LogPModel, step) = false
 
     @testset "copy particle container" begin

diff --git a/test/pgas.jl b/test/pgas.jl
@@ -5,19 +5,16 @@
         r::Float64
     end
 
-    mutable struct BaseModel <: AdvancedPS.AbstractStateSpaceModel
+    mutable struct BaseModel <: SSMProblems.AbstractStateSpaceModel
         X::Vector{Float64}
         θ::Params
         BaseModel(params::Params) = new(Vector{Float64}(), params)
     end
 
-    AdvancedPS.initialization(model::BaseModel) = Normal(0, model.θ.q)
-    function AdvancedPS.transition(model::BaseModel, state, step)
-        return Distributions.Normal(model.θ.a * state, model.θ.q)
-    end
-    function AdvancedPS.observation(model::BaseModel, state, step)
-        return Distributions.logpdf(Distributions.Normal(state, model.θ.r), 0)
-    end
+    SSMProblems.transition!!(rng::AbstractRNG, model::BaseModel) = rand(rng, Normal(0, model.θ.q))
+    SSMProblems.transition!!(rng::AbstractRNG, model::BaseModel, state, step) = rand(rng, Normal(model.θ.a * state, model.θ.q))
+    SSMProblems.emission_logdensity(model::BaseModel, state, step) = logpdf(Distributions.Normal(state, model.θ.r), 0)
+    SSMProblems.transition_logdensity(model::BaseModel, prev_state, current_state, step) = logpdf(Normal(model.θ.a * prev_state, model.θ.q), current_state)
 
     AdvancedPS.isdone(::BaseModel, step) = step > 3
 

diff --git a/test/runtests.jl b/test/runtests.jl
@@ -4,6 +4,7 @@ using Distributions
 using Libtask
 using Random
 using Test
+using SSMProblems
 
 @testset "AdvancedPS.jl" begin
     @testset "Resampling tests" begin