Added signal definition

models/autoencoder.py

@@ -8,6 +8,7 @@ from tensorflow.keras.models import Model
 from functools import partial
 import misc
 import defs
+from models import basic
 
 latent_dim = 64
 
@@ -19,7 +20,7 @@ class AutoencoderMod(defs.Modulator):
         super().__init__(2**autoencoder.N)
         self.autoencoder = autoencoder
 
-    def forward(self, binary: np.ndarray) -> np.ndarray:
+    def forward(self, binary: np.ndarray) -> defs.Signal:
         reshaped = binary.reshape((-1, self.N))
         reshaped_ho = misc.bit_matrix2one_hot(reshaped)
         encoded = self.autoencoder.encoder(reshaped_ho)
@@ -28,7 +29,7 @@ class AutoencoderMod(defs.Modulator):
 
         f = np.zeros(x2.shape[0])
         x3 = misc.rect2polar(np.c_[x2[:, 0], x2[:, 1], f])
-        return x3
+        return defs.Signal(x3)
 
 
 class AutoencoderDemod(defs.Demodulator):
@@ -36,9 +37,8 @@ class AutoencoderDemod(defs.Demodulator):
         super().__init__(2**autoencoder.N)
         self.autoencoder = autoencoder
 
-    def forward(self, values: np.ndarray) -> np.ndarray:
-        rect = misc.polar2rect(values[:, [0, 1]])
-        decoded = self.autoencoder.decoder(rect).numpy()
+    def forward(self, values: defs.Signal) -> np.ndarray:
+        decoded = self.autoencoder.decoder(values.rect).numpy()
         result = misc.int2bit_array(decoded.argmax(axis=1), self.N)
         return result.reshape(-1, )
 
@@ -69,7 +69,7 @@ class Autoencoder(Model):
         self.compiled = False
 
         # Divide by 2 because encoder outputs values between 0 and 1 instead of -1 and 1
-        self.noise = 10 ** (noise / 10)  # / 2
+        self.noise = noise #10 ** (noise / 10)  # / 2
 
         # self.decoder.add(layers.Softmax(units=4, dtype=bool))
 
@@ -85,13 +85,16 @@ class Autoencoder(Model):
         # ])
 
     def call(self, x, **kwargs):
-        encoded = self.encoder(x)
-        encoded = encoded * 2 - 1
+        chan = basic.AWGNChannel(self.noise)
+        signal = self.encoder(x)
+        signal = signal * 2 - 1
+        signal = chan.forward_tensor(signal)
+        # encoded = encoded * 2 - 1
         # encoded = tf.clip_by_value(encoded, clip_value_min=0, clip_value_max=1, name=None)
         # noise = self.randomiser(shape=(-1, 2), dtype=tf.float32)
-        noise = np.random.normal(0, 1, (1, 2)) * self.noise
-        noisy = tf.convert_to_tensor(noise, dtype=tf.float32)
-        decoded = self.decoder(encoded + noisy)
+        # noise = np.random.normal(0, 1, (1, 2)) * self.noise
+        # noisy = tf.convert_to_tensor(noise, dtype=tf.float32)
+        decoded = self.decoder(signal)
         return decoded
 
     def train(self, samples=1e6):