ba-thesis/sw/utility/simulations.py

73 lines
2.4 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(x: np.array,
decoder: typing.Any,
SNRs: typing.Sequence[float] = np.linspace(1, 7, 7),
target_frame_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 assumes the all-zeros assumption holds. Progress is printed to stdout.
:param x: Codeword to be sent (Element of [0, 1]^n)
:param decoder: Instance of the decoder to be tested
:param SNRs: List of SNRs for which the BER should be calculated
:param target_frame_errors: Number of frame 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_bpsk = 1 - 2 * x # Map x from [0, 1]^n to [-1, 1]^n
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
total_frame_errors = 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}"):
# Simulate channel
y = noise.add_awgn(x_bpsk, SNR, signal_amp=np.sqrt(2))
# Decode received frame
x_hat = decoder.decode(y)
# Calculate statistics
total_bit_errors += count_bit_errors(x, x_hat)
total_bits += x.size
total_frame_errors += 1 if total_bit_errors > 0 else 0
if total_frame_errors >= target_frame_errors:
break
BERs.append(total_bit_errors / total_bits)
return np.array(SNRs), np.array(BERs)