{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "8482bf05",
   "metadata": {},
   "source": [
    "# Comparing the Polar Coding results with Matlab\n",
    "Applying Polar encoding/decoding on random transport blocks and comparing the results with the equivalent Matlab code \"MatlabFiles/Polar.mlx\". [Here](MatlabFiles/Polar.html) is the execution results of this code in Matlab."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "e51dea53",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.io\n",
    "import time\n",
    "\n",
    "import neoradium as nr\n",
    "\n",
    "matlabFilesPath = \"./MatlabFiles\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "b90cb6b4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Code Rate:      0.45\n",
      "SNR:            0.3424250943932492 db\n",
      "Noise Variance: 0.9241819678918566\n",
      "Noise STD:      0.9613438343755353\n"
     ]
    }
   ],
   "source": [
    "payloadLen = 30           # A (Must be no larger than 1706)\n",
    "rateMatchedLen = 120      \n",
    "\n",
    "ebNo = 0.8\n",
    "codeRate = (payloadLen+24)/rateMatchedLen  # Effective code rate\n",
    "\n",
    "bps = 2;                             # bits per symbol = 2 for QPSK\n",
    "esNo = ebNo + 10*np.log10(bps)       # energy per symbol over noise\n",
    "snrdB = esNo + 10*np.log10(codeRate) # SNR in dB\n",
    "noiseVar = 1/(10**(snrdB/10))\n",
    "noiseStd = np.sqrt(noiseVar)\n",
    "print(\"Code Rate:     \", codeRate)\n",
    "print(\"SNR:           \", snrdB, \"db\")\n",
    "print(\"Noise Variance:\", noiseVar)\n",
    "print(\"Noise STD:     \", noiseStd)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "64b5dd72",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Polar Encoder Properties:\n",
      "  dataType ......................: DCI\n",
      "  payloadSize (A) ...............: 30\n",
      "  rateMatchedLen (E) ............: 120\n",
      "  codeBlockSize (K) .............: 54\n",
      "  polarCodeSize (N) .............: 128\n",
      "  Max Log2(N) (nMax) ............: 9\n",
      "  Segmentation (iSeg) ...........: Disabled\n",
      "  Code Block CRC (crcPoly) ......: 24C\n",
      "  Input Interleaving (iIL) ......: Enabled\n",
      "  Coded bit Interleaving (iBIL)..: Disabled\n",
      "  Num Parity checks (nPC, nPCwm) : 0,0\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a Polar Encoder\n",
    "polarEncoder = nr.PolarEncoder(payloadLen, rateMatchedLen, 'dci')\n",
    "polarEncoder.print()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8838c829",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transport Block Shape:  (30,)\n",
      "Transport Block:        100100110110000110100110000010\n",
      "Code Blocks Shape:      (1, 54)\n",
      "Code Blocks:            100100110110000110100110000010111100001010001110100010\n"
     ]
    }
   ],
   "source": [
    "# Read input bits from Matlab-generated file\n",
    "txpBlock = scipy.io.loadmat(matlabFilesPath+'/msg.mat')['msg'].reshape(-1)\n",
    "print(\"Transport Block Shape: \", txpBlock.shape)\n",
    "print(\"Transport Block:       \", \"\".join(str(x) for x in txpBlock))\n",
    "\n",
    "# Segmentation\n",
    "codeBlocks = polarEncoder.doSegmentation(txpBlock)\n",
    "print(\"Code Blocks Shape:     \", codeBlocks.shape)\n",
    "print(\"Code Blocks:           \", \"\".join(str(x) for x in codeBlocks[0]))\n",
    "\n",
    "# Compare with Matlab\n",
    "codeBlocksMatlab = scipy.io.loadmat(matlabFilesPath+'/msgcrc.mat')['msgcrc'].T\n",
    "assert np.abs(codeBlocks-codeBlocksMatlab).sum()==0, \"MISMATCH WITH MATLAB!!!\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "d66ea1eb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Code Words Shape:         (1, 128)\n",
      "Code Words first 10 bits: 1001110011\n"
     ]
    }
   ],
   "source": [
    "# Encode\n",
    "codeWords = polarEncoder.encode(codeBlocks)\n",
    "print(\"Code Words Shape:        \", codeWords.shape)\n",
    "print(\"Code Words first 10 bits:\", \"\".join(str(x) for x in codeWords[0][:10]))\n",
    "\n",
    "# Compare with Matlab\n",
    "codeWordsMatlab = scipy.io.loadmat(matlabFilesPath+'/encOut.mat')['encOut'].T\n",
    "assert np.abs(codeWords-codeWordsMatlab).sum()==0, \"MISMATCH WITH MATLAB!!!\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "c18f8fd4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rate-Matched Shape:   (1, 120)\n",
      "First 10 bits:        1001110011\n"
     ]
    }
   ],
   "source": [
    "# Rate Matching\n",
    "rateMatchedCodeWords = polarEncoder.rateMatch(codeWords)\n",
    "print(\"Rate-Matched Shape:  \", rateMatchedCodeWords.shape)\n",
    "print(\"First 10 bits:       \", \"\".join(str(x) for x in rateMatchedCodeWords[0][:10]))\n",
    "\n",
    "# Compare with Matlab\n",
    "rateMatchedCodeWordsMatlab = scipy.io.loadmat(matlabFilesPath+'/modIn.mat')['modIn'].T\n",
    "assert np.abs(rateMatchedCodeWords-rateMatchedCodeWordsMatlab).sum()==0, \"MISMATCH WITH MATLAB!!!\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "8fdc2d7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "modulated Shape:      (1, 60)\n",
      "Max Diff. with Matlab: 1.5700924586837752e-16\n",
      "First 5 symbols:\n",
      " [-0.70710678+0.70710678j  0.70710678-0.70710678j -0.70710678-0.70710678j\n",
      "  0.70710678+0.70710678j -0.70710678-0.70710678j]\n"
     ]
    }
   ],
   "source": [
    "# QPSK Modulation\n",
    "modulated = nr.Modem('QPSK').modulate(rateMatchedCodeWords)\n",
    "print(\"modulated Shape:     \", modulated.shape)\n",
    "\n",
    "# Compare with Matlab\n",
    "modulatedMatlab = scipy.io.loadmat(matlabFilesPath+'/modOut.mat')['modOut'].T\n",
    "assert np.abs(modulated-modulatedMatlab).max()<1e-12, \"MISMATCH WITH MATLAB!!!\"\n",
    "print(\"Max Diff. with Matlab:\", np.abs(modulated-modulatedMatlab).max())\n",
    "print(\"First 5 symbols:\\n\", modulated[0,:5])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "87c1aa7b",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise Shape:      (1, 60)\n",
      "First 5 Rx Symbols:\n",
      " [-0.51312371+0.76995541j  0.88901799-1.41258547j -1.44741359-1.38917752j\n",
      "  0.49504718+0.34599014j -1.58021683-0.80163037j]\n"
     ]
    }
   ],
   "source": [
    "# Read the noise from the Matlab file\n",
    "complexNoise = scipy.io.loadmat(matlabFilesPath+'/chanNoise.mat')['chanNoise'].T\n",
    "print(\"Noise Shape:     \", complexNoise.shape)\n",
    "\n",
    "# Add noise to the modulated signal\n",
    "rxSymbols = modulated + complexNoise\n",
    "print(\"First 5 Rx Symbols:\\n\", rxSymbols[0,:5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "ffc77206",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "LLR Shape:      (1, 120)\n",
      "First 10 LLRs: \n",
      " [-1.57039747  2.35642204  2.72080897 -4.32316924 -4.42975952 -4.25153001\n",
      "  1.51507488  1.05889091 -4.83619924 -2.45336219]\n",
      "Matlab First 10 LLRs: \n",
      " [-1.57039747  2.35642204  2.72080897 -4.32316924 -4.42975952 -4.25153001\n",
      "  1.51507488  1.05889091 -4.83619924 -2.45336219]\n",
      "Max Diff. with Matlab: 3.552713678800501e-15\n"
     ]
    }
   ],
   "source": [
    "# Demodulation: Calculate LLR values from symbols:\n",
    "llrs = nr.Modem('QPSK').getLLRsFromSymbols(rxSymbols, noiseVar)\n",
    "print(\"LLR Shape:     \", llrs.shape)\n",
    "print(\"First 10 LLRs: \\n\", llrs[0,:10])\n",
    "\n",
    "# Compare with Matlab\n",
    "llrsMatlab = scipy.io.loadmat(matlabFilesPath+'/rxLLR.mat')['rxLLR'].T\n",
    "print(\"Matlab First 10 LLRs: \\n\", llrsMatlab[0,:10])\n",
    "\n",
    "assert np.abs(llrs-llrsMatlab).max()<1e-12, \"MISMATCH WITH MATLAB!!!\"\n",
    "print(\"Max Diff. with Matlab:\", np.abs(llrs-llrsMatlab).max())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "05809cf1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Hard Decision first 10 bits: 1001110011\n",
      "Org CodeWord first 10 bits:  1001110011\n"
     ]
    }
   ],
   "source": [
    "#Hard Decision:\n",
    "hardCodeWord = 1*(llrs<0)\n",
    "print(\"Hard Decision first 10 bits:\", \"\".join(str(x) for x in hardCodeWord[0][:10]))\n",
    "print(\"Org CodeWord first 10 bits: \", \"\".join(str(x) for x in codeWords[0][:10]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "5abb0032",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Polar Decoder Properties:\n",
      "  dataType ......................: DCI\n",
      "  payloadSize (A) ...............: 30\n",
      "  rateMatchedLen (E) ............: 120\n",
      "  codeBlockSize (K) .............: 54\n",
      "  polarCodeSize (N) .............: 128\n",
      "  Max Log2(N) (nMax) ............: 9\n",
      "  Segmentation (iSeg) ...........: Disabled\n",
      "  Code Block CRC (crcPoly) ......: 24C\n",
      "  Input Interleaving (iIL) ......: Enabled\n",
      "  Coded bit Interleaving (iBIL)..: Disabled\n",
      "  Num Parity checks (nPC, nPCwm) : 0,0\n",
      "  SCL List Size .................: 8\n",
      "  MinSum Approximation ..........: Enabled\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Create a Polar Decoder\n",
    "polarDecoder = nr.PolarDecoder(payloadLen, rateMatchedLen, 'dci',\n",
    "                               sclListSize=8, useMinsum=True)\n",
    "polarDecoder.print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "ce7ce92f",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Rate Recovered Shape:      (1, 128)\n",
      "First 10 LLRs: \n",
      " [-1.57039747  2.35642204  2.72080897 -4.32316924 -4.42975952 -4.25153001\n",
      "  1.51507488  1.05889091 -4.83619924 -2.45336219]\n",
      "Matlab First 10 LLRs: \n",
      " [-1.57039747  2.35642204  2.72080897 -4.32316924 -4.42975952 -4.25153001\n",
      "  1.51507488  1.05889091 -4.83619924 -2.45336219]\n",
      "Max Diff. with Matlab: 3.552713678800501e-15\n"
     ]
    }
   ],
   "source": [
    "# Recover rate\n",
    "rateRecoveredRxBlocks = polarDecoder.recoverRate(llrs)\n",
    "print(\"Rate Recovered Shape:     \", rateRecoveredRxBlocks.shape)\n",
    "print(\"First 10 LLRs: \\n\", rateRecoveredRxBlocks[0,:10])\n",
    "\n",
    "# Compare with Matlab\n",
    "rateRecoveredRxBlocksMatlab = scipy.io.loadmat(matlabFilesPath+'/decIn.mat')['decIn'].T\n",
    "print(\"Matlab First 10 LLRs: \\n\", rateRecoveredRxBlocksMatlab[0,:10])\n",
    "\n",
    "assert np.abs(rateRecoveredRxBlocks-rateRecoveredRxBlocksMatlab).max()<1e-12, \"MISMATCH WITH MATLAB!!!\"\n",
    "print(\"Max Diff. with Matlab:\", np.abs(rateRecoveredRxBlocks-rateRecoveredRxBlocksMatlab).max())\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "f7085147",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Number of CRC Errors:           0\n",
      "Decoded Transport Block Shape:  (30,)\n",
      "Decoded Transport Block:        100100110110000110100110000010\n",
      "Original Transport Block:       100100110110000110100110000010\n",
      "Matlab Decoded Transport Block: 100100110110000110100110000010\n"
     ]
    }
   ],
   "source": [
    "# Decode the rate-recovered transport blocks\n",
    "decTxBlock, numCrcErrors = polarDecoder.decode(rateRecoveredRxBlocks)\n",
    "print(\"Number of CRC Errors:          \", numCrcErrors)\n",
    "print(\"Decoded Transport Block Shape: \", decTxBlock.shape)\n",
    "print(\"Decoded Transport Block:       \", \"\".join(str(x) for x in decTxBlock))\n",
    "print(\"Original Transport Block:      \", \"\".join(str(x) for x in txpBlock))\n",
    "\n",
    "# Compare with Matlab\n",
    "decTxBlockMatlab = scipy.io.loadmat(matlabFilesPath+'/decBits.mat')['decBits'].reshape(-1)[:payloadLen]\n",
    "print(\"Matlab Decoded Transport Block:\", \"\".join(str(x) for x in decTxBlockMatlab))\n",
    "\n",
    "assert np.abs(txpBlock-decTxBlock).sum()==0, \"MISMATCH WITH MATLAB!!!\"\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73bd2b57",
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a12f08ab",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}