Add Gaussian channel model

.gitignore

@@ -1 +1,2 @@
 __pycache__
+.ipynb_checkpoints

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)
 

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

constellation/__init__.py

@@ -1,2 +1,3 @@
 from constellation.ConstellationNet import ConstellationNet
+from constellation.GaussianChannel import GaussianChannel
 import constellation.util

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():

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)
-
-model = constellation.ConstellationNet(order=order)
+output_file = 'output/constellation-order-{}.pth'.format(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()