AI for Positioning Accuracy Enhancement

simParameters.Mode            = "Simulate"; % Simulate with the pretrained network or train the network
simParameters.TrainRatio      = 0.95; % Define the ratio of training data relative to validation and test data

simParameters.Scenario        = "InF-DH"; % "InF-SL", "InF-DL", "InF-SH", "InF-DH", "InF-HH"
simParameters.HallSize        = [120 60 10]; % Dimensions of hall for InF scenarios: [length width height] in meters
simParameters.ClutterSize     = 2; % Clutter size (m)
simParameters.ClutterHeight   = 6; % Clutter height (m)
simParameters.ClutterDensity  = 0.6; % Clutter density, between 0 and 1

simParameters.BSSeparation    = [20 20]; 
simParameters.BSAntennaHeight = 8;
simParameters.BSNoiseFigure   = 5;

simParameters.UEDrop          = "random"; % UE drop type can be either 'random' or 'grid'
simParameters.NumUEs          = 15; % Number of UE positions; applicable if ueDropType is 'random'
simParameters.UESeparation    = [30 30]; % UE separation in x and y axes in meters; applicable if ueDropType is 'grid'
simParameters.UEAntennaHeight = 1.5; % UE antenna height on the z axis in meters

% Set a random number seed for repeatability
simParameters.Seed = 23;
rng(simParameters.Seed,"twister");

% Configure and visualize the scenario
simParameters = h5GIndoorFactoryScenario(simParameters);

simParameters.Carrier = nrCarrierConfig;
simParameters.Carrier.NSizeGrid         = 270;
simParameters.Carrier.SubcarrierSpacing = 30; % in kHz

simParameters.SRS = nrSRSConfig;
simParameters.SRS.NumSRSSymbols = 12;  % Number of SRS OFDM symbols
simParameters.SRS.KTC           = 8 ; % Transmission comb number (subcarrier sparsity)
simParameters.SRS = hSRSConfig(simParameters);

simParameters.Algorithmic = struct();
simParameters.Algorithmic.PerfectChannelEstimator = true;
simParameters.Algorithmic.PerfectTimingEstimator = true;
simParameters.Algorithmic.FrequencyDomainChannelFiltering = false;

% Generate propagation channels based on TR 38.901
channels = h38901ChannelSetup(simParameters);

% Generate training and validation data sets if in the training workflow
if strcmp(simParameters.Mode,"Train")
    [dataTrain,dataValidate,inputSize,outputSize] = channels2ImageDataset(simParameters,channels);
end

if strcmp(simParameters.Mode,"Train")
    % Generate the ResNet architecture
    [dlnet,dlparam] = customResNetLayers(inputSize,outputSize,[64 128 256],[1 1 1]);
    
    % Training options
    miniBatchSize = 128;

    options = trainingOptions("adam",...
        InitialLearnRate=1e-3,...
        LearnRateSchedule="piecewise",...
        LearnRateDropFactor=0.1,...
        MiniBatchSize=miniBatchSize,...
        MaxEpochs=70,...
        Shuffle="every-epoch",...
        ValidationData=dataValidate,...
        Metrics="rmse",...
        ExecutionEnvironment="auto",...
        Plots="training-progress",...
        Verbose=false);
    lossFcn = "mse";
    
    % Display the network summary
    hDisplayNetworkSummary(dlnet,dlparam);
    
    % Train the network
    net = trainnet(dataTrain,dlnet,lossFcn,options);
else % Simulate
    % Load the pre-trained network
    if exist("trained_positioning.mat","file")
        loadNet = load("trained_positioning.mat");
    else
        loadNet = load("resnet_positioning.mat");
    end
    net = loadNet.net;
    dlparam = loadNet.dlparam;

    % Display the loaded network summary
    hDisplayNetworkSummary(net,dlparam);
end

if strcmp(simParameters.Mode,"Train")
    save("trained_positioning.mat","net","dlparam")
else % Simulate
    results = positioningSimulation(simParameters,channels,net);
end

if strcmp(simParameters.Mode,"Train")
    msgbox("Network is trained and saved in the current folder. You can now run the example in simulation mode.")
else % Simulate
    positioningSimulationResults(results)
end

function [dlnet,dlparam] = customResNetLayers(inputSize,outputSize,resFilters,resRepetitions)
%customResNetLayers Generate custom ResNet architecture.
%   [DLNET, DLPARAM] = customResNetLayers(INPUTSIZE, OUTPUTSIZE, RESFILTERS, RESREPETITIONS)
%   creates and returns the custom residual dlnetwork (DLNET) and the
%   network parameters (DLPARAM). The network has an input size of
%   INPUTSIZE and an output size of OUTPUTSIZE. The number of filters for
%   the 3x3 convolutional layers within each residual block is specified by
%   RESFILTERS, and the number of repetitions of each 3x3 convolution layer
%   is given by RESREPETITIONS.
%
%   The function generates a ResNet architecture without bottleneck
%   residual blocks. The architecture for an 18-layer network is specified
%   by (resFilters = [64 128 256 512]) and (resRepetitions = [2 2 2 2]), and
%   for a 34-layer network by (resFilters = [64 128 256 512]) and
%   (resRepetitions = [3 4 6 3]).
%
%   The architecture follows the pattern:
%
%   Input -> Conv2D -> BN -> ReLU -> Conv2D -> BN -> Addition -> ReLU -> Output    
%     |_________________________________________________|
%                              |
%         Identity or projection shortcut (using 1x1 filters, with the same
%         number of filters as the main path output, and a stride of 2 if
%         downsampling is required)
%
%   Note: This function does not support the 50-layer, 101-layer, and
%   152-layer ResNet architectures.

    dlparam = struct('idx',1,'conv',0,'relu',0,'bn',0,'pool',0,'fc',0,'add',0,'shortcut',0);
    dlparam.numChannels = 'auto';
    
    % Step 1: Generate initial layers before residual blocks
    dlparam.conv = dlparam.conv + 1;
    dlparam.bn = dlparam.bn + 1;
    dlparam.relu = dlparam.relu + 1;
    initLayers = [imageInputLayer(inputSize,Normalization='zerocenter',Name='input')
                  convolution2dLayer(7,64,Stride=2,Padding='same',Name='conv')
                  batchNormalizationLayer(Name='bn')
                  reluLayer(Name='relu')];
    dlnet = dlnetwork(initLayers);

    % Step 2: Generate residual blocks (intermediate layers)
    for filter = 1:length(resFilters)
        for repeat = 1:resRepetitions(filter)
            mainPath = [];
            shortcutPath = [];
      
            % Step 2.1: Create the main path
            % Reduce dimensions (use stride 2) if the block spans across the feature map

            % Begin with activation for the previous layer
            if dlparam.idx == 1
                dlparam.pool = dlparam.pool + 1; % Plus one to account for the batch norm from the previous layer
                mainPath = [mainPath
                            maxPooling2dLayer(3,Stride=2,Padding='same',Name='maxpool')];
                            sourceName = 'relu'; % Only for the first residual block
            else
                dlparam.relu = dlparam.relu + 1; % Plus one to account for the activation from the previous layer
                mainPath = [mainPath
                            reluLayer(Name=['relu_',num2str(dlparam.relu)])];
            end

            % First convolution of the main path (conv + bn + activation)
            dlparam.conv = dlparam.conv + 1;
            dlparam.bn = dlparam.bn + 1;
            dlparam.relu = dlparam.relu + 1;
            mainPath = [mainPath
                        convolution2dLayer(3,resFilters(filter),NumChannels=dlparam.numChannels,Stride=((dlparam.idx ~= 1) && (repeat == 1))+1,...
                            Padding='same',Name=['conv_',num2str(dlparam.conv)])
                        batchNormalizationLayer(Name=['bn_',num2str(dlparam.bn)])
                        reluLayer(Name=['relu_',num2str(dlparam.relu)])]; % Stride is 2 if the block represents a new feature map size, otherwise 1
      
            % Second convolution in the main path, excluding the activation
            % function (conv + bn)
            dlparam.conv = dlparam.conv + 1;
            dlparam.bn = dlparam.bn + 1;
            mainPath = [mainPath
                            convolution2dLayer(3,resFilters(filter),NumChannels=dlparam.numChannels,Stride=1,Padding="same",Name=['conv_',num2str(dlparam.conv)])
                            batchNormalizationLayer(Name=['bn_',num2str(dlparam.bn)])];
      
            % Step 2.2: Perform addition operation in the main path
            % followed by batch normalization
            dlparam.add = dlparam.add + 1;
            mainPath = [mainPath
                        additionLayer(2,Name=['add_',num2str(dlparam.add)])];

                dlnet = addLayers(dlnet,mainPath); % Add mainPath

            % Step 2.3: Create the shortcut path
            % Use identity shortcut if the block spans across the feature map;
            % otherwise, use projection to reduce dimensions
            if (dlparam.idx ~= 1) && (repeat == 1) % Projection shortcut
                dlparam.shortcut = dlparam.shortcut + 1;
                shortcutPath = [shortcutPath
                                convolution2dLayer(1,resFilters(filter),Stride=(dlparam.idx ~= 1)+1,Name=['shortcut_conv_',num2str(dlparam.shortcut)])
                                batchNormalizationLayer(Name=['shortcut_bn_',num2str(dlparam.shortcut)])]; % Stride is 1 in the initial projection shortcut

                dlnet = addLayers(dlnet,shortcutPath);
                dlnet = connectLayers(dlnet,mainPath(1).Name,['shortcut_conv_',num2str(dlparam.shortcut)]);
                dlnet = connectLayers(dlnet,['shortcut_bn_',num2str(dlparam.shortcut)],['add_',num2str(dlparam.add),'/in2']);
            else % Identity shortcut
                dlparam.shortcut = dlparam.shortcut + 1;
                shortcutPath = [shortcutPath
                                functionLayer(@(X) X,Name=['shortcut_',num2str(dlparam.shortcut)])];

                dlnet = addLayers(dlnet,shortcutPath); % Add shortcutPath
                dlnet = connectLayers(dlnet,mainPath(1).Name,['shortcut_',num2str(dlparam.shortcut)]);
                dlnet = connectLayers(dlnet,['shortcut_',num2str(dlparam.shortcut)],['add_',num2str(dlparam.add),'/in2']);
            end
      
            % Step 2.4: Interconnect the generated blocks
            destName = mainPath(1).Name;
            dlnet = connectLayers(dlnet,sourceName,destName);
            sourceName = mainPath(end).Name; % Store the last layer's name
            dlparam.idx = dlparam.idx + 1;
        end
    end

    % Step 3: Generate final layers after residual blocks
    dlparam.relu = dlparam.relu + 1;
    dlnet = addLayers(dlnet,reluLayer(Name=['relu_',num2str(dlparam.relu)]));
    dlnet = connectLayers(dlnet,sourceName,['relu_',num2str(dlparam.relu)]);
    dlparam.pool = dlparam.pool + 1;
    dlnet = addLayers(dlnet,globalAveragePooling2dLayer(Name='gloavgpool'));
    dlnet = connectLayers(dlnet,['relu_',num2str(dlparam.relu)],'gloavgpool');
    dlparam.fc = dlparam.fc + 1;
    dlnet = addLayers(dlnet,fullyConnectedLayer(outputSize,Name='fc'));
    dlnet = connectLayers(dlnet,'gloavgpool','fc');
end

function results = positioningSimulation(cfg,channels,net)
% Estimate 2D UE positions for both AI-based and TDoA techniques using UL
% SRS estimated channels

    algParameters = cfg.Algorithmic;
    carrier = cfg.Carrier;
    srs = cfg.SRS;
    Pue = cfg.Pos3UE;
    Pbs = cfg.Pos3BS;
    numUE = size(Pue,1);
    ueTxPower = cfg.PowerBSs;
    bsNoiseFigure = cfg.BSNoiseFigure;
    DisplaySimulationInformation = true;
    
    % Enable channel filtering in all channels
    for i=1:numel(channels)
        release(channels(i).SmallScale);
        channels(i).SmallScale.ChannelFiltering = ~algParameters.FrequencyDomainChannelFiltering;
    end
    
    if ~algParameters.PerfectTimingEstimator && algParameters.FrequencyDomainChannelFiltering
        algParameters.PerfectTimingEstimator = true;
        warning('For frequency domain channel filtering, perfect timing estimation is required.')
    end
    
    LOS = arrayfun(@(x) x.SmallScale.HasLOSCluster, channels);

    for ue = 1:numUE
    
        % This is a single-slot simulation. Set slot number to 0
        carrier.NSlot = 0;
      
        % Transmit SRS from current UE to all BSs
        txUL = hSRSTransmit(carrier,srs,ueTxPower,algParameters);
      
        % Apply fading channels to the signal
        [channels(:,ue),rxUL] = hChannel(carrier,channels(:,ue),txUL);
      
        % Apply AWGN
        rxUL = hAWGN(rxUL,bsNoiseFigure);
      
        % Estimate delays and channel responses from current UE to all BSs
        [H,~,delay] = hSRSReceive(carrier,srs,rxUL,algParameters);
      
        % Estimate the position of the UEs based on TDoA
        estPue = hEstimatePositionTDOA(Pbs,LOS(:,ue),delay.*1/rxUL(1).ofdmInfo.SampleRate);

        % Calculate the positioning estimation error
        results.TDoA.EstimatedPositions(ue,:) = estPue(1:2);
        results.TDoA.EstimationError(ue) = norm(estPue(1:2)-Pue(ue,1:2));

        % Estimate the positions of the UE based on AI
        estPue = hEstimatePositionAI(cfg,H,net);
      
        % Calculate the positioning estimation error
        results.AI.EstimatedPositions(ue,:) = estPue;
        results.AI.EstimationError(ue) = norm(estPue-Pue(ue,1:2));
      
        % Display progress if diagnostics are enabled
        if DisplaySimulationInformation
            fprintf('UE %3g/%3g (%5.1f%%): AI/ML positioning error: %3.3f m. | TDoA positioning error: %3.3f m. \n',...
              ue,numUE,ue/numUE*100,results.AI.EstimationError(ue),results.TDoA.EstimationError(ue));
        end
    
    end
    
    % Calculate the 10th, 50th, and 90th percentiles of the positioning estimation error vector
    results.TDoA.Percentiles = prctile(results.TDoA.EstimationError,[10 50 90]);
    results.AI.Percentiles = prctile(results.AI.EstimationError,[10 50 90]);
end

function txUL = hSRSTransmit(carrier,srs,txPower,algParameters)
% Map the SRS to the resource grid and transmit the waveform

    % Determine active UEs in the current slot by identifying non-empty SRS indices
    indices = nrSRSIndices(carrier,srs);
    active = ~isempty(indices);
    srs = srs(active);
    
    % Create an empty output structure
    txUL = struct('ulWaveform',[],'ulGrid',[],'ofdmInfo',[]);
    
    % Create the transmit resource grid
    ulGrid = nrResourceGrid(carrier,srs.NumSRSPorts,'OutputDataType','single');
    
    % Create SRS and map it to the resource grid
    symbols = nrSRS(carrier,srs);
    ulGrid(indices) = symbols;

    % Scale the amplitude of the resource grid and waveform to account for the
    % configured overall transmitted power
    amplitudeScaling = @(Pref,Pt) 1/sqrt(Pref)*sqrt(10^((Pt-30)/10));
    
    % Perform OFDM modulation on the signal
    if algParameters.FrequencyDomainChannelFiltering
        txUL.ofdmInfo = nrOFDMInfo(carrier);
        pref = sum(rms(ulGrid,[1 2]).^2)*(carrier.NSizeGrid*12/txUL.ofdmInfo.Nfft^2);
        txUL.ulGrid = ulGrid*amplitudeScaling(pref,txPower);
    else
        [ulWaveform,txUL.ofdmInfo] = nrOFDMModulate(carrier,ulGrid);
        pref = sum(rms(ulWaveform).^2);
        txUL.ulWaveform = ulWaveform*amplitudeScaling(pref,txPower);
    end
end

function rx = hAWGN(rx,noiseFigure)
% Add noise to the received waveform

  persistent kBoltz;
  if isempty(kBoltz)
      kBoltz = physconst('Boltzmann');
  end
  
  % For each BS
  numBS = numel(rx);
  for bs = 1:numBS
  
      % Establish dimensionality based on:
      % - the received waveform, if ChannelFiltering is true, or
      % - the received grid, if ChannelFiltering is false
      ofdmInfo = rx(bs).ofdmInfo;
      if (~isempty(rx(bs).rxWaveform))
          [T,Nr] = size(rx(bs).rxWaveform);
      else
          [K,L,R] = size(rx(bs).rxGrid);
      end
    
      % Calculate the required noise power spectral density
      NF = 10^(noiseFigure/10);
      N0 = sqrt(kBoltz*ofdmInfo.SampleRate*290*NF);
    
      % Create noise for the first UE and record it in the output. This
      % will be used if perfect channel estimation is configured
      if ~isempty(rx(bs).rxWaveform)
          noise = N0*randn([T Nr],'like',1i);
          rx(bs).noise = noise;
      else
          noiseGrid = N0*sqrt(ofdmInfo.Nfft)*randn(K,L,R,'like',1i);
          rx(bs).noiseGrid = noiseGrid;
      end
    
      % Add noise to the received waveform or received grid. The choice
      % depends on the ChannelFiltering setting (true or false)
      if (~isempty(rx(bs).rxWaveform))
          rx(bs).rxWaveform = rx(bs).rxWaveform + noise;
      else
          rx(bs).rxGrid = rx(bs).rxGrid + noiseGrid;
      end  
  end
end

function [H,nVar,delay] = hSRSReceive(carrier,srs,rx,alg)
% Generate received signal demodulation, operations on the received signal

    numBS = size(rx,1);
    H = cell(numBS,1);
    nVar = cell(numBS,1);
    delay = NaN(numBS,1);
    
   % Create SRS symbols and indices
    indices = nrSRSIndices(carrier,srs);
    symbols = nrSRS(carrier,srs);

    for bs = 1:numBS
    
        % Perform OFDM demodulation on the received signal
        if (alg.PerfectChannelEstimator)
            [offset,corr] = nrPerfectTimingEstimate(rx(bs).pathGains,rx(bs).pathFilters.');
        else
            [offset,corr] = nrTimingEstimate(carrier,rx(bs).rxWaveform,indices,symbols);
        end
      
        delay(bs) = offset;
      
        % Receiver aspects related to TDoA-based positioning
        scorr = sum(corr,2);
        if length(scorr) >= (2+offset)
            % Find the correlation peak
            scorr = scorr/max(scorr);
            minPeak = max(0.3+std(scorr));
            [~,idx] = findpeaks(scorr,MinPeakHeight=minPeak,WidthReference = 'halfprom');
            offset = idx(1)-1;
        
            % Perform linear interpolation of the correlation peak to estimate fractional delay
            idx = [-1 0 1];
            scorr = scorr(1 + offset + idx).';
            scorr = scorr/sum(scorr);
            delay(bs) = scorr*(offset + idx.') - rx(bs).ChannelFilterDelay;
        end
      
        % Receiver aspects required for AI-based positioning
        rxGrid = rx(bs).rxGrid;
        if (isempty(rxGrid))
            rxWaveform = rx(bs).rxWaveform;
            rxWaveform = rxWaveform(1+offset:end,:);
            rxGrid = nrOFDMDemodulate(carrier,rxWaveform);
        end
      
        % Perform channel and noise estimation
        if (alg.PerfectChannelEstimator)
            HL = nrPerfectChannelEstimate(carrier,rx(bs).pathGains,rx(bs).pathFilters.',offset);
            nVarL = var(rx(bs).noiseGrid(:));
        else
            % Perform channel estimation for each antenna layer
            [HL,nVarL] = nrChannelEstimate(rxGrid,indices,symbols,'CDMLengths',hSRSCDMLengths(srs));
        end
      
        % Create output for channel estimate and noise estimate
        H{bs} = HL;
        nVar{bs} = nVarL;
    end
end

function estPue = hEstimatePositionTDOA(Pbs,los,delay)
% Estimate the UE position using TDoA

    persistent lightSpeed;
    if isempty(lightSpeed)
        lightSpeed = physconst('LightSpeed');
    end

    if sum(los)>3
        idx = find(los);
        Pbs = Pbs(idx,:);
        delay = delay(idx);
    end
    
    d = delay*lightSpeed;
    
    idx = 2:size(Pbs,1);
    A = 2*Pbs(idx,1)./d(idx) - 2*Pbs(1,1)/d(1);
    B = 2*Pbs(idx,2)./d(idx) - 2*Pbs(1,2)/d(1);
    C = 2*Pbs(idx,3)./d(idx) - 2*Pbs(1,3)/d(1);
    D = d(idx) - d(1) - sum(Pbs(idx,:).^2,2)./d(idx) + sum(Pbs(1,:).^2,2)/d(1);
    
    M = [A B C];
    estPue = (-pinv(M)*D).';
end

function srs = hSRSConfig(cfg)
% Create and configure the SRS object

    carrier = cfg.Carrier;
    
    % Create SRS configuration object and set its parameters
    srs = nrSRSConfig;
    srs.NumSRSSymbols = cfg.SRS.NumSRSSymbols;
    srs.SymbolStart = carrier.SymbolsPerSlot-cfg.SRS.NumSRSSymbols;
    
    % Configure SRS for the widest bandwidth that fits within the carrier
    CSRS = srs.BandwidthConfigurationTable{:,"C_SRS"};
    mSRS0 = srs.BandwidthConfigurationTable{:,"m_SRS_0"};
    srs.CSRS = CSRS(find(mSRS0 <= carrier.NSizeGrid,1,'last'));
    
    srs.BHop = 3;           % Disable frequency hopping
    srs.KTC = cfg.SRS.KTC;
    srs.SRSPositioning = 1;
    
    srs.KBarTC = 0;
    srs.NumSRSPorts = cfg.SRS.NumSRSPorts;
    srs.NSRSID = 0;
end

function [channels,rxUL] = hChannel(carrier,channels,txUL)
% Pass the waveform through the channel in the UL direction

    % Set the sample rate and initial time, then calculate the maximum
    % delay
    ofdmInfo = txUL(1).ofdmInfo;
    [channels,maxDelay] = setupChannels(channels,ofdmInfo);
    
    L = carrier.SymbolsPerSlot;
    symbolLengths = ofdmInfo.SymbolLengths(mod(carrier.NSlot,carrier.SlotsPerSubframe)*L + (1:L));
    samplesPerSlot = sum(symbolLengths);
    
    numBS = size(channels,1);
    rxUL = repmat(newOutputStruct(),numBS,1);
    
    for bs = 1:numBS % For each BS
    
        % Select the channel for processing
        cdl = channels(bs).SmallScale;
      
        % Extract the uplink waveform and grid for the current UE
        ulWaveform = txUL.ulWaveform;
        ulGrid = txUL.ulGrid;
      
        % Pad the waveform to ensure the channel filter is fully flushed
        nT = size(ulWaveform,2);
        ulWaveform = [ulWaveform; zeros(maxDelay,nT)]; %#ok<AGROW>
      
        % Pass the uplink waveform through the channel
        ofdmInfo = txUL.ofdmInfo;
        if (cdl.ChannelFiltering)
            s = applyChannel(channels(bs),ulWaveform,ofdmInfo);
        else
            cdl.NumTimeSamples = samplesPerSlot;
            s = applyFrequencyDomainChannel(carrier,channels(bs),ulGrid,ofdmInfo);
        end
      
        % Record the output from the channel
        rxUL(bs) = s;
    end
end

function s = applyChannel(channel,txWaveform,ofdmInfo)
% Pass the waveform through the channel, record the output and related
% information in the structure 's'

    s = newOutputStruct();
    s.ofdmInfo = ofdmInfo;
    [rxWaveform,pathGains] = channel.SmallScale(txWaveform);
    s.rxWaveform = rxWaveform*channel.LargeScale;
    
    % Apply beamforming (TXRU virtualization) to the channel output
    pathGains = h38901Channel.applyBeamforming(channel.SmallScale,pathGains,channel.TXRUVirtualization);
    
    % Combine large-scale channel effects with path gains
    s.pathGains = pathGains * channel.LargeScale;
    
    s.ChannelFilterDelay = channel.chInfo.ChannelFilterDelay;
    s.pathFilters = channel.PathFilters;
end

function [channels,maxChDelay] = setupChannels(channels,ofdmInfo)
% Set the sample rate and initial time, then calculate the maximum delay
    
    numBS = size(channels,1);
    maxChDelay = 0;
    
    for bs = 1:numBS
        % Set the channel sample rate
        cdl = channels(bs).SmallScale;
        cdl.SampleRate = ofdmInfo.SampleRate;
      
        % Obtain the maximum delay of the channel
        chInfo = info(cdl);
        channels(bs).chInfo = chInfo;
        maxChDelay = max(maxChDelay,chInfo.MaximumChannelDelay);
    end
end

function s = newOutputStruct()

    s = struct('rxWaveform',[],'rxGrid',[],'ofdmInfo',[],'pathGains',[],'noise',[],...
      'noiseGrid',[],'pathFilters',[],'ChannelFilterDelay',[]);
end

function s = applyFrequencyDomainChannel(carrier,channel,txGrid,ofdmInfo)
% Pass the waveform through the channel in frequency domain

    s = newOutputStruct();
    s.ofdmInfo = ofdmInfo;
    s.pathGains = channel.SmallScale();
    s.pathFilters = channel.PathFilters;

    % Apply beamforming (TXRU virtualization) to the channel output
    pathGains = channel.SmallScale();
    pathGains = h38901Channel.applyBeamforming(channel.SmallScale,pathGains,channel.TXRUVirtualization);
    
    % Combine large-scale channel effects with path gains
    s.pathGains = pathGains * channel.LargeScale;
    
    % Calculate perfect channel estimates and apply them to the resource grid
    offset = nrPerfectTimingEstimate(s.pathGains,s.pathFilters.');
    
    H = nrPerfectChannelEstimate(carrier,s.pathGains,s.pathFilters.',offset);
    H = H(:,1:size(txGrid,2),:,:);
    s.rxGrid = applyChannelMatrices(txGrid,H);
    
    s.ChannelFilterDelay = channel.chInfo.ChannelFilterDelay;
end

function out = applyChannelMatrices(in,H)

    [K,L,R,P] = size(H);
    out = zeros([K L R],'like',H);
    a = reshape(in,K*L,P);
    b = permute(H,[1 2 4 3]);
    b = reshape(b,K*L,P,R);
    for r = 1:R
        out(:,:,r) = sum(reshape(a.*b(:,:,r),K,L,P),3);
    end
end

function channels = h38901ChannelSetup(cfg)
% Generate channels compliant with TR 38.901 using the scenario parameters

    % Create a carrier configuration and retrieve OFDM information
    carrier = cfg.Carrier;
    ofdmInfo = nrOFDMInfo(carrier);
    
    % Set up the channel configuration structure for BS -> UE direction;
    % the channel must be configured as downlink due to TXRU virtualization
    % on the transmit side
    cfg.AbsoluteTOA = true;
    cfg.SampleRate = ofdmInfo.SampleRate;
    cfg.HasSmallScale = true;
    cfg.FastFading = true;
    cfg.SpatialConsistency = false;
    
    % Create the path loss configuration
    plc = nrPathLossConfig(Scenario=cfg.Scenario);
    
    % Loop over each BS
    for i = 1:cfg.NumBSs
    
        % Set BS position
        cfg.BSPosition = cfg.Pos3BS(i,:);
      
        % Loop over each each UE
        for j = 1:cfg.NumUEs
      
            % -----------------------------------------------------------------
            % Channel generator setup
        
            if (i==1 && j==1)
                % SitePositions should be set to an N-by-3 matrix (each row
                % being the [x y z] position in metres of a BS) when creating
                % the first link of a set of links, to initialize state related
                % to spatially-correlated LSPs and spatially-consistent SSPs if
                % enabled.
                cfg.SitePositions = cfg.Pos3BS;
            else
                % SitePositions should be set to [] when creating subsequent
                % links in the same environment, to re-use the state described
                % above
                cfg.SitePositions = [];
            end
        
            % -----------------------------------------------------------------
            % Create the channel for a UE-BS link
        
            % Select a random UE position
            cfg.UEPosition = cfg.Pos3UE(j,:);
        
            % Set up subscripts for the current nodes and one sector
            cfg.NodeSubs = [i 1 j];
        
            % Create channel for this UE
            channel = h38901Channel.createChannelLink(cfg);
        
            % Apply it to large scale channel. Note that in the SLS, for
            % efficiency the h38901Channel object maintains a single
            % nrPathLossConfig and applies it to each channel
            channel.LargeScale = channel.LargeScale(plc);
        
            % Swap transmit and receive to configure for UE -> BS channel
            channel.LargeScale.swapTransmitAndReceive();
            swapTransmitAndReceive(channel.SmallScale);
        
            % Evaluate large scale channel
            PLdB = channel.LargeScale.execute(cfg.UEPosition,cfg.BSPosition,cfg.CenterFrequency);
            channel.LargeScale = db2mag(-PLdB);
        
            channel.chInfo = [];
        
            % -----------------------------------------------------------------
            % Record the channel for a UE-BS link
        
            channels(i,j) = channel; %#ok<AGROW>
      
        end
    end
end

function hDisplayNetworkSummary(net,details)
% Display the details of the deep neural network and plot its architecture

    numLearnables = sum(cellfun(@(x) numel(x),net.Learnables.Value))/1e6; % millions
    numLayers = details.conv+details.fc;
    
    disp('Neural Network Summary')
    disp(repmat('_',1,30))
    fprintf(' - Number of layers: %i\n',numLayers)
    if net.Initialized
        fprintf(' - Complexity: %.4f million\n',numLearnables)
    end
    fprintf(' - Input size: [ %s ]\n',num2str(net.Layers(1).InputSize))
    fprintf(' - Output size: [ %s ]\n',num2str(net.Layers(end).OutputSize))
    figure;
    plot(net);
    axis off
    box off
    axis square
    title(['Custom ',num2str(numLayers),'-Layer ResNet Architecture'])
    drawnow
end

function [trainData,valData,inputSize,outputSize] = channels2ImageDataset(cfg,channels)
% Generate an image-like data set of UL SRS estimated channels using TR 38.901 channel models

    algParameters = cfg.Algorithmic;
    carrier = cfg.Carrier;
    ofdmInfo = nrOFDMInfo(carrier);
    srs = cfg.SRS;
    Pue = cfg.Pos3UE;
    numUE = size(Pue,1);
    ueTxPower = cfg.PowerBSs;
    bsNoiseFigure = cfg.BSNoiseFigure;
    DisplaySimulationInformation = true;
    numDisp = 10; % Number of display text items shown
    deltaDisp = (numUE-3)/(numDisp-1);
    progressIdx = [1 floor(deltaDisp):floor(deltaDisp):numUE-floor(deltaDisp) numUE];
    
    X = []; % features array
    if algParameters.PerfectChannelEstimator % perfect channel estimation
        for j=1:size(channels,2)
            HArray = [];
            dispProgress = DisplaySimulationInformation;
            for i=1:size(channels,1)
                channel = channels(i,j); % (numBSs x numUEs)
                % Set up any other aspects of the small-scale channel that might vary
                % depending on the context of the calling code
                channel.SmallScale.NumTimeSamples = sum(ofdmInfo.SymbolLengths(1:ofdmInfo.SymbolsPerSlot));
          
                % Execute the small-scale channel to obtain CDL path gains
                % and sample times
                cdl = channel.SmallScale;
                [pathGains,sampleTimes] = cdl();
          
                % Retrieve precalculated path filters
                pathFilters = channel.PathFilters.';
          
                % Apply beamforming (TXRU virtualization) to the channel output
                pathGains = h38901Channel.applyBeamforming(channel.SmallScale,pathGains,channel.TXRUVirtualization);
          
                % Combine large-scale channel effects with path gains
                pathGains = pathGains * channel.LargeScale;
          
                % Perform perfect timing synchronization and channel estimation
                offset = nrPerfectTimingEstimate(pathGains,pathFilters);
          
                % H is a complex 3-D array with dimensions: 
                % (NumSC x NumSym x NumPorts)
                H = nrPerfectChannelEstimate(carrier,pathGains,pathFilters,offset,sampleTimes);
              
                % Perform time-domain dimensionality reduction due to lack
                % of UE mobility
                H = squeeze(H(:,1,:));

                % Store the reduced-dimension channels
                % (NumPorts x NumPorts x NumBaseStations)
                HArray = cat(3,HArray,H);
          
                % Display progress at custom intervals if diagnostics are enabled
                if (DisplaySimulationInformation) && dispProgress && any(j == progressIdx)
                    fprintf('Generating Data Set: %3.2f%% complete.\n',100*j/numUE);
                    dispProgress = false;
                end
            end
            % Image-like channel data set
            % (Subcarriers x OFDM Symbols x NumPorts x NumBaseStations)
            X = cat(4,X,hChannel2Image(cfg,HArray));
        end
    else % practical channel estimation

        % Enable channel filtering in all channels
        for i=1:numel(channels)
            release(channels(i).SmallScale);
            channels(i).SmallScale.ChannelFiltering = true;
        end

        for ue = 1:numUE
        
            % This is a single-slot simulation. Set slot number to 0
            carrier.NSlot = 0;
          
            % Transmit SRS from current UE to all BSs
            txUL = hSRSTransmit(carrier,srs,ueTxPower,algParameters);
          
            % Apply fading channels to the signal
            [channels(:,ue),rxUL] = hChannel(carrier,channels(:,ue),txUL);
          
            % Apply AWGN
            rxUL = hAWGN(rxUL,bsNoiseFigure);
          
            % For all UEs scheduled for SRS, estimate the Channel State
            % Information (CSI) and record it
            [H,~,~] = hSRSReceive(carrier,srs,rxUL,algParameters);
    
            % Perform time-domain dimensionality reduction due to lack of UE
            % mobility
            H = cellfun(@(x) squeeze(x(:,1,:)),H,'UniformOutput',false);
          
            % Generate a collection of estimated CSI
            % HArray dimensions:
            % (NumPorts x NumPorts x NumBaseStations)
            HArray = cat(3,H{:},[]);
    
            % Image-like channel data set
            % (Subcarriers x OFDM Symbols x NumPorts x NumBaseStations)
            X = cat(4,X,hChannel2Image(cfg,HArray));
          
            % Display progress at custom intervals if diagnostics are enabled
            if (DisplaySimulationInformation) && any(ue == progressIdx)
                fprintf('Generating Data Set: %3.2f%% complete.\n',100*ue/numUE);
            end
        end
    end
    disp(repmat('_',1,30))
    clearvars("channels"); % Remove for memory efficiency

    % Generate training and validation data sets
    % Data set dimensions: [NumPorts, NumPorts, NumBSs, NumUEs]
    % where the last dimension is the 'batch' dimension
    inputSize = size(X,1:3);
    outputSize = size(cfg.Pos2UE,2);
    y = cfg.Pos2UE; % labels
    
    % Split the data set into training, validation, and test sets
    [trainInd,valInd,~] = ...
      dividerand(size(X,4),cfg.TrainRatio,1-cfg.TrainRatio,0);
    
    % Training data
    XTrain = X(:,:,:,trainInd);
    yTrain = y(trainInd,:);
    XTrain = arrayDatastore(XTrain,"IterationDimension",ndims(X));
    yTrain = arrayDatastore(yTrain,"IterationDimension",1);
    trainData = combine(XTrain,yTrain);
    
    % Validation data
    XVal = X(:,:,:,valInd);
    yVal = y(valInd,:);
    
    % Display data set information
    fprintf('Number of training images: %i\n',numpartitions(trainData))
    if ~isempty(valInd)
        XVal = arrayDatastore(XVal,"IterationDimension",ndims(X));
        yVal = arrayDatastore(yVal,"IterationDimension",1);
        valData = combine(XVal,yVal);
        fprintf('Number of validation images: %i\n',numpartitions(valData))
    else
        valData = []; % No validation data
    end
end

function estPue = hEstimatePositionAI(cfg,H,net)
% Estimate the UE position using AI

    % AI-based CSI preprocessing. Perform time-domain dimensionality
    % reduction due to lack of UE mobility
    H = cellfun(@(x) squeeze(x(:,1,:)),H,'UniformOutput',false);

    % Generate a collection of estimated CSI
    % HArray dimensions: [Subcarriers, OFDM Symbols, NumPorts, NumBaseStations]
    HArray = cat(3,H{:},[]);

    % Convert channel estimates to an image-like format [NumPorts, NumPorts, NumBaseStations]
    csiImages = hChannel2Image(cfg,HArray);
  
    % Estimate the positions of the UEs using AI-based techniques
    estPue = predict(net,csiImages);   
end

function positioningSimulationResults(results)
% Display and visualize the results from AI and TDoA positioning simulations

    t = table(results.AI.Percentiles(:),results.TDoA.Percentiles(:));
    t.Properties.VariableNames = ["AI/ML (ResNet)", "TDoA"];
    t.Properties.RowNames = ["10th percentile","50th percentile","90th percentile"];
    disp(t)
    
    plotCDF = @(x) stairs(sort(x),(1:length(x))/length(x),'LineWidth',1.5);
    
    figure
    plotCDF(results.AI.EstimationError)
    hold on
    plotCDF(results.TDoA.EstimationError)
    hold off
    title('CDF of 2D Positioning Error');
    xlabel('Positioning Error (m)')
    ylabel('Cumulative Probability')
    legend('AI/ML (ResNet)','TDoA','Location','best')
    grid on
    drawnow
end

function csiImages = hChannel2Image(cfg,HArray)
% Convert the channel matrices into a (NumAntennaPorts x NumAntennaPorts x NumBSs)
% image-like array
    csiImages = abs(HArray);
    csiImages = imresize(csiImages,repmat(prod(cfg.TransmitAntennaArray.Size),1,2));
    csiImages = rescale(csiImages,0,1);
end