Split trajectory

Author: FredericWantiez

Diff for: src/algorithms/ffbs.jl

@@ -25,6 +25,43 @@ function (callback::WeightedParticleRecorderCallback)(
     return nothing
 end
 
+function gen_trajectory(
+    rng::Random.AbstractRNG, 
+    model::StateSpaceModel, 
+    particles::AbstractMatrix{T},  # Need better container
+    log_weights::AbstractMatrix{WT}, 
+    forward_state,
+    n_timestep::Int;
+    kwargs...
+) where {T,WT}
+    trajectory = Vector{T}(undef, n_timestep)
+    trajectory[end] = forward_state
+    for step in (n_timestep - 1):-1:1
+        backward_weights = backward(
+            model,
+            step,
+            trajectory[step + 1],
+            particles[step, :],
+            log_weights[step, :];
+            kwargs...,
+        )
+        ancestor = rand(rng, Categorical(softmax(backward_weights)))
+        trajectory[step] = particles[step, ancestor]
+    end
+    return trajectory
+end
+
+
+function backward(
+    model::StateSpaceModel, step::Integer, state, particles::T, log_weights::WT; kwargs...
+) where {T,WT}
+    transitions = map(particles) do prev_state
+        SSMProblems.logdensity(model.dyn, step, prev_state, state; kwargs...)
+    end
+    return log_weights + transitions
+end
+
+
 function sample(
     rng::Random.AbstractRNG,
     model::StateSpaceModel{T,LDT},
@@ -34,6 +71,7 @@ function sample(
     callback=nothing,
     kwargs...,
 ) where {T,LDT,N}
+
     n_timestep = length(obs)
     recorder = WeightedParticleRecorderCallback(
         Array{eltype(model.dyn)}(undef, n_timestep, N), Array{T}(undef, n_timestep, N)
@@ -46,28 +84,8 @@ function sample(
     trajectories = Array{eltype(model.dyn)}(undef, n_timestep, M)
 
     trajectories[end, :] = particles.filtered[idx_ref]
-    for step in (n_timestep - 1):-1:1
-        for j in 1:M
-            backward_weights = backward(
-                model::StateSpaceModel,
-                step,
-                trajectories[step + 1],
-                recorder.particles[step, :],
-                recorder.log_weights[step, :];
-                kwargs...,
-            )
-            ancestor = rand(rng, Categorical(softmax(backward_weights)))
-            trajectories[step, j] = recorder.particles[step, ancestor]
-        end
+    for j in 1:M
+        trajectories[:, j] = gen_trajectory(rng, model, recorder.particles, recorder.log_weights, trajectories[end, j], n_timestep)
     end
     return trajectories
 end
-
-function backward(
-    model::StateSpaceModel, step::Integer, state, particles::T, log_weights::WT; kwargs...
-) where {T,WT}
-    transitions = map(
-        x -> SSMProblems.logdensity(model.dyn, step, x, state; kwargs...), particles
-    )
-    return log_weights + transitions
-end