Refactor code for 5G standard-compliant cyclic prefix length

Author: danielschaeufele

sionna/nr/carrier_config.py

@@ -6,6 +6,8 @@
 """
 # pylint: disable=line-too-long
 
+import numpy as np
+
 from .config import Config
 
 class CarrierConfig(Config):
@@ -244,27 +246,22 @@ def kappa(self):
     @property
     def cyclic_prefix_length(self):
         r"""
-        float, read-only : Cyclic prefix length of all symbols except for the
-            first symbol in each half subframe
-            :math:`N_{\text{CP},l}^{\mu} \cdot T_{\text{c}}` [s]
-        """
-        if self.cyclic_prefix=="extended":
-            cp = 512*self.kappa*2**(-self.mu)
-        else:
-            cp = 144*self.kappa*2**(-self.mu)
-        return cp*self.t_c
-
-    @property
-    def cyclic_prefix_length_first_symbol(self):
-        r"""
-        float, read-only : Cyclic prefix length of first symbol in each
-            half subframe
+        np.ndarray[float], read-only : Vector of cyclic prefix length of all
+            symbols in the current slot as defined in Section 5.3.1 [3GPP38211]_
             :math:`N_{\text{CP},l}^{\mu} \cdot T_{\text{c}}` [s]
         """
+        cp = np.zeros(self.num_symbols_per_slot)
         if self.cyclic_prefix=="extended":
-            cp = 512*self.kappa*2**(-self.mu)
+            cp[:] = 512*self.kappa*2**(-self.mu)
         else:
-            cp = 144*self.kappa*2**(-self.mu) + 16*self.kappa
+            cp[:] = 144*self.kappa*2**(-self.mu)
+
+            # Extend cyclic prefix for l=0 or l=7*2^\mu
+            long_cp_period = 7 * 2 ** self.mu
+            l_start = self.slot_number * self.num_symbols_per_slot
+            for i in range(l_start % long_cp_period,
+                           self.num_symbols_per_slot, long_cp_period):
+                cp[i] += 16*self.kappa
         return cp*self.t_c
 
     #-------------------#

sionna/nr/pusch_config.py

@@ -1050,11 +1050,8 @@ def check_pusch_configs(pusch_configs):
         "num_cdm_groups_without_data" : pc.dmrs.num_cdm_groups_without_data
     }
     params["bandwidth"] = params["num_subcarriers"]*params["subcarrier_spacing"]
-    params["cyclic_prefix_length"] = int(np.ceil(carrier.cyclic_prefix_length *
-                                             params["bandwidth"]))
-    params["cyclic_prefix_length_first_symbol"] =\
-        int(np.ceil(carrier.cyclic_prefix_length_first_symbol
-                    * params["bandwidth"]))
+    params["cyclic_prefix_length"] = np.ceil(carrier.cyclic_prefix_length *
+                                             params["bandwidth"]).astype(int)
 
     for pusch_config in pusch_configs:
         if params["precoding"]=="codebook":

sionna/nr/pusch_receiver.py

@@ -156,16 +156,10 @@ def __init__(self,
             assert l_min is not None, \
                 "l_min must be provided for input_domain==time"
             self._l_min = l_min
-            symbols_per_block = (
-                pusch_transmitter._carrier_config.num_slots_per_subframe *
-                pusch_transmitter._carrier_config.num_symbols_per_slot // 2)
             self._ofdm_demodulator = OFDMDemodulator(
                 fft_size=pusch_transmitter._num_subcarriers,
                 l_min=self._l_min,
-                cyclic_prefix_length=pusch_transmitter._cyclic_prefix_length,
-                cyclic_prefix_length_first_symbol=
-                    pusch_transmitter._cyclic_prefix_length_first_symbol,
-                symbols_per_block=symbols_per_block)
+                cyclic_prefix_length=pusch_transmitter._cyclic_prefix_length)
 
         # Use or create default ChannelEstimator
         self._perfect_csi = False

sionna/nr/pusch_transmitter.py

@@ -168,7 +168,9 @@ def __init__(self,
                             subcarrier_spacing=self._subcarrier_spacing,
                             num_tx=self._num_tx,
                             num_streams_per_tx=self._num_layers,
-                            cyclic_prefix_length=self._cyclic_prefix_length,
+                            # TODO: pass vector of cyclic prefix lengths
+                            # (requires rewrite of channel simulation code)
+                            cyclic_prefix_length=self._cyclic_prefix_length[1],
                             pilot_pattern=self._pilot_pattern,
                             dtype=dtype)
 
@@ -183,13 +185,7 @@ def __init__(self,
 
         # (Optionally) Create OFDMModulator
         if self._output_domain=="time":
-            symbols_per_block = (self._carrier_config.num_slots_per_subframe *
-                             self._carrier_config.num_symbols_per_slot // 2)
-            self._ofdm_modulator = OFDMModulator(
-                cyclic_prefix_length=self._cyclic_prefix_length,
-                cyclic_prefix_length_first_symbol=
-                    self._cyclic_prefix_length_first_symbol,
-                symbols_per_block=symbols_per_block)
+            self._ofdm_modulator = OFDMModulator(cyclic_prefix_length=self._cyclic_prefix_length)
 
     #########################################
     # Public methods and properties

sionna/ofdm/demodulator.py

@@ -15,20 +15,11 @@
 class OFDMDemodulator(Layer):
     # pylint: disable=line-too-long
     r"""
-    OFDMDemodulator(fft_size, l_min, cyclic_prefix_length=0, cyclic_prefix_length_first_symbol=None, symbols_per_block=1, **kwargs)
+    OFDMDemodulator(fft_size, l_min, cyclic_prefix_length=0, **kwargs)
 
     Computes the frequency-domain representation of an OFDM waveform
     with cyclic prefix removal.
 
-    When only `cyclic_prefix_length` is given then a cyclic prefix of length
-    `cyclic_prefix_length` is removed from each symbol. When additionally
-    `cyclic_prefix_length_first_symbol` and `symbols_per_block` are given then
-    the length of the cyclic prefix is `cyclic_prefix_length_first_symbol` for
-    the first symbol of each block and `cyclic_prefix_length` for the
-    remaining symbols. For LTE one block corresponds to one slot (i.e., 7
-    symbols). For 5G NR one block corresponds to one half subframe and the
-    number of symbols depends on the numerology.
-
     The demodulator assumes that the input sequence is generated by the
     :class:`~sionna.channel.TimeChannel`. For a single pair of antennas,
     the received signal sequence is given as:
@@ -72,17 +63,9 @@ class OFDMDemodulator(Layer):
         impulse response. It should be the same value as that used by the
         `cir_to_time_channel` function.
 
-    cyclic_prefix_length : int
-        Integer indicating the length of the cyclic prefix that it prepended
-        to each OFDM symbol (except for the first symbol of each block if
-        `cyclic_prefix_length_first_symbol` and `symbols per block` is given).
-
-    cyclic_prefix_length_first_symbol : int
-        Integer indicating the length of the cyclic prefix that it prepended
-        to the first OFDM symbol of each block.
-
-    symbols_per_block : int
-        Integer indicating the number of symbols per block.
+    cyclic_prefix_length : int or list[int] or np.ndarray[int]
+        Integer or vector of integers indicating the length of the cyclic
+        prefix that it prepended to each OFDM symbol.
 
     Input
     -----
@@ -98,15 +81,11 @@ class OFDMDemodulator(Layer):
     """
 
     def __init__(self, fft_size, l_min, cyclic_prefix_length=0,
-                 cyclic_prefix_length_first_symbol=None, symbols_per_block=1,
                  **kwargs):
         super().__init__(**kwargs)
         self.fft_size = fft_size
         self.l_min = l_min
         self.cyclic_prefix_length = cyclic_prefix_length
-        self.cyclic_prefix_length_first_symbol =(
-            cyclic_prefix_length_first_symbol)
-        self.symbols_per_block = symbols_per_block
 
     @property
     def fft_size(self):
@@ -133,34 +112,18 @@ def cyclic_prefix_length(self):
 
     @cyclic_prefix_length.setter
     def cyclic_prefix_length(self, value):
-        assert isinstance(value, int) and value >=0,\
-            "`cyclic_prefix_length` must be a nonnegative integer."
-        self._cyclic_prefix_length = int(value)
-
-    @property
-    def cyclic_prefix_length_first_symbol(self):
-        if self._cyclic_prefix_length_first_symbol is None:
-            return self._cyclic_prefix_length
+        if isinstance(value, list):
+            value = np.array(value)
+        if isinstance(value, np.ndarray):
+            assert (np.issubdtype(value.dtype, np.integer) and
+                    value.ndim == 1 and np.all(value >= 0)),\
+                ("`cyclic_prefix_length` must be a 1D array with"
+                 " only nonnegative integers.")
         else:
-            return self._cyclic_prefix_length_first_symbol
+            assert isinstance(value, int) and value >=0,\
+                "`cyclic_prefix_length` must be a nonnegative integer."
+        self._cyclic_prefix_length = value
 
-    @cyclic_prefix_length_first_symbol.setter
-    def cyclic_prefix_length_first_symbol(self, value):
-        assert (value is None or isinstance(value, int) and
-                value >= self._cyclic_prefix_length),\
-            ("`cyclic_prefix_length_first_symbol` must be integer and " +
-             "larger or equal to `cyclic_prefix_length`.")
-        self._cyclic_prefix_length_first_symbol = value
-
-    @property
-    def symbols_per_block(self):
-        return self._symbols_per_block
-
-    @symbols_per_block.setter
-    def symbols_per_block(self, value):
-        assert isinstance(value, int) and  value >= 1,\
-            "`symbols_per_block` must be a positive integer."
-        self._symbols_per_block = value
 
     def build(self, input_shape):
         num_samples = input_shape[-1]
@@ -170,24 +133,23 @@ def build(self, input_shape):
               * tf.range(self.fft_size, dtype=tf.float32)
         self._phase_compensation = tf.exp(tf.complex(0., tmp))
 
-        self._samples_per_block = (self.cyclic_prefix_length_first_symbol +
-            (self.symbols_per_block - 1) * self.cyclic_prefix_length +
-            self.symbols_per_block * self.fft_size)
-
-        # Compute number of elements that will be truncated and number of
-        # symbols for padding
-        self._rest = num_samples % self._samples_per_block
-        samples_first_symbol = (self.cyclic_prefix_length_first_symbol +
-                                self.fft_size)
-        samples_other_symbols = (self.cyclic_prefix_length + self.fft_size)
-        if self._rest > samples_first_symbol:
-            self._rest -= samples_first_symbol
-            excess_symbols = self._rest // samples_other_symbols
-            self._rest -= excess_symbols * samples_other_symbols
-            excess_symbols += 1  # Because of first symbol in block
-            self._num_pad_symbols = self.symbols_per_block - excess_symbols
+        if isinstance(self.cyclic_prefix_length, int):
+            self._num_ofdm_symbols = (input_shape[-1] //
+                          (self.fft_size + self.cyclic_prefix_length))
+            self.cyclic_prefix_length = np.full(self._num_ofdm_symbols,
+                                                self.cyclic_prefix_length)
         else:
-            self._num_pad_symbols = 0
+            self._num_ofdm_symbols = self.cyclic_prefix_length.shape[0]
+
+        symbol_ends = tf.math.cumsum(self.cyclic_prefix_length + self.fft_size)
+        assert num_samples >= symbol_ends[-1],\
+            "shape(inputs)[-1] must be larger or equal than samples per slot"
+
+        gather_idx = []
+        for i in range(self._num_ofdm_symbols):
+            gather_idx.append(tf.range(symbol_ends[i] - self.fft_size,
+                                       symbol_ends[i]))
+        self._gather_idx = tf.concat(gather_idx, 0)
 
     def call(self, inputs):
         """Demodulate OFDM waveform onto a resource grid.
@@ -202,39 +164,14 @@ def call(self, inputs):
         """
         batch_dims = tf.shape(inputs)[:-1]
 
-        # Cut last samples that do not fit into an OFDM symbol
-        x = inputs if self._rest == 0 else inputs[..., :-self._rest]
-
-        if self._num_pad_symbols > 0:
-            pad_samples = self._num_pad_symbols * (self.fft_size +
-                                                   self.cyclic_prefix_length)
-            padding_shape = tf.concat([batch_dims, [pad_samples]], axis=0)
-            padding = tf.zeros(padding_shape, dtype=x.dtype)
-            x = tf.concat([x, padding], axis=-1)
-
-        # Reshape input to blocks
-        num_blocks = tf.shape(x)[-1] // self._samples_per_block
-        new_shape = tf.concat([batch_dims,
-                               [num_blocks, self._samples_per_block]], 0)
-        x = tf.reshape(x, new_shape)
-
-        # Remove extra cyclic prefix from first symbol
-        x = x[...,(self.cyclic_prefix_length_first_symbol -
-                   self.cyclic_prefix_length):]
+        # Gather samples that do not belong to cyclic prefix
+        x = tf.gather(inputs, self._gather_idx, axis=-1)
 
         # Reshape input to separate OFDM symbols
         new_shape = tf.concat([batch_dims,
-                               [num_blocks * self.symbols_per_block],
-                               [self.fft_size + self.cyclic_prefix_length]], 0)
+                               [self._num_ofdm_symbols, self.fft_size]], 0)
         x = tf.reshape(x, new_shape)
 
-        # Remove padding
-        if self._num_pad_symbols > 0:
-            x = x[..., :-self._num_pad_symbols, :]
-
-        # Remove cyclic prefix
-        x = x[...,self.cyclic_prefix_length:]
-
         # Compute FFT
         x = fft(x)