ofdm_channel.py (5663B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 """ 6 Layer for implementing an ideal OFDM channel response, i.e., single-tap 7 channel response in the frequency domain 8 """ 9 10 import tensorflow as tf 11 from tensorflow.keras.layers import Layer 12 13 from . import GenerateOFDMChannel, ApplyOFDMChannel 14 15 class OFDMChannel(Layer): 16 # pylint: disable=line-too-long 17 r"""OFDMChannel(channel_model, resource_grid, add_awgn=True, normalize_channel=False, return_channel=False, dtype=tf.complex64, **kwargs) 18 19 Generate channel frequency responses and apply them to channel inputs 20 assuming an OFDM waveform with no ICI nor ISI. 21 22 This class inherits from the Keras `Layer` class and can be used as layer 23 in a Keras model. 24 25 For each OFDM symbol :math:`s` and subcarrier :math:`n`, the channel output is computed as follows: 26 27 .. math:: 28 y_{s,n} = \widehat{h}_{s, n} x_{s,n} + w_{s,n} 29 30 where :math:`y_{s,n}` is the channel output computed by this layer, 31 :math:`\widehat{h}_{s, n}` the frequency channel response, 32 :math:`x_{s,n}` the channel input ``x``, and :math:`w_{s,n}` the additive noise. 33 34 For multiple-input multiple-output (MIMO) links, the channel output is computed for each antenna 35 of each receiver and by summing over all the antennas of all transmitters. 36 37 The channel frequency response for the :math:`s^{th}` OFDM symbol and 38 :math:`n^{th}` subcarrier is computed from a given channel impulse response 39 :math:`(a_{m}(t), \tau_{m}), 0 \leq m \leq M-1` generated by the ``channel_model`` 40 as follows: 41 42 .. math:: 43 \widehat{h}_{s, n} = \sum_{m=0}^{M-1} a_{m}(s) e^{-j2\pi n \Delta_f \tau_{m}} 44 45 where :math:`\Delta_f` is the subcarrier spacing, and :math:`s` is used as time 46 step to indicate that the channel impulse response can change from one OFDM symbol to the 47 next in the event of mobility, even if it is assumed static over the duration 48 of an OFDM symbol. 49 50 Parameters 51 ---------- 52 channel_model : :class:`~sionna.channel.ChannelModel` object 53 An instance of a :class:`~sionna.channel.ChannelModel` object, such as 54 :class:`~sionna.channel.RayleighBlockFading` or 55 :class:`~sionna.channel.tr38901.UMi`. 56 57 resource_grid : :class:`~sionna.ofdm.ResourceGrid` 58 Resource grid 59 60 add_awgn : bool 61 If set to `False`, no white Gaussian noise is added. 62 Defaults to `True`. 63 64 normalize_channel : bool 65 If set to `True`, the channel is normalized over the resource grid 66 to ensure unit average energy per resource element. Defaults to `False`. 67 68 return_channel : bool 69 If set to `True`, the channel response is returned in addition to the 70 channel output. Defaults to `False`. 71 72 dtype : tf.DType 73 Complex datatype to use for internal processing and output. 74 Defaults to tf.complex64. 75 76 Input 77 ----- 78 79 (x, no) or x: 80 Tuple or Tensor: 81 82 x : [batch size, num_tx, num_tx_ant, num_ofdm_symbols, fft_size], tf.complex 83 Channel inputs 84 85 no : Scalar or Tensor, tf.float 86 Scalar or tensor whose shape can be broadcast to the shape of the 87 channel outputs: 88 [batch size, num_rx, num_rx_ant, num_ofdm_symbols, fft_size]. 89 Only required if ``add_awgn`` is set to `True`. 90 The noise power ``no`` is per complex dimension. If ``no`` is a scalar, 91 noise of the same variance will be added to the outputs. 92 If ``no`` is a tensor, it must have a shape that can be broadcast to 93 the shape of the channel outputs. This allows, e.g., adding noise of 94 different variance to each example in a batch. If ``no`` has a lower 95 rank than the channel outputs, then ``no`` will be broadcast to the 96 shape of the channel outputs by adding dummy dimensions after the last 97 axis. 98 99 Output 100 ------- 101 y : [batch size, num_rx, num_rx_ant, num_ofdm_symbols, fft_size], tf.complex 102 Channel outputs 103 h_freq : [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_ofdm_symbols, fft_size], tf.complex 104 (Optional) Channel frequency responses. Returned only if 105 ``return_channel`` is set to `True`. 106 """ 107 108 def __init__(self, channel_model, resource_grid, add_awgn=True, 109 normalize_channel=False, return_channel=False, 110 dtype=tf.complex64, **kwargs): 111 super().__init__(trainable=False, dtype=dtype, **kwargs) 112 113 self._cir_sampler = channel_model 114 self._rg = resource_grid 115 self._add_awgn = add_awgn 116 self._normalize_channel = normalize_channel 117 self._return_channel = return_channel 118 119 def build(self, input_shape): #pylint: disable=unused-argument 120 121 self._generate_channel = GenerateOFDMChannel(self._cir_sampler, 122 self._rg, 123 self._normalize_channel, 124 tf.as_dtype(self.dtype)) 125 self._apply_channel = ApplyOFDMChannel( self._add_awgn, 126 tf.as_dtype(self.dtype)) 127 128 def call(self, inputs): 129 130 if self._add_awgn: 131 x, no = inputs 132 else: 133 x = inputs 134 135 h_freq = self._generate_channel(tf.shape(x)[0]) 136 if self._add_awgn: 137 y = self._apply_channel([x, h_freq, no]) 138 else: 139 y = self._apply_channel([x, h_freq]) 140 141 if self._return_channel: 142 return y, h_freq 143 else: 144 return y