Просмотр исходного кода

Merge branch 'photonics'

# Conflicts:
#	main.py
Min 5 лет назад
Родитель
Сommit
6fa2f7d82c
8 измененных файлов с 321 добавлено и 217 удалено
  1. 2 1
      .gitignore
  2. 33 6
      defs.py
  3. 88 0
      graphs.py
  4. 29 91
      main.py
  5. 3 0
      misc.py
  6. 110 34
      models/autoencoder.py
  7. 49 83
      models/basic.py
  8. 7 2
      models/optical_channel.py

+ 2 - 1
.gitignore

@@ -4,4 +4,5 @@ __pycache__
 *.pyo
 
 # Environments
-venv/
+venv/
+tests/local_test.py

+ 33 - 6
defs.py

@@ -1,5 +1,30 @@
 import math
 import numpy as np
+import tensorflow as tf
+
+class Signal:
+
+    @property
+    def rect_x(self) -> np.ndarray:
+        return self.rect[:, 0]
+
+    @property
+    def rect_y(self) -> np.ndarray:
+        return self.rect[:, 1]
+
+    @property
+    def rect(self) -> np.ndarray:
+        raise NotImplemented("Not implemented")
+
+    def set_rect_xy(self, x_mat: np.ndarray, y_mat: np.ndarray):
+        raise NotImplemented("Not implemented")
+
+    def set_rect(self, mat: np.ndarray):
+        raise NotImplemented("Not implemented")
+
+    @property
+    def apf(self):
+        raise NotImplemented("Not implemented")
 
 
 class COMComponent:
@@ -13,12 +38,14 @@ class Channel(COMComponent):
     This model is just empty therefore just bypasses any input to output
     """
 
-    def forward(self, values: np.ndarray) -> np.ndarray:
+    def forward(self, values: Signal) -> Signal:
+        raise NotImplemented("Need to define forward function")
+
+    def forward_tensor(self, tensor: tf.Tensor) -> tf.Tensor:
         """
-        :param values: value generator, each iteration returns tuple of (amplitude, phase, frequency)
-        :return: affected tuple of (amplitude, phase, frequency)
+        Forward operation optimised for tensorflow tensors
         """
-        raise NotImplemented("Need to define forward function")
+        raise NotImplemented("Need to define forward_tensor function")
 
 
 class ModComponent(COMComponent):
@@ -30,7 +57,7 @@ class ModComponent(COMComponent):
 
 class Modulator(ModComponent):
 
-    def forward(self, binary: np.ndarray) -> np.ndarray:
+    def forward(self, binary: np.ndarray) -> Signal:
         """
         :param binary: raw bytes as input (most be dtype=bool)
         :return: amplitude, phase, frequency
@@ -40,7 +67,7 @@ class Modulator(ModComponent):
 
 class Demodulator(ModComponent):
 
-    def forward(self, values: np.ndarray) -> np.ndarray:
+    def forward(self, values: Signal) -> np.ndarray:
         """
         :param values: value generator, each iteration returns tuple of (amplitude, phase, frequency)
         :return: binary resulting values (dtype=bool)

+ 88 - 0
graphs.py

@@ -0,0 +1,88 @@
+import math
+import os
+from multiprocessing import Pool
+
+from sklearn.metrics import accuracy_score
+
+from defs import Modulator, Demodulator, Channel
+from models.basic import AWGNChannel
+from misc import generate_random_bit_array
+from models.optical_channel import OpticalChannel
+import matplotlib.pyplot as plt
+import numpy as np
+
+CPU_COUNT = os.environ.get("CPU_COUNT", os.cpu_count())
+
+
+def show_constellation(mod: Modulator, chan: Channel, demod: Demodulator, samples=1000):
+    x = generate_random_bit_array(samples)
+    x_mod = mod.forward(x)
+    x_chan = chan.forward(x_mod)
+    x_demod = demod.forward(x_chan)
+
+    plt.plot(x_chan.rect_x[x], x_chan.rect_y[x], '+')
+    plt.plot(x_chan.rect_x[:, 0][~x], x_chan.rect_y[:, 1][~x], '+')
+    plt.plot(x_mod.rect_x[:, 0], x_mod.rect_y[:, 1], 'ro')
+    axes = plt.gca()
+    axes.set_xlim([-2, +2])
+    axes.set_ylim([-2, +2])
+    plt.grid()
+    plt.show()
+    print('Accuracy : ' + str())
+
+
+def get_ber(mod, chan, demod, samples=1000):
+    if samples % mod.N:
+        samples += mod.N - (samples % mod.N)
+    x = generate_random_bit_array(samples)
+    x_mod = mod.forward(x)
+    x_chan = chan.forward(x_mod)
+    x_demod = demod.forward(x_chan)
+    return 1 - accuracy_score(x, x_demod)
+
+
+def get_AWGN_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30):
+    ber_x = np.linspace(start, stop, steps)
+    ber_y = []
+    for noise in ber_x:
+        ber_y.append(get_ber(mod, AWGNChannel(noise), demod, samples=samples))
+    return ber_x, ber_y
+
+
+def __calc_ber(packed):
+    # This function has to be outside get_Optical_ber in order to be pickled by pool
+    mod, demod, noise, length, pulse_shape, samples = packed
+    tx_channel = OpticalChannel(noise_level=noise, dispersion=-21.7, symbol_rate=10e9, sample_rate=400e9,
+                                length=length, pulse_shape=pulse_shape, sqrt_out=True)
+    return get_ber(mod, tx_channel, demod, samples=samples)
+
+
+def get_Optical_ber(mod, demod, samples=1000, start=-8., stop=5., steps=30, length=100, pulse_shape='rect'):
+    ber_x = np.linspace(start, stop, steps)
+    ber_y = []
+    print(f"Computing Optical BER.. 0/{len(ber_x)}", end='')
+    with Pool(CPU_COUNT) as pool:
+        packed_args = [(mod, demod, noise, length, pulse_shape, samples) for noise in ber_x]
+        for i, ber in enumerate(pool.imap(__calc_ber, packed_args)):
+            ber_y.append(ber)
+            i += 1  # just offset by 1
+            print(f"\rComputing Optical BER.. {i}/{len(ber_x)} ({i * 100 / len(ber_x):6.2f}%)", end='')
+    print()
+    return ber_x, ber_y
+
+
+def get_SNR(mod, demod, ber_func=get_Optical_ber, samples=1000, start=-5, stop=15, **ber_kwargs):
+    """
+    SNR for optics and RF should be calculated the same, that is A^2
+    Because P∝V² and P∝I²
+    """
+    x_mod = mod.forward(generate_random_bit_array(samples * mod.N))
+    sig_power = [A ** 2 for A in x_mod.amplitude]
+    av_sig_pow = np.mean(sig_power)
+    av_sig_pow = math.log(av_sig_pow, 10)
+
+    noise_start = -start + av_sig_pow
+    noise_stop = -stop + av_sig_pow
+    ber_x, ber_y = ber_func(mod, demod, samples, noise_start, noise_stop, **ber_kwargs)
+    SNR = -ber_x + av_sig_pow
+    return SNR, ber_y

+ 29 - 91
main.py

@@ -1,72 +1,9 @@
 import matplotlib.pyplot as plt
 
-import numpy as np
-from sklearn.metrics import accuracy_score
-from models import basic
+import graphs
 from models.basic import AWGNChannel, BPSKDemod, BPSKMod, BypassChannel, AlphabetMod, AlphabetDemod
-import misc
-import math
-from models.autoencoder import Autoencoder, view_encoder
-from models.optical_channel import OpticalChannel
 
 
-def show_constellation(mod, chan, demod, samples=1000):
-    x = misc.generate_random_bit_array(samples)
-    x_mod = mod.forward(x)
-    x_chan = chan.forward(x_mod)
-    x_demod = demod.forward(x_chan)
-
-    x_mod_rect = misc.polar2rect(x_mod)
-    x_chan_rect = misc.polar2rect(x_chan)
-    plt.plot(x_chan_rect[:, 0][x], x_chan_rect[:, 1][x], '+')
-    plt.plot(x_chan_rect[:, 0][~x], x_chan_rect[:, 1][~x], '+')
-    plt.plot(x_mod_rect[:, 0], x_mod_rect[:, 1], 'ro')
-    axes = plt.gca()
-    axes.set_xlim([-2, +2])
-    axes.set_ylim([-2, +2])
-    plt.grid()
-    plt.show()
-    print('Accuracy : ' + str())
-
-
-def get_ber(mod, chan, demod, samples=1000):
-    if samples % mod.N:
-        samples += mod.N - (samples % mod.N)
-    x = misc.generate_random_bit_array(samples)
-    x_mod = mod.forward(x)
-    x_chan = chan.forward(x_mod)
-    x_demod = demod.forward(x_chan)
-    return 1 - accuracy_score(x, x_demod)
-
-
-def get_AWGN_ber(mod, demod, samples=1000, start=-8, stop=5, steps=30):
-    ber_x = np.linspace(start, stop, steps)
-    ber_y = []
-    for noise in ber_x:
-        ber_y.append(get_ber(mod, AWGNChannel(noise), demod, samples=samples))
-    return ber_x, ber_y
-
-
-def get_SNR(mod, demod, samples=1000, start=-8, stop=5, steps=30):
-    ber_x, ber_y = get_AWGN_ber(mod, demod, samples, start, stop, steps)
-    x_mod = mod.forward(misc.generate_random_bit_array(samples*mod.N))
-    sig_amp = x_mod[:, 0]
-    sig_power = [A ** 2 for A in sig_amp]
-    av_sig_pow = np.mean(sig_power)
-    av_sig_pow = math.log(av_sig_pow, 10)
-    SNR = (ber_x * -1) + av_sig_pow
-    return SNR, ber_y
-
-def get_Optical_ber(mod, demod, samples=1000, start=-8, stop=5, steps=30, length=100, pulse_shape='rect'):
-    ber_x = np.linspace(start, stop, steps)
-    ber_y = []
-
-    for noise in ber_x:
-        tx_channel = OpticalChannel(noise_level=noise, dispersion=-21.7, symbol_rate=10e9, sample_rate=400e9,
-                                    length=length, pulse_shape=pulse_shape, sqrt_out=True)
-        ber_y.append(get_ber(mod, tx_channel, demod, samples=samples))
-    return ber_x, ber_y
-
 if __name__ == '__main__':
     # show_constellation(BPSKMod(10e6), AWGNChannel(-1), BPSKDemod(10e6, 10e3))
 
@@ -147,44 +84,45 @@ if __name__ == '__main__':
     # plt.legend()
     # plt.show()
 
-    for l in np.logspace(start=0, stop=3, num=5):
-        plt.plot(*get_Optical_ber(
-            AlphabetMod('4pam', 10e6),
-            AlphabetDemod('4pam', 10e6),
-            samples=1000,
-            steps=40,
-            start=-15,
-            length=l,
-            pulse_shape='rcos'
-        ), '-', label=(str(int(l))+'km'))
-
-    plt.yscale('log')
-    plt.gca().invert_xaxis()
-    plt.grid()
-    plt.xlabel('Noise dB')
-    plt.ylabel('BER')
-    plt.title("BER against Fiber length")
-    plt.legend()
-    plt.show()
+    # for l in np.logspace(start=0, stop=3, num=6):
+    #     plt.plot(*misc.get_SNR(
+    #         AlphabetMod('4pam', 10e6),
+    #         AlphabetDemod('4pam', 10e6),
+    #         samples=2000,
+    #         steps=200,
+    #         start=-5,
+    #         stop=20,
+    #         length=l,
+    #         pulse_shape='rcos'
+    #     ), '-', label=(str(int(l))+'km'))
+    #
+    # plt.yscale('log')
+    # # plt.gca().invert_xaxis()
+    # plt.grid()
+    # plt.xlabel('SNR dB')
+    # # plt.ylabel('BER')
+    # plt.title("BER against Fiber length")
+    # plt.legend()
+    # plt.show()
 
     for ps in ['rect', 'rcos', 'rrcos']:
-        plt.plot(*get_Optical_ber(
+        plt.plot(*graphs.get_SNR(
             AlphabetMod('4pam', 10e6),
             AlphabetDemod('4pam', 10e6),
-            samples=1000,
-            steps=40,
-            start=-15,
-            length=10,
+            samples=30000,
+            steps=100,
+            start=-5,
+            stop=20,
+            length=1,
             pulse_shape=ps
         ), '-', label=ps)
 
     plt.yscale('log')
-    plt.gca().invert_xaxis()
     plt.grid()
-    plt.xlabel('Noise dB')
+    plt.xlabel('SNR dB')
     plt.ylabel('BER')
     plt.title("BER for different pulse shapes")
     plt.legend()
     plt.show()
 
-    pass
+    pass

+ 3 - 0
misc.py

@@ -1,3 +1,4 @@
+
 import numpy as np
 import math
 import matplotlib.pyplot as plt
@@ -102,3 +103,5 @@ def generate_random_bit_array(size):
     arr = np.concatenate(p)
     np.random.shuffle(arr)
     return arr
+
+

+ 110 - 34
models/autoencoder.py

@@ -3,11 +3,16 @@ import numpy as np
 import tensorflow as tf
 
 from sklearn.metrics import accuracy_score
+from sklearn.model_selection import train_test_split
 from tensorflow.keras import layers, losses
 from tensorflow.keras.models import Model
+from tensorflow.python.keras.layers import LeakyReLU, ReLU
+
 from functools import partial
 import misc
 import defs
+from models import basic
+import os
 
 latent_dim = 64
 
@@ -16,10 +21,10 @@ print("# GPUs Available: ", len(tf.config.experimental.list_physical_devices('GP
 
 class AutoencoderMod(defs.Modulator):
     def __init__(self, autoencoder):
-        super().__init__(2**autoencoder.N)
+        super().__init__(2 ** autoencoder.N)
         self.autoencoder = autoencoder
 
-    def forward(self, binary: np.ndarray) -> np.ndarray:
+    def forward(self, binary: np.ndarray):
         reshaped = binary.reshape((-1, self.N))
         reshaped_ho = misc.bit_matrix2one_hot(reshaped)
         encoded = self.autoencoder.encoder(reshaped_ho)
@@ -28,38 +33,39 @@ class AutoencoderMod(defs.Modulator):
 
         f = np.zeros(x2.shape[0])
         x3 = misc.rect2polar(np.c_[x2[:, 0], x2[:, 1], f])
-        return x3
+        return basic.RFSignal(x3)
 
 
 class AutoencoderDemod(defs.Demodulator):
     def __init__(self, autoencoder):
-        super().__init__(2**autoencoder.N)
+        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, )
 
 
 class Autoencoder(Model):
-    def __init__(self, N, noise):
+    def __init__(self, N, channel, signal_dim=2):
         super(Autoencoder, self).__init__()
         self.N = N
         self.encoder = tf.keras.Sequential()
         self.encoder.add(tf.keras.Input(shape=(2 ** N,), dtype=bool))
         self.encoder.add(layers.Dense(units=2 ** (N + 1)))
+        self.encoder.add(LeakyReLU(alpha=0.001))
         # self.encoder.add(layers.Dropout(0.2))
         self.encoder.add(layers.Dense(units=2 ** (N + 1)))
-        self.encoder.add(layers.Dense(units=2, activation="sigmoid"))
+        self.encoder.add(LeakyReLU(alpha=0.001))
+        self.encoder.add(layers.Dense(units=signal_dim, activation="tanh"))
         # self.encoder.add(layers.ReLU(max_value=1.0))
 
         self.decoder = tf.keras.Sequential()
-        self.decoder.add(tf.keras.Input(shape=(2,)))
-        self.decoder.add(layers.Dense(units=2 ** (N + 1)))
-        # self.decoder.add(layers.Dropout(0.2))
+        self.decoder.add(tf.keras.Input(shape=(signal_dim,)))
         self.decoder.add(layers.Dense(units=2 ** (N + 1)))
+        # leaky relu with alpha=1 gives by far best results
+        self.decoder.add(LeakyReLU(alpha=1))
         self.decoder.add(layers.Dense(units=2 ** N, activation="softmax"))
 
         # self.randomiser = tf.random_normal_initializer(mean=0.0, stddev=0.1, seed=None)
@@ -68,8 +74,14 @@ class Autoencoder(Model):
         self.demod = None
         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
+        if isinstance(channel, int) or isinstance(channel, float):
+            self.channel = basic.AWGNChannel(channel)
+        else:
+            if not hasattr(channel, 'forward_tensor'):
+                raise ValueError("Channel has no forward_tensor function")
+            if not callable(channel.forward_tensor):
+                raise ValueError("Channel.forward_tensor is not callable")
+            self.channel = channel
 
         # self.decoder.add(layers.Softmax(units=4, dtype=bool))
 
@@ -85,15 +97,61 @@ class Autoencoder(Model):
         # ])
 
     def call(self, x, **kwargs):
-        encoded = self.encoder(x)
-        encoded = encoded * 2 - 1
+        signal = self.encoder(x)
+        signal = signal * 2 - 1
+        signal = self.channel.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 fit_encoder(self, modulation, sample_size, train_size=0.8, epochs=1, batch_size=1, shuffle=False):
+        alphabet = basic.load_alphabet(modulation, polar=False)
+
+        if not alphabet.shape[0] == self.N ** 2:
+            raise Exception("Cardinality of modulation scheme is different from cardinality of autoencoder!")
+
+        x_train = np.random.randint(self.N ** 2, size=int(sample_size * train_size))
+        y_train = alphabet[x_train]
+        x_train_ho = np.zeros((int(sample_size * train_size), self.N ** 2))
+        for idx, x in np.ndenumerate(x_train):
+            x_train_ho[idx, x] = 1
+
+        x_test = np.random.randint(self.N ** 2, size=int(sample_size * (1 - train_size)))
+        y_test = alphabet[x_test]
+        x_test_ho = np.zeros((int(sample_size * (1 - train_size)), self.N ** 2))
+        for idx, x in np.ndenumerate(x_test):
+            x_test_ho[idx, x] = 1
+
+        self.encoder.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
+        self.encoder.fit(x_train_ho, y_train,
+                         epochs=epochs,
+                         batch_size=batch_size,
+                         shuffle=shuffle,
+                         validation_data=(x_test_ho, y_test))
+
+    def fit_decoder(self, modulation, samples):
+        samples = int(samples * 1.3)
+        demod = basic.AlphabetDemod(modulation, 0)
+        x = np.random.rand(samples, 2) * 2 - 1
+        x = x.reshape((-1, 2))
+        f = np.zeros(x.shape[0])
+        xf = np.c_[x[:, 0], x[:, 1], f]
+        y = demod.forward(basic.RFSignal(misc.rect2polar(xf)))
+        y_ho = misc.bit_matrix2one_hot(y.reshape((-1, 4)))
+
+        X_train, X_test, y_train, y_test = train_test_split(x, y_ho)
+        self.decoder.compile(optimizer='adam', loss=tf.keras.losses.MeanSquaredError())
+        self.decoder.fit(X_train, y_train, shuffle=False, validation_data=(X_test, y_test))
+        y_pred = self.decoder(X_test).numpy()
+        y_pred2 = np.zeros(y_test.shape, dtype=bool)
+        y_pred2[np.arange(y_pred2.shape[0]), np.argmax(y_pred, axis=1)] = True
+
+        print("Decoder accuracy: %.4f" % accuracy_score(y_pred2, y_test))
+
     def train(self, samples=1e6):
         if samples % self.N:
             samples += self.N - (samples % self.N)
@@ -157,26 +215,44 @@ if __name__ == '__main__':
 
     n = 4
 
-    samples = 1e6
-    x_train = misc.generate_random_bit_array(samples).reshape((-1, n))
-    x_train_ho = misc.bit_matrix2one_hot(x_train)
-    x_test_array = misc.generate_random_bit_array(samples * 0.3)
-    x_test = x_test_array.reshape((-1, n))
-    x_test_ho = misc.bit_matrix2one_hot(x_test)
+    # samples = 1e6
+    # x_train = misc.generate_random_bit_array(samples).reshape((-1, n))
+    # x_train_ho = misc.bit_matrix2one_hot(x_train)
+    # x_test_array = misc.generate_random_bit_array(samples * 0.3)
+    # x_test = x_test_array.reshape((-1, n))
+    # x_test_ho = misc.bit_matrix2one_hot(x_test)
 
     autoencoder = Autoencoder(n, -8)
     autoencoder.compile(optimizer='adam', loss=losses.MeanSquaredError())
 
-    autoencoder.fit(x_train_ho, x_train_ho,
-                    epochs=1,
-                    shuffle=False,
-                    validation_data=(x_test_ho, x_test_ho))
+    autoencoder.fit_encoder(modulation='16qam',
+                            sample_size=2e6,
+                            train_size=0.8,
+                            epochs=1,
+                            batch_size=256,
+                            shuffle=True)
 
-    encoded_data = autoencoder.encoder(x_test_ho)
-    decoded_data = autoencoder.decoder(encoded_data).numpy()
-
-    result = misc.int2bit_array(decoded_data.argmax(axis=1), n)
-    print("Accuracy: %.4f" % accuracy_score(x_test_array, result.reshape(-1, )))
     view_encoder(autoencoder.encoder, n)
+    autoencoder.fit_decoder(modulation='16qam', samples=2e6)
+    autoencoder.train()
+    view_encoder(autoencoder.encoder, n)
+
+    # view_encoder(autoencoder.encoder, n)
+    # view_encoder(autoencoder.encoder, n)
+
+
+    # autoencoder.compile(optimizer='adam', loss=losses.MeanSquaredError())
+    #
+    # autoencoder.fit(x_train_ho, x_train_ho,
+    #                 epochs=1,
+    #                 shuffle=False,
+    #                 validation_data=(x_test_ho, x_test_ho))
+    #
+    # encoded_data = autoencoder.encoder(x_test_ho)
+    # decoded_data = autoencoder.decoder(encoded_data).numpy()
+    #
+    # result = misc.int2bit_array(decoded_data.argmax(axis=1), n)
+    # print("Accuracy: %.4f" % accuracy_score(x_test_array, result.reshape(-1, )))
+    # view_encoder(autoencoder.encoder, n)
 
     pass

+ 49 - 83
models/basic.py

@@ -4,72 +4,9 @@ import math
 import misc
 from scipy.spatial import cKDTree
 from os import path
+import tensorflow as tf
 
 ALPHABET_DIR = "./alphabets"
-# def _make_gray(n):
-#     if n <= 0:
-#         return []
-#     arr = ['0', '1']
-#     i = 2
-#     while True:
-#         if i >= 1 << n:
-#             break
-#         for j in range(i - 1, -1, -1):
-#             arr.append(arr[j])
-#         for j in range(i):
-#             arr[j] = "0" + arr[j]
-#         for j in range(i, 2 * i):
-#             arr[j] = "1" + arr[j]
-#         i = i << 1
-#     return list(map(lambda x: int(x, 2), arr))
-#
-#
-# def _gen_mary_alphabet(size, gray=True, polar=True):
-#     alphabet = np.zeros((size, 2))
-#     N = math.ceil(math.sqrt(size))
-#
-#     # if sqrt(size) != size^2 (not a perfect square),
-#     # skip defines how many corners to cut off.
-#     skip = 0
-#     if N ** 2 > size:
-#         skip = int(math.sqrt((N ** 2 - size) // 4))
-#
-#     step = 2 / (N - 1)
-#     skipped = 0
-#     for x in range(N):
-#         for y in range(N):
-#             i = x * N + y - skipped
-#             if i >= size:
-#                 break
-#             # Reverse y every odd column
-#             if x % 2 == 0 and N < 4:
-#                 y = N - y - 1
-#             if skip > 0:
-#                 if (x < skip or x + 1 > N - skip) and \
-#                         (y < skip or y + 1 > N - skip):
-#                     skipped += 1
-#                     continue
-#             # Exception for 3-ary alphabet, skip centre point
-#             if size == 8 and x == 1 and y == 1:
-#                 skipped += 1
-#                 continue
-#             alphabet[i, :] = [step * x - 1, step * y - 1]
-#     if gray:
-#         shape = alphabet.shape
-#         d1 = 4 if N > 4 else 2 ** N // 4
-#         g1 = np.array([0, 1, 3, 2])
-#         g2 = g1[:d1]
-#         hypershape = (d1, 4, 2)
-#         if N > 4:
-#             hypercube = alphabet.reshape(hypershape + (N-4, ))
-#             hypercube = hypercube[:, g1, :, :][g2, :, :, :]
-#         else:
-#             hypercube = alphabet.reshape(hypershape)
-#             hypercube = hypercube[:, g1, :][g2, :, :]
-#         alphabet = hypercube.reshape(shape)
-#     if polar:
-#         alphabet = misc.rect2polar(alphabet)
-#     return alphabet
 
 
 def load_alphabet(name, polar=True):
@@ -102,12 +39,12 @@ def load_alphabet(name, polar=True):
                 else:
                     raise ValueError()
                 if 'd' in header:
-                    p = y*math.pi/180
+                    p = y * math.pi / 180
                     y = math.sin(p) * x
                     x = math.cos(p) * x
                 data.append((x, y))
             except ValueError:
-                raise ValueError(f"Alphabet {name} line {i+1}: '{row}' has invalid values")
+                raise ValueError(f"Alphabet {name} line {i + 1}: '{row}' has invalid values")
 
     data2 = [None] * len(data)
     for i, d in enumerate(data):
@@ -118,6 +55,30 @@ def load_alphabet(name, polar=True):
     return arr
 
 
+class RFSignal(defs.Signal):
+    def __init__(self, array: np.ndarray):
+        self.amplitude = array[:, 0]
+        self.phase = array[:, 1]
+        self.frequency = array[:, 2]
+        self.symbols = array.shape[0]
+
+    @property
+    def rect(self) -> np.ndarray:
+        return misc.polar2rect(np.c_[self.amplitude, self.phase])
+
+    def set_rect_xy(self, x_mat: np.ndarray, y_mat: np.ndarray):
+        self.set_rect(np.c_[x_mat, y_mat])
+
+    def set_rect(self, mat: np.ndarray):
+        polar = misc.rect2polar(mat)
+        self.amplitude = polar[:, 0]
+        self.phase = polar[:, 1]
+
+    @property
+    def apf(self):
+        return np.c_[self.amplitude, self.phase, self.frequency]
+
+
 class BypassChannel(defs.Channel):
     def forward(self, values):
         return values
@@ -131,12 +92,17 @@ class AWGNChannel(defs.Channel):
         super().__init__(**kwargs)
         self.noise = 10 ** (noise_level / 10)
 
-    def forward(self, values):
-        a = np.random.normal(0, 1, values.shape[0]) * self.noise
-        p = np.random.normal(0, 1, values.shape[0]) * self.noise
-        f = np.zeros(values.shape[0])
-        noise_mat = np.c_[a, p, f]
-        return values + noise_mat
+    def forward(self, values: RFSignal) -> RFSignal:
+        values.set_rect_xy(
+            values.rect_x + np.random.normal(0, 1, values.symbols) * self.noise,
+            values.rect_y + np.random.normal(0, 1, values.symbols) * self.noise,
+        )
+        return values
+
+    def forward_tensor(self, tensor: tf.Tensor) -> tf.Tensor:
+        noise = tf.random.normal([2], mean=0.0, stddev=1.0, dtype=tf.dtypes.float32, seed=None, name=None)
+        tensor += noise * self.noise
+        return tensor
 
 
 class BPSKMod(defs.Modulator):
@@ -145,12 +111,12 @@ class BPSKMod(defs.Modulator):
         super().__init__(2, **kwargs)
         self.f = carrier_f
 
-    def forward(self, binary: np.ndarray):
+    def forward(self, binary):
         a = np.ones(binary.shape[0])
         p = np.zeros(binary.shape[0])
         p[binary == True] = np.pi
         f = np.zeros(binary.shape[0]) + self.f
-        return np.c_[a, p, f]
+        return RFSignal(np.c_[a, p, f])
 
 
 class BPSKDemod(defs.Demodulator):
@@ -167,11 +133,11 @@ class BPSKDemod(defs.Demodulator):
     def forward(self, values):
         # TODO: Channel noise simulator for frequency component?
         # for now we only care about amplitude and phase
-        ap = np.delete(values, 2, 1)
-        ap = misc.polar2rect(ap)
+        # ap = np.delete(values, 2, 1)
+        # ap = misc.polar2rect(ap)
 
-        result = np.ones(values.shape[0], dtype=bool)
-        result[ap[:, 0] > 0] = False
+        result = np.ones(values.symbols, dtype=bool)
+        result[values.rect_x[:, 0] > 0] = False
         return result
 
 
@@ -196,7 +162,7 @@ class AlphabetMod(defs.Modulator):
         a = values[:, 0]
         p = values[:, 1]
         f = np.zeros(reshaped.shape[0]) + self.f
-        return np.c_[a, p, f]  #, indices
+        return RFSignal(np.c_[a, p, f])  # , indices
 
 
 class AlphabetDemod(defs.Demodulator):
@@ -211,11 +177,11 @@ class AlphabetDemod(defs.Demodulator):
         self.ktree = cKDTree(self.alphabet)
 
     def forward(self, binary):
-        binary = binary[:, :2]  # ignore frequency
-        rbin = misc.polar2rect(binary)
-        indices = self.ktree.query(rbin)[1]
+        # binary = binary[:, :2]  # ignore frequency
+        # rbin = misc.polar2rect(binary)
+        indices = self.ktree.query(binary.rect)[1]
 
         # Converting indices to bite array
         # FIXME: unpackbits requires 8bit inputs, thus largest demodulation is 256-QAM
         values = np.unpackbits(np.array([indices], dtype=np.uint8).T, bitorder='little', axis=1)
-        return values[:, :self.N].reshape((-1,)).astype(bool)  #, indices
+        return values[:, :self.N].reshape((-1,)).astype(bool)  # , indices

+ 7 - 2
models/optical_channel.py

@@ -6,6 +6,9 @@ import math
 from numpy.fft import fft, fftfreq, ifft
 from commpy.filters import rrcosfilter, rcosfilter, rectfilter
 
+from models import basic
+
+
 class OpticalChannel(defs.Channel):
     def __init__(self, noise_level, dispersion, symbol_rate, sample_rate, length, pulse_shape='rect',
                  sqrt_out=False, show_graphs=False, **kwargs):
@@ -98,6 +101,8 @@ class OpticalChannel(defs.Channel):
         return t, val_t
 
     def forward(self, values):
+        if hasattr(values, 'apf'):
+            values = values.apf
         # Converting APF representation to time-series
         t, val_t = self.__get_time_domain(values)
 
@@ -127,7 +132,7 @@ class OpticalChannel(defs.Channel):
 
         # Symbol Decisions
         idx = np.arange(self.filter_samples/2, t.shape[0] - (self.filter_samples/2),
-                        self.symbol_period/self.sample_period, dtype='int16')
+                        self.symbol_period/self.sample_period, dtype='int64')
         t_descision = self.sample_period * idx
 
         if self.show_graphs:
@@ -152,7 +157,7 @@ class OpticalChannel(defs.Channel):
         if self.sqrt_out:
             out[:, 0] = np.sqrt(out[:, 0])
 
-        return out
+        return basic.RFSignal(out)
 
 
 if __name__ == '__main__':