simulations/main.py

import matplotlib.pyplot as plt

import numpy as np
from sklearn.metrics import accuracy_score
from models import basic
from models.basic import AWGNChannel, BPSKDemod, BPSKMod, BypassChannel, AlphabetMod, AlphabetDemod
import misc
import math
import os
from models.autoencoder import Autoencoder, view_encoder
from models.optical_channel import OpticalChannel
from multiprocessing import Pool

CPU_COUNT = os.environ.get("CPU_COUNT", os.cpu_count())


def show_constellation(mod, chan, demod, samples=1000):
    x = misc.generate_random_bit_array(samples)
    x_mod = mod.forward(x)
    x_chan = chan.forward(x_mod)
    x_demod = demod.forward(x_chan)

    x_mod_rect = misc.polar2rect(x_mod)
    x_chan_rect = misc.polar2rect(x_chan)
    plt.plot(x_chan_rect[:, 0][x], x_chan_rect[:, 1][x], '+')
    plt.plot(x_chan_rect[:, 0][~x], x_chan_rect[:, 1][~x], '+')
    plt.plot(x_mod_rect[:, 0], x_mod_rect[:, 1], 'ro')
    axes = plt.gca()
    axes.set_xlim([-2, +2])
    axes.set_ylim([-2, +2])
    plt.grid()
    plt.show()
    print('Accuracy : ' + str())


def get_ber(mod, chan, demod, samples=1000):
    if samples % mod.N:
        samples += mod.N - (samples % mod.N)
    x = misc.generate_random_bit_array(samples)
    x_mod = mod.forward(x)
    x_chan = chan.forward(x_mod)
    x_demod = demod.forward(x_chan)
    return 1 - accuracy_score(x, x_demod)


def get_AWGN_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30):
    ber_x = np.linspace(start, stop, steps)
    ber_y = []
    for noise in ber_x:
        ber_y.append(get_ber(mod, AWGNChannel(noise), demod, samples=samples))
    return ber_x, ber_y


def __calc_ber(packed):
    # This function has to be outside get_Optical_ber in order to be pickled by pool
    mod, demod, noise, length, pulse_shape, samples = packed
    tx_channel = OpticalChannel(noise_level=noise, dispersion=-21.7, symbol_rate=10e9, sample_rate=400e9,
                                length=length, pulse_shape=pulse_shape, sqrt_out=True)
    return get_ber(mod, tx_channel, demod, samples=samples)


def get_Optical_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30, length=100, pulse_shape='rect'):
    ber_x = np.linspace(start, stop, steps)
    ber_y = []
    print(f"Computing Optical BER.. 0/{len(ber_x)}", end='')
    with Pool(CPU_COUNT) as pool:
        packed_args = [(mod, demod, noise, length, pulse_shape, samples) for noise in ber_x]
        for i, ber in enumerate(pool.imap(__calc_ber, packed_args)):
            ber_y.append(ber)
            i += 1  # just offset by 1
            print(f"\rComputing Optical BER.. {i}/{len(ber_x)} ({i*100/len(ber_x):6.2f}%)", end='')
    print()
    return ber_x, ber_y


def get_SNR(mod, demod, ber_func=get_Optical_ber, samples=1000, start=-5, stop=15, **ber_kwargs):
    """
    SNR for optics and RF should be calculated the same, that is A^2
    Because P∝V² and P∝I²
    """
    x_mod = mod.forward(misc.generate_random_bit_array(samples * mod.N))
    sig_amp = x_mod[:, 0]
    sig_power = [A ** 2 for A in sig_amp]
    av_sig_pow = np.mean(sig_power)
    av_sig_pow = math.log(av_sig_pow, 10)

    noise_start = -start + av_sig_pow
    noise_stop = -stop + av_sig_pow
    ber_x, ber_y = ber_func(mod, demod, samples, noise_start, noise_stop, **ber_kwargs)
    SNR = -ber_x + av_sig_pow
    return SNR, ber_y



if __name__ == '__main__':
    # show_constellation(BPSKMod(10e6), AWGNChannel(-1), BPSKDemod(10e6, 10e3))

    # get_ber(BPSKMod(10e6), AWGNChannel(-20), BPSKDemod(10e6, 10e3))
    # mod = MaryMod('8psk', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('qpsk', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('16qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # mod = MaryMod('64qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # aenc = Autoencoder(4, -25)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -25dB')

    # aenc = Autoencoder(5, -25)
    # aenc.train(samples=2e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 5bit -25dB')

    # view_encoder(aenc.encoder, 5)
    # plt.plot(*get_AWGN_ber(AlphabetMod('32qam', 10e6), AlphabetDemod('32qam', 10e6), samples=12000, start=-15), '-',
    #          label='32-QAM')
    # show_constellation(AlphabetMod('32qam', 10e6), AWGNChannel(-1), AlphabetDemod('32qam', 10e6))
    # mod = AlphabetMod('32qam', 10e6)
    # misc.display_alphabet(mod.alphabet, a_vals=True)
    # pass

    # aenc = Autoencoder(5, -15)
    # aenc.train(samples=2e6)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 5bit -15dB')
    #
    # aenc = Autoencoder(4, -25)
    # aenc.train(samples=6e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -20dB')
    #
    # aenc = Autoencoder(4, -15)
    # aenc.train(samples=6e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -15dB')

    # aenc = Autoencoder(2, -20)
    # aenc.train(samples=6e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 2bit -20dB')
    #
    # aenc = Autoencoder(2, -15)
    # aenc.train(samples=6e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 2bit -15dB')

    # aenc = Autoencoder(4, -10)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -10dB')
    #
    # aenc = Autoencoder(4, -8)
    # aenc.train(samples=5e5)
    # plt.plot(*get_AWGN_ber(aenc.get_modulator(), aenc.get_demodulator(), samples=12000, start=-15), '-',
    #          label='AE 4bit -8dB')

    # for scheme in ['64qam', '32qam', '16qam', 'qpsk', '8psk']:
    #     plt.plot(*get_SNR(
    #         AlphabetMod(scheme, 10e6),
    #         AlphabetDemod(scheme, 10e6),
    #         samples=100e3,
    #         steps=40,
    #         start=-15
    #     ), '-', label=scheme.upper())
    # plt.yscale('log')
    # plt.grid()
    # plt.xlabel('SNR dB')
    # plt.ylabel('BER')
    # plt.legend()
    # plt.show()

    for l in np.logspace(start=0, stop=3, num=5):
        plt.plot(*get_SNR(
            AlphabetMod('4pam', 10e6),
            AlphabetDemod('4pam', 10e6),
            samples=2000,
            steps=200,
            start=-5,
            stop=20,
            length=l,
            pulse_shape='rcos'
        ), '-', label=(str(int(l))+'km'))

    plt.yscale('log')
    # plt.gca().invert_xaxis()
    plt.grid()
    plt.xlabel('SNR dB')
    # plt.ylabel('BER')
    plt.title("BER against Fiber length")
    plt.legend()
    plt.show()
    # FIXME: Exit for now
    exit()

    for ps in ['rect']: #, 'rcos', 'rrcos']:
        plt.plot(*get_Optical_ber(
            AlphabetMod('4pam', 10e6),
            AlphabetDemod('4pam', 10e6),
            samples=30000,
            steps=40,
            start=-35,
            # stop=10,
            length=1,
            pulse_shape=ps
        ), '-', label=ps)

    plt.yscale('log')
    plt.grid()
    plt.xlabel('SNR dB')
    plt.ylabel('BER')
    plt.title("BER for different pulse shapes")
    plt.legend()
    plt.show()

    pass