{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "c15a52e1-8ae5-441f-8c60-362e851f4de3",
   "metadata": {},
   "source": [
    "# SNR Calculations in NeoRadium\n",
    "This notebook demonstrates how NeoRadium calculates the noise standard deviation from a given SNR value \n",
    "in both the time and frequency domains."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "2415601e",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import scipy.io\n",
    "import time\n",
    "\n",
    "from neoradium import Carrier, PDSCH, CdlChannel, AntennaPanel, LdpcEncoder, Grid, random, Waveform\n",
    "from neoradium.utils import toDb, toLinear, getNmse\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "45896ed9-ebac-4b14-a6e2-55a6f1844f96",
   "metadata": {},
   "source": [
    "First let's try a simple case without a channel model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "224b0f9d-0fa6-4d2a-a67c-35932eaa6378",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Bandwidth Part Properties:\n",
      "  Resource Blocks:    52 RBs starting at 0 (624 subcarriers)\n",
      "  Subcarrier Spacing: 30 kHz\n",
      "  CP Type:            normal\n",
      "  Bandwidth:          18.72 MHz\n",
      "  symbolsPerSlot:     14\n",
      "  slotsPerSubFrame:   2\n",
      "  nFFT:               1024\n",
      "  frameNo:            0\n",
      "  slotNo:             0\n",
      "\n",
      "Shape of Resource Grid:         (1, 14, 624)\n",
      "Shape of Precoded Resource Grid:(2, 14, 624)\n",
      "Shape of txWaveform:            (2, 15360)\n",
      "Shape of rxWaveform:            (2, 15360)\n",
      "Shape of rxGrid:                (2, 14, 624)\n"
     ]
    }
   ],
   "source": [
    "# NOTE: You can find a similar Matlab example at:\n",
    "# https://www.mathworks.com/help/5g/ug/snr-definition-used-in-link-simulations.html\n",
    "snrDb = 0\n",
    "snr = toLinear(snrDb)   # SNR in linear scale\n",
    "\n",
    "carrier = Carrier(numRbs=52, spacing=30)\n",
    "bwp = carrier.curBwp\n",
    "nr, nt = 2, 2  # receiver and transmitter antennas\n",
    "print(bwp)\n",
    "\n",
    "pdsch = PDSCH(bwp, interleavingBundleSize=0, numLayers=1, modulation='16QAM', nID=carrier.cellId)\n",
    "pdsch.setDMRS(prgSize=0, configType=2, additionalPos=2)\n",
    "\n",
    "codeRate = 490/1024\n",
    "ldpcEncoder = LdpcEncoder(baseGraphNo=1, modulation=pdsch.modems[0].modulation, \n",
    "                          txLayers=pdsch.numLayers, targetRate=codeRate)\n",
    "\n",
    "\n",
    "\n",
    "random.setSeed(123)                          # Make the results reproducible.\n",
    "grid = pdsch.getGrid()                       # Create a resource grid populated with DMRS\n",
    "txBlockSize = pdsch.getTxBlockSize(codeRate) # Calculate the transport block size based on 3GPP TS 38.214\n",
    "txBlock = random.bits(txBlockSize[0])        # Create random binary data\n",
    "        \n",
    "numBits = pdsch.getBitSizes(grid)            # Actual number of bits available in the resource grid\n",
    "\n",
    "# Now perform the segmentation, rate-matching, and encoding in one call:\n",
    "rateMatchedCodeBlocks = ldpcEncoder.getRateMatchedCodeBlocks(txBlock, numBits[0])\n",
    "\n",
    "# Now populate the resource grid with coded data. This includes QAM modulation and resource mapping.\n",
    "pdsch.populateGrid(grid, rateMatchedCodeBlocks)\n",
    "\n",
    "# Store the indexes of the PDSCH data in pdschIndexes to be used later.\n",
    "pdschIndexes = pdsch.getReIndexes(grid, \"PDSCH\")   \n",
    "\n",
    "# Getting the Precoding Matrix, and precoding the resource grid\n",
    "precoder = np.ones((nt, pdsch.numLayers))/np.sqrt(pdsch.numLayers)      # Get the precoder matrix from the PDSCH object\n",
    "precodedGrid = grid.precode(precoder)                                   # Perform the precoding\n",
    "\n",
    "print(f\"Shape of Resource Grid:         {grid.shape}\")\n",
    "print(f\"Shape of Precoded Resource Grid:{precodedGrid.shape}\")\n",
    "\n",
    "txWaveform = precodedGrid.ofdmModulate()\n",
    "print(f\"Shape of txWaveform:            {txWaveform.shape}\")\n",
    "\n",
    "# Applying channel to the transmitted signal\n",
    "rxWaveform = Waveform(txWaveform.waveform/np.sqrt(nr))\n",
    "print(f\"Shape of rxWaveform:            {rxWaveform.shape}\")\n",
    "rxGrid = rxWaveform.ofdmDemodulate(bwp)\n",
    "print(f\"Shape of rxGrid:                {rxGrid.shape}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a76db02d-d05f-46b5-9b35-7cbdfb9d4b2c",
   "metadata": {},
   "source": [
    "The standard deviation of the noise in time domain is:\n",
    "$$\\sigma_T = \\sqrt{\\frac {N_{FFT} \\sigma_x^2} {K.{SNR}}} $$\n",
    "where $\\sigma_x^2$ is the variance of the time-domain signal (the Waveform object), $K$ is the number of subcarriers in the current bandwidth part, and $N_{FFT}$ is the FFT size (``bwp.nFFT``). This value can be used to generate Gaussian noise applied to a Waveform object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "13e80b5e-5845-47e2-94b7-5346f57dbc7c",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Time): 0.0220441589451537\n"
     ]
    }
   ],
   "source": [
    "noiseStdTime = rxWaveform.getNoiseStd(snr, bwp)\n",
    "print(f\"Noise STD (Time): {noiseStdTime}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "fb2327da-5121-4e68-b50a-a98b29857200",
   "metadata": {},
   "source": [
    "The standard deviation of the noise in Frequency domain is:\n",
    "$$\\sigma_F = \\sqrt{\\frac {\\sigma_X^2} {SNR}} $$\n",
    "where $\\sigma_X^2$ is the variance of the frequency-domain signal (the Grid object). This can be used to generate Gaussian noise applied to a Grid object:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "05384ed4-7584-408f-8b60-d62e2e6a004d",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Freq): 0.705375297931587\n"
     ]
    }
   ],
   "source": [
    "noiseStdFreq = rxGrid.getNoiseStd(snr)\n",
    "print(f\"Noise STD (Freq): {noiseStdFreq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2c4debc3-c2d0-4c4f-ab1d-3cdd430ca2a2",
   "metadata": {},
   "source": [
    "Verifying Equation:\n",
    "$$\\sigma_T = \\sqrt{{\\frac 1 {N_{FFT}}}} \\sigma_F$$"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "7ebcbb4d-132e-41e6-9eb4-e84e066f2d00",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Freq)/sqrt(1024): 0.022042978060362095 = Noise STD (Time)\n"
     ]
    }
   ],
   "source": [
    "print(f\"Noise STD (Freq)/sqrt({bwp.nFFT}): {noiseStdFreq/np.sqrt(bwp.nFFT)} = Noise STD (Time)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "609e924a-135a-4801-859b-65ef8200e668",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise Grid STD: 0.7020339000063331\n",
      "Noise STD (Freq): 0.705375297931587\n"
     ]
    }
   ],
   "source": [
    "# Apply noise in time domain and measure noise in frequency domain:\n",
    "noisyRxWaveform = rxWaveform.addNoise(noiseStd=noiseStdTime)\n",
    "noisyRxGrid = noisyRxWaveform.ofdmDemodulate(bwp)\n",
    "noiseGrid = noisyRxGrid.grid - rxGrid.grid\n",
    "print(f\"Noise Grid STD: {noiseGrid.std()}\")\n",
    "print(f\"Noise STD (Freq): {noiseStdFreq}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e55514f3-5ddc-4dac-90a0-834c930d6e30",
   "metadata": {},
   "source": [
    "## Matlab's Assumption\n",
    "Matlab assumes $\\sigma_{X}^2 = \\frac{1}{N_r}$, where $N_r$ is the number of receiver antennas. With that assumption, we have:\n",
    "\n",
    "$\\sigma_F = \\frac{1}{\\sqrt{N_r \\cdot SNR}}$ and $\\sigma_T = \\frac{1}{\\sqrt{N_r \\cdot N_{FFT} \\cdot SNR}}$\n",
    "\n",
    "Note that in this case, these values are the same since we used $\\sigma_{X}^2 = \\frac{1}{N_r}$. This is not always true in practice."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "af0ebbca-cc11-409a-a1d3-17ccb61401a8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Time Domain - RxPower=1/Nr): 0.022097086912079608\n",
      "Noise STD (Freq Domain - RxPower=1/Nr): 0.7071067811865475\n"
     ]
    }
   ],
   "source": [
    "noiseStdTime1 = 1/np.sqrt(nr*bwp.nFFT*snr)\n",
    "noiseStdFreq1 = 1/np.sqrt(nr*snr)\n",
    "print(f\"Noise STD (Time Domain - RxPower=1/Nr): {noiseStdTime1}\")\n",
    "print(f\"Noise STD (Freq Domain - RxPower=1/Nr): {noiseStdFreq1}\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07999e33-b4b8-45d4-abc8-c4c03e2c902a",
   "metadata": {},
   "source": [
    "Now we try the same steps with a CDL channel model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "06b2af6d-72c2-41c4-83a0-840fc51e77ac",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Shape of Resource Grid:         (2, 14, 624)\n",
      "Shape of Precoded Resource Grid:(4, 14, 624)\n",
      "Shape of txWaveform:            (4, 15360)\n",
      "Shape of rxWaveform:            (2, 15360)\n",
      "Timing Offset:                  13\n",
      "Shape of rxGrid:                (2, 14, 624)\n",
      "Shape of rxGridF:               (2, 14, 624)\n",
      "NMSE(rxGrid,rxGridF):           6.457297490367846e-07\n"
     ]
    }
   ],
   "source": [
    "# Now repeat the same process using a CDL channel model with two layers:\n",
    "snrDb = 0\n",
    "snr = toLinear(snrDb)\n",
    "\n",
    "carrier = Carrier(numRbs=52, spacing=30)\n",
    "bwp = carrier.curBwp\n",
    "\n",
    "pdsch = PDSCH(bwp, interleavingBundleSize=0, numLayers=2, modulation='16QAM', nID=carrier.cellId)\n",
    "pdsch.setDMRS(prgSize=0, configType=2, additionalPos=2)\n",
    "\n",
    "codeRate = 490/1024\n",
    "ldpcEncoder = LdpcEncoder(baseGraphNo=1, modulation=pdsch.modems[0].modulation, \n",
    "                          txLayers=pdsch.numLayers, targetRate=codeRate)\n",
    "\n",
    "random.setSeed(123)                          # Making the results reproducible.\n",
    "grid = pdsch.getGrid()                       # Create a resource grid already populated with DMRS \n",
    "txBlockSize = pdsch.getTxBlockSize(codeRate) # Calculate the Transport Block Size based on 3GPP TS 38.214 \n",
    "txBlock = random.bits(txBlockSize[0])        # Create random binary data\n",
    "        \n",
    "numBits = pdsch.getBitSizes(grid)            # Actual number of bits available in the resource grid\n",
    "\n",
    "# Now perform the segmentation, rate-matching, and encoding in one call:\n",
    "rateMatchedCodeBlocks = ldpcEncoder.getRateMatchedCodeBlocks(txBlock, numBits[0])\n",
    "\n",
    "# Now populate the resource grid with coded data. This includes QAM modulation and resource mapping.\n",
    "pdsch.populateGrid(grid, rateMatchedCodeBlocks)\n",
    "\n",
    "# Store the indexes of the PDSCH data in pdschIndexes to be used later.\n",
    "pdschIndexes = pdsch.getReIndexes(grid, \"PDSCH\")   \n",
    "\n",
    "# Creating a CdlChannel object:\n",
    "f0 = 4e9  # Carrier Frequency\n",
    "channel = CdlChannel(bwp, 'C', delaySpread=300, carrierFreq=f0, dopplerShift=5,\n",
    "                     txAntenna = AntennaPanel([1,4], polarization=\"|\"),  # 4 TX antennas\n",
    "                     rxAntenna = AntennaPanel([1,2], polarization=\"|\"),  # 2 RX antennas\n",
    "                     normalizeGains = True, normalizeOutput=True)\n",
    "nr,nt = channel.nrNt\n",
    "\n",
    "# Getting the Precoding Matrix, and precoding the resource grid\n",
    "channelMatrix = channel.getChannelMatrix()              # Get the channel matrix\n",
    "precoder = pdsch.getPrecodingMatrix(channelMatrix)      # Get the precoder matrix from the PDSCH object\n",
    "precodedGrid = grid.precode(precoder)                   # Perform the precoding\n",
    "\n",
    "print(f\"Shape of Resource Grid:         {grid.shape}\")\n",
    "print(f\"Shape of Precoded Resource Grid:{precodedGrid.shape}\")\n",
    "\n",
    "txWaveform = precodedGrid.ofdmModulate()\n",
    "print(f\"Shape of txWaveform:            {txWaveform.shape}\")\n",
    "\n",
    "rxWaveform = channel.applyToSignal(txWaveform)\n",
    "print(f\"Shape of rxWaveform:            {rxWaveform.shape}\")\n",
    "\n",
    "offset = channel.getTimingOffset()\n",
    "syncedWaveform = rxWaveform.sync(offset)\n",
    "print(f\"Timing Offset:                  {offset}\")\n",
    "\n",
    "rxGrid = syncedWaveform.ofdmDemodulate(bwp)\n",
    "print(f\"Shape of rxGrid:                {rxGrid.shape}\")\n",
    "\n",
    "rxGridF = channel.applyToGrid(precodedGrid)\n",
    "print(f\"Shape of rxGridF:               {rxGridF.shape}\")\n",
    "print(f\"NMSE(rxGrid,rxGridF):           {getNmse(rxGrid.grid, rxGridF.grid)}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "91b01a1b-1c53-4748-9b05-909e89929659",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Time): 0.06558965831573078\n"
     ]
    }
   ],
   "source": [
    "noiseStdTime = syncedWaveform.getNoiseStd(snr, bwp)\n",
    "print(f\"Noise STD (Time): {noiseStdTime}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b267a539-9ad7-45f7-a10e-56be2d04717e",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Freq): 2.0989851489946787\n"
     ]
    }
   ],
   "source": [
    "noiseStdFreq = rxGrid.getNoiseStd(snr)\n",
    "print(f\"Noise STD (Freq): {noiseStdFreq}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "11f16d51-78b1-47c1-8026-6d5d9f20f852",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Freq)/sqrt(1024): 0.06559328590608371 = Noise STD (Time)\n"
     ]
    }
   ],
   "source": [
    "print(f\"Noise STD (Freq)/sqrt({bwp.nFFT}): {noiseStdFreq/np.sqrt(bwp.nFFT)} = Noise STD (Time)\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "54580eb7-e35a-4fa8-9e35-4ff2be54daee",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise STD (Time Domain - RxPower=1/Nr): 0.022097086912079608\n",
      "Noise STD (Freq Domain - RxPower=1/Nr): 0.7071067811865475\n"
     ]
    }
   ],
   "source": [
    "# In this case, Matlab's assumption is not valid; these values do not match the calculated STD values above.\n",
    "noiseStdTime1 = 1/np.sqrt(nr*bwp.nFFT*snr)\n",
    "noiseStdFreq1 = 1/np.sqrt(nr*snr)\n",
    "print(f\"Noise STD (Time Domain - RxPower=1/Nr): {noiseStdTime1}\")\n",
    "print(f\"Noise STD (Freq Domain - RxPower=1/Nr): {noiseStdFreq1}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6bc3f49e-bb36-4a0f-995c-aaa6cbbaadf8",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Noise Grid STD: 2.082620653246523\n",
      "Noise STD (Freq): 2.0989851489946787\n"
     ]
    }
   ],
   "source": [
    "# Apply noise in the time domain and measure the resulting noise in the frequency domain:\n",
    "noisyRxWaveform = rxWaveform.addNoise(noiseStd=noiseStdTime)\n",
    "noisySyncedWaveform = noisyRxWaveform.sync(offset)\n",
    "noisyRxGrid = noisySyncedWaveform.ofdmDemodulate(bwp)\n",
    "noiseGrid = noisyRxGrid.grid - rxGrid.grid\n",
    "print(f\"Noise Grid STD: {noiseGrid.std()}\")\n",
    "print(f\"Noise STD (Freq): {noiseStdFreq}\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7154c02a-e752-4c2e-ae0c-a4b29c5e5f59",
   "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
}