From 3b40e27070086f9d6fd2b6bdd2511bff664832fa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Matt=C3=A9o=20Delabre?= Date: Sat, 14 Dec 2019 23:04:35 -0500 Subject: [PATCH] Add Gaussian channel model --- .gitignore | 1 + constellation/ConstellationNet.py | 14 +++--- constellation/GaussianChannel.py | 74 +++++++++++++++++++++++++++++++ constellation/__init__.py | 1 + plot.py | 14 ++++-- train.py | 9 +++- 6 files changed, 101 insertions(+), 12 deletions(-) create mode 100644 constellation/GaussianChannel.py diff --git a/.gitignore b/.gitignore index bee8a64..eeb514c 100644 --- a/.gitignore +++ b/.gitignore @@ -1 +1,2 @@ __pycache__ +.ipynb_checkpoints diff --git a/constellation/ConstellationNet.py b/constellation/ConstellationNet.py index 7bcfa4b..299bcb7 100644 --- a/constellation/ConstellationNet.py +++ b/constellation/ConstellationNet.py @@ -1,4 +1,5 @@ import torch.nn as nn +from .GaussianChannel import GaussianChannel class ConstellationNet(nn.Module): @@ -6,7 +7,8 @@ class ConstellationNet(nn.Module): self, order=2, encoder_layers_sizes=(), - decoder_layers_sizes=() + decoder_layers_sizes=(), + channel_model=GaussianChannel() ): """ Create an encoder-decoder network to automatically shape a @@ -14,13 +16,15 @@ class ConstellationNet(nn.Module): fiber channel. :param order: Order of the constellation, i.e. the number of messages - that are to be transmitted or equivalently the number of symbols whose + that are to be transmitted, or equivalently the number of symbols whose placements in the constellation have to be learned. :param encoder_layers_sizes: Shape of the encoder’s hidden layers. The size of this sequence is the number of hidden layers, with each element being a number which specifies the number of neurons in its channel. :param decoder_layers_sizes: Shape of the decoder’s hidden layers. Uses the same convention as `encoder_layers_sizes` above. + :param channel_model: Instance of the channel model to use between the + encoder and decoder network. """ super().__init__() @@ -41,9 +45,7 @@ class ConstellationNet(nn.Module): ] self.encoder = nn.Sequential(*encoder_layers) - - # TODO: Add real channel model - self.channel = nn.Identity() + self.channel = channel_model # Build the decoder network taking the noisy I/Q vector received from # the channel as input and outputting a probability vector for each @@ -60,7 +62,7 @@ class ConstellationNet(nn.Module): # Softmax is not used at the end of the network because the # CrossEntropyLoss criterion is used for training, which includes # LogSoftmax - decoder_layers.append(nn.Linear(prev_layer_size, order),) + decoder_layers.append(nn.Linear(prev_layer_size, order)) self.decoder = nn.Sequential(*decoder_layers) diff --git a/constellation/GaussianChannel.py b/constellation/GaussianChannel.py new file mode 100644 index 0000000..ab37672 --- /dev/null +++ b/constellation/GaussianChannel.py @@ -0,0 +1,74 @@ +import torch.nn as nn +import torch +import numpy as np + + +def channel_OSNR(): + Sys_rate = 32e9 + r = 0.05 + Dispersion = 16.48e-6 + B_2 = Dispersion + Non_linear_index = 1.3e3 + Gam = Non_linear_index + Loss = 10**20 + Alpha = Loss + Span_count = 20 + N_s = Span_count + Span_length = 10e5 # (km) + L_s = Span_length + Noise_figure = 10**0.5 # (dB) + h = 6.6261e-34 + v = 299792458 + B_WDM = Sys_rate*(1+r) + B_N = 0.1 + + P_ASE_1 = h*v*B_N*(Loss*Span_length*Noise_figure-1) + P_ASE = P_ASE_1 * Span_count + L_eff = 1-np.exp(-Loss*Span_length)/2/Alpha + + eps = 0.3*np.log(1+(6/L_s)*(L_eff/np.arcsinh((np.pi**2/3)*B_2*L_eff*B_WDM**2))) + b = P_ASE_1/(2*(N_s**eps)*B_N*(Gam**2)*L_eff*np.arcsinh((np.pi**2/3)*B_2*L_eff*B_WDM**2)) + P_ch = Sys_rate*(((27*np.pi*B_2/16)*b)**(1/3)) + OSNR = (2*P_ch/3/P_ASE) + + OSNR_dB = 10*np.log10(OSNR) + return OSNR_dB + + +def Const_Points_Pow(IQ): + """ + Calculate the average power of a set of vectors. + """ + p_enc_avg = (torch.norm(IQ, dim=1) ** 2).mean() + p_enc_avg_dB = 10 * torch.log10(p_enc_avg) + return p_enc_avg_dB + + +def Pow_Noise(CH_OSNR, CPP): + """ + Calculate the power of channel noise. + """ + P_N_dB = CPP - CH_OSNR + p_N_watt = 10**(P_N_dB/10) + Var_Noise = p_N_watt + return Var_Noise + + +def Channel_Noise_Model(NV, S): + """ + Compute the Gaussian noise to be added to each vector to simulate passing + through a channel. + """ + return torch.distributions.normal.Normal( + 0, torch.sqrt(NV*5000) + ).sample(S) + + +class GaussianChannel(nn.Module): + def __init__(self): + super().__init__() + + def forward(self, x): + Noise_Variance = Pow_Noise(channel_OSNR(), Const_Points_Pow(x)) + Noise_Volts = Channel_Noise_Model(Noise_Variance, [len(x), 2]) + return x + Noise_Volts diff --git a/constellation/__init__.py b/constellation/__init__.py index e461218..587d83d 100644 --- a/constellation/__init__.py +++ b/constellation/__init__.py @@ -1,2 +1,3 @@ from constellation.ConstellationNet import ConstellationNet +from constellation.GaussianChannel import GaussianChannel import constellation.util diff --git a/plot.py b/plot.py index 65d9041..6e495f2 100644 --- a/plot.py +++ b/plot.py @@ -3,17 +3,23 @@ from constellation import util import torch from matplotlib import pyplot from mpl_toolkits.axisartist.axislines import SubplotZero -import math -import numpy # Number learned symbols order = 4 # File in which the trained model is saved -input_file = 'output/constellation-order-{}.tc'.format(order) +input_file = 'output/constellation-order-{}.pth'.format(order) + +# Restore model from file +model = constellation.ConstellationNet( + order=order, + encoder_layers_sizes=(4,), + decoder_layers_sizes=(4,), + channel_model=constellation.GaussianChannel() +) -model = constellation.ConstellationNet(order=order) model.load_state_dict(torch.load(input_file)) +model.eval() # Compute encoded vectors with torch.no_grad(): diff --git a/train.py b/train.py index 3d1eacd..edb1ad2 100644 --- a/train.py +++ b/train.py @@ -15,9 +15,14 @@ num_epochs = 20000 loss_report_epoch_skip = 500 # File in which the trained model is saved -output_file = 'output/constellation-order-{}.tc'.format(order) +output_file = 'output/constellation-order-{}.pth'.format(order) -model = constellation.ConstellationNet(order=order) +model = constellation.ConstellationNet( + order=order, + encoder_layers_sizes=(4,), + decoder_layers_sizes=(4,), + channel_model=constellation.GaussianChannel() +) # Train the model with random data model.train()