From 1324398231f588a7de3d5937306e8cdc2d07053f Mon Sep 17 00:00:00 2001
From: Moritz Klischat <moritz.klischat@tum.de>
Date: Mon, 28 Oct 2019 15:46:03 +0100
Subject: [PATCH] fixed viterbi algorithm #1126

---
 probability.py            | 27 +++++++++++++++++++++------
 tests/test_probability.py |  6 ++++--
 2 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/probability.py b/probability.py
index c503084c4..e3fe6cddb 100644
--- a/probability.py
+++ b/probability.py
@@ -11,6 +11,7 @@
 import random
 from collections import defaultdict
 from functools import reduce
+import numpy as np
 
 
 # ______________________________________________________________________________
@@ -687,28 +688,42 @@ def forward_backward(HMM, ev):
 
 def viterbi(HMM, ev):
     """[Equation 15.11]
-    Viterbi algorithm to find the most likely sequence. Computes the best path,
+    Viterbi algorithm to find the most likely sequence. Computes the best path and the corresponding probabilities,
     given an HMM model and a sequence of observations."""
     t = len(ev)
+    ev = ev.copy()
     ev.insert(0, None)
 
     m = [[0.0, 0.0] for _ in range(len(ev) - 1)]
 
     # the recursion is initialized with m1 = forward(P(X0), e1)
     m[0] = forward(HMM, HMM.prior, ev[1])
+    # keep track of maximizing predecessors
+    backtracking_graph = []
 
     for i in range(1, t):
         m[i] = element_wise_product(HMM.sensor_dist(ev[i + 1]),
                                     [max(element_wise_product(HMM.transition_model[0], m[i - 1])),
                                      max(element_wise_product(HMM.transition_model[1], m[i - 1]))])
+        backtracking_graph.append([np.argmax(element_wise_product(HMM.transition_model[0], m[i - 1])),
+                                   np.argmax(element_wise_product(HMM.transition_model[1], m[i - 1]))])
+
+    # computed probabilities
+    ml_probabilities = [0.0] * (len(ev) - 1)
+    # most likely sequence
+    ml_path = [True] * (len(ev) - 1)
 
-    path = [0.0] * (len(ev) - 1)
     # the construction of the most likely sequence starts in the final state with the largest probability,
-    # and runs backwards; the algorithm needs to store for each xt its best predecessor xt-1
-    for i in range(t, -1, -1):
-        path[i - 1] = max(m[i - 1])
+    # and runs backwards; the algorithm needs to store for each xt its predecessor xt-1 maximizing its probability
+    i_max = np.argmax(m[-1])
+
+    for i in range(t - 1, -1, -1):
+        ml_probabilities[i] = m[i][i_max]
+        ml_path[i] = True if i_max == 0 else False
+        if i > 0:
+            i_max = backtracking_graph[i - 1][i_max]
 
-    return path
+    return ml_path, ml_probabilities
 
 
 # _________________________________________________________________________
diff --git a/tests/test_probability.py b/tests/test_probability.py
index 5acd862bc..b38052894 100644
--- a/tests/test_probability.py
+++ b/tests/test_probability.py
@@ -288,10 +288,12 @@ def test_viterbi():
     umbrellaHMM = HiddenMarkovModel(umbrella_transition, umbrella_sensor)
 
     umbrella_evidence = [T, T, F, T, T]
-    assert rounder(viterbi(umbrellaHMM, umbrella_evidence)) == [0.8182, 0.5155, 0.1237, 0.0334, 0.0210]
+    assert viterbi(umbrellaHMM, umbrella_evidence)[0] == [T, T, F, T, T]
+    assert rounder(viterbi(umbrellaHMM, umbrella_evidence)[1]) == [0.8182, 0.5155, 0.1237, 0.0334, 0.0210]
 
     umbrella_evidence = [T, F, T, F, T]
-    assert rounder(viterbi(umbrellaHMM, umbrella_evidence)) == [0.8182, 0.1964, 0.053, 0.0154, 0.0042]
+    assert viterbi(umbrellaHMM, umbrella_evidence)[0] == [T, F, F, F, T]
+    assert rounder(viterbi(umbrellaHMM, umbrella_evidence)[1]) == [0.8182, 0.1964, 0.0275, 0.0154, 0.0042]
 
 
 def test_fixed_lag_smoothing():