constellationnet/constellation/GaussianChannel.py

76 lines
1.8 KiB
Python

import torch.nn as nn
import torch
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__()
def forward(self, x):
Noise_Variance = Pow_Noise(channel_OSNR())
Noise_Volts = Channel_Noise_Model(Noise_Variance, [len(x), 2])
return x + Noise_Volts