apply_ofdm_channel.py (3510B)
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 applying OFDM channel: single-tap channel response in the frequency 7 domain 8 """ 9 10 import tensorflow as tf 11 12 from sionna.utils import expand_to_rank 13 from .awgn import AWGN 14 15 class ApplyOFDMChannel(tf.keras.layers.Layer): 16 # pylint: disable=line-too-long 17 r"""ApplyOFDMChannel(add_awgn=True, dtype=tf.complex64, **kwargs) 18 19 Apply single-tap channel frequency responses to channel inputs. 20 21 This class inherits from the Keras `Layer` class and can be used as layer 22 in a Keras model. 23 24 For each OFDM symbol :math:`s` and subcarrier :math:`n`, the single-tap channel 25 is applied 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 (``h_freq``), 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 Parameters 38 ---------- 39 40 add_awgn : bool 41 If set to `False`, no white Gaussian noise is added. 42 Defaults to `True`. 43 44 dtype : tf.DType 45 Complex datatype to use for internal processing and output. Defaults to 46 `tf.complex64`. 47 48 Input 49 ----- 50 51 (x, h_freq, no) or (x, h_freq): 52 Tuple: 53 54 x : [batch size, num_tx, num_tx_ant, num_ofdm_symbols, fft_size], tf.complex 55 Channel inputs 56 57 h_freq : [batch size, num_rx, num_rx_ant, num_tx, num_tx_ant, num_ofdm_symbols, fft_size], tf.complex 58 Channel frequency responses 59 60 no : Scalar or Tensor, tf.float 61 Scalar or tensor whose shape can be broadcast to the shape of the 62 channel outputs: 63 [batch size, num_rx, num_rx_ant, num_ofdm_symbols, fft_size]. 64 Only required if ``add_awgn`` is set to `True`. 65 The noise power ``no`` is per complex dimension. If ``no`` is a 66 scalar, noise of the same variance will be added to the outputs. 67 If ``no`` is a tensor, it must have a shape that can be broadcast to 68 the shape of the channel outputs. This allows, e.g., adding noise of 69 different variance to each example in a batch. If ``no`` has a lower 70 rank than the channel outputs, then ``no`` will be broadcast to the 71 shape of the channel outputs by adding dummy dimensions after the 72 last axis. 73 74 Output 75 ------- 76 y : [batch size, num_rx, num_rx_ant, num_ofdm_symbols, fft_size], tf.complex 77 Channel outputs 78 """ 79 80 def __init__(self, add_awgn=True, dtype=tf.complex64, **kwargs): 81 82 super().__init__(trainable=False, dtype=dtype, **kwargs) 83 84 self._add_awgn = add_awgn 85 86 def build(self, input_shape): #pylint: disable=unused-argument 87 88 if self._add_awgn: 89 self._awgn = AWGN(dtype=self.dtype) 90 91 def call(self, inputs): 92 93 if self._add_awgn: 94 x, h_freq, no = inputs 95 else: 96 x, h_freq = inputs 97 98 # Apply the channel response 99 x = expand_to_rank(x, h_freq.shape.rank, axis=1) 100 y = tf.reduce_sum(tf.reduce_sum(h_freq*x, axis=4), axis=3) 101 102 # Add AWGN if requested 103 if self._add_awgn: 104 y = self._awgn((y, no)) 105 106 return y