Cleanup channel model

constellation/GaussianChannel.py

@@ -1,75 +1,78 @@
 import torch.nn as nn
-import torch
+from torch.distributions.normal import Normal
 import math
-import numpy as np
-
-
-def channel_OSNR():
-    Sys_rate = 32e9
-    r = 0.05
-    Dispersion = 16.48e-6
-    B_2 = Dispersion
-    Non_linear_index = 1.3e3
-    Gam = Non_linear_index
-    Loss = 10**20
-    Alpha = Loss
-    Span_count = 20
-    N_s = Span_count
-    Span_length = 10e5  # (km)
-    L_s = Span_length
-    Noise_figure = 10**0.5  # (dB)
-    h = 6.6261e-34
-    v = 299792458
-    B_WDM = Sys_rate*(1+r)
-    B_N = 0.1
-
-    P_ASE_1 = h*v*B_N*(Loss*Span_length*Noise_figure-1)
-    P_ASE = P_ASE_1 * Span_count
-    L_eff = 1-np.exp(-Loss*Span_length)/2/Alpha
-
-    eps = 0.3*np.log(1+(6/L_s)*(L_eff/np.arcsinh((np.pi**2/3)*B_2*L_eff*B_WDM**2)))
-    b = P_ASE_1/(2*(N_s**eps)*B_N*(Gam**2)*L_eff*np.arcsinh((np.pi**2/3)*B_2*L_eff*B_WDM**2))
-    P_ch = Sys_rate*(((27*np.pi*B_2/16)*b)**(1/3))
-    OSNR = (2*P_ch/3/P_ASE)
-
-    OSNR_dB = 10*np.log10(OSNR)
-    return OSNR_dB
-
-
-def Const_Points_Pow(IQ):
-    """
-    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)
-    return p_enc_avg_dB
-
-
-def Pow_Noise(CH_OSNR):
-    """
-    Calculate the power of channel noise.
-    """
-    P_N_dB = -CH_OSNR
-    p_N_watt = 10**(P_N_dB/10)
-    Var_Noise = p_N_watt
-    return Var_Noise
-
-
-def Channel_Noise_Model(variance, shape):
-    """
-    Compute the Gaussian noise to be added to each vector to simulate passing
-    through a channel.
-    """
-    return torch.distributions.normal.Normal(
-        0, math.sqrt(variance * 5000)
-    ).sample(shape)
 
 
 class GaussianChannel(nn.Module):
     def __init__(self):
         super().__init__()
 
+        # Initialize channel parameters
+        sys_rate = 32e9
+        r = 0.05
+        dispersion = 16.48e-6
+        B_2 = dispersion
+        non_linear_index = 1.3e3
+        gam = non_linear_index
+        loss = 10**20
+        alpha = loss
+        span_count = 20
+        N_s = span_count
+        span_length = 10e5  # (km)
+        L_s = span_length
+        noise_figure = 10 ** 0.5  # (dB)
+        h = 6.6261e-34
+        v = 299792458
+        B_WDM = sys_rate * (1 + r)
+        B_N = 0.1
+
+        P_ASE_1 = h * v * B_N * (loss * span_length * noise_figure - 1)
+        P_ASE = P_ASE_1 * span_count
+        L_eff = 1 - math.exp(-loss * span_length) / 2 / alpha
+
+        eps = 0.3 * math.log(
+            1 + (6 / L_s) * (
+                L_eff / math.asinh(
+                    (math.pi ** 2)
+                    / 3
+                    * B_2
+                    * L_eff
+                    * (B_WDM ** 2)
+                )
+            )
+        )
+
+        b = P_ASE_1 / (
+            2
+            * (N_s ** eps)
+            * B_N
+            * (gam ** 2)
+            * L_eff * math.asinh(
+                (math.pi ** 2) / 3
+                * B_2
+                * L_eff
+                * (B_WDM ** 2)
+            )
+        )
+
+        P_ch = sys_rate * (((27 * math.pi * B_2 / 16) * b) ** (1 / 3))
+        OSNR = (2 * P_ch / 3 / P_ASE)
+
+        OSNR_dB = 10 * math.log10(OSNR)
+
+        p_N_dB = -OSNR_dB
+        p_N_watt = 10**(p_N_dB/10)
+
+        self.noise_std = math.sqrt(p_N_watt * 5000)
+
+    def get_noise(self, rows):
+        """
+        Generate Gaussian random noise according to the channel’s parameters.
+
+        :param rows: Number of noise vectors to generate.
+        :return: Matrix of shape `rows` × 2.
+        """
+        return Normal(0, self.noise_std).sample((rows, 2))
+
     def forward(self, x):
-        Noise_Variance = Pow_Noise(channel_OSNR())
-        Noise_Volts = Channel_Noise_Model(Noise_Variance, [len(x), 2])
-        return x + Noise_Volts
+        return x + self.get_noise(len(x))