Functional API for applying transition matrix

gdalle · Nov 15, 2024 · 0947ce1 · 0947ce1
1 parent a370596
commit 0947ce1
Show file tree

Hide file tree

Showing 2 changed files with 15 additions and 3 deletions.
diff --git a/src/inference/forward_backward.jl b/src/inference/forward_backward.jl
@@ -47,10 +47,11 @@ function _forward_backward!(
     # Backward
     β[:, t2] .= c[t2]
     for t in (t2 - 1):-1:t1
-        trans = transition_matrix(hmm, control_seq[t])
         Bβ[:, t + 1] .= view(B, :, t + 1) .* view(β, :, t + 1)
-        mul!(view(β, :, t), trans, view(Bβ, :, t + 1))
-        lmul!(c[t], view(β, :, t))
+        βₜ = view(β, :, t)
+        Bβₜ₊₁ = view(Bβ, :, t + 1)
+        predict_previous_state!(βₜ, hmm, Bβₜ₊₁, control_seq[t])
+        lmul!(c[t], βₜ)
     end
     Bβ[:, t1] .= view(B, :, t1) .* view(β, :, t1)
 

diff --git a/src/inference/predict.jl b/src/inference/predict.jl
@@ -8,3 +8,14 @@ function predict_next_state!(
     mul!(next_state_marginals, transpose(trans), current_state_marginals)
     return next_state_marginals
 end
+
+function predict_previous_state!(
+    previous_state_marginals::AbstractVector{<:Real},
+    hmm::AbstractHMM,
+    current_state_marginals::AbstractVector{<:Real},
+    control=nothing,
+)
+    trans = transition_matrix(hmm, control)
+    mul!(previous_state_marginals, trans, current_state_marginals)
+    return previous_state_marginals
+end