Merge remote-tracking branch 'origin/chromatic_dispersion'

models/optical_channel.py

@@ -0,0 +1,127 @@
+import matplotlib.pyplot as plt
+
+import defs
+import numpy as np
+import math
+from scipy.fft import fft, ifft
+
+
+class OpticalChannel(defs.Channel):
+    def __init__(self, noise_level, dispersion, symbol_rate, sample_rate, length, show_graphs=False, **kwargs):
+        """
+        :param noise_level: Noise level in dB
+        :param dispersion: dispersion coefficient is ps^2/km
+        :param symbol_rate: Symbol rate of modulated signal in Hz
+        :param sample_rate: Sample rate of time-domain model (time steps in simulation) in Hz
+        :param length: fibre length in km
+        :param show_graphs: if graphs should be displayed or not
+
+        Optical Channel class constructor
+        """
+        super().__init__(**kwargs)
+        self.noise = 10 ** (noise_level / 10)
+
+        self.dispersion = dispersion # * 1e-24  # Converting from ps^2/km to s^2/km
+        self.symbol_rate = symbol_rate
+        self.symbol_period = 1 / self.symbol_rate
+        self.sample_rate = sample_rate
+        self.sample_period = 1 / self.sample_rate
+        self.length = length
+        self.show_graphs = show_graphs
+
+    def __get_time_domain(self, symbol_vals):
+        samples_per_symbol = int(self.sample_rate / self.symbol_rate)
+        samples = int(symbol_vals.shape[0] * samples_per_symbol)
+
+        symbol_vals_a = np.repeat(symbol_vals, repeats=samples_per_symbol, axis=0)
+        t = np.linspace(start=0, stop=samples * self.sample_period, num=samples)
+        val_t = symbol_vals_a[:, 0] * np.cos(2 * math.pi * symbol_vals_a[:, 2] * t + symbol_vals_a[:, 1])
+
+        return t, val_t
+
+    def __time_to_frequency(self, values):
+        val_f = fft(values)
+        f = np.linspace(0.0, 1 / (2 * self.sample_period), (values.size // 2))
+        f_neg = -1 * np.flip(f)
+        f = np.concatenate((f, f_neg), axis=0)
+        return f, val_f
+
+    def __frequency_to_time(self, values):
+        val_t = ifft(values)
+        t = np.linspace(start=0, stop=values.size * self.sample_period, num=values.size)
+        return t, val_t
+
+    def __apply_dispersion(self, values):
+        # Obtain fft
+        f, val_f = self.__time_to_frequency(values)
+
+        if self.show_graphs:
+            plt.plot(f, val_f)
+            plt.title('frequency domain (pre-distortion)')
+            plt.show()
+
+        # Apply distortion
+        dist_val_f = val_f * np.exp(0.5j * self.dispersion * self.length * np.power(2 * math.pi * f, 2))
+
+        if self.show_graphs:
+            plt.plot(f, dist_val_f)
+            plt.title('frequency domain (post-distortion)')
+            plt.show()
+
+        # Inverse fft
+        t, val_t = self.__frequency_to_time(dist_val_f)
+
+        return t, val_t
+
+    def __photodiode_detection(self, values):
+        t = np.linspace(start=0, stop=values.size * self.sample_period, num=values.size)
+        val_t = np.power(np.absolute(values), 2)
+        return t, val_t
+
+    def forward(self, values):
+        # Converting APF representation to time-series
+        t, val_t = self.__get_time_domain(values)
+
+        if self.show_graphs:
+            plt.plot(t, val_t)
+            plt.title('time domain (raw)')
+            plt.show()
+
+        # Adding AWGN
+        val_t += np.random.normal(0, 1, val_t.shape) * self.noise
+
+        if self.show_graphs:
+            plt.plot(t, val_t)
+            plt.title('time domain (AWGN)')
+            plt.show()
+
+        # Applying chromatic dispersion
+        t, val_t = self.__apply_dispersion(val_t)
+
+        if self.show_graphs:
+            plt.plot(t, val_t)
+            plt.title('time domain (post-distortion)')
+            plt.show()
+
+        # Photodiode Detection
+        t, val_t = self.__photodiode_detection(val_t)
+
+        if self.show_graphs:
+            plt.plot(t, val_t)
+            plt.title('time domain (post-detection)')
+            plt.show()
+
+        return t, val_t
+
+
+if __name__ == '__main__':
+    # Simple OOK modulation
+    num_of_symbols = 10
+    symbol_vals = np.zeros((num_of_symbols, 3))
+
+    symbol_vals[:, 0] = np.random.randint(2, size=symbol_vals.shape[0])
+    symbol_vals[:, 2] = 10e6
+
+    channel = OpticalChannel(noise_level=-20, dispersion=-21.7, symbol_rate=100e3,
+                             sample_rate=500e6, length=100, show_graphs=True)
+    time, v = channel.forward(symbol_vals)