Generating the Dataset for Channel Estimation Training

The first step in this project is generating the datasets used for training, validation, and testing. In this notebook, we construct an OFDM communication pipeline and generate input/output pairs for supervised learning. The inputs consist of received resource grids together with transmitted DMRS information, while the outputs are the corresponding ground-truth channel labels.

The experiments assume a two-layer MIMO configuration with 16 transmit antennas and 4 receive antennas. Under this configuration, each slot may contain up to four code blocks. Each dataset sample is represented as a 6 × 14 × 540 tensor, and the corresponding ground-truth label is a 4 × 14 × 540 tensor.

The figure below illustrates the dataset generation process. Communication over a single slot produces four dataset samples, one for each receive antenna. In this work, we use:

  • \(N_L = 2\): number of layers

  • \(L = 14\): number of OFDM symbols per slot

  • \(K = 540\): number of subcarriers

  • \(N_r = 4\): number of receive antennas

Because DMRS-based channel estimation is used throughout the pipeline, the effective channel includes the impact of precoding.

Data Generation Pipeline

We begin by importing the required modules from NeoRadium.

[1]:
import numpy as np
import time, os

from neoradium import BandwidthPart, PDSCH, CdlChannel, AntennaPanel, Grid, random
from ChEstUtils import getRandomPilotInfo, getModelIn, getLabels

The makeDataset function below runs the pipeline and creates samples with 0, 1, 2, or 3 successfully decoded code blocks. When 0 code blocks are used, the sample contains DMRS only. Please refer to the ChEstUtils.py file for the implementation details of the functions getRandomPilotInfo, getModelIn, and getLabels.

[2]:
def makeDataset(numSlots, snrDbs, seed, pdsch, fileName=None, freqDomain=True):
    bwp = pdsch.bwp                       # The only bandwidth part in the carrier

    # We don't need to use channel coding in the pipeline during dataset generation. However, we need
    # to know the code-block sizes to be able to find the corresponding REs in the grid. The following
    # code calculates the code-block sizes for each of the four code blocks.
    ldpc = pdsch.getLdpcCodec(coderates=490/1024)

    pdsch.initGrid()                                            # Create and initialize PDSCH's internal grid
    numBits = pdsch.getBitCapacity()                            # number of bits available in the resource grid
    txBlock = random.bits(ldpc.txBlockSizes[0])                 # Create random binary data for txBlock
    rateMatchedCBs = ldpc.encode(txBlock, numBits[0], concatCBs=False)
    cbSizes = [len(cb) for cb in rateMatchedCBs]                # Code block sizes

    # Create a random CDL channel-matrix generator
    chanGen = CdlChannel.getChanGen(numSlots, bwp,              # Number of channels and bandwidth part
                                    profiles="ABCDE",           # Randomly pick a CDL profile
                                    delaySpread=300,            # 300 ns
                                    ueSpeed=0.5,                # 0.5 m/s ≈ 6.7 Hz Doppler
                                    carrierFreq=4e9,            # Carrier frequency
                                    txAntenna=AntennaPanel([2,4], polarization="x"),  # 16 TX antennas
                                    rxAntenna=AntennaPanel([1,2], polarization="x"),  # 4 RX antennas
                                    seed=seed)

    samples, labels = [], []                        # Initialize samples and labels
    t0 = time.monotonic()                           # Start time for time estimation
    random.setSeed(seed)
    print(f"Making dataset for {numSlots:,} slots")
    for s, channelMatrix in enumerate(chanGen):
        pdsch.initGrid()                            # Create and initialize PDSCH's internal grid
        numBits = pdsch.getBitCapacity()[0]         # Number of bits available in the resource grid
        txBits = random.bits(numBits)               # Create random binary data

        pdsch.setPdschData(txBits)                  # Map/modulate the data to the resource grid

        precoder = pdsch.getPrecodingMatrix(channelMatrix)      # Get the precoder matrix
        txGrid = bwp.createGrid(channelMatrix.shape[3])         # Create the transmitted resource grid
        pdsch.precodeTo(txGrid, precoder)                       # Precode PDSCH data into the txGrid

        snrDb = snrDbs[s%len(snrDbs)]               # Get next SNR value

        if freqDomain:
            rxGrid = txGrid.applyChannel(channelMatrix)         # Apply the channel in the frequency domain
            noisyRxGrid = rxGrid.addNoise(snrDb=snrDb)          # Add noise
        else:
            channel = chanGen.curChan                           # Get the channel model
            txWaveform = txGrid.ofdmModulate()                  # OFDM modulation
            maxDelay = channel.getMaxDelay()                    # Calculate the maximum channel delay
            txWaveform = txWaveform.pad(maxDelay)               # Pad the waveform with zeros
            rxWaveform = channel.applyToSignal(txWaveform)      # Apply the channel to the waveform
            noisyRxWaveform = rxWaveform.addNoise(snrDb=snrDb, bwp=bwp) # Add noise
            offset = channel.getTimingOffset()                  # Get the timing offset for synchronization
            syncedWaveform = noisyRxWaveform.sync(offset)       # Synchronize the received waveform
            noisyRxGrid = syncedWaveform.ofdmDemodulate(bwp)    # OFDM demodulate the synchronized waveform

        # pilotIdx contains the inices in the txGrid with known values. Known values could include DMRS
        # and a number of successfully decoded code blocks (0..numCBs-1)
        pilotIdx = getRandomPilotInfo(pdsch, cbSizes)
        newSamples = getModelIn(pilotIdx, pdsch.grid, noisyRxGrid, 10000)   # rr x 2*(pp+1) x ll x kk
        effChannel = CdlChannel.getEffChannel(channelMatrix, precoder)      # Get ground-truth effective channel
        newLabels = getLabels( effChannel )                                 # rr x 2*pp x ll x kk

        samples += [ newSamples ]
        labels += [ newLabels ]

        dt = time.monotonic()-t0                                # Calculate the elapsed time
        percentDone = (s+1)*100/numSlots                        # Calculate the completion percentage

        # Print progress messages
        remainTime = int(np.round(100*dt/percentDone-dt))       # Estimated remaining time
        print(f"  {int(percentDone)}% done in {int(np.round(dt)):,} Sec., Remaining time: {remainTime:,} Sec.",
              end='            \r')

    # n: Number of samples in the dataset, pp: Number of ports, ll: Number of OFDM symbols, kk: Number of subcarriers
    samples = np.concatenate(samples, axis=0)       # n x 2*(pp+1) x ll x kk
    labels = np.concatenate(labels, axis=0)         # n x 2*pp x ll x kk

    if fileName is not None:
        np.save(fileName, np.concatenate([samples,labels],axis=1))   # Save the dataset to the specified file
        print(f"\r  Done. ({dt:.2f} Sec.) Saved to \"{fileName}\".                       ")
    else:
        print(f"\r  Done. ({dt:.2f} Sec.)                                                ")

    return samples, labels
[3]:

bwp = BandwidthPart(numRbs=45, spacing=30) # Create a BandwidthPart object # Create a 2-layer PDSCH object with type-2 DMRS on symbols 2 and 11 pdsch = PDSCH(bwp, numLayers=2, modulation="16QAM") pdsch.setDMRS(configType=2, additionalPos=1) dataPath = "/data/datasets/SelfRefine/" # Replace with the location of your data files os.makedirs(dataPath, exist_ok=True) # Create the data folder if it does not exist snrDbs = np.arange(-15,-9,.5) # Set the range of SNR values (in dB) # Create training, validation, and test dataset files (About 30 GB of disk space needed): makeDataset(17500, snrDbs, seed=123, pdsch=pdsch, fileName=os.path.join(dataPath,"Train.npy")) makeDataset(2500, snrDbs, seed=456, pdsch=pdsch, fileName=os.path.join(dataPath,"Valid.npy")) makeDataset(5000, snrDbs, seed=789, pdsch=pdsch, fileName=os.path.join(dataPath,"Test.npy"));
Making dataset for 17,500 slots
  Done. (3135.66 Sec.) Saved to "/data/datasets/SelfRefine/Train.npy".
Making dataset for 2,500 slots
  Done. (445.14 Sec.) Saved to "/data/datasets/SelfRefine/Valid.npy".
Making dataset for 5,000 slots
  Done. (892.34 Sec.) Saved to "/data/datasets/SelfRefine/Test.npy".
[ ]: