fiber.py (17495B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 6 """ 7 This module defines the split-step Fourier method to approximate the solution of 8 the nonlinear Schroedinger equation. 9 """ 10 11 import tensorflow as tf 12 from tensorflow.keras.layers import Layer 13 import sionna 14 from sionna import config 15 from sionna.channel import utils 16 17 18 class SSFM(Layer): 19 # pylint: disable=line-too-long 20 r"""SSFM(alpha=0.046,beta_2=-21.67,f_c=193.55e12,gamma=1.27,half_window_length=0,length=80,n_ssfm=1,n_sp=1.0,sample_duration=1.0,t_norm=1e-12,with_amplification=False,with_attenuation=True,with_dispersion=True,with_manakov=False,with_nonlinearity=True,swap_memory=True,dtype=tf.complex64,**kwargs) 21 22 Layer implementing the split-step Fourier method (SSFM) 23 24 The SSFM (first mentioned in [HT1973]_) numerically solves the generalized 25 nonlinear Schrödinger equation (NLSE) 26 27 .. math:: 28 29 \frac{\partial E(t,z)}{\partial z}=-\frac{\alpha}{2} E(t,z)+j\frac{\beta_2}{2}\frac{\partial^2 E(t,z)}{\partial t^2}-j\gamma |E(t,z)|^2 E(t,z) + n(n_{\text{sp}};\,t,\,z) 30 31 for an unpolarized (or single polarized) optical signal; 32 or the Manakov equation (according to [WMC1991]_) 33 34 .. math:: 35 36 \frac{\partial \mathbf{E}(t,z)}{\partial z}=-\frac{\alpha}{2} \mathbf{E}(t,z)+j\frac{\beta_2}{2}\frac{\partial^2 \mathbf{E}(t,z)}{\partial t^2}-j\gamma \frac{8}{9}||\mathbf{E}(t,z)||_2^2 \mathbf{E}(t,z) + \mathbf{n}(n_{\text{sp}};\,t,\,z) 37 38 for dual polarization, with attenuation coefficient :math:`\alpha`, group 39 velocity dispersion parameters :math:`\beta_2`, and nonlinearity 40 coefficient :math:`\gamma`. The noise terms :math:`n(n_{\text{sp}};\,t,\,z)` 41 and :math:`\mathbf{n}(n_{\text{sp}};\,t,\,z)`, respectively, stem from 42 an (optional) ideally distributed Raman amplification with 43 spontaneous emission factor :math:`n_\text{sp}`. The optical signal 44 :math:`E(t,\,z)` has the unit :math:`\sqrt{\text{W}}`. For the dual 45 polarized case, :math:`\mathbf{E}(t,\,z)=(E_x(t,\,z), E_y(t,\,z))` 46 is a vector consisting of the signal components of both polarizations. 47 48 The symmetrized SSFM is applied according to Eq. (7) of [FMF1976]_ that 49 can be written as 50 51 .. math:: 52 53 E(z+\Delta_z,t) \approx \exp\left(\frac{\Delta_z}{2}\hat{D}\right)\exp\left(\int^{z+\Delta_z}_z \hat{N}(z')dz'\right)\exp\left(\frac{\Delta_z}{2}\hat{D}\right)E(z,\,t) 54 55 where only the single-polarized case is shown. The integral is 56 approximated by :math:`\Delta_z\hat{N}` with :math:`\hat{D}` and 57 :math:`\hat{N}` denoting the linear and nonlinear SSFM operator, 58 respectively [A2012]_. 59 60 Additionally, ideally distributed Raman amplification may be applied, which 61 is implemented as in [MFFP2009]_. Please note that the implemented 62 Raman amplification currently results in a transparent fiber link. Hence, 63 the introduced gain cannot be parametrized. 64 65 The SSFM operates on normalized time :math:`T_\text{norm}` 66 (e.g., :math:`T_\text{norm}=1\,\text{ps}=1\cdot 10^{-12}\,\text{s}`) and 67 distance units :math:`L_\text{norm}` 68 (e.g., :math:`L_\text{norm}=1\,\text{km}=1\cdot 10^{3}\,\text{m}`). 69 Hence, all parameters as well as the signal itself have to be given with the 70 same unit prefix for the 71 same unit (e.g., always pico for time, or kilo for distance). Despite the normalization, 72 the SSFM is implemented with physical 73 units, which is different from the normalization, e.g., used for the 74 nonlinear Fourier transform. For simulations, only :math:`T_\text{norm}` has to be 75 provided. 76 77 To avoid reflections at the signal boundaries during simulation, a Hamming 78 window can be applied in each SSFM-step, whose length can be 79 defined by ``half_window_length``. 80 81 Example 82 -------- 83 84 Setting-up: 85 86 >>> ssfm = SSFM( 87 >>> alpha=0.046, 88 >>> beta_2=-21.67, 89 >>> f_c=193.55e12, 90 >>> gamma=1.27, 91 >>> half_window_length=100, 92 >>> length=80, 93 >>> n_ssfm=200, 94 >>> n_sp=1.0, 95 >>> t_norm=1e-12, 96 >>> with_amplification=False, 97 >>> with_attenuation=True, 98 >>> with_dispersion=True, 99 >>> with_manakov=False, 100 >>> with_nonlinearity=True) 101 102 Running: 103 104 >>> # x is the optical input signal 105 >>> y = ssfm(x) 106 107 Parameters 108 ---------- 109 alpha : float 110 Attenuation coefficient :math:`\alpha` in :math:`(1/L_\text{norm})`. 111 Defaults to 0.046. 112 113 beta_2 : float 114 Group velocity dispersion coefficient :math:`\beta_2` in :math:`(T_\text{norm}^2/L_\text{norm})`. 115 Defaults to -21.67 116 117 f_c : float 118 Carrier frequency :math:`f_\mathrm{c}` in :math:`(\text{Hz})`. 119 Defaults to 193.55e12. 120 121 gamma : float 122 Nonlinearity coefficient :math:`\gamma` in :math:`(1/L_\text{norm}/\text{W})`. 123 Defaults to 1.27. 124 125 half_window_length : int 126 Half of the Hamming window length. Defaults to 0. 127 128 length : float 129 Fiber length :math:`\ell` in :math:`(L_\text{norm})`. 130 Defaults to 80.0. 131 132 n_ssfm : int | "adaptive" 133 Number of steps :math:`N_\mathrm{SSFM}`. 134 Set to "adaptive" to use nonlinear-phase rotation to calculate 135 the step widths adaptively (maxmimum rotation can be set in phase_inc). 136 Defaults to 1. 137 138 n_sp : float 139 Spontaneous emission factor :math:`n_\mathrm{sp}` of Raman amplification. 140 Defaults to 1.0. 141 142 sample_duration : float 143 Normalized time step :math:`\Delta_t` in :math:`(T_\text{norm})`. 144 Defaults to 1.0. 145 146 t_norm : float 147 Time normalization :math:`T_\text{norm}` in :math:`(\text{s})`. 148 Defaults to 1e-12. 149 150 with_amplification : bool 151 Enable ideal inline amplification and corresponding 152 noise. Defaults to `False`. 153 154 with_attenuation : bool 155 Enable attenuation. Defaults to `True`. 156 157 with_dispersion : bool 158 Apply chromatic dispersion. Defaults to `True`. 159 160 with_manakov : bool 161 Considers axis [-2] as x- and y-polarization and calculates the 162 nonlinear step as given by the Manakov equation. Defaults to `False.` 163 164 with_nonlinearity : bool 165 Apply Kerr nonlinearity. Defaults to `True`. 166 167 phase_inc: float 168 Maximum nonlinear-phase rotation in rad allowed during simulation. 169 To be used with ``n_ssfm`` = "adaptive". 170 Defaults to 1e-4. 171 172 swap_memory : bool 173 Use CPU memory for while loop. Defaults to `True`. 174 175 dtype : tf.complex 176 Defines the datatype for internal calculations and the output 177 dtype. Defaults to `tf.complex64`. 178 179 Input 180 ----- 181 x : [...,n] or [...,2,n], tf.complex 182 Input signal in :math:`(\sqrt{\text{W}})`. If ``with_manakov`` is `True`, 183 the second last dimension is interpreted as x- and y-polarization, 184 respectively. 185 186 Output 187 ------ 188 y : Tensor with same shape as ``x``, `tf.complex` 189 Channel output 190 """ 191 def __init__(self, 192 alpha=0.046, 193 beta_2=-21.67, 194 f_c=193.55e12, 195 gamma=1.27, 196 half_window_length=0, 197 length=80, 198 n_ssfm=1, 199 n_sp=1.0, 200 sample_duration=1.0, 201 t_norm=1e-12, 202 with_amplification=False, 203 with_attenuation=True, 204 with_dispersion=True, 205 with_manakov=False, 206 with_nonlinearity=True, 207 phase_inc=1e-4, 208 swap_memory=True, 209 dtype=tf.complex64, 210 **kwargs): 211 212 super().__init__(dtype=dtype, **kwargs) 213 214 self._dtype = dtype 215 self._cdtype = tf.as_dtype(dtype) 216 self._rdtype = tf.as_dtype(dtype).real_dtype 217 218 self._alpha = tf.cast(alpha, dtype=self._rdtype) 219 self._beta_2 = tf.cast(beta_2, dtype=self._rdtype) 220 self._f_c = tf.cast(f_c, dtype=self._rdtype) 221 self._gamma = tf.cast(gamma, dtype=self._rdtype) 222 self._half_window_length = half_window_length 223 self._length = tf.cast(length, dtype=self._rdtype) 224 self._phase_inc = tf.cast(phase_inc, dtype=self._rdtype) 225 226 if n_ssfm == "adaptive": 227 self._n_ssfm = tf.cast(-1, dtype=tf.int32) # adaptive == -1 228 elif isinstance(n_ssfm, int): 229 self._n_ssfm = tf.cast(n_ssfm, dtype=tf.int32) 230 # Precalculate uniform step size 231 tf.assert_greater(self._n_ssfm, 0) 232 else: 233 raise ValueError("Unsupported parameter for n_ssfm. \ 234 Either an integer or 'adaptive'.") 235 236 # only used for constant step width -> negative value calculated 237 # with adaptive step widths can be ignored 238 self._dz = self._length / tf.cast(self._n_ssfm, dtype=self._rdtype) 239 self._n_sp = tf.cast(n_sp, dtype=self._rdtype) 240 self._swap_memory = swap_memory 241 self._t_norm = tf.cast(t_norm, dtype=self._rdtype) 242 self._sample_duration = tf.cast(sample_duration, dtype=self._rdtype) 243 244 # Booleans are not casted to avoid branches in the graph 245 self._with_amplification = with_amplification 246 self._with_attenuation = with_attenuation 247 self._with_dispersion = with_dispersion 248 self._with_manakov = with_manakov 249 self._with_nonlinearity = with_nonlinearity 250 251 self._rho_n = \ 252 sionna.constants.H * self._f_c * self._alpha * self._length * \ 253 self._n_sp # in (W/Hz) 254 255 # Calculate noise power depending on simulation bandwidth 256 self._p_n_ase = self._rho_n / self._sample_duration / self._t_norm 257 # in (Ws) 258 if self._with_manakov: 259 self._p_n_ase = self._p_n_ase / 2.0 260 261 self._window = tf.complex( 262 tf.signal.hamming_window( 263 window_length=2*self._half_window_length, 264 dtype=self._rdtype 265 ), 266 tf.zeros( 267 2*self._half_window_length, 268 dtype=self._rdtype 269 ) 270 ) 271 272 def _apply_linear_operator(self, q, dz, zeros, frequency_vector): 273 # Chromatic dispersion 274 if self._with_dispersion: 275 dispersion = tf.exp( 276 tf.complex( 277 zeros, 278 -self._beta_2 / tf.cast(2.0, self._rdtype) * dz * 279 ( 280 tf.cast(2.0, self._rdtype) * 281 tf.cast(sionna.constants.PI, self._rdtype) * 282 frequency_vector 283 ) ** tf.cast(2.0, self._rdtype) 284 ) 285 ) 286 dispersion = tf.signal.fftshift(dispersion, axes=-1) 287 q = tf.signal.ifft(tf.signal.fft(q) * dispersion) 288 289 # Attenuation 290 if self._with_attenuation: 291 q = q * tf.cast(tf.exp(-self._alpha / 2.0 * dz), self._cdtype) 292 293 # Amplification (Raman) 294 if self._with_amplification: 295 q = q * tf.cast(tf.exp(self._alpha / 2.0 * dz), self._cdtype) 296 297 return q 298 299 def _apply_noise(self, q, dz): 300 # Noise due to Amplification (Raman) 301 if self._with_amplification: 302 step_noise = self._p_n_ase * tf.cast(dz, self._rdtype) \ 303 / tf.cast(self._length, self._rdtype) \ 304 / tf.cast(2.0, self._rdtype) 305 q_n = tf.complex( 306 config.tf_rng.normal( 307 q.shape, 308 tf.cast(0.0, self._rdtype), 309 tf.sqrt(step_noise), 310 self._rdtype), 311 config.tf_rng.normal( 312 q.shape, 313 tf.cast(0.0, self._rdtype), 314 tf.sqrt(step_noise), 315 self._rdtype) 316 ) 317 q = q + q_n 318 319 return q 320 321 def _apply_nonlinear_operator(self, q, dz, zeros): 322 if self._with_nonlinearity: 323 if self._with_manakov: 324 q = q * tf.exp( 325 tf.complex( 326 zeros, 327 tf.cast(8.0/9.0, self._rdtype) * tf.reduce_sum( 328 tf.abs(q) ** tf.cast(2.0, self._rdtype), 329 axis=-2, 330 keepdims=True 331 ) * self._gamma * tf.negative(tf.math.real(dz))) 332 ) 333 else: 334 q = q * tf.exp( 335 tf.complex( 336 zeros, 337 tf.abs(q) ** tf.cast(2.0, self._rdtype) * self._gamma * 338 tf.negative(tf.math.real(dz))) 339 ) 340 341 return q 342 343 344 def _calculate_step_width(self, q, remaining_length): 345 max_power = tf.math.reduce_max(tf.math.pow(tf.math.abs(q),2.0),axis=None) 346 # ensure that the exact length is reached in the end 347 dz = tf.math.minimum(self._phase_inc / self._gamma / max_power,remaining_length) 348 return dz 349 350 def _adaptive_step(self,q, precalculations, remaining_length, step_counter): 351 352 (window, _, zeros, f) = precalculations 353 354 dz = self._calculate_step_width(q,remaining_length) 355 356 # Apply window-function 357 q = self._apply_window(q, window) 358 q = self._apply_linear_operator(q, dz, zeros, f) # D 359 q = self._apply_nonlinear_operator(q, dz, zeros) # N 360 q = self._apply_noise(q, dz) 361 remaining_length = remaining_length - dz 362 363 precalculations = (window, dz, zeros, f) 364 step_counter = step_counter + 1 365 return q, precalculations, remaining_length, step_counter 366 367 def _cond_adaptive(self, q, precalculations,remaining_length,step_counter): 368 # pylint: disable=unused-argument 369 return tf.greater_equal(remaining_length, 1e-3) # avoid numerical issues for 0 370 371 372 def _apply_window(self, q, window): 373 return q * window 374 375 def _step(self, q, precalculations, n_steps, step_counter): 376 (window, dz, zeros, f) = precalculations 377 378 # Apply window-function 379 q = self._apply_window(q, window) 380 q = self._apply_nonlinear_operator(q, dz, zeros) # N 381 q = self._apply_noise(q, dz) 382 q = self._apply_linear_operator(q, dz, zeros, f) # D 383 384 step_counter = step_counter + 1 385 386 return q, precalculations, n_steps, step_counter 387 388 def _cond(self, q, precalculations, n_steps, step_counter): 389 # pylint: disable=unused-argument 390 return tf.less(step_counter, n_steps) 391 392 def call(self, inputs): 393 if self._with_manakov: 394 tf.assert_equal(tf.shape(inputs)[-2], 2) 395 396 x = inputs 397 398 # Calculate support parameters 399 input_shape = x.shape 400 401 # Generate frequency vectors 402 _, f = utils.time_frequency_vector( 403 input_shape[-1], self._sample_duration, dtype=self._rdtype) 404 405 # Window function calculation (depends on length of the signal) 406 window = tf.concat( 407 [ 408 self._window[0:self._half_window_length], 409 tf.complex( 410 tf.ones( 411 [input_shape[-1] - 2*self._half_window_length], 412 dtype=self._rdtype 413 ), 414 tf.zeros( 415 [input_shape[-1] - 2*self._half_window_length], 416 dtype=self._rdtype 417 ) 418 ), 419 self._window[self._half_window_length::] 420 ], 421 axis=0 422 ) 423 424 # All-zero vector 425 zeros = tf.zeros(input_shape, dtype=self._rdtype) 426 # SSFM step counter 427 iterator = tf.constant(0, dtype=tf.int32, name="step_counter") 428 429 if self._n_ssfm == -1: # adaptive step width 430 431 x, _, _, _ = tf.while_loop( 432 self._cond_adaptive, 433 self._adaptive_step, 434 (x, (window, tf.cast(0.,self._rdtype), zeros, f), self._length, iterator), 435 swap_memory=self._swap_memory, 436 parallel_iterations=1 437 ) 438 439 # constant step size 440 else: 441 # Spatial step size 442 dz = tf.cast(self._dz, dtype=self._rdtype) 443 444 dz_half = dz/tf.cast(2.0, self._rdtype) 445 446 # Symmetric SSFM 447 # Start with half linear propagation 448 x = self._apply_linear_operator(x, dz_half, zeros, f) 449 # Proceed with N_SSFM-1 steps applying nonlinear and linear operator 450 x, _, _, _ = tf.while_loop( 451 self._cond, 452 self._step, 453 (x, (window, dz, zeros, f), self._n_ssfm-1, iterator), 454 swap_memory=self._swap_memory, 455 parallel_iterations=1 456 ) 457 # Final nonlinear operator 458 x = self._apply_nonlinear_operator(x, dz, zeros) 459 # Final noise application 460 x = self._apply_noise(x, dz) 461 # End with half linear propagation 462 x = self._apply_linear_operator(x, dz_half, zeros, f) 463 464 return x