encoding.py (8840B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 """Layers for encoding of linear codes.""" 6 7 import tensorflow as tf 8 from tensorflow.keras.layers import Layer 9 from sionna.fec.utils import pcm2gm 10 import numbers # to check if n, k are numbers 11 12 class LinearEncoder(Layer): 13 # pylint: disable=line-too-long 14 r"""LinearEncoder(enc_mat, is_pcm=False, dtype=tf.float32, **kwargs) 15 16 Linear binary encoder for a given generator or parity-check matrix ``enc_mat``. 17 18 If ``is_pcm`` is True, ``enc_mat`` is interpreted as parity-check 19 matrix and internally converted to a corresponding generator matrix. 20 21 The class inherits from the Keras layer class and can be used as layer in a 22 Keras model. 23 24 Parameters 25 ---------- 26 enc_mat : [k, n] or [n-k, n], ndarray 27 Binary generator matrix of shape `[k, n]`. If ``is_pcm`` is 28 True, ``enc_mat`` is interpreted as parity-check matrix of shape 29 `[n-k, n]`. 30 31 dtype: tf.DType 32 Defaults to `tf.float32`. Defines the datatype for the output dtype. 33 34 Input 35 ----- 36 inputs: [...,k], tf.float32 37 2+D tensor containing information bits. 38 39 Output 40 ------ 41 : [...,n], tf.float32 42 2+D tensor containing codewords with same shape as inputs, except the 43 last dimension changes to `[...,n]`. 44 45 Raises 46 ------ 47 AssertionError 48 If the encoding matrix is not a valid binary 2-D matrix. 49 50 Note 51 ---- 52 If ``is_pcm`` is True, this layer uses 53 :class:`~sionna.fec.utils.pcm2gm` to find the generator matrix for 54 encoding. Please note that this imposes a few constraints on the 55 provided parity-check matrix such as full rank and it must be binary. 56 57 Note that this encoder is generic for all binary linear block codes 58 and, thus, cannot implement any code specific optimizations. As a 59 result, the encoding complexity is :math:`O(k^2)`. Please consider code 60 specific encoders such as the 61 :class:`~sionna.fec.polar.encoding.Polar5GEncoder` or 62 :class:`~sionna.fec.ldpc.encoding.LDPC5GEncoder` for an improved 63 encoding performance. 64 """ 65 66 def __init__(self, 67 enc_mat, 68 is_pcm=False, 69 dtype=tf.float32, 70 **kwargs): 71 72 super().__init__(dtype=dtype, **kwargs) 73 74 # tf.int8 currently not supported by tf.matmult 75 assert (dtype in 76 (tf.float16, tf.float32, tf.float64, tf.int32, tf.int64)), \ 77 "Unsupported dtype." 78 79 # check input values for consistency 80 assert isinstance(is_pcm, bool), \ 81 'is_parity_check must be bool.' 82 83 # verify that enc_mat is binary 84 assert ((enc_mat==0) | (enc_mat==1)).all(), "enc_mat is not binary." 85 assert (len(enc_mat.shape)==2), "enc_mat must be 2-D array." 86 87 # in case parity-check matrix is provided, convert to generator matrix 88 if is_pcm: 89 self._gm = pcm2gm(enc_mat, verify_results=True) 90 else: 91 self._gm = enc_mat 92 93 self._k = self._gm.shape[0] 94 self._n = self._gm.shape[1] 95 self._coderate = self._k / self._n 96 97 assert (self._k<=self._n), "Invalid matrix dimensions." 98 99 self._gm = tf.cast(self._gm, dtype=self.dtype) 100 101 ######################################### 102 # Public methods and properties 103 ######################################### 104 105 @property 106 def k(self): 107 """Number of information bits per codeword.""" 108 return self._k 109 110 @property 111 def n(self): 112 "Codeword length." 113 return self._n 114 115 @property 116 def gm(self): 117 "Generator matrix used for encoding." 118 return self._gm 119 120 @property 121 def coderate(self): 122 """Coderate of the code.""" 123 return self._coderate 124 125 ######################### 126 # Keras layer functions 127 ######################### 128 129 def build(self, input_shape): 130 """Nothing to build, but check for valid shapes.""" 131 assert input_shape[-1]==self._k, "Invalid input shape." 132 assert (len(input_shape)>=2), 'The inputs must have at least rank 2.' 133 134 def call(self, inputs): 135 """Generic encoding function based on generator matrix multiplication. 136 """ 137 138 c = tf.linalg.matmul(inputs, self._gm) 139 140 # faster implementation of tf.math.mod(c, 2) 141 c_uint8 = tf.cast(c, tf.uint8) 142 c_bin = tf.bitwise.bitwise_and(c_uint8, tf.constant(1, tf.uint8)) 143 c = tf.cast(c_bin, self.dtype) 144 145 return c 146 147 class AllZeroEncoder(Layer): 148 r"""AllZeroEncoder(k, n, dtype=tf.float32, **kwargs) 149 Dummy encoder that always outputs the all-zero codeword of length ``n``. 150 151 Note that this encoder is a dummy encoder and does NOT perform real 152 encoding! 153 154 The class inherits from the Keras layer class and can be used as layer in a 155 Keras model. 156 157 Parameters 158 ---------- 159 k: int 160 Defining the number of information bit per codeword. 161 162 n: int 163 Defining the desired codeword length. 164 165 dtype: tf.DType 166 Defaults to `tf.float32`. Defines the datatype for internal 167 calculations and the output dtype. 168 169 Input 170 ----- 171 inputs: [...,k], tf.float32 172 2+D tensor containing arbitrary values (not used!). 173 174 Output 175 ------ 176 : [...,n], tf.float32 177 2+D tensor containing all-zero codewords. 178 179 Raises 180 ------ 181 AssertionError 182 ``k`` and ``n`` must be positive integers and ``k`` must be smaller 183 (or equal) than ``n``. 184 185 Note 186 ---- 187 As the all-zero codeword is part of any linear code, it is often used 188 to simulate BER curves of arbitrary (LDPC) codes without the need of 189 having access to the actual generator matrix. However, this `"all-zero 190 codeword trick"` requires symmetric channels (such as BPSK), otherwise 191 scrambling is required (cf. [Pfister]_ for further details). 192 193 This encoder is a dummy encoder that is needed for some all-zero 194 codeword simulations independent of the input. It does NOT perform 195 real encoding although the information bits are taken as input. 196 This is just to ensure compatibility with other encoding layers. 197 """ 198 199 def __init__(self, 200 k, 201 n, 202 dtype=tf.float32, 203 **kwargs): 204 205 super().__init__(dtype=dtype, **kwargs) 206 207 #assert error if r>1 or k,n are negativ 208 assert isinstance(k, numbers.Number), "k must be a number." 209 assert isinstance(n, numbers.Number), "n must be a number." 210 k = int(k) # k or n can be float (e.g. as result of n=k*r) 211 n = int(n) # k or n can be float (e.g. as result of n=k*r) 212 assert k>-1, "k cannot be negative." 213 assert n>-1, "n cannot be negative." 214 assert n>=k, "Invalid coderate (>1)." 215 # init encoder parameters 216 self._k = k 217 self._n = n 218 self._coderate = k / n 219 220 ######################################### 221 # Public methods and properties 222 ######################################### 223 224 @property 225 def k(self): 226 """Number of information bits per codeword.""" 227 return self._k 228 229 @property 230 def n(self): 231 "Codeword length." 232 return self._n 233 234 @property 235 def coderate(self): 236 """Coderate of the LDPC code.""" 237 return self._coderate 238 239 ######################### 240 # Keras layer functions 241 ######################### 242 243 def build(self, input_shape): 244 """Nothing to build.""" 245 pass 246 247 def call(self, inputs): 248 """Encoding function that outputs the all-zero codeword. 249 250 This function returns the all-zero codeword of shape `[..., n]`. 251 Note that this encoder is a dummy encoder and does NOT perform real 252 encoding! 253 254 Args: 255 inputs (tf.float32): Tensor of arbitrary shape. 256 257 Returns: 258 `tf.float32`: Tensor of shape `[...,n]`. 259 260 Note: 261 This encoder is a dummy encoder that is needed for some all-zero 262 codeword simulations independent of the input. It does NOT perform 263 real encoding although the information bits are taken as input. 264 This is just to ensure compatibility with other encoding layers. 265 """ 266 # keep shape of first dimensions 267 # return an all-zero tensor of shape [..., n] 268 output_shape = tf.concat([tf.shape(inputs)[:-1], 269 tf.constant(self._n, shape=[1])], 270 0) 271 c = tf.zeros(output_shape, dtype=super().dtype) 272 return c