Refactor DAG shortest path algorithm using Viterbi

Source/Engine/gramambular2/reading_grid.cpp

@@ -110,90 +110,6 @@ std::optional<ReadingGrid::NodePtr> ReadingGrid::findInSpan(
 
 namespace {
 
-// Defines a vertex of a DAG. This is a mutable data structure used for both
-// DAG construction and single-source shortest-path computation.
-struct Vertex {
-  explicit Vertex(ReadingGrid::NodePtr nodePtr) : node(std::move(nodePtr)) {
-    edges.reserve(ReadingGrid::kMaximumSpanLength);
-  }
-  ReadingGrid::NodePtr node;
-  std::vector<Vertex*> edges;
-
-  // Used during topological-sort.
-  bool topologicallySorted = false;
-
-  // Used during shortest-path computation. We are actually computing the
-  // path with the *largest* weight, hence distance's initial value being
-  // negative infinity. If we were to compute the *shortest* weight/distance,
-  // we would have initialized this to infinity.
-  double distance = -std::numeric_limits<double>::infinity();
-  Vertex* prev = nullptr;
-};
-
-// Cormen et al. 2001 explains the historical origin of the term "relax."
-void Relax(Vertex* u, Vertex* v) {
-  // The distance from u to w is simply v's score.
-  double w = v->node->score();
-
-  // Since we are computing the largest weight, we update v's distance and prev
-  // if the current distance to v is *less* than that of u's plus the distance
-  // to v (which is represented by w).
-  if (v->distance < u->distance + w) {
-    v->distance = u->distance + w;
-    v->prev = u;
-  }
-}
-
-using VertexSpan = std::vector<Vertex>;
-
-// Topological-sorts a DAG that has a single root and returns the vertices in
-// topological order. Here, a non-recursive version is implemented using our own
-// stack and state definitions, so that we are not constrained by the current
-// thread's stack size. This is the equivalent to this recursive version:
-//
-//  void TopologicalSort(Vertex* v) {
-//    for (Vertex* nv : v->edges) {
-//      if (!nv->topologicallySorted) {
-//        dfs(nv, result);
-//      }
-//    }
-//    v->topologicallySorted = true;
-//    result.push_back(v);
-//  }
-//
-// The recursive version is similar to the TOPOLOGICAL-SORT algorithm found in
-// Cormen et al. 2001.
-std::vector<Vertex*> TopologicalSort(Vertex* root) {
-  std::vector<Vertex*> result;
-  struct State {
-    explicit State(Vertex* vv) : v(vv), edgeIter(v->edges.begin()) {}
-    Vertex* v;
-    std::vector<Vertex*>::iterator edgeIter;
-  };
-  std::stack<State> stack;
-  stack.emplace(root);
-
-  while (!stack.empty()) {
-    State& state = stack.top();
-    Vertex* v = state.v;
-
-    if (state.edgeIter != v->edges.end()) {
-      Vertex* nv = *state.edgeIter;
-      ++state.edgeIter;
-      if (!nv->topologicallySorted) {
-        stack.emplace(nv);
-        continue;
-      }
-    }
-
-    v->topologicallySorted = true;
-    result.push_back(v);
-    stack.pop();
-  }
-
-  return result;
-}
-
 int64_t GetEpochNowInMicroseconds() {
   auto now = std::chrono::system_clock::now();
   int64_t timestamp =
@@ -206,8 +122,8 @@ int64_t GetEpochNowInMicroseconds() {
 }  // namespace
 
 // Find the weightiest path in the grid graph. The path represents the most
-// likely hidden chain of events from the observations. We use the
-// DAG-SHORTEST-PATHS algorithm in Cormen et al. 2001 to compute such path.
+// likely hidden chain of events from the observations.
+// We use the Viterbi algorithm to compute such path.
 // Instead of computing the path with the shortest distance, though, we compute
 // the path with the longest distance (so the weightiest), since with log
 // probability a larger value means a larger probability. The algorithm runs in
@@ -220,67 +136,68 @@ ReadingGrid::WalkResult ReadingGrid::walk() {
   }
   int64_t start = GetEpochNowInMicroseconds();
 
-  std::vector<VertexSpan> vspans(spans_.size(), VertexSpan());
-  size_t vertices = 0;
-  size_t edges = 0;
-  for (size_t i = 0, len = spans_.size(); i < len; ++i) {
+  // Defines a state in the DP table. This structure tracks the maximum
+  // accumulated score and the back-pointer required for path reconstruction in
+  // the Viterbi algorithm.
+  struct State {
+    size_t fromIndex = 0;
+    ReadingGrid::NodePtr fromNode = nullptr;
+    double maxScore = -std::numeric_limits<double>::infinity();
+  };
+
+  const size_t readingLen = readings_.size();
+  std::vector<State> viterbi(readingLen + 1);
+  viterbi[0].maxScore = 0.0;
+
+  // Iterate through the grid and compute the maximum accumulated score for each
+  // reachable position. Since the grid is a lattice where edges only point
+  // forward, processing nodes in index order is equivalent to processing them
+  // in topological order.
+  size_t reachableStates = 0;
+  size_t evaluatedEdges = 0;
+  for (size_t i = 0; i < readingLen; ++i) {
+    ++reachableStates;
+
     const ReadingGrid::Span& span = spans_[i];
-    for (size_t j = 1, maxSpanLen = span.maxLength(); j <= maxSpanLen; ++j) {
-      NodePtr p = span.nodeOf(j);
-      if (p != nullptr) {
-        vspans[i].emplace_back(std::move(p));
-        ++vertices;
-      }
-    }
-  }
-  result.vertices = vertices;
+    const size_t maxSpanLen = span.maxLength();
 
-  Vertex terminal(std::make_shared<ReadingGrid::Node>(
-      "_TERMINAL_", 0, std::vector<LanguageModel::Unigram>()));
-  for (size_t i = 0, vspansLen = vspans.size(); i < vspansLen; ++i) {
-    for (Vertex& v : vspans[i]) {
-      size_t nextVertexPos = i + v.node->spanningLength();
-      if (nextVertexPos == vspansLen) {
-        v.edges.push_back(&terminal);
+    for (size_t spanLen = 1; spanLen <= maxSpanLen; ++spanLen) {
+      const ReadingGrid::NodePtr& node = span.nodeOf(spanLen);
+      if (node == nullptr) {
         continue;
       }
-
-      for (Vertex& nv : vspans[nextVertexPos]) {
-        v.edges.push_back(&nv);
-        ++edges;
+      ++evaluatedEdges;
+
+      // Performs a relaxation on a transition. This updates the destination
+      // state if the path through the current node yields a higher score than
+      // the previously known best path. This is the core operation of the
+      // Viterbi algorithm, adapted for finding the maximum likelihood path.
+      double score = viterbi[i].maxScore + node->score();
+      State& target = viterbi[i + spanLen];
+      if (score > target.maxScore) {
+        target.maxScore = score;
+        target.fromNode = node;
+        target.fromIndex = i;
       }
     }
   }
-  result.edges = edges;
+  // Vertices are the reachable states
+  // Edges are the candidate word transitions
+  result.vertices = reachableStates;
+  result.edges = evaluatedEdges;
 
-  Vertex root(std::make_shared<ReadingGrid::Node>(
-      "_ROOT_", 0, std::vector<LanguageModel::Unigram>()));
-  root.distance = 0;
-  for (Vertex& v : vspans[0]) {
-    root.edges.push_back(&v);
-  }
-
-  std::vector<Vertex*> ordered = TopologicalSort(&root);
-  for (auto it = ordered.rbegin(), rend = ordered.rend(); it != rend; ++it) {
-    Vertex* u = *it;
-    for (Vertex* v : u->edges) {
-      Relax(u, v);
-    }
-  }
-
-  std::vector<NodePtr> walked;
+  // Reconstruct the most likely path by tracing back from the end of the grid
+  // to the root using the back-pointers
   size_t totalReadingLen = 0;
-  Vertex* it = &terminal;
-  while (it->prev != nullptr) {
-    walked.push_back(it->prev->node);
-    it = it->prev;
-    totalReadingLen += it->node->spanningLength();
+  for (size_t curr = readingLen; curr > 0; curr = viterbi[curr].fromIndex) {
+    assert(viterbi[curr].fromNode != nullptr);
+    totalReadingLen += viterbi[curr].fromNode->spanningLength();
+    result.nodes.emplace_back(std::move(viterbi[curr].fromNode));
   }
-
-  assert(totalReadingLen == readings_.size());
-  assert(walked.size() >= 2);
+  std::reverse(result.nodes.begin(), result.nodes.end());
+  assert(totalReadingLen == readingLen);
   result.totalReadings = totalReadingLen;
-  result.nodes = std::vector<NodePtr>(walked.rbegin() + 1, walked.rend());
+
   result.elapsedMicroseconds = GetEpochNowInMicroseconds() - start;
   return result;
 }