utils.py (7256B)
1 # 2 # SPDX-FileCopyrightText: Copyright (c) 2021-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved. 3 # SPDX-License-Identifier: Apache-2.0 4 # 5 """Layer for utility functions needed for Convolutional Codes.""" 6 7 import numpy as np 8 import tensorflow as tf 9 from sionna.fec.utils import int2bin, bin2int 10 11 def polynomial_selector(rate, constraint_length): 12 """Returns generator polynomials for given code parameters. The 13 polynomials are chosen from [Moon]_ which are tabulated by searching 14 for polynomials with best free distances for a given rate and 15 constraint length. 16 17 Input 18 ----- 19 rate: float 20 Desired rate of the code. 21 Currently, only r=1/3 and r=1/2 are supported. 22 23 constraint_length: int 24 Desired constraint length of the encoder 25 26 Output 27 ------ 28 : tuple 29 Tuple of strings with each string being a 0,1 sequence where 30 each polynomial is represented in binary form. 31 32 """ 33 34 assert(isinstance(constraint_length, int)),\ 35 "constraint_length must be int." 36 assert(2 < constraint_length < 9),\ 37 "Unsupported constraint_length." 38 39 assert(rate in (1/2, 1/3)), "Unsupported rate." 40 41 rate_half_dict = { 42 3: ('101', '111'), # (5,7) 43 4: ('1101', '1011'), # (15, 13) 44 5: ('10011', '11011'), # (23, 33) # taken from GSM05.03, 4.1.3 45 6: ('110101', '101111'), # (65, 57) 46 7: ('1011011', '1111001'), # (133, 171) 47 8: ('11100101', '10011111'), # (345, 237) 48 } 49 rate_third_dict = { 50 3: ('101', '111', '111'), # (5,7,7) 51 4: ('1011', '1101', '1111'),# (54, 64, 74) 52 5: ('10101', '11011', '11111'), # (52, 66, 76) 53 6: ('100111', '101011', '111101'), # (47,53,75) 54 7: ('1111001','1100101','1011011'), # (554, 744) 55 8: ('10010101', '11011001', '11110111') # (452, 662, 756) 56 } 57 58 gen_poly_dict = { 59 1/2: rate_half_dict, 60 1/3: rate_third_dict 61 } 62 gen_poly = gen_poly_dict[rate][constraint_length] 63 return gen_poly 64 65 66 class Trellis(object): 67 """Trellis(gen_poly, rsc=True) 68 69 Trellis structure for a given generator polynomial. Defines 70 state transitions and output symbols (and bits) for each current 71 state and input. 72 73 Parameters 74 ---------- 75 gen_poly: tuple 76 Sequence of strings with each string being a 0,1 sequence. 77 If `None`, ``rate`` and ``constraint_length`` must be provided. If 78 `rsc` is True, then first polynomial will act as denominator for 79 the remaining generator polynomials. For e.g., ``rsc`` = `True` and 80 ``gen_poly`` = (`111`, `101`, `011`) implies generator matrix equals 81 :math:`G(D)=[\\frac{1+D^2}{1+D+D^2}, \\frac{D+D^2}{1+D+D^2}]`. 82 Currently Trellis is only implemented for generator matrices of 83 size :math:`\\frac{1}{n}`. 84 85 rsc: boolean 86 Boolean flag indicating whether the Trellis is recursive systematic 87 or not. If `True`, the encoder is recursive systematic in which 88 case first polynomial in ``gen_poly`` is used as the feedback 89 polynomial. Default is `True`. 90 91 """ 92 def __init__(self, gen_poly, rsc=True): 93 self.rsc=rsc 94 self.gen_poly = gen_poly 95 self.constraint_length = len(self.gen_poly[0]) 96 97 self.conv_k = 1 98 self.conv_n = len(self.gen_poly) 99 self.ni = 2**self.conv_k 100 self.ns = 2**(self.constraint_length-1) 101 self._mu = len(gen_poly[0])-1 102 103 if self.rsc: 104 self.fb_poly = [int(x) for x in self.gen_poly[0]] 105 assert self.fb_poly[0]==1 106 assert self.conv_k==1 107 108 #For current state i and input j, state transitions i->to_nodes[i][j] 109 self.to_nodes = None 110 111 #For current state i, valid state transitions are from_nodes[i][:]-> i 112 self.from_nodes = None 113 114 # Given states i and j, Trellis emits ops[i][j] symbol if neq -1 115 self.op_mat = None 116 117 # Given next state as i, trellis emits op_by_tonode[i][:] symbols 118 self.op_by_tonode = None 119 120 # Given ip_by_tonode[i][:] bits as input, trellis transitions to State i 121 self.ip_by_tonode = None 122 123 self._generate_transitions() 124 125 def _binary_matmul(self, st): 126 """ 127 For a given state st, this method multiplies each generator 128 polynomial with st and returns the sum modulo 2 bit as output 129 """ 130 op = np.zeros(self.conv_n, int) 131 assert len(st) == len(self.gen_poly[0]) 132 for i, poly in enumerate(self.gen_poly): 133 op_int = sum(int(char)*int(poly[idx]) for idx,char in enumerate(st)) 134 op[i] = int2bin(op_int % 2, 1)[0] 135 return op 136 137 def _binary_vecmul(self, v1, v2): 138 """ 139 For given vectors v1, v2, this method multiplies the two binary vectors 140 with each other and returns binary output i.e. sum modulo 2. 141 """ 142 assert len(v1) == len(v2) 143 op_int = sum(x*int(v2[idx]) for idx,x in enumerate(v1)) 144 op = int2bin(op_int, 1)[0] 145 return op 146 147 def _generate_transitions(self): 148 """Utility method that generates state transitions for different 149 input symbols. This depends only on constraint_length and independent 150 of the generator polynomials. 151 """ 152 to_nodes = np.full((self.ns, self.ni), -1, int) 153 from_nodes = np.full((self.ns, self.ni), -1, int) 154 op_mat = np.full((self.ns, self.ns), -1, int) 155 ip_by_tonode = np.full((self.ns, self.ni), -1, int) 156 op_by_tonode = np.full((self.ns, self.ni), -1, int) 157 op_by_fromnode = np.full((self.ns, self.ni), -1, int) 158 159 from_nodes_ctr = np.zeros(self.ns, int) 160 for i in range(self.ni): 161 ip_bit = int2bin(i, self.conv_k)[0] 162 for j in range(self.ns): 163 curr_st_bits = int2bin(j, self.constraint_length-1) 164 if self.rsc: 165 fb_bit = self._binary_vecmul(curr_st_bits, self.fb_poly[1:]) 166 new_bit = int2bin(ip_bit + fb_bit, 1)[0] 167 else: 168 new_bit = ip_bit 169 state_bits = [new_bit] + curr_st_bits 170 j_to = bin2int(state_bits[:-1]) 171 172 to_nodes[j][i] = j_to 173 from_nodes[j_to][from_nodes_ctr[j_to]] = j 174 175 op_bits = self._binary_matmul(state_bits) 176 op_sym = bin2int(op_bits) 177 op_mat[j, j_to] = op_sym 178 op_by_tonode[j_to, from_nodes_ctr[j_to]] = op_sym 179 ip_by_tonode[j_to, from_nodes_ctr[j_to]] = i 180 op_by_fromnode[j][i] = op_sym 181 from_nodes_ctr[j_to] += 1 182 183 self.to_nodes = tf.convert_to_tensor(to_nodes, dtype=tf.int32) 184 self.from_nodes = tf.convert_to_tensor(from_nodes, dtype=tf.int32) 185 self.op_mat = tf.convert_to_tensor(op_mat, dtype=tf.int32) 186 self.ip_by_tonode = tf.convert_to_tensor(ip_by_tonode, dtype=tf.int32) 187 self.op_by_tonode = tf.convert_to_tensor(op_by_tonode, dtype=tf.int32) 188 self.op_by_fromnode = tf.convert_to_tensor( 189 op_by_fromnode, dtype=tf.int32)