Viterbi encoder and decoder example

viterbi.py

# Viterbi algorithm
# See example at https://www.youtube.com/watch?v=dKIf6mQUfnY
# -- Lars.Rasmusson@gmail.com (2018-10-04)

import numpy as np

# Define the coder by its internal state transition and output
states = [(0,0), (0,1), (1,0), (1,1)]
def output(s, b): return (b^s[1], b^s[1]^s[0]) # xor binary values
def next_state(s, b): return b, s[0]

# Create an encoder function
def encode(v, s=(0,0)):
    for b in v:
        yield output(s, b)
        s = next_state(s, b)

# helper function:
#   compute which states and input can lead to state s
def incoming_states(s, B=[0, 1]):
    return [(u, b) for u in states for b in B
            if next_state(u, b) == s]

# Hamming distance
def hdist(x, y): return sum([a != b for a, b in zip(x, y)])

# Update state probability to maximally likely probability
# under Viterbi decoding assumptions (= as few errors as possible).
def update(state, c):
    n = {}
    for s in states:
        best, best_w = np.inf, []
        for u, b in incoming_states(s):
            q = hdist(c, output(u, b))
            p, w = state[u]
            if p + q < best:
                best = p + q
                best_w = w + [b]
        n[s] = (best, best_w)
    return n

def decode(code):
    # initialize - we start in state (0,0)
    s = dict((k, (np.inf, [])) for k in states)
    s[(0,0)] = (0, [])

    # update state probabilities based on incoming symbols
    for c in code:
        s = update(s, c)

    # print possible symbols
    best = min(v for (v,_) in s.values())
    return [v for (p, v) in s.values() if p == best]

# Example from youtube
# print(decode([(0,0), (0,1), (0,1), (1,0), (1,0)]))

def flip(l, n):
    l = np.array(l)
    s = l.shape
    l = np.array(l.reshape(-1))
    f = np.random.permutation(np.arange(len(l)))[:n]
    l[f] = 1 - l[f]
    return l.reshape(s)

def test():
    for v, n in [([1,1,0,1,0], 0), ([0,1,0,1,0,0,1,0,1], 1),
                 ([0,1,0,1,0,0,1,0,1], 3), ([0,0,1,1,0],2)]:
        e = flip(list(encode(v)),n)
        d = decode(e)
        print('encoding', v, '\n code with', n, 'bit flips:', e.tolist())
        print(v in d, 'decoding as', d, '\n')

test()