simulations/tests/ofdm_test.py

import itertools
import math
import numpy as np
from models.custom_layers import OpticalChannel
from matplotlib import pyplot as plt

# mapping_table = np.asarray([-3 - 3j, -3 - 1j, -3 + 3j, -3 + 1j,
#                             -1 - 3j, -1 - 1j, -1 + 3j, -1 + 1j,
#                             3 - 3j, 3 - 1j, 3 + 3j, 3 + 1j,
#                             1 - 3j, 1 - 1j, 1 + 3j, 1 + 1j])

mapping_table = np.asarray([-1 - 1j, -1 + 1j, 1 - 1j, 1 + 1j])

BANDWIDTH = 50e9

CARDINALITY = mapping_table.shape[0]
DISPERSION_FACTOR = -21.7 * 1e-24
FIBER_LENGTH = 5
FIBER_LENGTH_STDDEV = 0
RX_STDDEV = 0.01
SIG_AVG = 0.5
ENOB = 10

# Number of OFDM symbols to simulate
OFDM_N = 100
# Number of OFDM subcarriers
K = 32
# length of the cyclic prefix: 25% of the block
CP = K // 4
# number of pilot carriers per OFDM block
P = 0
# The known value each pilot transmits
pilotValue = 3 + 3j
# DC offset used to ensure all values are positive
DC_OFFSET = 50
DC_OFFSET = np.asarray([DC_OFFSET])

SHOW_PLOTS = True


bits_per_symbol = int(math.log(CARDINALITY, 2))
bits_lst = [list(i) for i in itertools.product([0, 1], repeat=bits_per_symbol)]

# Set true to view plot of symbols as IQ plot
if False:
    for idx, sym in enumerate(mapping_table):
        plt.plot(sym.real, sym.imag, 'bo')
        plt.text(sym.real, sym.imag + 0.2, str(bits_lst[idx])[1:-1], ha='center')

    plt.title('Symbols used in each subcarrier')
    plt.xlabel('I')
    plt.ylabel('Q')
    plt.xlim(-4, 4)
    plt.ylim(-4, 4)
    plt.show()

tx_stream = []
tx_syms = []

for ofdm_sym in range(OFDM_N):
    print(ofdm_sym)
    # All subcarriers used
    allCarriers = np.arange(K)

    # Identifying pilot carriers (and adding final subcarrier as a pilot for convenience)
    if P == 0:
        pilotCarriers = np.asarray([-1])
    else:
        pilotCarriers = allCarriers[::K // P]
        pilotCarriers = np.hstack([pilotCarriers, np.array([allCarriers[-1]])])
        # P = P + 1

    # Removing pilot carriers to obtain data carriers
    dataCarriers = np.delete(allCarriers, pilotCarriers)

    # if SHOW_PLOTS:
    #     f = np.fft.fftfreq(allCarriers.shape[0], d=1 / BANDWIDTH)
    #     plt.plot(f[1]*pilotCarriers, np.zeros_like(pilotCarriers), 'bo', label='pilot')
    #     plt.plot(f[1]*dataCarriers, np.zeros_like(dataCarriers), 'ro', label='data')
    #     plt.show()

    # Generate random symbols as integers and then map symbol values onto them
    input_syms = np.random.randint(CARDINALITY, size=len(dataCarriers))
    tx_syms.append(input_syms)
    mapped_syms = mapping_table[input_syms]

    # Generate the upper sideband of the OFDM symbol
    enc_stream_upper_f = np.zeros(K, dtype=complex)
    enc_stream_upper_f[pilotCarriers] = pilotValue
    enc_stream_upper_f[dataCarriers] = mapped_syms

    # Generate the lower sideband of the OFDM symbol
    enc_stream_lower_f = np.conjugate(np.flip(enc_stream_upper_f))

    # Combine the two sidebands with a DC offset to ensure values are always positive
    enc_stream_f = np.concatenate((DC_OFFSET, enc_stream_upper_f, enc_stream_lower_f), axis=None)

    # if SHOW_PLOTS:
    #     f = np.fft.fftfreq(enc_stream_f.shape[0], d=1/BANDWIDTH)
    #     plt.plot(f, np.real(enc_stream_f), 'x')
    #     plt.plot(f, np.imag(enc_stream_f), 'x')
    #     # plt.xlim(-0.5e10, 0.5e10)
    #     plt.show()

    # Take the inverse fourier transform
    enc_stream_t = np.fft.ifft(enc_stream_f)

    # if SHOW_PLOTS:
    #     t = np.arange(len(enc_stream_t))*(1/BANDWIDTH)
    #     plt.plot(t, np.real(enc_stream_t))
    #     plt.plot(t, np.imag(enc_stream_t))
    #     plt.show()

    # Take the real part to be transmitted via the channel
    tx_stream.append(np.real(enc_stream_t))


tx_stream = np.asarray(tx_stream).flatten()
tx_syms = np.asarray(tx_syms).flatten()

optical_channel = OpticalChannel(fs=BANDWIDTH,
                                 dispersion_factor=DISPERSION_FACTOR,
                                 fiber_length=FIBER_LENGTH,
                                 fiber_length_stddev=FIBER_LENGTH_STDDEV,
                                 lpf_cutoff=BANDWIDTH / 2,
                                 rx_stddev=RX_STDDEV,
                                 sig_avg=SIG_AVG,
                                 enob=ENOB,
                                 num_of_samples=len(tx_stream))

rx_stream = optical_channel(tx_stream)

rx_stream_f = np.fft.fft(np.sqrt(rx_stream))
freq = np.fft.fftfreq(rx_stream_f.shape[0], d=1/BANDWIDTH)
# compensation = np.exp(-0.5j*DISPERSION_FACTOR*FIBER_LENGTH*np.square(2*math.pi*freq))

# rx_compensated_f = rx_stream_f * compensation

rx_compensated_t = np.fft.ifft(rx_stream_f)

rx_reshaped_t = rx_compensated_t.reshape((OFDM_N, -1))

if False:
    t = np.arange(len(tx_stream))*(1/BANDWIDTH)
    plt.plot(t, tx_stream)
    plt.plot(t, 2+np.sqrt(rx_stream))
    plt.plot(t, rx_compensated_t-2)
    plt.show()

rx_syms = []

# chan_num = 2

for chan_num in range(K):
    i = 0
    for sym in rx_reshaped_t:
        sym_f = np.fft.fft(sym)

        freq = np.fft.fftfreq(sym_f.shape[0], d=1 / BANDWIDTH)
        # compensation = np.exp(-0.5j * DISPERSION_FACTOR * FIBER_LENGTH * np.square(2 * math.pi * freq))

        sym_f = sym_f

        upper_f = sym_f[1:1 + K]

        if SHOW_PLOTS:
            # for x in upper_f:
            if tx_syms[i+chan_num] == 0:
                plt.plot(np.real(upper_f[chan_num]), np.imag(upper_f[chan_num]), 'bo')
            elif tx_syms[i+chan_num] == 1:
                plt.plot(np.real(upper_f[chan_num]), np.imag(upper_f[chan_num]), 'ro')
            elif tx_syms[i+chan_num] == 2:
                plt.plot(np.real(upper_f[chan_num]), np.imag(upper_f[chan_num]), 'go')
            elif tx_syms[i+chan_num] == 3:
                plt.plot(np.real(upper_f[chan_num]), np.imag(upper_f[chan_num]), 'mo')
            i += K
    plt.title((str(chan_num) + " rot"))
    plt.show()

# for sym in rx_reshaped_t:
#     sym_f = np.fft.fft(sym)
#     upper_f = sym_f[1:1 + K]
#
#     if SHOW_PLOTS:
#         for x in upper_f:
#             if tx_syms[i] == 0:
#                 plt.plot(np.real(x), np.imag(x), 'bo')
#             elif tx_syms[i] == 1:
#                 plt.plot(np.real(x), np.imag(x), 'ro')
#             elif tx_syms[i] == 2:
#                 plt.plot(np.real(x), np.imag(x), 'go')
#             elif tx_syms[i] == 3:
#                 plt.plot(np.real(x), np.imag(x), 'mo')
#             # plt.plot(np.real(mapping_table[tx_syms[i]]), np.imag(mapping_table[tx_syms[i]]), 'ro')
#             # plt.text(np.real(sym), np.imag(sym), str(i), ha='center')
#             # plt.text(np.real(mapping_table[tx_syms[i]]), np.imag(mapping_table[tx_syms[i]]), str(i), ha='center')
#             i += 1
plt.show()