ba-thesis/sw/utility/simulations.py

72 lines
2.3 KiB
Python

"""This file contains utility functions relating to tests and simulations of the decoders."""
import numpy as np
import typing
from tqdm import tqdm
from utility import noise
def count_bit_errors(d: np.array, d_hat: np.array) -> int:
"""Count the number of wrong bits in a decoded codeword.
:param d: Originally sent data
:param d_hat: Received data
:return: Number of bit errors
"""
return np.sum(d != d_hat)
def test_decoder(encoder: typing.Any,
decoder: typing.Any,
d: np.array,
SNRs: typing.Sequence[float] = np.linspace(1, 4, 7),
target_bit_errors: int = 100,
N_max: int = 10000) \
-> typing.Tuple[np.array, np.array]:
"""Calculate the Bit Error Rate (BER) for a given decoder for a number of SNRs.
This function prints its progress to stdout.
:param encoder: Instance of the encoder used to generate the codeword to transmit
:param decoder: Instance of the decoder to be tested
:param d: Dataword (element of [0, 1]^n)
:param SNRs: List of SNRs for which the BER should be calculated
:param target_bit_errors: Number of bit errors after which to stop the simulation
:param N_max: Maximum number of iterations to perform for each SNR
:return: Tuple of numpy arrays of the form (SNRs, BERs)
"""
x = encoder.encode(d)
BERs = []
for SNR in tqdm(SNRs, desc="Calculating Bit-Error-Rates",
position=0,
leave=False,
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
total_bit_errors = 0
total_bits = 0
for n in tqdm(range(N_max), desc=f"Simulating for SNR = {SNR} dB",
position=1,
leave=False,
bar_format="{l_bar}{bar}| {n_fmt}/{total_fmt}"):
# TODO: Is this a valid simulation? Can we just add AWGN to the codeword,
# ignoring and modulation and (e.g. matched) filtering?
y = noise.add_awgn(x, SNR, signal_amp=np.sqrt(2))
y_hat = decoder.decode(y)
total_bit_errors += count_bit_errors(d, y_hat)
total_bits += d.size
if total_bit_errors >= target_bit_errors:
break
BERs.append(total_bit_errors / total_bits)
return np.array(SNRs), np.array(BERs)