model.py

import torch
import torch.nn as nn
import numpy as np

from outer_code import OuterCode
from utils import PositionalEncoder_fixed, Power_reallocate, power_constraint, generate_LUT, \
    convert_class_to_bitwise_matmul, convert_prior_to_probabilities
from layers import BERT, Identity, LinearLayers

class TransCoder(nn.Module):
    """
        TransCoder class for communication systems.
        This class implements the transmitter (encoder), receiver (decoder), and outer code integration as a part of TransCoder.
        It supports various configurations such as Turbo codes, Polar codes, BCH codes and LDPC codes.

        Attributes:
            configs (dict): Configuration dictionary containing model parameters.
            bs (int): Batch size.
            padd_symbol_len (int): Length of padding symbols.
            K_code (int): Codeword length.
            device (torch.device): Device to run the model on (CPU/GPU).
            outer_code_type (str): Type of outer code (e.g., Turbo, Polar_usc, LDPC_nn, BCH_nn).
            M (int): Modulation order.
            block_rate (int): Block rate for encoding.
            K_in (int): Information message length for the outer code.
            model_num_iters (int): Number of decoding iterations.
            power_norm (bool): Whether to additionally normalize power.
            no_inner (bool): Whether to disable the TransCoder.
            no_Tmodel (bool): Whether to disable the TransCoder encoder model at the transmitter.
            no_Rmodel (bool): Whether to disable the TransCoder decoder model at the receiver.
        """
    def __init__(self, configs):
        """
            Initialize the TransCoder model with the given configurations.

            Args:
                configs (dict): Configuration dictionary containing model parameters.
            """
        super(TransCoder, self).__init__()
        self.configs = configs
        self.bs = configs["batch_size"]
        self.padd_symbol_len = configs["padd_symbol_len"]
        self.K_code = configs["K"]
        self.device = configs["device"]
        self.if_complex = True
        self.outer_code_type = configs["outer_code_type"]
        self.M = configs["M"]
        self.block_rate = configs["block_rate"]
        self.K_in = configs["K_in"]
        self.batch_size = configs["batch_size"]
        self.rayleigh = configs["rayleigh"]
        padd_bit = 0

        self.K_transmit = self.K_code

        self.model_num_iters = configs["num_iters"]
        self.backprop = configs["backprop_outer"]
        self.llrs_norm_loss = configs["llrs_norm_loss"]

        self.K_extended = self.K_code
        if self.padd_symbol_len > 0:
            self.K_extended += self.padd_symbol_len
        self.num_blocks = self.K_extended // self.M
        self.NS_model = configs["NS_model"]
        self.ones_symbols_len = 0
        self.bp_iters_arr = configs["bp_iters_arr"]
        self.no_Tmodel = configs["no_Tmodel"]
        self.no_Rmodel = configs["no_Rmodel"]
        self.power_norm = configs["power_norm"]

        block_length = self.num_blocks * self.block_rate - self.padd_symbol_len

        if self.padd_symbol_len > 0:
            self.ones_symbols = torch.ones(self.bs, self.padd_symbol_len,
                                             requires_grad=False).to(self.device)

            self.added_symb = torch.zeros(self.bs, self.padd_symbol_len,
                                          requires_grad=False).to(self.device)
            print(f"Inner coded sequence now has length {self.padd_symbol_len + self.K_code} by employing "
                  f"{self.padd_symbol_len} padded with ones symbols")
        else:
            self.ones_symbols = torch.empty((self.bs, ),
                                           requires_grad=False).to(self.device)

        # Positional encoder will add bias to each index of any sequence according to indices of the elements
        max_seq_len = 200
        if configs['K'] > 500:
            max_seq_len = 300
        if configs['K'] > 1500:
            max_seq_len = 400
        if configs['K'] == 1536:
            max_seq_len = 512
        self.pe = PositionalEncoder_fixed(configs["d_model_trx"], max_seq_len=max_seq_len)

        self.power_reallocator = Power_reallocate(block_length=block_length,
                                                  padd_bit=padd_bit)



        # OUTER CODE INITIALIZATION
        self.outer_code = OuterCode(device=self.device, k=self.K_in, n=self.K_code,
                                    belief_propagation_iterations=configs["bp_iters"],
                                    hard_out=False, outer_code_type=self.outer_code_type,
                                    loss_type=configs["loss_type"], bs=self.batch_size, method_bp=configs["method_bp"],
                                    train_regime=bool(configs["train"]), ldpc_type=configs["ldpc_type"],
                                    multilabel=configs["multilabel"], weights=configs["weights"],
                                    w_n=configs["w_n"], protograph=configs["protograph"],
                                    label_smoothing=configs["label_smoothing"])

        multclass = True
        # The autoencoder will encode K bits with a single iteration since there is no feedback
        self.no_inner = configs["no_inner"]
        self.table = generate_LUT(configs['M']).to(self.device)
        if self.no_inner:
            self.Tmodel = Identity()
            self.Rmodel = Identity()

        elif configs["linear_layers"]:
            self.Tmodel = LinearLayers(mod="trx", input_size=self.M,
                               m=self.M, block_rate=self.block_rate, multclass=multclass,)
            self.Rmodel = LinearLayers(mod="rec",input_size=self.M,
                               m=self.M, block_rate=self.block_rate, multclass=multclass,)

        else:
            self.Tmodel = BERT(mod="trx", input_size=self.M,
                               m=self.M, block_rate=self.block_rate,
                               d_model=configs["d_model_trx"],
                               N=configs["N_trx"],
                               heads=configs["heads_trx"],
                               dropout=configs["dropout"],
                               custom_attn=configs["custom_attn"],
                               multclass=multclass,
                               NS_model=self.NS_model)

            self.Rmodel = BERT(mod="rec", input_size=self.block_rate, m=self.M,
                               block_rate=self.block_rate,
                               d_model=configs["d_model_trx"],
                               N=configs["N_rec"], heads=configs["heads_trx"],
                               dropout=configs["dropout"], custom_attn=configs["custom_attn"],
                               multclass=multclass, NS_model=self.NS_model,
                               af_module=configs["af_module"])
        if self.no_Tmodel:
            self.Tmodel = Identity()

        if self.outer_code_type == 'Turbo':
            self.Rmodelf = BERT(mod='rec', input_size=3, m=1,
                                block_rate=self.block_rate,
                                d_model=configs["d_model_trx"],
                                N=configs["N_rec"], heads=configs["heads_trx_rf"],
                                dropout=configs["dropout"],
                                custom_attn=configs["custom_attn"], multclass=multclass,
                                NS_model=self.NS_model, af_module=configs["af_module"])
        else:
            self.Rmodelf = BERT(mod='rec', input_size=self.block_rate + configs["heads_trx_rf"] * self.M,
                                m=self.M, block_rate=self.block_rate,
                                d_model=configs["d_model_trx"], N=configs["N_rec"], heads=configs["heads_trx_rf"], dropout=configs["dropout"],
                                custom_attn=configs["custom_attn"], multclass=multclass,
                                NS_model=self.NS_model, af_module=configs["af_module"])

        print(f'Tmodel # of Parameters: {np.sum([np.prod(p.shape) for p in self.Tmodel.parameters()])}')
        print(f'Rmodel # of Parameters: {np.sum([np.prod(p.shape) for p in self.Rmodel.parameters()])}')
        print(f'Rfmodel # of Parameters: {np.sum([np.prod(p.shape) for p in self.Rmodelf.parameters()])}')

    def encoder(self, each_batch, bVec_md, isTraining=1, normalization_type="soft"):
        """
            Perform the TransCoder encoding process using the transmitter model.

            Args:
                each_batch (torch.Tensor): Number of data batch.
                bVec_md (torch.Tensor): Input modulated codeword vector.
                isTraining (int): Flag indicating whether the model is in training mode.
                normalization_type (str): Type of power normalization to apply.

            Returns:
                torch.Tensor: Encoded blocks.
            """
        if not self.no_inner and not self.no_Tmodel:
            generated_blocks = self.transmitter_nn(each_batch, bVec_md,
                                                                   isTraining=isTraining,
                                                                   normalization_type=normalization_type)
        else:
            generated_blocks = bVec_md.clone()
        generated_blocks = generated_blocks.reshape(self.bs, -1)
        return generated_blocks

    def transmitter_nn(self, each_batch, bVec_md,
                         isTraining=1, normalization_type="soft"):
        """
            Transmitter neural network for encoding the input bit vector.

            Args:
                each_batch (torch.Tensor): Number of data batch.
                bVec_md (torch.Tensor): Input modulated codeword vector.
                isTraining (int): Flag indicating whether the model is in training mode.
                normalization_type (str): Type of power normalization to apply.

            Returns:
                torch.Tensor: Encoded blocks after power constraint and reallocation.
            """
        if self.padd_symbol_len > 0:
            bVec_md = torch.cat((bVec_md, self.ones_symbols), dim=1)

        bVec_md = bVec_md.reshape(self.bs, self.num_blocks, self.M)
        coded_blocks = self.Tmodel(bVec_md, None, self.pe)

        coded_blocks = coded_blocks.reshape(self.bs, self.num_blocks * self.block_rate)
        if self.padd_symbol_len > 0:
            coded_blocks = coded_blocks[:, :-self.padd_symbol_len]

        coded_blocks = power_constraint(coded_blocks, isTraining, each_batch,
                                        normalization_type=normalization_type,
                                        snr=self.configs["snr"], path=self.model_path,
                                        save_first=self.configs["save_first"])

        coded_blocks = self.power_reallocator(coded_blocks)

        generated_blocks = coded_blocks.clone()
        return generated_blocks

    def decoder(self, coded_blocks, current_sigma):
        """
            Perform the TransCoder decoding process using the receiver model.

            Args:
                coded_blocks (torch.Tensor): Encoded blocks with noise added.
                current_sigma (float): Noise standard deviation.

            Returns:
                tuple: Decoded probabilities and intermediate decoded blocks.
            """
        if self.no_inner or self.no_Rmodel:
            if self.padd_symbol_len > 0:
                decoded_probs_initial = torch.cat((coded_blocks,
                                               self.ones_symbols), dim=1).reshape(self.bs,
                                                                                  self.num_blocks,
                                                                                  self.block_rate)
            else:
                decoded_probs_initial = coded_blocks.reshape(self.bs, self.num_blocks,
                                                             self.block_rate)

        else:
            if self.padd_symbol_len > 0:
                added_symb = self.ones_symbols
                coded_blocks = torch.cat((coded_blocks, added_symb), dim=1)
            coded_blocks = coded_blocks.reshape(self.bs, self.num_blocks, self.block_rate)
            decoded_probs = self.Rmodel(coded_blocks, None, self.pe, std=current_sigma)

            decoded_probs_initial = decoded_probs.clone()
        return  coded_blocks, decoded_probs_initial

    def compute_llrs(self, coded_blocks, current_sigma, k, table, decoded_probs):
        """
           Compute Log-Likelihood Ratios (LLRs) for the given coded blocks.

           Args:
               coded_blocks (torch.Tensor): Encoded blocks.
               current_sigma (float): Noise standard deviation.
               k (int): Current decoding iteration.
               table (torch.Tensor): Lookup table for TransCoder.
               decoded_probs (torch.Tensor): Decoded probabilities.

           Returns:
               torch.Tensor: Computed LLRs.
           """
        if (self.no_inner and k == 0) or self.no_Rmodel:
            llrs = (2 * coded_blocks / torch.pow(current_sigma, 2)).to(torch.float32)
        elif self.outer_code_type != 'Turbo' or k == 0:
            bit_wise_one_probs, bit_wise_zero_probs = convert_class_to_bitwise_matmul(decoded_probs, table)
            # Getting the LLRs (it is log(p(x=1)/p(x=0))
            llrs = torch.log(bit_wise_one_probs) - torch.log(bit_wise_zero_probs)

            llrs = llrs.reshape(self.bs, self.K_extended)
            if self.padd_symbol_len > 0 and (self.outer_code_type != 'Turbo' or self.model_num_iters == 0):
                llrs = llrs[:, :-self.padd_symbol_len]
        elif self.outer_code_type == 'Turbo' and k != 0:
            llrs = (torch.log(decoded_probs.view(self.bs, -1, 2)[:, :, 0]) - torch.log(
                decoded_probs.view(self.bs, -1, 2)[:, :, 1])).view(self.bs, -1)

        return llrs

    def compute_priors(self, logits):
        """
            Compute priors from logits for iterative decoding.

            Args:
                logits (torch.Tensor): Logits from the outer code.

            Returns:
                torch.Tensor: Computed priors.
        """
        logits_orig = convert_prior_to_probabilities(logits.clone(), modulate=False)
        logits = convert_prior_to_probabilities(logits, modulate=True)
        if self.padd_symbol_len > 0:
            added_symb = self.ones_symbols
            logits = torch.cat((logits, added_symb), dim=1)
            logits_orig = torch.cat((logits_orig, added_symb), dim=1).reshape(self.bs, self.num_blocks, self.M).to(
                self.device)
            priors = logits.reshape(self.bs, self.num_blocks, self.M).to(self.device)
        else:
            priors = logits.reshape(self.bs, self.num_blocks, self.M).to(self.device)
            logits_orig = logits_orig.reshape(self.bs, self.num_blocks, self.M).to(self.device)
        return priors

    def encode_only(self, each_batch, bVec_md, noise_samples, table=None, isTraining=1,
                normalization_type="soft", bVec=None,
                current_sigma=None):
        """
            Forward pass of the TransCoder model.

            Args:
                each_batch (torch.Tensor): Number of data batch.
                bVec_md (torch.Tensor): Input modulated codeword vector.
                noise_samples (torch.Tensor): Noise samples to add to simulate transmission.
                table (torch.Tensor, optional): Lookup table for TransCoder.
                isTraining (int): Flag indicating whether the model is in training mode.
                normalization_type (str): Type of power normalization to apply (default: "average").
                current_sigma (torch.Tensor, optional): Noise standard deviation.

            Returns:
                tuple: Decoded probabilities, initial decoded probabilities, generated blocks,
                       priors from the first iteration, and priors history.
            """
        generated_blocks = self.encoder(each_batch, bVec_md, isTraining=isTraining,
                                        normalization_type=normalization_type)
        return generated_blocks.clone()


    def forward(self, each_batch, bVec_md, noise_samples, table=None, isTraining=1,
                normalization_type="soft",
                current_sigma=None):
        """
            Forward pass of the TransCoder model.

            Args:
                each_batch (torch.Tensor): Number of data batch.
                bVec_md (torch.Tensor): Input modulated codeword vector.
                noise_samples (torch.Tensor): Noise samples to add to simulate transmission.
                table (torch.Tensor, optional): Lookup table for TransCoder.
                isTraining (int): Flag indicating whether the model is in training mode.
                normalization_type (str): Type of power normalization to apply (default: "average").
                current_sigma (torch.Tensor, optional): Noise standard deviation.

            Returns:
                tuple: Decoded probabilities, initial decoded probabilities, generated blocks,
                       priors from the first iteration, and priors history.
            """
        generated_blocks = self.encoder(each_batch, bVec_md, isTraining=isTraining,
                                        normalization_type=normalization_type)
        coded_blocks = generated_blocks.clone()

        if self.rayleigh:
            uni_vec = torch.rand_like(coded_blocks.reshape(self.bs, -1))
            sigma = 1 #/ np.sqrt(2)
            gains = sigma * torch.sqrt(-2 * torch.log(1 - uni_vec))
            coded_blocks = gains*coded_blocks.reshape(self.bs, -1) + noise_samples
        else:
            coded_blocks = coded_blocks.reshape(self.bs, -1) + noise_samples

        coded_blocks, decoded_probs_initial = self.decoder(coded_blocks, current_sigma)
        decoded_probs = decoded_probs_initial.clone()

        decoded_probs_history = []
        decoded_probs_history.append(decoded_probs_initial.clone())
        priors_history = []
        for k in range(self.model_num_iters):
            # Getting bitwise probabilities
            llrs = self.compute_llrs(coded_blocks, current_sigma, k, table, decoded_probs)

            priors, logits = self.outer_code.decode(llrs, bp_iters=self.bp_iters_arr[k], backprop=self.backprop)

            if k == 0:
                priors_first = priors.clone()

            priors_history.append(priors.clone())

            if self.outer_code_type == 'Turbo':
                belief_input = self.outer_code.belief_inputs
            else:
                priors = self.compute_priors(logits)
                coded_blocks = coded_blocks.reshape(self.bs, self.num_blocks, self.block_rate)
                belief_input = torch.cat([coded_blocks, priors], dim=2)

            decoded_probs = self.Rmodelf(belief_input, None, self.pe, std=current_sigma)
            if self.outer_code_type != 'Turbo':
                decoded_probs_history.append(decoded_probs.clone())

        llrs = self.compute_llrs(coded_blocks, current_sigma, self.model_num_iters, table, decoded_probs)

        # Obtaining the outer code predictions
        priors_final, _ = self.outer_code.decode(llrs,
                                                 bp_iters=self.bp_iters_arr[self.model_num_iters],
                                                 backprop=self.backprop)

        if self.model_num_iters == 0:
            priors_first = priors_final
        priors_history.append(priors_final.clone())
        return decoded_probs_history, decoded_probs_initial, generated_blocks, priors_first, priors_history