#
# SPDX-FileCopyrightText: Copyright (c) 2021-2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
"""Layer for Turbo Code Encoding."""
import math
import tensorflow as tf
from tensorflow.keras.layers import Layer
from sionna.fec import interleaving
from sionna.fec.utils import bin2int_tf, int2bin_tf
from sionna.fec.conv.encoding import ConvEncoder
from sionna.fec.conv.utils import Trellis
from sionna.fec.turbo.utils import polynomial_selector, puncture_pattern, TurboTermination
[docs]class TurboEncoder(Layer):
# pylint: disable=line-too-long
r"""TurboEncoder(gen_poly=None, constraint_length=3, rate=1/3, terminate=False, interleaver_type='3GPP', output_dtype=tf.float32, **kwargs)
Performs encoding of information bits to a Turbo code codeword [Berrou]_.
Implements the standard Turbo code framework [Berrou]_: Two identical
rate-1/2 convolutional encoders :class:`~sionna.fec.conv.encoding.ConvEncoder`
are combined to produce a rate-1/3 Turbo code. Further,
puncturing to attain a rate-1/2 Turbo code is supported.
The class inherits from the Keras layer class and can be used as layer in a
Keras model.
Parameters
----------
gen_poly: tuple
Tuple of strings with each string being a 0,1 sequence. If
`None`, ``constraint_length`` must be provided.
constraint_length: int
Valid values are between 3 and 6 inclusive. Only required if
``gen_poly`` is `None`.
rate: float
Valid values are 1/3 and 1/2. Note that ``rate`` here denotes
the `design` rate of the Turbo code. If ``terminate`` is `True`, a
small rate-loss occurs.
terminate: boolean
Underlying convolutional encoders are terminated to all zero state
if `True`. If terminated, the true rate of the code is slightly lower
than ``rate``.
interleaver_type: str
Valid values are `"3GPP"` or `"random"`. Determines the choice of
the interleaver to interleave the message bits before input to the
second convolutional encoder. If `"3GPP"`, the Turbo code interleaver
from the 3GPP LTE standard [3GPPTS36212_Turbo]_ is used. If `"random"`,
a random interleaver is used.
output_dtype: tf.DType
Defaults to `tf.float32`. Defines the output datatype of the layer.
Input
-----
inputs : [...,k], tf.float32
2+D tensor of information bits where `k` is the information length
Output
------
: `[...,k/rate]`, tf.float32
2+D tensor where `rate` is provided as input
parameter. The output is the encoded codeword for the input
information tensor. When ``terminate`` is `True`, the effective rate
of the Turbo code is slightly less than ``rate``.
Note
----
Various notations are used in literature to represent the generator
polynomials for convolutional codes. For simplicity
:class:`~sionna.fec.turbo.encoding.TurboEncoder` only
accepts the binary format, i.e., `10011`, for the ``gen_poly`` argument
which corresponds to the polynomial :math:`1 + D^3 + D^4`.
Note that Turbo codes require the underlying convolutional encoders
to be recursive systematic encoders. Only then the channel output
from the systematic part of the first encoder can be used to decode
the second encoder.
Also note that ``constraint_length`` and ``memory`` are two different
terms often used to denote the strength of the convolutional code. In
this sub-package we use ``constraint_length``. For example, the polynomial
`10011` has a ``constraint_length`` of 5, however its ``memory`` is
only 4.
When ``terminate`` is `True`, the true rate of the Turbo code is
slightly lower than ``rate``. It can be computed as
:math:`\frac{k}{\frac{k}{r}+\frac{4\mu}{3r}}` where `r` denotes
``rate`` and :math:`\mu` is the ``constraint_length`` - 1. For example, in
3GPP, ``constraint_length`` = 4, ``terminate`` = `True`, for
``rate`` = 1/3, true rate is equal to :math:`\frac{k}{3k+12}` .
"""
def __init__(self,
gen_poly=None,
constraint_length=3,
rate=1/3,
terminate=False,
interleaver_type='3GPP',
output_dtype=tf.float32,
**kwargs):
super().__init__(**kwargs)
if gen_poly is not None:
assert all(isinstance(poly, str) for poly in gen_poly), \
"Each element of gen_poly must be a string."
assert all(len(poly)==len(gen_poly[0]) for poly in gen_poly), \
"Each polynomial must be of same length."
assert all(all(
char in ['0','1'] for char in poly) for poly in gen_poly),\
"Each Polynomial must be a string of 0/1 s."
assert len(gen_poly)==2, \
"Generator polynomials need to be of Rate-1/2 "
self._gen_poly = gen_poly
else:
valid_constraint_length = (3, 4, 5, 6)
assert constraint_length in valid_constraint_length, \
"Constraint length must be between 3 and 6."
self._gen_poly = polynomial_selector(constraint_length)
valid_rates = (1/2, 1/3)
assert rate in valid_rates, "Invalid coderate."
assert isinstance(terminate, bool), "terminate must be bool."
assert interleaver_type in ('3GPP', 'random'),\
"Invalid interleaver_type."
self._coderate_desired = rate
self._coderate = self._coderate_desired
self._terminate = terminate
self._interleaver_type = interleaver_type
self.output_dtype = output_dtype
# Underlying convolutional encoders to be rsc or not
rsc = True
self._coderate_conv = 1/len(self.gen_poly)
self._punct_pattern = puncture_pattern(rate, self._coderate_conv)
self._trellis = Trellis(self.gen_poly, rsc=rsc)
self._mu = self.trellis._mu
# conv_n denotes number of output bits for conv_k input bits.
self._conv_k = self._trellis.conv_k
self._conv_n = self._trellis.conv_n
self._ni = 2**self._conv_k
self._no = 2**self._conv_n
self._ns = self._trellis.ns
self._k = None
self._n = None
if self.terminate:
self.turbo_term = TurboTermination(self._mu+1, conv_n=self._conv_n)
if self._interleaver_type == '3GPP':
self.internal_interleaver = interleaving.Turbo3GPPInterleaver()
else:
self.internal_interleaver = interleaving.RandomInterleaver(
keep_batch_constant=True,
keep_state=True,
axis=-1)
if self.punct_pattern is not None:
self.punct_idx = tf.where(self.punct_pattern)
self.convencoder = ConvEncoder(gen_poly=self._gen_poly,
rsc=rsc,
terminate=self._terminate)
#########################################
# Public methods and properties
#########################################
@property
def gen_poly(self):
"""Generator polynomial used by the encoder"""
return self._gen_poly
@property
def constraint_length(self):
"""Constraint length of the encoder"""
return self._mu + 1
@property
def coderate(self):
"""Rate of the code used in the encoder"""
if self.terminate and self._k is None:
print("Note that, due to termination, the true coderate is lower "\
"than the returned design rate. "\
"The exact true rate is dependent on the value of k and "\
"hence cannot be computed before the first call().")
elif self.terminate and self._k is not None:
term_factor = 1+math.ceil(4*self._mu/3)/self._k
self._coderate = self._coderate_desired/term_factor
return self._coderate
@property
def trellis(self):
"""Trellis object used during encoding"""
return self._trellis
@property
def terminate(self):
"""Indicates if the convolutional encoders are terminated"""
return self._terminate
@property
def punct_pattern(self):
"""Puncturing pattern for the Turbo codeword"""
return self._punct_pattern
@property
def k(self):
"""Number of information bits per codeword"""
if self._k is None:
print("Note: The value of k cannot be computed before the first " \
"call().")
return self._k
@property
def n(self):
"""Number of codeword bits"""
if self._n is None:
print("Note: The value of n cannot be computed before the first " \
"call().")
return self._n
def _conv_enc(self, info_vec, terminate):
"""
This method encodes the information tensor info_vec using the
underlying convolutional encoder. Returns the encoded codeword tensor
array ta, and the tensor array containing termination bits ta_term.
If the terminate variable is False, ta_term is array of length 0.
"""
msg = tf.cast(info_vec, tf.int32)
msg_reshaped = tf.reshape(msg, [-1, self._k])
term_syms = int(self._mu) if terminate else 0
prev_st = tf.zeros([tf.shape(msg_reshaped)[0]], tf.int32)
ta = tf.TensorArray(tf.int32, size=self.num_syms, dynamic_size=False)
idx_offset = range(0, self._conv_k)
for idx in tf.range(0, self._k, self._conv_k):
msg_bits_idx = tf.gather(msg_reshaped,
idx + idx_offset,
axis=-1)
#msg_bits_idx = tf.experimental.numpy.take_along_axis(msg_reshaped)
msg_idx = bin2int_tf(msg_bits_idx)
indices = tf.stack([prev_st, msg_idx], -1)
new_st = tf.gather_nd(self._trellis.to_nodes, indices=indices)
idx_syms = tf.gather_nd(self._trellis.op_mat,
tf.stack([prev_st, new_st], -1))
idx_bits = int2bin_tf(idx_syms, self._conv_n)
ta = ta.write(idx//self._conv_k, idx_bits)
prev_st = new_st
ta_term = tf.TensorArray(tf.int32, size=term_syms, dynamic_size=False)
# Termination
if terminate:
fb_poly = tf.constant([int(x) for x in self.gen_poly[0][1:]])
fb_poly_tiled = tf.tile(
tf.expand_dims(fb_poly,0),[tf.shape(prev_st)[0],1])
for idx in tf.range(0, term_syms, self._conv_k):
prev_st_bits = int2bin_tf(prev_st, self._mu)
msg_idx = tf.math.reduce_sum(
tf.multiply(fb_poly_tiled, prev_st_bits),-1)
msg_idx = tf.squeeze(int2bin_tf(msg_idx,1),-1)
indices = tf.stack([prev_st, msg_idx], -1)
new_st = tf.gather_nd(self._trellis.to_nodes, indices=indices)
idx_syms = tf.gather_nd(self._trellis.op_mat,
tf.stack([prev_st, new_st], -1))
idx_bits = int2bin_tf(idx_syms, self._conv_n)
ta_term = ta_term.write(idx//self._conv_k, idx_bits)
prev_st = new_st
return ta, ta_term
def _puncture_cw(self, cw):
"""
Given the codeword ``cw``, this method punctures ``cw`` using the
puncturing pattern defined in self.punct_pattern. A simple tile
operation of self.punct_pattern followed by tf.boolean_mask(cw, mask_)
works. However this fails in XLA mode as the dimension of the above
operation is unknown.
Hence, idx is obtained from `tf.where(self.punct_pattern)` during
initialization. This way the dimension of idx is known during graph
creation. Then during the call(), idx is tiled followed by row offset
addition to idx (the indices tensor) will achieve the same result as
applying a tiled boolean_mask.
"""
# cw shape: (bs, n, 3)- transpose to (n, 3, bs)
cw = tf.transpose(cw, perm=[1, 2, 0])
cw_n = cw.get_shape()[0]
punct_period = self.punct_pattern.shape[0]
mask_reps = cw_n//punct_period
idx = tf.tile(self.punct_idx, [mask_reps, 1])
idx_per_period = self.punct_idx.shape[0]
idx_per_time = idx_per_period/punct_period
# When tiling punct_pattern doesn't cover cw, delta_times > 0
delta_times = cw_n - (mask_reps * punct_period)
delta_idx_rows = int(delta_times*idx_per_time)
time_offset = punct_period * tf.range(mask_reps)[None,:]
row_idx = tf.transpose(tf.tile(time_offset,[idx_per_period,1]))
row_idx = tf.reshape(row_idx, (-1, 1))
total_indices = mask_reps*idx_per_period + delta_idx_rows
col_idx = tf.zeros((total_indices,1), tf.int32)
if delta_times > 0:
idx = tf.concat([idx, self.punct_idx[:delta_idx_rows]], axis=0)
# Additional index row offsets if delta_times > 0
time_n = punct_period*mask_reps
row_idx_delta = tf.tile(
tf.range(time_n, time_n+delta_times)[None, :],
[delta_idx_rows, 1])
row_idx = tf.concat([row_idx, row_idx_delta], axis=0)
idx_offset = tf.cast(tf.concat([row_idx, col_idx], axis=1), tf.int64)
idx = tf.add(idx, idx_offset)
cw = tf.gather_nd(cw, idx)
cw = tf.transpose(cw)
return cw
#########################
# Keras layer functions
#########################
def build(self, input_shape):
"""Build layer and check dimensions.
Args:
input_shape: shape of input tensor (...,k).
"""
self._k = input_shape[-1]
self._n = int(self._k/self._coderate_desired)
if self._interleaver_type == '3GPP':
assert self._k <= 6144, '3GPP Turbo Codes define Interleavers only\
upto frame lengths of 6144'
# Num. of encoding periods/state transitions.
# Not equal to _k if_conv_k>1.
self.num_syms = int(self._k//self._conv_k)
def call(self, inputs):
"""Turbo code encoding function.
Args:
inputs (tf.float32): Information tensor of shape `[...,k]`.
Returns:
`tf.float32`: Encoded codeword tensor of shape `[...,n]`.
"""
tf.debugging.assert_greater(tf.rank(inputs), 1)
if inputs.shape[-1] != self._k:
self.build(inputs.shape)
if self._terminate:
num_term_bits_ = int(
self.turbo_term.get_num_term_syms()/self._coderate_conv)
num_term_bits_punct = int(
num_term_bits_*self._coderate_conv/self._coderate_desired)
else:
num_term_bits_ = 0
num_term_bits_punct = 0
output_shape = inputs.get_shape().as_list()
output_shape[0] = -1
output_shape[-1] = self._n + num_term_bits_punct
preterm_n = int(self._k/self._coderate_conv)
msg = tf.cast(tf.reshape(inputs, [-1, self._k]), tf.int32)
msg2 = self.internal_interleaver(msg)
cw1_ = self.convencoder(msg)
cw2_ = self.convencoder(msg2)
cw1, term1 = cw1_[:, :preterm_n], cw1_[:, preterm_n:]
cw2, term2 = cw2_[:, :preterm_n], cw2_[:, preterm_n:]
# Gather parity stream from 2nd enc
par_idx = tf.range(1, preterm_n, delta=self._conv_n)
cw2_par = tf.gather(cw2, indices=par_idx, axis=-1)
cw1 = tf.reshape(cw1,(-1, self._k, self._conv_n))
cw2_par = tf.reshape(cw2_par, (-1, self._k, 1))
# Concatenate 2nd enc parity to _conv_n streams from first encoder
cw = tf.concat([cw1, cw2_par], axis=-1)
if self.terminate:
term_syms_turbo = self.turbo_term.termbits_conv2turbo(term1, term2)
term_syms_turbo = tf.reshape(
term_syms_turbo, (-1, num_term_bits_//2, 3))
cw = tf.concat([cw, term_syms_turbo], axis=-2)
if self.punct_pattern is not None:
cw = self._puncture_cw(cw)
cw = tf.cast(cw, self.output_dtype)
cw_reshaped = tf.reshape(cw, output_shape)
return cw_reshaped