Utils module

TuringLang · Sep 13, 2023 · 09f1687 · 09f1687
1 parent 5bcb4f1
commit 09f1687
Show file tree

Hide file tree

Showing 3 changed files with 17 additions and 6 deletions.
diff --git a/examples/smc.jl b/examples/smc.jl
@@ -17,7 +17,7 @@ end
 function sweep!(
     rng::AbstractRNG,
     model::AbstractStateSpaceModel,
-    particles::SSMProblems.ParticleContainer,
+    particles::SSMProblems.Utils.ParticleContainer,
     observations::AbstractArray,
     resampling=systematic_resampling,
     threshold=0.5,
@@ -37,7 +37,7 @@ function sweep!(
         # Mutation step
         for i in eachindex(particles)
             latent_state = transition!!(rng, model, particles[i].state, timestep)
-            particles[i] = SSMProblems.Particle(particles[i], latent_state)
+            particles[i] = SSMProblems.Utils.Particle(particles[i], latent_state)
             logweights[i] += emission_logdensity(
                 model, particles[i].state, observation, timestep
             )
@@ -60,7 +60,7 @@ function sample(
 )
     particles = map(1:N) do i
         state = transition!!(rng, model)
-        SSMProblems.Particle(state)
+        SSMProblems.Utils.Particle(state)
     end
     samples = sweep!(rng, model, particles, observations, resampling, threshold)
     return samples
@@ -109,7 +109,7 @@ end
 
 # Sample latent state trajectories
 samples = sample(rng, LinearSSM(), N, observations)
-traces = reverse(hcat(map(SSMProblems.linearize, samples)...))
+traces = reverse(hcat(map(SSMProblems.Utils.linearize, samples)...))
 
 scatter(traces; color=:black, opacity=0.3, label=false)
 plot!(x; label="True state")

diff --git a/src/SSMProblems.jl b/src/SSMProblems.jl
@@ -10,7 +10,7 @@ abstract type AbstractStateSpaceModel end
 abstract type AbstractParticleCache end
 
 """
-    transition!!(rng, model[, timestep, state, cache])
+    transition!!(rng, model[, state, timestep, cache])
 
 Simulate the particle for the next time step from the forward dynamics.
 """

diff --git a/src/utils/particles.jl b/src/utils/particles.jl
@@ -1,10 +1,19 @@
-# Concrete Particles as LinkedLists
+"""
+    Common concrete implementations of Particle types for Particle Filter kernels.
+"""
+module Utils
+
 abstract type Node{T} end
 
 struct Root{T} <: Node{T} end
 Root(T) = Root{T}()
 Root() = Root(Any)
 
+"""
+    Particle{T}
+
+Particle as immutable LinkedList. 
+"""
 struct Particle{T} <: Node{T}
     parent::Node{T}
     state::T
@@ -30,3 +39,5 @@ function linearize(particle::Particle{T}) where {T}
     end
     return trace
 end
+
+end