Viterbi Algorithm › Implementing Viterbi Efficiently

import numpy as np


class EfficientViterbi:
    """
    Optimized Viterbi implementation using NumPy.

    Uses matrix operations for faster computation on larger state spaces.
    """

    def __init__(self, states, trans_matrix, emit_matrix, start_probs):
        """
        Args:
            states: List of state names
            trans_matrix: (n_states, n_states) transition probabilities
            emit_matrix: (n_states, n_observations) emission probabilities
            start_probs: (n_states,) initial state probabilities
        """
        self.states = states
        self.n_states = len(states)

        # Store log probabilities
        self.log_trans = np.log(trans_matrix + 1e-10)
        self.log_emit = emit_matrix  # Will be indexed per observation
        self.log_start = np.log(start_probs + 1e-10)

    def decode(self, observation_indices):
        """
        Find the optimal state sequence for a sequence of observations.

        Args:
            observation_indices: List of observation indices

        Returns:
            path: List of state indices
            score: Log probability of the path
        """
        n = len(observation_indices)

        # Viterbi table: (sequence_length, n_states)
        V = np.full((n, self.n_states), float("-inf"))
        backpointers = np.zeros((n, self.n_states), dtype=int)

        # Base case
        obs_0 = observation_indices[0]
        V[0] = self.log_start + self.log_emit[:, obs_0]

        # Fill the trellis
        for i in range(1, n):
            obs_i = observation_indices[i]
            emit_i = self.log_emit[:, obs_i]

            # For each current state, find best predecessor
            # V[i-1] is (n_states,), log_trans is (n_states, n_states)
            # We want: for each current state s, max over s' of V[i-1, s'] + log_trans[s', s]

            # Broadcast: V[i-1, :, np.newaxis] + log_trans = (n_states, n_states)
            # where entry [s', s] = V[i-1, s'] + log_trans[s', s]
            scores = V[i - 1, :, np.newaxis] + self.log_trans

            # Max over predecessors (axis 0) gives best score for each current state
            best_scores = scores.max(axis=0)
            best_prevs = scores.argmax(axis=0)

            V[i] = best_scores + emit_i
            backpointers[i] = best_prevs

        # Find best final state
        best_final = V[n - 1].argmax()
        best_score = V[n - 1, best_final]

        # Backtrack
        path = [best_final]
        current = best_final
        for i in range(n - 1, 0, -1):
            current = backpointers[i, current]
            path.append(current)

        path.reverse()
        return path, best_score

# Demonstrate the efficient implementation

# Create a small HMM
n_states = 3
n_obs = 4

# Random transition and emission matrices
np.random.seed(42)
trans_matrix = np.random.dirichlet(np.ones(n_states), size=n_states)
emit_matrix = np.random.dirichlet(np.ones(n_obs), size=n_states)
start_probs = np.ones(n_states) / n_states

# Convert emission to log space
log_emit = np.log(emit_matrix + 1e-10)

# Create decoder
states = ["State0", "State1", "State2"]
decoder = EfficientViterbi(states, trans_matrix, log_emit, start_probs)

# Decode a sequence
observations = [0, 1, 2, 1, 0, 3, 2]
path, score = decoder.decode(observations)
path_names = [states[i] for i in path]

def viterbi_sparse(
    observations, states, start_prob, valid_transitions, emit_prob
):
    """
    Viterbi with sparse transition handling.

    Args:
        valid_transitions: Dict mapping state -> list of valid predecessor states
                          with their log probabilities
    """
    n = len(observations)
    V = [{} for _ in range(n)]
    backpointer = [{} for _ in range(n)]

    # Base case
    for s in states:
        emit = emit_prob.get((s, observations[0]), float("-inf"))
        V[0][s] = start_prob.get(s, float("-inf")) + emit
        backpointer[0][s] = None

    # Recursion - only consider valid predecessors
    for i in range(1, n):
        for s in states:
            best_prev = None
            best_score = float("-inf")
            emit = emit_prob.get((s, observations[i]), float("-inf"))

            # Only iterate over valid predecessors
            for s_prev, log_trans in valid_transitions.get(s, []):
                score = V[i - 1][s_prev] + log_trans + emit
                if score > best_score:
                    best_score = score
                    best_prev = s_prev

            if best_score > float("-inf"):
                V[i][s] = best_score
                backpointer[i][s] = best_prev

    # Find best path
    valid_finals = [s for s in states if s in V[n - 1]]
    if not valid_finals:
        return None, float("-inf")

    best_final = max(valid_finals, key=lambda s: V[n - 1][s])
    best_score = V[n - 1][best_final]

    path = [best_final]
    current = best_final
    for i in range(n - 1, 0, -1):
        current = backpointer[i][current]
        if current is None:
            break
        path.append(current)

    path.reverse()
    return path, best_score