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