# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python 3
#     name: python3
# ---

# + [markdown] id="view-in-github" colab_type="text"
# <a href="https://colab.research.google.com/github/BenYavor/Wireless_encoding_with_MI_estimation/blob/master/Communication_Autoencoder_with_MI_channel_estimation.ipynb" target="_parent"><img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/></a>

# + id="EkWq8Zt_isjR" colab_type="code" outputId="8f3cc779-e404-4025-d29a-7c569caa2d15" colab={"base_uri": "https://localhost:8080/", "height": 644}
# !pip install tensorflow==2.0.0 
import sys
assert sys.version_info >= (3, 5)
import numpy as np
# %matplotlib inline
import matplotlib as mpl
import matplotlib.pyplot as plt
mpl.rc('axes', labelsize=14)
mpl.rc('xtick', labelsize=12)
mpl.rc('ytick', labelsize=12)
import tensorflow as tf
from tensorflow import keras
import pandas as pd
from scipy import special
from tensorflow.keras import layers
np.random.seed(42)
tf.random.set_seed(42)
import time


# + id="MXJRVGze4Lgb" colab_type="code" colab={}
M = 16
k = int(np.log2(M))
n = 1
TRAINING_SNR = 7


# + id="oUypWuZ9jAcV" colab_type="code" colab={}
def EbNo_to_noise(ebnodb):
    '''Transform EbNo[dB]/snr to noise power'''
    ebno = 10**(ebnodb/10)
    noise_std = 1/np.sqrt(2*(k/n)*ebno) 
    return noise_std

def SNR_to_noise(snrdb):
    '''Transform EbNo[dB]/snr to noise power'''
    snr = 10**(snrdb/10)
    noise_std = 1/np.sqrt(2*snr)
    return noise_std


# + id="x81YbgfUjDLE" colab_type="code" colab={}
randN_05 = keras.initializers.RandomNormal(mean=0.0, stddev=0.05, seed=None)
bias_init = keras.initializers.Constant(0.01)

input_A = keras.layers.Input(shape=[4*n])
input_B = keras.layers.Input(shape=[4*n])
    
transform = keras.models.Sequential([
layers.Dense(128, bias_initializer=bias_init, kernel_initializer=randN_05, activation="relu"),   #256
#keras.layers.Dropout(rate=0.3), # To regularize higher dimensionality
layers.Dense(128, bias_initializer=bias_init, kernel_initializer=randN_05, activation="relu"),    #256
#keras.layers.Dropout(rate=0.3), # To regularize higher dimensionality
layers.Dense(1, bias_initializer=bias_init, kernel_initializer=randN_05, activation=None)])

output_A = transform(input_A)
output_B = transform(input_B)
output_C = tf.reduce_mean(output_A) - tf.math.log(tf.reduce_mean(tf.exp(output_B))) # MINE
MI_mod = keras.models.Model(inputs=[input_A, input_B], outputs=output_C)

#print(MI_mod.inputs)
#MI_mod.summary()
#keras.utils.plot_model(MI_mod, 'Structure_of_MI_estimation.png', show_shapes=True)


# + id="E6C69Kzzi4cZ" colab_type="code" colab={}
noise_std = EbNo_to_noise(TRAINING_SNR)
# custom functions / layers without weights
norm_layer = keras.layers.Lambda(lambda x: tf.divide(x,tf.sqrt(2*tf.reduce_mean(tf.square(x)))))
shape_layer = keras.layers.Lambda(lambda x: tf.reshape(x, shape=[-1,2,n]))
shape_layer2 = keras.layers.Lambda(lambda x: tf.reshape(x, shape=[-1,2*n]))
channel_layer = keras.layers.Lambda(lambda x: 
                    x + tf.random.normal(tf.shape(x), mean=0.0, stddev=noise_std))

encoder = keras.models.Sequential([
            keras.layers.Embedding(M, M, embeddings_initializer='glorot_normal'),
            keras.layers.Dense(M*4, activation="elu"),                                 # M
            keras.layers.Dense(M*4, activation="elu"),                                 #optional
            keras.layers.Dense(2*n, activation=None),
            shape_layer,
            norm_layer])

channel = keras.models.Sequential([channel_layer])

decoder = keras.models.Sequential([
                keras.layers.InputLayer(input_shape=[2,n]),
                shape_layer2,
                keras.layers.Dense(M*4, activation="elu"),                              #M
                keras.layers.Dense(M*2, activation="elu"),                              #optional
                keras.layers.Dense(M, activation="softmax")
                ])

autoencoder = keras.models.Sequential([encoder, channel, decoder])


# + id="PA9p9oXcjIa4" colab_type="code" colab={}
def B_Ber_m(input_msg, msg):
    '''Calculate the Batch Bit Error Rate'''
    pred_error = tf.not_equal(input_msg, tf.argmax(msg, 1))      
    bber = tf.reduce_mean(tf.cast(pred_error, tf.float32))
    return bber


# + id="l75z0TB6jKOV" colab_type="code" colab={}
def random_sample(batch_size=32):
    msg = np.random.randint(M, size=batch_size)
    return msg


# + id="9yqE4rYojNw_" colab_type="code" colab={}
def test_encoding(M=16, n=1):
    inp = np.arange(0,M)
    coding = encoder.predict(inp)
    fig = plt.figure(figsize=(4,4))
    plt.plot(coding[:,0], coding[:, 1], "b.")
    plt.xlabel("$x_1$", fontsize=18)
    plt.ylabel("$x_2$", fontsize=18, rotation=0)
    plt.grid(True)
    plt.gca().set_ylim(-2, 2)
    plt.gca().set_xlim(-2, 2)
    plt.show()

def test_noisy_codeword(data):
    rcvd_word = data[1:2000]
    fig = plt.figure(figsize=(4,4))
    plt.plot(rcvd_word[:,0], rcvd_word[:, 1], "b.")
    plt.xlabel("$x_1$", fontsize=18)
    plt.ylabel("$x_2$", fontsize=18, rotation=0)
    plt.grid(True)
    plt.gca().set_ylim(-2, 2)
    plt.gca().set_xlim(-2, 2)
    plt.show()


# + id="YMwTXx4qjQgY" colab_type="code" colab={}
loss_fn = keras.losses.SparseCategoricalCrossentropy()
mean_loss = keras.metrics.Mean()


# + id="afmRKoJ8jTGx" colab_type="code" colab={}
def plot_loss(step, epoch, mean_loss, X_batch, y_pred, plot_encoding):
    template = 'Iteration: {}, Epoch: {}, Loss: {:.5f}, Batch_BER: {:.5f}'
    if step % 10 == 0:
        print(template.format(step, epoch, mean_loss.result(), B_Ber_m(X_batch, y_pred)))
        if plot_encoding:
            test_encoding()


# + id="MUxXl-Hy4sCQ" colab_type="code" colab={}
def plot_batch_loss(epoch, mean_loss, X_batch, y_pred):
        template_outer_loop = 'Interim result for Epoch: {}, Loss: {:.5f}, Batch_BER: {:.5f}'
        print(template_outer_loop.format(epoch, mean_loss.result(), B_Ber_m(X_batch, y_pred)))


# + id="2SI75bbcJ0bg" colab_type="code" colab={}
def split_train_data(x_sample, y_sample):
    x_shaped = tf.reshape(x_sample, shape=[-1,2*n])
    y_shaped = tf.reshape(y_sample, shape=[-1,2*n])
    x_sample1, x_sample2 = tf.split(x_shaped, num_or_size_splits=2)
    y_sample1, y_sample2 = tf.split(y_shaped, num_or_size_splits=2)
    joint_sample = tf.concat([x_sample1, y_sample1], axis=1)
    marg_sample = tf.concat([x_sample2, y_sample1], axis=1)
    return joint_sample, marg_sample


# + id="Tuf3BQ1dhl1F" colab_type="code" colab={}
def train_mi(n_epochs=5, n_steps=20, batch_size=200):
    i = 0
    for epoch in range(1, n_epochs + 1): 
        print("Training in Epoch {}/{}".format(epoch, n_epochs)) 
        for step in range(1, n_steps + 1):
            i +=1
            losses =[]
            X_batch = random_sample(batch_size*2)
            loss = mi_train_step(X_batch)
            mi_avg = -mean_loss(loss)
            losses.append(loss)
            if i%20==0:
              MINE_losses.append(np.mean(losses))
        print('Epoch: {}, Mi is {}'.format(epoch, mi_avg))
        mean_loss.reset_states()

@tf.function  
def mi_train_step(X_batch):
  with tf.GradientTape() as tape:
    x_enc = encoder(X_batch, training=True)
    y_recv = channel(x_enc)
    joint_marg_s = split_train_data(x_enc, y_recv)
    loss = -MI_mod(joint_marg_s)
    gradients = tape.gradient(loss, MI_mod.trainable_variables) 
    optimizer_mi.apply_gradients(zip(gradients, MI_mod.trainable_variables))
  return loss


# + [markdown] id="zB73RkszjT9j" colab_type="text"
# ## Training Function

# + id="_sUbcMOpjUGr" colab_type="code" colab={}
def train_decoder(n_epochs=5, n_steps=20, batch_size=200, plot_encoding=True):
    i = 0
    for epoch in range(1, n_epochs + 1): 
        print("Training Bob in Epoch {}/{}".format(epoch, n_epochs)) 
        for step in range(1, n_steps + 1):
            i+=1
            losses =[]
            X_batch  = random_sample(batch_size)
            loss, y_pred = dec_train_step(X_batch)
            mean_loss(loss)
            plot_loss(step, epoch, mean_loss, X_batch, y_pred, plot_encoding)
            losses.append(loss)
            if i%20==0:
              Dec_losses.append(np.mean(losses))
        plot_batch_loss(epoch, mean_loss, X_batch, y_pred) 
        mean_loss.reset_states()

@tf.function
def dec_train_step(X_batch):
  with tf.GradientTape() as tape:
    y_pred = autoencoder(X_batch, training=True)
    loss = tf.reduce_mean(loss_fn(X_batch, y_pred))
    gradients = tape.gradient(loss, decoder.trainable_variables) 
    optimizer_ae.apply_gradients(zip(gradients, decoder.trainable_variables)) 
  return loss, y_pred



# + id="gofvRVrTk_9L" colab_type="code" colab={}
def train_encoder(n_epochs=5, n_steps=20, batch_size=200):
    i=0
    for epoch in range(1, n_epochs + 1): 
        print("Training Bob in Epoch {}/{}".format(epoch, n_epochs)) 
        for step in range(1, n_steps + 1):
            i+=1
            losses =[]
            X_batch  = random_sample(batch_size)
            loss = enc_train_step1(X_batch,optimizer_ae)               
            mi_avg = -mean_loss(loss)
            losses.append(loss)
            if i%20==0:
              Enc_losses.append(np.mean(losses))
        X_batch  = random_sample(batch_size) 
        enc_train_step2(X_batch)
        print('Epoch: {}, Mi is {}'.format(epoch, mi_avg))
       # test_encoding(M, 1)

@tf.function
def enc_train_step1(X_batch,optimizer_ae):
  with tf.GradientTape() as tape:
    x_enc = encoder(X_batch, training=True)
    y_recv = channel(x_enc)
    joint_marg_s = split_train_data(x_enc, y_recv)
    loss = -MI_mod(joint_marg_s)
    gradients = tape.gradient(loss, encoder.trainable_variables) 
    optimizer_ae.apply_gradients(zip(gradients, encoder.trainable_variables))
  return loss

@tf.function
def enc_train_step2(X_batch):
  with tf.GradientTape() as tape:
    x_enc = encoder(X_batch, training=True)
    y_recv = channel(x_enc)
    joint_marg_s = split_train_data(x_enc, y_recv)
    loss = -MI_mod(joint_marg_s)
    gradients = tape.gradient(loss, MI_mod.trainable_variables) 
    optimizer_mi.apply_gradients(zip(gradients, MI_mod.trainable_variables))


# + colab_type="code" id="eHGEqRx5Q7ML" colab={}
def train_encoder2(n_epochs=5, n_steps=20, batch_size=200):
    i=0
    for epoch in range(1, n_epochs + 1): 
        print("Training Bob in Epoch {}/{}".format(epoch, n_epochs)) 
        for step in range(1, n_steps + 1):
            losses =[]
            i +=1
            X_batch  = random_sample(batch_size)
            loss = enc_train_step12(X_batch)
            mi_avg = -mean_loss(loss)
            losses.append(loss)
            if i%20 == 0:
              Enc_losses.append(np.mean(losses))
        X_batch  = random_sample(batch_size) 
        enc_train_step2(X_batch)
        print('Epoch: {}, Mi is {}'.format(epoch, mi_avg))
       # test_encoding(M, 1)

@tf.function
def enc_train_step12(X_batch):
  with tf.GradientTape() as tape:
    x_enc = encoder(X_batch, training=True)
    y_recv = channel(x_enc)
    joint_marg_s = split_train_data(x_enc, y_recv)
    loss = -MI_mod(joint_marg_s)
    gradients = tape.gradient(loss, encoder.trainable_variables) 
    optimizer_ae2.apply_gradients(zip(gradients, encoder.trainable_variables))
  return loss



# + id="cfFUT6DwlmvJ" colab_type="code" colab={}
# test msg sequence for normal encoding
N_test = 500000
test_msg = np.random.randint(M, size=N_test)


# + id="xVXMlZzTlz8i" colab_type="code" colab={}
def Test_AE(data):
    '''Calculate Bit Error for varying SNRs'''
    snr_range = np.linspace(0, 15, 31)
    bber_vec = [None] * len(snr_range)
        
    for db in range(len(snr_range)):           
        noise_std = EbNo_to_noise(snr_range[db])
        code_word = encoder(data)
        rcvd_word = code_word + tf.random.normal(tf.shape(code_word), mean=0.0, stddev=noise_std)
        dcoded_msg = decoder(rcvd_word)
        bber_vec[db] = B_Ber_m(data, dcoded_msg)
        if (db % 6 == 0) & (db > 0):
            print(f'Progress: {db} of {30} parts')

    return (snr_range, bber_vec)


# + id="vFKf3kxv6h1l" colab_type="code" colab={}
time_to_train_MI = 0
time_to_train_enc = 0
time_to_train_dec = 0

samples_for_training = 0

# + id="c36sWVydq0g6" colab_type="code" colab={}
MINE_losses = []
Enc_losses = []
Dec_losses = []

# + id="wZQLuwHHl1Gj" colab_type="code" outputId="492687dc-7bb3-43d6-b9f0-d390fdf389e1" colab={"base_uri": "https://localhost:8080/", "height": 1000}
# %%time
start_mi = time.time()
optimizer_mi = keras.optimizers.Nadam(lr=0.005)
train_mi(n_epochs=1, n_steps=500, batch_size=128)
samples_for_training += 2*(1*128*500)                                      #could be smaller if used the same data for the n_steps
time_to_train_MI += time.time()-start_mi

start_enc1 = time.time()
optimizer_ae = keras.optimizers.Nadam(lr=0.005)
train_encoder(n_epochs=5, n_steps=400, batch_size=500)
samples_for_training += (5*500) + (5*500*400)
time_to_train_enc += time.time() - start_enc1
test_encoding(M, 1)

start_enc2 = time.time()
optimizer_ae2 = keras.optimizers.Nadam(lr=0.0005)
train_encoder(n_epochs=5, n_steps=400, batch_size=500)
samples_for_training += (5*500) + (5*500*400)
time_to_train_enc += time.time() - start_enc2
test_encoding(M, 1)

start_dec = time.time()
optimizer_ae = keras.optimizers.Nadam(lr=0.005)
train_decoder(n_epochs=5, n_steps=400, batch_size=500, plot_encoding=False)
samples_for_training += 5*400*500
time_to_train_dec += time.time() - start_dec

# + id="N9dYDWne8UQN" colab_type="code" outputId="778af068-0fa3-4e50-cb0d-69a99a345073" colab={"base_uri": "https://localhost:8080/", "height": 106}
print('MI time',time_to_train_MI)
print('Enc time',time_to_train_enc)
print('Dec time',time_to_train_dec)
print('time for training',(time_to_train_dec +time_to_train_enc+ time_to_train_MI))
print('Samples used for training', samples_for_training )

# + id="z869ARDCr4wg" colab_type="code" colab={"base_uri": "https://localhost:8080/", "height": 395} outputId="508e6fb6-c624-49f6-cdf2-ef929f7f45aa"
fig = plt.figure(figsize=(10, 6))
#plt.plot(Enc_losses,'^-') 
#plt.plot(Dec_losses,'*-') 
iterations = np.linspace(0,(len(MINE_losses)*20),len(MINE_losses))
plt.plot(iterations, MINE_losses ,'^-') 
plt.ylabel("Losses", fontsize=14, rotation=90)
plt.xlabel("Iterations", fontsize=14)
plt.legend(['MINE training loss'],
           prop={'size': 14}, loc='upper right');
plt.grid(True, which="both")

# + colab_type="code" outputId="a3edf718-896b-4a2e-8853-e4b151124528" id="8XsoDMdaw1vo" colab={"base_uri": "https://localhost:8080/", "height": 395}
fig = plt.figure(figsize=(10, 6))
#plt.plot(MINE_losses,'o-') 
iterations = np.linspace(0,(len(Enc_losses)*20),len(Enc_losses))
plt.plot(iterations, Enc_losses ,'^-') 
plt.ylabel("Losses", fontsize=14, rotation=90)
plt.xlabel("Iterations", fontsize=14)
plt.legend(['Enc training loss'],
           prop={'size': 14}, loc='upper right');
plt.grid(True, which="both")

# + colab_type="code" outputId="c10f2efa-c3fc-4d70-88a2-837cf571a0a5" id="Cj0MS69Dw2NJ" colab={"base_uri": "https://localhost:8080/", "height": 395}
fig = plt.figure(figsize=(10, 6))
#plt.plot(MINE_losses,'o-') 
#plt.plot(Enc_losses,'^-') 
iterations = np.linspace(0,2000,100)
plt.plot(iterations,Dec_losses,'*-') 
plt.ylabel("Losses", fontsize=14, rotation=90)
plt.xlabel("Iterations", fontsize=14)
plt.legend(['Dec training loss'],
           prop={'size': 14}, loc='upper right');
plt.grid(True, which="both")

# + id="z-ELSw_ZnCXT" colab_type="code" outputId="b7b9614c-b7ac-4104-fe0d-4e512938ee97" colab={"base_uri": "https://localhost:8080/", "height": 106}
bber_data = Test_AE(test_msg)


# + id="eN-jOLPqTvrz" colab_type="code" outputId="015c6ead-0461-4d9e-dab2-c53b2c46dcfa" colab={"base_uri": "https://localhost:8080/", "height": 406}
# Approximate 16 QAM Error
def SIXT_QAM_sim(ebno):
    return (3.0/2)*special.erfc(np.sqrt((4.0/10)*10.**(ebno/10)))

ebnodbs = np.linspace(0,15,16)
fig = plt.figure(figsize=(10, 6))
plt.semilogy(bber_data[0], bber_data[1], 'o-')

plt.semilogy(ebnodbs, SIXT_QAM_sim(ebnodbs), '^-');
plt.gca().set_ylim(1e-5, 1)
plt.gca().set_xlim(0, 15)
plt.ylabel("Batch Symbol Error Rate", fontsize=14, rotation=90)
plt.xlabel("SNR [dB]", fontsize=18)
plt.legend(['AE with MINE', '16QAM'],
           prop={'size': 14}, loc='upper right');
plt.grid(True, which="both")

# + id="7Ai0b10W8XJl" colab_type="code" colab={}


StarcoderdataJupyterScriptsDedupFiltered