Add example for convolutional codes

Author: BastienTr

commpy/examples/conv_encode_decode.py

@@ -0,0 +1,121 @@
+# Authors: CommPy contributors
+# License: BSD 3-Clause
+
+import math
+
+import matplotlib.pyplot as plt
+import numpy as np
+
+import commpy.channelcoding.convcode as cc
+import commpy.channels as chan
+import commpy.links as lk
+import commpy.modulation as mod
+import commpy.utilities as util
+
+# =============================================================================
+# Convolutional Code 1: G(D) = [1+D^2, 1+D+D^2]
+# Standard code with rate 1/2
+# =============================================================================
+
+# Number of delay elements in the convolutional encoder
+memory1 = np.array(2, ndmin=1)
+
+# Generator matrix
+g_matrix1 = np.array((0o5, 0o7), ndmin=2)
+
+# Create trellis data structure
+trellis1 = cc.Trellis(memory1, g_matrix1)
+
+# =============================================================================
+# Convolutional Code 2: G(D) = [[1, 0, 0], [0, 1, 1+D]]; F(D) = [[D, D], [1+D, 1]]
+# RSC with rate 2/3
+# =============================================================================
+
+# Number of delay elements in the convolutional encoder
+memory2 = np.array((1, 1))
+
+# Generator matrix & feedback matrix
+g_matrix2 = np.array(((1, 0, 0), (0, 1, 3)))
+feedback = np.array(((2, 2), (3, 1)))
+
+# Create trellis data structure
+trellis2 = cc.Trellis(memory2, g_matrix2, feedback, 'rsc')
+
+# =============================================================================
+# Basic example using homemade counting and hard decoding
+# =============================================================================
+
+# Traceback depth of the decoder
+tb_depth = None  # Default value is 5 times the number or memories
+
+for trellis in (trellis1, trellis2):
+    for i in range(10):
+        # Generate random message bits to be encoded
+        message_bits = np.random.randint(0, 2, 1000)
+
+        # Encode message bits
+        coded_bits = cc.conv_encode(message_bits, trellis)
+
+        # Introduce bit errors (channel)
+        coded_bits[np.random.randint(0, 1000)] = 0
+        coded_bits[np.random.randint(0, 1000)] = 0
+        coded_bits[np.random.randint(0, 1000)] = 1
+        coded_bits[np.random.randint(0, 1000)] = 1
+
+        # Decode the received bits
+        decoded_bits = cc.viterbi_decode(coded_bits.astype(float), trellis, tb_depth)
+
+        num_bit_errors = util.hamming_dist(message_bits, decoded_bits[:len(message_bits)])
+
+        if num_bit_errors != 0:
+            print(num_bit_errors, "Bit Errors found!")
+        elif i == 9:
+            print("No Bit Errors :)")
+
+# =============================================================================
+# Complete example using commpy features. Example with code 1
+# =============================================================================
+
+# Modem : QPSK
+modem = mod.QAMModem(4)
+
+# AWGN channel
+channels = chan.SISOFlatChannel(None, (1+0j, 0j))
+
+# SNR range to test
+SNRs = np.arange(0, 6) + 10 * math.log10(modem.num_bits_symbol)
+
+
+# Modulation function
+def modulate(bits):
+    return modem.modulate(cc.conv_encode(bits, trellis1, 'cont'))
+
+
+# Demodulation function
+def demode(msg):
+    return modem.demodulate(msg, 'hard')
+
+
+# Receiver function (no process required as there are no fading)
+def receiver(y, h, constellation, noise_var):
+    return modem.demodulate(y, 'hard')
+
+
+# Decoder function
+def decoder(msg):
+    return cc.viterbi_decode(msg, trellis1)
+
+
+# Build model from parameters
+code_rate = trellis1.k / trellis1.n
+model = lk.LinkModel(modulate, channels, receiver,
+                     modem.num_bits_symbol, modem.constellation, modem.Es,
+                     decoder, code_rate)
+
+# Test
+BERs = lk.link_performance(model, SNRs, 10000, 600, 5000, code_rate)
+plt.semilogy(SNRs, BERs, 'o-')
+plt.grid()
+plt.xlabel('Signal to Noise Ration (dB)')
+plt.ylabel('Bit Error Rate')
+plt.show()