From 3365603614a852d3dfa6cc0959b1413ce0c9fb31 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Matt=C3=A9o=20Delabre?= <spam@delab.re>
Date: Sun, 15 Dec 2019 09:42:33 -0500
Subject: [PATCH] Add normalization layer

---
 constellation/ConstellationNet.py | 14 ++++++++++----
 constellation/GaussianChannel.py  |  2 +-
 constellation/NormalizePower.py   | 20 ++++++++++++++++++++
 3 files changed, 31 insertions(+), 5 deletions(-)
 create mode 100644 constellation/NormalizePower.py

diff --git a/constellation/ConstellationNet.py b/constellation/ConstellationNet.py
index 6bacd07..946a729 100644
--- a/constellation/ConstellationNet.py
+++ b/constellation/ConstellationNet.py
@@ -1,8 +1,13 @@
 import torch.nn as nn
 from .GaussianChannel import GaussianChannel
+from .NormalizePower import NormalizePower
 
 
 class ConstellationNet(nn.Module):
+    """
+    Autoencoder network to automatically shape a constellation of symbols for
+    efficient communication over an optical fiber channel.
+    """
     def __init__(
         self,
         order=2,
@@ -11,9 +16,7 @@ class ConstellationNet(nn.Module):
         channel_model=GaussianChannel()
     ):
         """
-        Create an encoder-decoder network to automatically shape a
-        constellation of symbols for efficient communication over an optical
-        fiber channel.
+        Create an autoencoder.
 
         :param order: Order of the constellation, i.e. the number of messages
         that are to be transmitted, or equivalently the number of symbols whose
@@ -39,7 +42,10 @@ class ConstellationNet(nn.Module):
             encoder_layers.append(nn.SELU())
             prev_layer_size = layer_size
 
-        encoder_layers.append(nn.Linear(prev_layer_size, 2))
+        encoder_layers += [
+            nn.Linear(prev_layer_size, 2),
+            NormalizePower(),
+        ]
 
         self.encoder = nn.Sequential(*encoder_layers)
         self.channel = channel_model
diff --git a/constellation/GaussianChannel.py b/constellation/GaussianChannel.py
index af2425b..09638d7 100644
--- a/constellation/GaussianChannel.py
+++ b/constellation/GaussianChannel.py
@@ -37,7 +37,7 @@ def channel_OSNR():
 
 def Const_Points_Pow(IQ):
     """
-    Calculate the average power of a set of vectors.
+    Calculate the average power of a constellation.
     """
     p_enc_avg = (torch.norm(IQ, dim=1) ** 2).mean()
     p_enc_avg_dB = 10 * torch.log10(p_enc_avg)
diff --git a/constellation/NormalizePower.py b/constellation/NormalizePower.py
new file mode 100644
index 0000000..cede8e6
--- /dev/null
+++ b/constellation/NormalizePower.py
@@ -0,0 +1,20 @@
+import torch.nn as nn
+import torch
+
+
+class NormalizePower(nn.Module):
+    """
+    Layer for normalizing a batch of vectors so that their average length is 1.
+
+    :attr epsilon: Minimum mean length to avoid division by zero.
+    """
+    epsilon = 1e-12
+
+    def forward(self, x):
+        average_power = (torch.norm(x, dim=1) ** 2).mean()
+        average_power = torch.max(torch.tensor([
+            NormalizePower.epsilon,
+            average_power
+        ]))
+
+        return x * torch.rsqrt(average_power)