DeepLearningDataSynthesisFor5GChannelEstimationExample.m

%% Deep Learning Data Synthesis for 5G Channel Estimation
%
% This example shows how to train a convolutional neural network (CNN) for
% channel estimation using Deep Learning Toolbox(TM) and data generated
% with 5G Toolbox(TM). Using the trained CNN, you perform channel
% estimation in single-input single-output (SISO) mode, utilizing the
% physical downlink shared channel (PDSCH) demodulation reference signal
% (DM-RS).
% 
% Copyright 2019-2020 The MathWorks, Inc.

%% Introduction
%
% The general approach to channel estimation is to insert known reference
% pilot symbols into the transmission and then interpolate the rest of the
% channel response by using these pilot symbols.
%
% <<../DeepLearningDataSynthesis5G_ChEstimationOverview.png>>
%
% 
% For an example showing how to use this channel estimation approach, see
% <docid:5g_ug#mw_new-radio-pdsch-throughput NR PDSCH
% Throughput>.
%
% You can also use deep learning techniques to perform channel estimation.
% For example, by viewing the PDSCH resource grid as a 2-D image, you can
% turn the problem of channel estimation into an image processing problem,
% similar to denoising or super-resolution, where CNNs are effective.
%
% Using 5G Toolbox, you can customize and generate standard-compliant
% waveforms and channel models to use as training data. Using Deep Learning
% Toolbox, you can use this training data to train a channel estimation
% CNN. This example shows how to generate such training data and how to
% train a channel estimation CNN. The example also shows how to use the
% channel estimation CNN to process images that contain linearly
% interpolated received pilot symbols. The example concludes by visualizing
% the results of the neural network channel estimator in comparison to
% practical and perfect estimators.
%
% <<../DeepLearningDataSynthesis5G_ExampleOverview.png>>
% 
%% Neural Network Training
%
% Neural network training consists of these steps:
%
% * Data generation
% * Splitting the generated data into training and validation sets
% * Defining the CNN architecture
% * Specifying the training options, optimizer, and learning rate
% * Training the network
%
% Due to the large number of signals and possible scenarios, training can
% take several minutes. By default, training is disabled, a pretrained
% model is used. You can enable training by setting |trainModel| to true.

trainModel = false;
%% 
% If you have Parallel Computing Toolbox(TM) installed and a supported
% CUDA-enabled NVIDIA(R) GPU set up, the network training uses GPU
% acceleration by default. The <docid:nnet_ref#bu6sn4c trainNetwork>
% function allows you to override this default behaviour. For a list of
% supported GPUs, see <docid:distcomp_ug#mw_57e04559-0b60-42d5-ad55-e77ec5f5865f GPU Support by Release>.
%
% Data generation is set to produce 256 training examples or training data
% sets. This amount of data is sufficient to train a functional channel
% estimation network on a CPU in a reasonable time. For comparison, the
% pretrained model is based on 16,384 training examples.
%
% Training data of the CNN model has a fixed size dimensionality, the
% network can only accept 612-by-14-by-1 grids, i.e. 612 subcarriers, 14
% OFDM symbols and 1 antenna. Therefore, the model can only operate on a
% fixed bandwidth allocation, cyclic prefix length, and a single receive
% antenna.
%
% The CNN treats the resource grids as 2-D images, hence each element of
% the grid must be a real number. In a channel estimation scenario, the
% resource grids have complex data. Therefore, the real and imaginary parts
% of these grids are input separately to the CNN. In this example, the
% training data is converted from a complex 612-by-14 matrix into a
% real-valued 612-by-14-by-2 matrix, where the third dimension denotes the
% real and imaginary components. Because you have to input the real and
% imaginary grids into the neural network separately when making
% predictions, the example converts the training data into 4-D arrays of
% the form 612-by-14-by-1-by-2N, where N is the number of training
% examples.
%
% To ensure that the CNN does not overfit the training data, the training
% data is split into validation and training sets. The validation data is
% used for monitoring the performance of the trained neural network at
% regular intervals, as defined by |valFrequency|, approximately 5 per
% epoch. Stop training when the validation loss stops improving. In this
% instance, the validation data size is the same as the size of a single
% mini-batch due to the small size of the data set.
%
% The returned channel estimation CNN is trained on various channel
% configurations based on different delay spreads, doppler shifts, and SNR
% ranges between 0 and 10 dB.

% Set the random seed for reproducibility (this has no effect if a GPU is
% used)
rng(42)

if trainModel
    % Generate the training data
    [trainData,trainLabels] = hGenerateTrainingData(256);

    % Set the number of examples per mini-batch
    batchSize = 32;

    % Split real and imaginary grids into 2 image sets, then concatenate
    trainData = cat(4,trainData(:,:,1,:),trainData(:,:,2,:));
    trainLabels = cat(4,trainLabels(:,:,1,:),trainLabels(:,:,2,:));

    % Split into training and validation sets
    valData = trainData(:,:,:,1:batchSize);
    valLabels = trainLabels(:,:,:,1:batchSize);

    trainData = trainData(:,:,:,batchSize+1:end);
    trainLabels = trainLabels(:,:,:,batchSize+1:end);

    % Validate roughly 5 times every epoch
    valFrequency = round(size(trainData,4)/batchSize/5); 

    % Define the CNN structure
    layers = [ ...
        imageInputLayer([612 14 1],'Normalization','none')
        convolution2dLayer(9,64,'Padding',4)
        reluLayer
        convolution2dLayer(5,64,'Padding',2,'NumChannels',64)
        reluLayer
        convolution2dLayer(5,64,'Padding',2,'NumChannels',64)
        reluLayer
        convolution2dLayer(5,32,'Padding',2,'NumChannels',64)
        reluLayer
        convolution2dLayer(5,1,'Padding',2,'NumChannels',32)
        regressionLayer
    ];

    % Set up a training policy
    options = trainingOptions('adam', ...
        'InitialLearnRate',3e-4, ...
        'MaxEpochs',5, ...
        'Shuffle','every-epoch', ...
        'Verbose',false, ...
        'Plots','training-progress', ...
        'MiniBatchSize',batchSize, ...
        'ValidationData',{valData, valLabels}, ...
        'ValidationFrequency',valFrequency, ...
        'ValidationPatience',5);

    % Train the network. The saved structure trainingInfo contains the
    % training progress for later inspection. This structure is useful for
    % comparing optimal convergence speeds of different optimization
    % methods.
    [channelEstimationCNN,trainingInfo] = trainNetwork(trainData, ...
        trainLabels,layers,options);

else
    % Load pretrained network if trainModel is set to false
    load('trainedChannelEstimationNetwork.mat')
end
%%
% Inspect the composition and individual layers of the model. The model has
% 5 convolutional layers. The input layer expects matrices of size
% 612-by-14, where 612 is the number of subcarriers and 14 is the number of
% OFDM symbols. Each element is a real number, since the real and imaginary
% parts of the complex grids are input separately.
channelEstimationCNN.Layers

%% Create Channel Model for Simulation
%
% Set the simulation noise level in dB.
SNRdB = 10;

%% 
% Load the predefined simulation parameters, including the PDSCH
% parameters and DM-RS configuration. The returned object |carrier| is a
% valid carrier configuration object and |pdsch| is a PDSCH configuration
% structure set for a SISO transmission.
[gnb,carrier,pdsch] = hDeepLearningChanEstSimParameters();

%%
% Create a TDL channel model and set channel parameters. To compare
% different channel responses of the estimators, you can change these
% parameters later.

channel = nrTDLChannel;
channel.Seed = 0;
channel.DelayProfile = 'TDL-A';
channel.DelaySpread = 3e-7;
channel.MaximumDopplerShift = 50;

% This example supports only SISO configuration
channel.NumTransmitAntennas = 1;
channel.NumReceiveAntennas = 1;

waveformInfo = nrOFDMInfo(carrier);
channel.SampleRate = waveformInfo.SampleRate;

%% 
% Get the maximum number of delayed samples by a channel multipath
% component. This number is calculated from the channel path with the
% largest delay and the implementation delay of the channel filter. This
% number is needed to flush the channel filter when obtaining the received
% signal.
chInfo = info(channel);
maxChDelay = ceil(max(chInfo.PathDelays*channel.SampleRate))+chInfo.ChannelFilterDelay;

%% Simulate PDSCH Transmission
%
% Simulate a PDSCH transmission by performing these steps:
%
% * Generate PDSCH resource grid
% * Insert DM-RS symbols
% * Perform OFDM modulation
% * Send modulated waveform through the channel model
% * Add white Gaussian noise
% * Perform perfect timing synchronization
% * Perform OFDM demodulation 
% 

% Generate DM-RS indices and symbols
[~,dmrsIndices,dmrsSymbols,pdschIndicesInfo] = hPDSCHResources(gnb,pdsch);

% Create PDSCH resource grid
pdschGrid = nrResourceGrid(carrier);

% Map PDSCH DM-RS symbols to the grid
pdschGrid(dmrsIndices) = pdschGrid(dmrsIndices)+dmrsSymbols;

% OFDM-modulate associated resource elements
txWaveform = nrOFDMModulate(carrier,pdschGrid);
%% 
% To flush the channel content, append zeros at the end of the transmitted
% waveform. These zeros take into account any delay introduced in the
% channel, such as multipath and implementation delay. The number of
% zeros depends on the sampling rate, delay profile, and delay spread.
txWaveform = [txWaveform; zeros(maxChDelay,size(txWaveform,2))];
%% 
% Send data through the TDL channel model. 
[rxWaveform,pathGains,sampleTimes] = channel(txWaveform);
%% 
% Add additive white Gaussian noise (AWGN) to the received time-domain
% waveform. To take into account sampling rate, normalize the noise power.
% The SNR is defined per resource element (RE) for each receive antenna
% (3GPP TS 38.101-4).
SNR = 10^(SNRdB/20); % Calculate linear noise gain
N0 = 1/(sqrt(2.0*gnb.NRxAnts*double(waveformInfo.Nfft))*SNR);
noise = N0*complex(randn(size(rxWaveform)),randn(size(rxWaveform)));
rxWaveform = rxWaveform + noise;
%% 
% Perform perfect synchronization. To find the strongest multipath
% component, use the information provided by the channel.

% Get path filters for perfect channel estimation
pathFilters = getPathFilters(channel); 
[offset,~] = nrPerfectTimingEstimate(pathGains,pathFilters);

rxWaveform = rxWaveform(1+offset:end, :);
%% 
% OFDM-demodulate the received data to recreate the resource grid.

rxGrid = nrOFDMDemodulate(carrier,rxWaveform);

% Pad the grid with zeros in case an incomplete slot has been demodulated
[K,L,R] = size(rxGrid);
if (L < carrier.SymbolsPerSlot)
    rxGrid = cat(2,rxGrid,zeros(K,carrier.SymbolsPerSlot-L,R));
end

%% Compare and Visualize Various Channel Estimations
%
% You can perform and compare the results of perfect, practical, and neural
% network estimations of the same channel model. 
% 
% To perform perfect channel estimation, use the
% <docid:5g_ref#mw_function_nrPerfectChannelEstimate
% nrPerfectChannelEstimate> function using the value of the path gains
% provided by the channel.
estChannelGridPerfect = nrPerfectChannelEstimate(carrier,pathGains, ...
    pathFilters,offset,sampleTimes);
%% 
% To perform practical channel estimation, use the
% <docid:5g_ref#mw_function_nrChannelEstimate nrChannelEstimate> function.
[estChannelGrid,~] = nrChannelEstimate(carrier,rxGrid,dmrsIndices, ...
    dmrsSymbols,'CDMLengths',pdschIndicesInfo.CDMLengths);
%%
% To perform channel estimation using the neural network, you must
% interpolate the received grid. Then split the interpolated image into its
% real and imaginary parts and input these images together into the neural
% network as a single batch. Use the
% <docid:nnet_ref#mw_0a51db93-cccf-4b2f-ae4c-6724cbf5ec46 predict> function
% to make predictions on the real and imaginary images. Finally,
% concatenate and transform the results back into complex data.

% Interpolate the received resource grid using pilot symbol locations
interpChannelGrid = hPreprocessInput(rxGrid,dmrsIndices,dmrsSymbols);

% Concatenate the real and imaginary grids along the batch dimension
nnInput = cat(4,real(interpChannelGrid),imag(interpChannelGrid));

% Use the neural network to estimate the channel
estChannelGridNN = predict(channelEstimationCNN,nnInput);

% Convert results to complex 
estChannelGridNN = complex(estChannelGridNN(:,:,:,1),estChannelGridNN(:,:,:,2));

%% 
% Calculate the mean squared error (MSE) of each estimation method.
neural_mse = mean(abs(estChannelGridPerfect(:) - estChannelGridNN(:)).^2);
interp_mse = mean(abs(estChannelGridPerfect(:) - interpChannelGrid(:)).^2);
practical_mse = mean(abs(estChannelGridPerfect(:) - estChannelGrid(:)).^2);

%%
% Plot the individual channel estimations and the actual channel
% realization obtained from the channel filter taps. Both the practical
% estimator and the neural network estimator outperform linear
% interpolation.

plotChEstimates(interpChannelGrid,estChannelGrid,estChannelGridNN,estChannelGridPerfect,...
    interp_mse,practical_mse,neural_mse);

%% References
% # van de Beek, Jan&ndash;Jaap, Ove Edfors, Magnus Sandell, Sarah Kate
% Wilson, and Per Ola Borjesson. &ldquo;On Channel Estimation in OFDM
% Systems.&rdquo; In 1995 IEEE 45th Vehicular Technology Conference.
% Countdown to the Wireless Twenty&ndash;First Century, 2:815&ndash;19,
% July 1995.
% # Ye, Hao, Geoffrey Ye Li, and Biing-Hwang Juang. &ldquo;Power of Deep
% Learning for Channel Estimation and Signal Detection in OFDM
% Systems.&rdquo; IEEE Wireless Communications Letters 7, no. 1 (February
% 2018): 114&ndash;17.
% # Soltani, Mehran, Vahid Pourahmadi, Ali Mirzaei, and Hamid Sheikhzadeh.
% &ldquo;Deep Learning&ndash;Based Channel Estimation.&rdquo; Preprint, submitted October 13,
% 2018.

%% Local Functions

function hest = hPreprocessInput(rxGrid,dmrsIndices,dmrsSymbols)
% Perform linear interpolation of the grid and input the result to the
% neural network This helper function extracts the DM-RS symbols from
% dmrsIndices locations in the received grid rxGrid and performs linear
% interpolation on the extracted pilots.

    % Obtain pilot symbol estimates
    dmrsRx = rxGrid(dmrsIndices);
    dmrsEsts = dmrsRx .* conj(dmrsSymbols);

    % Create empty grids to fill after linear interpolation
    [rxDMRSGrid, hest] = deal(zeros(size(rxGrid)));
    rxDMRSGrid(dmrsIndices) = dmrsSymbols;
    
    % Find the row and column coordinates for a given DMRS configuration
    [rows,cols] = find(rxDMRSGrid ~= 0);
    dmrsSubs = [rows,cols,ones(size(cols))];
    [l_hest,k_hest] = meshgrid(1:size(hest,2),1:size(hest,1));

    % Perform linear interpolation
    f = scatteredInterpolant(dmrsSubs(:,2),dmrsSubs(:,1),dmrsEsts);
    hest = f(l_hest,k_hest);

end

function [trainData,trainLabels] = hGenerateTrainingData(dataSize)
% Generate training data examples for channel estimation
% Run dataSize number of iterations to create random channel configurations
% and pass an OFDM-modulated fixed PDSCH grid with only the DM-RS symbols
% inserted. Perform perfect timing synchronization and OFDM demodulation,
% extracting the pilot symbols and performing linear interpolation at each
% iteration. Use perfect channel information to create the
% label data. The function returns 2 arrays - the training data and labels.

    fprintf('Starting data generation...\n')

    % List of possible channel profiles
    delayProfiles = {'TDL-A', 'TDL-B', 'TDL-C', 'TDL-D', 'TDL-E'};

    [simParameters, carrier, pdsch] = hDeepLearningChanEstSimParameters();

    % Create the channel model object
    nTxAnts = simParameters.NTxAnts;
    nRxAnts = simParameters.NRxAnts;

    channel = nrTDLChannel; % TDL channel object
    channel.NumTransmitAntennas = nTxAnts;
    channel.NumReceiveAntennas = nRxAnts;

    % Use the value returned from <matlab:edit('nrOFDMInfo') nrOFDMInfo> to
    % set the channel model sample rate
    waveformInfo = nrOFDMInfo(carrier);
    channel.SampleRate = waveformInfo.SampleRate;

    % Get the maximum number of delayed samples by a channel multipath
    % component. This number is calculated from the channel path with the largest
    % delay and the implementation delay of the channel filter, and is required
    % to flush the channel filter to obtain the received signal.
    chInfo = info(channel);
    maxChDelay = ceil(max(chInfo.PathDelays*channel.SampleRate)) + chInfo.ChannelFilterDelay;

    % Return DM-RS indices and symbols
    [~,dmrsIndices,dmrsSymbols,~] = hPDSCHResources(simParameters,pdsch);

    % PDSCH mapping in grid associated with PDSCH transmission period
    pdschGrid = nrResourceGrid(carrier,nTxAnts);

    % PDSCH DM-RS precoding and mapping
    [~,dmrsAntIndices] = nrExtractResources(dmrsIndices,pdschGrid);
    pdschGrid(dmrsAntIndices) = pdschGrid(dmrsAntIndices) + dmrsSymbols;

    % OFDM modulation of associated resource elements
    txWaveform_original = nrOFDMModulate(carrier,pdschGrid);

    % Acquire linear interpolator coordinates for neural net preprocessing
    [rows,cols] = find(pdschGrid ~= 0);
    dmrsSubs = [rows, cols, ones(size(cols))];
    hest = zeros(size(pdschGrid));
    [l_hest,k_hest] = meshgrid(1:size(hest,2),1:size(hest,1));

    % Preallocate memory for the training data and labels
    numExamples = dataSize;
    [trainData, trainLabels] = deal(zeros([612 14 2 numExamples]));

    % Main loop for data generation, iterating over the number of examples
    % specified in the function call. Each iteration of the loop produces a
    % new channel realization with a random delay spread, doppler shift,
    % and delay profile. Every perturbed version of the transmitted
    % waveform with the DM-RS symbols is stored in trainData, and the
    % perfect channel realization in trainLabels.
    for i = 1:numExamples
        % Release the channel to change nontunable properties
        channel.release

        % Pick a random seed to create different channel realizations
        channel.Seed = randi([1001 2000]);

        % Pick a random delay profile, delay spread, and maximum doppler shift
        channel.DelayProfile = string(delayProfiles(randi([1 numel(delayProfiles)])));
        channel.DelaySpread = randi([1 300])*1e-9;
        channel.MaximumDopplerShift = randi([5 400]);

        % Send data through the channel model. Append zeros at the end of
        % the transmitted waveform to flush channel content. These zeros
        % take into account any delay introduced in the channel, such as
        % multipath delay and implementation delay. This value depends on
        % the sampling rate, delay profile, and delay spread
        txWaveform = [txWaveform_original; zeros(maxChDelay, size(txWaveform_original,2))];
        [rxWaveform,pathGains,sampleTimes] = channel(txWaveform);

        % Add additive white Gaussian noise (AWGN) to the received time-domain
        % waveform. To take into account sampling rate, normalize the noise power.
        % The SNR is defined per RE for each receive antenna (3GPP TS 38.101-4).   
        SNRdB = randi([0 10]);  % Random SNR values between 0 and 10 dB
        SNR = 10^(SNRdB/20);    % Calculate linear noise gain
        N0 = 1/(sqrt(2.0*nRxAnts*double(waveformInfo.Nfft))*SNR);
        noise = N0*complex(randn(size(rxWaveform)),randn(size(rxWaveform)));
        rxWaveform = rxWaveform + noise;

        % Perfect synchronization. Use information provided by the channel
        % to find the strongest multipath component
        pathFilters = getPathFilters(channel); % Get path filters for perfect channel estimation
        [offset,~] = nrPerfectTimingEstimate(pathGains,pathFilters);

        rxWaveform = rxWaveform(1+offset:end, :);

        % Perform OFDM demodulation on the received data to recreate the
        % resource grid, including padding in case practical
        % synchronization results in an incomplete slot being demodulated
        rxGrid = nrOFDMDemodulate(carrier,rxWaveform);
        [K,L,R] = size(rxGrid);
        if (L < carrier.SymbolsPerSlot)
            rxGrid = cat(2,rxGrid,zeros(K,carrier.SymbolsPerSlot-L,R));
        end

        % Perfect channel estimation, using the value of the path gains
        % provided by the channel. This channel estimate does not
        % include the effect of transmitter precoding
        estChannelGridPerfect = nrPerfectChannelEstimate(carrier,pathGains, ...
            pathFilters,offset,sampleTimes);

        % Linear interpolation
        dmrsRx = rxGrid(dmrsIndices);
        dmrsEsts = dmrsRx .* conj(dmrsSymbols);
        f = scatteredInterpolant(dmrsSubs(:,2),dmrsSubs(:,1),dmrsEsts);
        hest = f(l_hest,k_hest);

        % Split interpolated grid into real and imaginary components and
        % concatenate them along the third dimension, as well as for the
        % true channel response
        rx_grid = cat(3, real(hest), imag(hest));
        est_grid = cat(3, real(estChannelGridPerfect), ...
            imag(estChannelGridPerfect));

        % Add generated training example and label to the respective arrays
        trainData(:,:,:,i) = rx_grid;
        trainLabels(:,:,:,i) = est_grid;

        % Data generation tracker
        if mod(i,round(numExamples/25)) == 0
            fprintf('%3.2f%% complete\n',i/numExamples*100);
        end
    end
    fprintf('Data generation complete!\n')
end

function plotChEstimates(interpChannelGrid,estChannelGrid,estChannelGridNN,estChannelGridPerfect,...
                         interp_mse,practical_mse,neural_mse)
% Plot the different channel estimates and display the measured MSE

    figure

    subplot(1,4,1)
    imagesc(abs(interpChannelGrid));
    xlabel('OFDM Symbol');
    ylabel('Subcarrier');
    title({'Linear Interpolation', ['MSE: ', num2str(interp_mse)]});

    subplot(1,4,2)
    imagesc(abs(estChannelGrid));
    xlabel('OFDM Symbol');
    ylabel('Subcarrier');
    title({'Practical Estimator', ['MSE: ', num2str(practical_mse)]});

    subplot(1,4,3)
    imagesc(abs(estChannelGridNN));
    xlabel('OFDM Symbol');
    ylabel('Subcarrier');
    title({'Neural Network', ['MSE: ', num2str(neural_mse)]});

    subplot(1,4,4)
    imagesc(abs(estChannelGridPerfect));
    xlabel('OFDM Symbol');
    ylabel('Subcarrier');
    title({'Actual Channel'});
end