<jupyter_start><jupyter_text>See [`tf_wdmSystem-learning.html`](tf_wdmSystem-learning.html) for version with `runWithNonlinear = True`<jupyter_code>if runWithNonlinear:  
    # load memory_saving_gradients.py
    # see https://github.com/cybertronai/gradient-checkpointing
    import urllib.request
    memSavingGradientsFilename = "memory_saving_gradients.py"
    if not os.path.exists(memSavingGradientsFilename):
        memSavingGradientsURL = "https://raw.githubusercontent.com/cybertronai/gradient-checkpointing/master/memory_saving_gradients.py"
        urllib.request.urlretrieve(memSavingGradientsURL, memSavingGradientsFilename)

        import memory_saving_gradients
        tf.__dict__["gradients"] = memory_saving_gradients.gradients_collection
# Tx
param = cfs.defaultParameters(precision='double')

if runWithNonlinear:    
    param.sps = 16
    param.nSamples = 1024
    param.optimizeP = False
    param.PdBm  = 1
else:    
    param.sps = 16
    param.nSamples = 1024
    param.PdBm  = -2 # linear regime such there is no nonlinear interference
    
param.M     = 64
param.Fs    = param.sps * param.Rs
param.nPilots = 128
param.pilotOrder = 4
print(param)
ssfParam = ssf.defaultParameters(precision='double')
ssfParam.Fs = param.Fs
ssfParam.N = param.sps * (param.nSamples+param.nPilots)
ssfParam.noiseEnabled = True
ssfParam.noiseFigure = 5
ssfParam.nSpans = 10

if runWithNonlinear:
    ssfParam.gamma = 1.3
    ssfParam.nSteps = 350
else:
    ssfParam.gamma = 0
    ssfParam.nSteps = 2

# ssfParam.stepSize = ssfParam.spanLength/ssfParam.nSteps
# ssfParam.stepSizeTemplate = ssfParam.stepSize * np.ones((ssfParam.nSteps,)) # constant stepsize

ssfParam.stepSizeTemplate = ssf.logStepSizes(ssfParam.spanLength, ssfParam.alpha, ssfParam.nSteps) # log stepsize

print(ssfParam)
aeParam = cu.AttrDict()
aeParam.seed    = 1337
aeParam.constellationOrder = param.M
aeParam.constellationDim = 2

if runWithNonlinear:
    aeParam.nLayers = 4
    aeParam.nHidden = 64
else:
    aeParam.nLayers = 2
    aeParam.nHidden = 32

aeParam.activation = tf.nn.relu

print(aeParam)
trainingParam = cu.AttrDict()
if runWithNonlinear:    
    trainingParam.batchSize     = 1
    trainingParam.nBatches      = 2
    trainingParam.nMiniBatches  = 4
    trainingParam.evalBatches   = 10
    trainingParam.removeSymbols = 32
    trainingParam.learningRate  = 0.0005
else:    
    trainingParam.batchSize     = 1
    trainingParam.nBatches      = 2
    trainingParam.nMiniBatches  = 4
    trainingParam.evalBatches   = 5
    trainingParam.removeSymbols = 32
    trainingParam.learningRate  = 0.005
    
trainingParam.nEpochs             = 100

trainingParam.displayStep         = 5
trainingParam.path                = 'results_SSFM'
trainingParam.filename            = 'M{:03d}_seed{:04d}'.format(param.M,aeParam.seed)
trainingParam.earlyStopping       = 10
trainingParam.earlyStoppingMetric = 'xentropy'
trainingParam.summaries           = True

# generalize this
if trainingParam.summaries:
    # tensorboard directory
    hyperParam = ['M']
    aeHyperParam = ['seed']
    trainingHyperParam = ['learningRate']

    trainingParam.summaryString = ','.join(  [ '{}={}'.format(item,param[item]) for item in hyperParam ]
                                            +[ '{}={}'.format(item,trainingParam[item]) for item in trainingHyperParam ]
                                            +[ '{}={}'.format(item,aeParam[item]) for item in aeHyperParam ] )

    print(trainingParam.summaryString,flush=True)
    
print(trainingParam)
# Constants
c = 299792458
tf.set_random_seed(aeParam.seed)
np.random.seed(aeParam.seed)
X = tf.placeholder(param.realType, shape=(None, param.nChannels, param.nPol, param.nSamples, param.M))
txSymbols, constellation = ae.encoder(X, aeParam, toComplex=True)
# pilot symbols
X_pilot = tf.placeholder(param.realType, shape=(None, param.nChannels, param.nPol, param.nPilots, param.pilotOrder))
pilot_constellation = tx.qammod(param.pilotOrder)
pilot_constellation_tf = tf.constant(pilot_constellation, param.complexType)
pilot_txSymbols = cfh.QAMencoder(X_pilot, pilot_constellation_tf, realOutput=False)
txSymbols = tf.concat([pilot_txSymbols,txSymbols], axis=-1)
txSymbols
param.nSamples = param.nSamples + param.nPilots
signal = cfs.wdmTransmitter(txSymbols, param)
ssfParam.dispersionCompensationEnabled = False
ssfParam.stepSize = tf.placeholder(param.realType, shape=(ssfParam.nSpans, ssfParam.nSteps))
signalOut = ssf.model(ssfParam, signal)

param.dispersionCompensation = True
param.beta2  = ssfParam.D * 1e-6 * ( c / ssfParam.Fc )**2 / ( 2 * np.pi * c )
param.dz     = ssfParam.spanLength * 1e3 * ssfParam.nSpans
rxSymbols = cfs.wdmReceiver(signalOut, param)
P0 = tf.get_default_graph().get_tensor_by_name("P0:0")
PdBm = tf.get_default_graph().get_tensor_by_name("PdBm:0")
normP0 = tf.get_default_graph().get_tensor_by_name("normP0:0")
rxSymbols = tf.cast( tf.rsqrt( normP0 ), param.complexType ) * rxSymbols
# simulate static phase rotation 
if not runWithNonlinear:
    rxSymbols = rxSymbols * tf.exp(1j * 0.5)
rxSymbols = cfh.staticPhaseRotationCompensation(rxSymbols, nPilots=param.nPilots)
rxSymbols = cfh.testPhases(pilot_constellation_tf, txSymbols, rxSymbols, 4, param.pilotOrder, nTestPhases=4, nPilots=param.nPilots)
txSymbols = txSymbols[..., param.nPilots:]
rxSymbols = rxSymbols[..., param.nPilots:]
rxSymbols, txSymbols, labels = cfh.truncate(trainingParam.removeSymbols, rxSymbols, txSymbols, tf.linalg.matrix_transpose(X))
labels = tf.linalg.matrix_transpose(labels)

decoder = ae.decoder( rxSymbols, aeParam, fromComplex=True)
softmax = tf.nn.softmax(decoder)

# loss
loss = tf.reduce_mean(tf.nn.softmax_cross_entropy_with_logits_v2(labels=labels, logits=decoder))
nDims = 4 # Dimension of X
reduce_axis = [0, 2, 3]
errorrate = cfh.symbolErrorrate(constellation, txSymbols, rxSymbols, nDims, param.M, reduce_axis)

one = tf.constant( 1, P0.dtype )
effSNR = cfh.effectiveSNR(txSymbols, rxSymbols, one, reduce_axis)

# gaussMI
gaussMI = []
for (txChannel, rxChannel) in zip(tf.split(txSymbols, param.nChannels, axis=1), tf.split(rxSymbols, param.nChannels, axis=1)):
    gaussMI_temp = cfh.gaussianMI(tf.reshape(txChannel,[-1]), tf.reshape(rxChannel,[-1]), constellation, param.M)
    gaussMI.append(gaussMI_temp)
gaussMI = tf.stack(gaussMI)

# nnMI
Px = tf.constant( 1/param.M, param.realType )
nnMI = []
for (sm, l) in zip(tf.split(softmax, param.nChannels, axis=1), tf.split(labels, param.nChannels, axis=1)):
    nnMI_temp = cfh.softmaxMI(sm, l, Px)
    nnMI.append(nnMI_temp)
nnMI = tf.stack(nnMI)

# only gaussMId channel
gaussMIdChannel = param.nChannels//2 +1
errorrate = errorrate[gaussMIdChannel]
gaussMI = gaussMI[gaussMIdChannel]
effSNR = effSNR[gaussMIdChannel]
metricsDict = {'xentropy':loss, 'errorrate':errorrate, 'gaussMI': gaussMI, 'nnMI': nnMI, 'effSNR': effSNR}
def feedDictFun(trainingParam):
    randomisedStepSize = ssf.randomizeSteps(ssfParam.stepSizeTemplate, ssfParam.spanLength, ssfParam.nSpans)
    x, idx, x_seed = cu.hotOnes((trainingParam.batchSize, param.nChannels, param.nPol, param.nSamples-param.nPilots), (1,2,3,4,0),param.M)
    x_pilot, _, _  = cu.hotOnes((trainingParam.batchSize, param.nChannels, param.nPol, param.nPilots), (1,2,3,4,0),param.pilotOrder)
    return { X:x, X_pilot:x_pilot, ssfParam.stepSize:randomisedStepSize }
sess = tf.compat.v1.Session()
sess, outMetrics = cft.train(sess, tf.train.AdamOptimizer, loss, metricsDict, trainingParam, feedDictFun, debug=False)
outPdBm = sess.run( PdBm )

np_constellation = sess.run(constellation)
plt.plot(np.real(np_constellation),np.imag(np_constellation),'.')
plt.axis('square')
lim_ = 1.5
plt.xlim(-lim_,lim_)
plt.ylim(-lim_,lim_)
plt.savefig('const.png')
import claude.models.NLIN as nlin
# System parameters
qamParam = nlin.defaultParameters()
qamParam.nSpans = ssfParam.nSpans

aseNoisePower, interConst, intraConst, interConstAdd, intraConstAdd = nlin.calcConstants(qamParam)
powerSweep = np.arange(-8, 4, 0.1)

# QAM
qam_constellation = np.squeeze(tx.qammod(param.M))
qamParam.kur, qamParam.kur3 = nlin.calcKur(qam_constellation)

qamEffSNR = powerSweep - nlin.calcNLIN(qamParam, powerSweep, aseNoisePower, interConst, intraConst, interConstAdd, intraConstAdd)

# learned constellation
aeNlinParam = qamParam
aeNlinParam.kur, aeNlinParam.kur3 = nlin.calcKur(np_constellation)

aeEffSNR = powerSweep - nlin.calcNLIN(aeNlinParam, powerSweep, aseNoisePower, interConst, intraConst, interConstAdd, intraConstAdd)
plt.plot(powerSweep, qamEffSNR, 'b')
plt.plot(powerSweep, aeEffSNR, 'r')
plt.plot(outPdBm, outMetrics['effSNR'], 'gx')
plt.title('Power [dBm] VS effective SNR')
plt.xlabel('Power [dBm]')
plt.ylabel('effective SNR [dB]')
plt.show()
plt.savefig('snr.png')
qamMI = np.zeros(np.shape(qamEffSNR))
aeMI = np.zeros(np.shape(qamEffSNR))
N = 10000

for ii,p in enumerate(powerSweep):
    qamMI[ii] = cu.SNRtoMI(N, qamEffSNR[ii], np.expand_dims(qam_constellation, 0))
    aeMI[ii] = cu.SNRtoMI(N, aeEffSNR[ii], np.expand_dims(np_constellation, 0))
plt.plot(powerSweep, qamMI, 'b')
plt.plot(powerSweep, aeMI, 'r')
plt.plot(outPdBm, outMetrics['gaussMI'], 'gx')
plt.plot(outPdBm, outMetrics['nnMI'], 'kx')
plt.title('Power [dBm] VS MI')
plt.xlabel('Power [dBm]')
plt.ylabel('MI [bits]')
plt.show()
plt.savefig('mi.png')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.12545728426942113, 'fraction_numerical': 0.015601463309662148, 'mean_word_length': 4.897842639593908, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 2, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '573286', 'license': 'MIT', 'n_tokens_mistral': 3705, 'n_tokens_neox': 3459, 'n_words': 581, 'path': 'examples/tf_wdmSystem-learning.ipynb', 'repo_name': 'meisme123/claude', 'size': '9294'}