Using pyAerial for PUSCH decoding on Aerial Data Lake data

This example shows how to use the pyAerial bindings to run cuPHY GPU accelerated PUSCH decoding for 5G NR PUSCH. The 5G NR PUSCH data is read from an example over the air captured PUSCH dataset collected and stored using Aerial Data Lake. Building a PUSCH receiver using pyAerial is demonstrated in two ways, first by using a fully fused, complete, PUSCH receiver called from Python using just a single function call. The same is then achieved by building the complete PUSCH receiver using individual separate Python function calls to individual PUSCH receiver components.

Note: This example requires that the clickhouse server is running and that the example data has been stored in the database. Refer to the Aerial Data Lake documentation on how to do this.

# Check platform.
import platform
if platform.machine() not in ['x86_64', 'aarch64']:
    raise SystemExit("Unsupported platform!")

Imports

import math
import os
os.environ["CUDA_VISIBLE_DEVICES"] = "0"

import numpy as np
import pandas as pd
from IPython.display import Markdown
from IPython.display import display

# Connecting to clickhouse on remote server
import clickhouse_connect

# Plotting with Matplotlib.
import matplotlib.pyplot as plt

# pyAerial imports
from aerial.phy5g.config import PuschConfig
from aerial.phy5g.config import PuschUeConfig
from aerial.phy5g.algorithms import ChannelEstimator
from aerial.phy5g.algorithms import ChannelEqualizer
from aerial.phy5g.algorithms import NoiseIntfEstimator
from aerial.phy5g.ldpc import LdpcDeRateMatch
from aerial.phy5g.ldpc import LdpcDecoder
from aerial.phy5g.ldpc import CrcChecker
from aerial.phy5g.pusch import PuschRx
from aerial.util.cuda import get_cuda_stream
from aerial.pycuphy.types import PuschLdpcKernelLaunch
from aerial.util.fapi import dmrs_fapi_to_bit_array

# Hide log10(10) warning
_ = np.seterr(divide='ignore', invalid='ignore')

Create the PUSCH pipelines

This is a PUSCH receiver pipeline made up of separately called pyAerial PUSCH receiver components.

# Whether to plot intermediate results within the PUSCH pipeline, such as channel estimates and equalized symbols.
plot_figures = True

num_ues = 1
num_tx_ant = 2             # UE antennas
num_rx_ant = 4             # gNB antennas
cell_id = 41               # Physical cell ID
enable_pusch_tdi = 0       # Enable time interpolation for equalizer coefficients
eq_coeff_algo = 1          # Equalizer algorithm

# The PUSCH receiver chain built from separately called pyAerial Python components is defined here.
class PuschRxSeparate:
    """PUSCH receiver class.

    This class encapsulates the whole PUSCH receiver chain built using
    pyAerial components.
    """

    def __init__(self,
                 num_rx_ant,
                 enable_pusch_tdi,
                 eq_coeff_algo,
                 plot_figures):
        """Initialize the PUSCH receiver."""
        self.cuda_stream = get_cuda_stream()

        # Build the components of the receiver.
        self.channel_estimator = ChannelEstimator(
            num_rx_ant=num_rx_ant,
            cuda_stream=self.cuda_stream)
        self.channel_equalizer = ChannelEqualizer(
            num_rx_ant=num_rx_ant,
            enable_pusch_tdi=enable_pusch_tdi,
            eq_coeff_algo=eq_coeff_algo,
            cuda_stream=self.cuda_stream)
        self.noise_intf_estimator = NoiseIntfEstimator(
            num_rx_ant=num_rx_ant,
            eq_coeff_algo=eq_coeff_algo,
            cuda_stream=self.cuda_stream)
        self.derate_match = LdpcDeRateMatch(
            enable_scrambling=True,
            cuda_stream=self.cuda_stream)
        self.decoder = LdpcDecoder(cuda_stream=self.cuda_stream)
        self.crc_checker = CrcChecker(cuda_stream=self.cuda_stream)

        # Whether to plot the intermediate results.
        self.plot_figures = plot_figures

    def run(
        self,
        rx_slot,
        slot,
        pusch_configs
    ):
        """Run the receiver."""
        # Channel estimation.
        ch_est = self.channel_estimator.estimate(
            rx_slot=rx_slot,
            slot=slot,
            pusch_configs=pusch_configs
        )

        # Noise and interference estimation.
        lw_inv, noise_var_pre_eq = self.noise_intf_estimator.estimate(
            rx_slot=rx_slot,
            channel_est=ch_est,
            slot=slot,
            pusch_configs=pusch_configs
        )

        # Channel equalization and soft demapping. The first return value are the LLRs,
        # second are the equalized symbols. We only want the LLRs now.
        llrs, sym = self.channel_equalizer.equalize(
            rx_slot=rx_slot,
            channel_est=ch_est,
            lw_inv=lw_inv,
            noise_var_pre_eq=noise_var_pre_eq,
            pusch_configs=pusch_configs
        )

        if self.plot_figures:
            fig, axs = plt.subplots(1,4)
            for ant in range(4):
                axs[ant].imshow(10*np.log10(np.abs(rx_slot[:, :, ant]**2)), aspect='auto')
                axs[ant].set_ylim([pusch_record.rbStart * 12, pusch_record.rbSize * 12])
                axs[ant].set_title('Ant ' + str(ant))
                axs[ant].set(xlabel='Symbol', ylabel='Resource Element')
                axs[ant].label_outer()
            fig.suptitle('Power in RU Antennas')

            fig, axs = plt.subplots(1,2)
            axs[0].scatter(rx_slot.reshape(-1).real, rx_slot.reshape(-1).imag)
            axs[0].set_title("Pre-Equalized samples")
            axs[0].set_aspect('equal')

            axs[1].scatter(np.array(sym).reshape(-1).real, np.array(sym).reshape(-1).imag)
            axs[1].set_title("Post-Equalized samples")
            axs[1].set_aspect('equal')

            fig, axs = plt.subplots(1)
            axs.set_title("Channel estimates from the PUSCH pipeline")
            for ant in range(4):
                axs.plot(np.abs(ch_est[0][ant, 0, :, 0]))
            axs.legend(["Rx antenna 0, estimate",
                        "Rx antenna 1, estimate",
                        "Rx antenna 2, estimate",
                        "Rx antenna 3, estimate"])
            axs.grid(True)
            plt.show()

        coded_blocks = self.derate_match.derate_match(
            input_llrs=llrs,
            pusch_configs=pusch_configs
        )

        code_blocks = self.decoder.decode(
            input_llrs=coded_blocks,
            pusch_configs=pusch_configs
        )

        decoded_tbs, tb_crcs = self.crc_checker.check_crc(
            input_bits=code_blocks,
            pusch_configs=pusch_configs
        )

        return decoded_tbs

pusch_rx_separate = PuschRxSeparate(
    num_rx_ant=num_rx_ant,
    enable_pusch_tdi=enable_pusch_tdi,
    eq_coeff_algo=eq_coeff_algo,
    plot_figures=plot_figures
)

# This is the fully fused PUSCH receiver chain.
pusch_rx = PuschRx(
    cell_id=cell_id,
    num_rx_ant=num_rx_ant,
    num_tx_ant=num_rx_ant,
    enable_pusch_tdi=enable_pusch_tdi,
    eq_coeff_algo=eq_coeff_algo,
    # To make this equal separate PUSCH Rx components configuration:
    ldpc_kernel_launch=PuschLdpcKernelLaunch.PUSCH_RX_LDPC_STREAM_SEQUENTIAL
)

Querying the database

Below shows how to connect to the clickhouse database and querying the data from it.

# Connect to the local database
client = clickhouse_connect.get_client(host='localhost')

# Pick a packet from the database
pusch_records = client.query_df('select * from fapi where mcsIndex != 0 order by TsTaiNs limit 10')

Extract the PUSCH parameters and run the pipelines

for index, pusch_record in pusch_records.iterrows():
    query = f"""select TsTaiNs,fhData from fh where
            TsTaiNs == {pusch_record.TsTaiNs.timestamp()} and
            CellId == 51
            """
    fh = client.query_df(query)

    display(Markdown("### Example {} - SFN.Slot {}.{} from time {}"
                     .format(index + 1, pusch_record.SFN, pusch_record.Slot, pusch_record.TsTaiNs
    )))

    # Make sure that the fronthaul database is complete for the SFN.Slot we've chosen.
    # Also make sure that PDU data exists for the entry.
    if fh.index.size < 1 or np.array(pusch_record.pduData).size == 0:
        pusch_records = pusch_records.drop(index)
        continue;

    fh_samp = np.array(fh['fhData'][0], dtype=np.float32)
    rx_slot = np.swapaxes(fh_samp.view(np.complex64).reshape(4, 14, 273 * 12), 2, 0)

    # Extract all the needed parameters from the PUSCH record and create the PuschConfig.
    pusch_ue_config = PuschUeConfig(
        scid=int(pusch_record.SCID),
        layers=pusch_record.nrOfLayers,
        dmrs_ports=pusch_record.dmrsPorts,
        rnti=pusch_record.rnti,
        data_scid=pusch_record.dataScramblingId,
        mcs_table=pusch_record.mcsTable,
        mcs_index=pusch_record.mcsIndex,
        code_rate=pusch_record.targetCodeRate,
        mod_order=pusch_record.qamModOrder,
        tb_size=pusch_record.TBSize
    )

    slot = int(pusch_record.Slot)
    tb_input = np.array(pusch_record.pduData)

    # Note that this is a list. One UE group only in this case.
    pusch_configs = [PuschConfig(
        ue_configs=[pusch_ue_config],
        num_dmrs_cdm_grps_no_data=pusch_record.numDmrsCdmGrpsNoData,
        dmrs_scrm_id=pusch_record.ulDmrsScramblingId,
        start_prb=pusch_record.rbStart,
        num_prbs=pusch_record.rbSize,
        dmrs_syms=dmrs_fapi_to_bit_array(int(pusch_record.ulDmrsSymbPos)),
        dmrs_max_len=1,
        dmrs_add_ln_pos=2,
        start_sym=pusch_record.StartSymbolIndex,
        num_symbols=pusch_record.NrOfSymbols
    )]

    # Run the receiver built from separately called components.
    tbs = pusch_rx_separate.run(
        slot=slot,
        rx_slot=rx_slot,
        pusch_configs=pusch_configs
    )

    if np.array_equal(tbs[0], tb_input):
        display(Markdown("**Separated kernels PUSCH decoding success** for SFN.Slot {}.{} from time {}"
                         .format(pusch_record.SFN, pusch_record.Slot, pusch_record.TsTaiNs)))
    else:
        display(Markdown("**Separated kernels PUSCH decoding failure**"))
        print("Output bytes:")
        print(tbs[0])
        print("Expected output:")
        print(tb_input)

    # Run the fused PUSCH receiver.
    # Note that this is where we set the dynamically changing parameters.
    tb_crcs, tbs = pusch_rx.run(
        rx_slot=rx_slot,
        slot=slot,
        pusch_configs=pusch_configs
    )

    if np.array_equal(tbs[0], tb_input):
        display(Markdown("**Fused PUSCH decoding success** for SFN.Slot {}.{} from time {}"
                         .format(pusch_record.SFN, pusch_record.Slot, pusch_record.TsTaiNs)))
    else:
        display(Markdown("**Fused PUSCH decoding failure**"))
        print("Output bytes:")
        print(tbs[0])
        print("Expected output:")
        print(tb_input)

Example 1 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 2 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 3 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 4 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

/tmp/ipykernel_43612/1271835907.py:85: UserWarning: Attempting to set identical low and high ylims makes transformation singular; automatically expanding.
  axs[ant].set_ylim([pusch_record.rbStart * 12, pusch_record.rbSize * 12])

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 5 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 6 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 7 - SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Separated kernels PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Fused PUSCH decoding success for SFN.Slot 391.4 from time 2024-07-19 10:42:46.272000

Example 8 - SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Example 9 - SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Separated kernels PUSCH decoding success for SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Fused PUSCH decoding success for SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Example 10 - SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Separated kernels PUSCH decoding success for SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000

Fused PUSCH decoding success for SFN.Slot 493.4 from time 2024-07-19 10:42:47.292000