snr fix2

models/layers.py

@@ -112,7 +112,9 @@ class DigitizationLayer(layers.Layer):
         """
         super(DigitizationLayer, self).__init__()
 
-        stddev = 3 * (sig_avg ** 2) * (10 ** ((-6.02 * enob + 1.76) / 10))
+        avg_power = (2*sig_avg) / 3
+
+        stddev = math.sqrt(3 * (avg_power ** 2) * (10 ** ((-6.02 * enob + 1.76) / 10)))
 
         self.noise_layer = layers.GaussianNoise(stddev)
         freq = np.fft.fftfreq(num_of_samples, d=1 / fs)
@@ -169,7 +171,7 @@ class OpticalChannel(layers.Layer):
         self.rx_stddev = rx_stddev
         self.sig_avg = sig_avg
         self.enob = enob
-        self.total_atten = tf.cast(10**(-0.1*atten*fiber_length), dtype=tf.float32)
+        self.total_atten = tf.cast(10**(-0.1*atten*fiber_length), dtype=tf.complex64)
 
         self.noise_layer = layers.GaussianNoise(self.rx_stddev)
         self.fiber_length_noise = layers.GaussianNoise(self.fiber_length_stddev)
@@ -187,7 +189,7 @@ class OpticalChannel(layers.Layer):
         average_optical_power = (2*self.sig_avg*self.total_atten) / 3
         signal_var = self.rx_stddev**2
 
-        return 10*math.log(average_optical_power/signal_var)
+        return 10*math.log(average_optical_power/signal_var, 10)
 
     def call(self, inputs, **kwargs):
         # DAC LPF and noise