Python - QAM

import numpy as np
import matplotlib.pyplot as plt

class Quadrature_Amplitude_Modulation:
    def __init__(self, mod_order: int) -> None:
        """! Constructor for quadrature amplitude modulation (QAM)
        @param mod_order Modulation order, e.g., 4 for QPSK.
        """
        self.mod_order = mod_order
        self.bit_num_sym = int(np.log2(mod_order))
        m = np.sqrt(mod_order)
        v = np.linspace(1 - m, m - 1, int(m))
        x, y = np.meshgrid(v, v)
        self.const = np.reshape(x + 1j * y, -1) / np.sqrt((mod_order - 1) * 2 / 3)
        self.name = "QPSK" if mod_order == 4 else f"{mod_order}QAM"
    def calc_mi(self, sinr: float, sample_num: int=1000) -> float:
        """! Calculate the per-symbol mutual information, a.k.a. RBIR, for the specified SINR.
        @param sinr SINR in dB
        @param sample_num Number of random samples
        @return Per-bit information
        """
        n_pwr = np.power(10, -0.1 * sinr)
        n = np.random.randn(sample_num) + 1j * np.random.randn(sample_num)
        n *= np.sqrt(n_pwr / 2)
        si = np.zeros((self.mod_order, sample_num))
        for sym_idx in range(self.mod_order):
            for sample_idx in range(sample_num):
                x = np.power(np.abs(self.const - self.const[sym_idx] + n[sample_idx]), 2) - np.power(np.abs(n[sample_idx]), 2)
                si[sym_idx][sample_idx] = np.sum(np.exp(-1 * x / n_pwr))
        return self.bit_num_sym - np.log2(si).mean()
    def visualize(self, ax: plt.Axes = None, **kwargs) -> None:
        """! Visualize the constellation
        @param ax An axis object
        @param kwargs Keyword argument for further customization
        """
        if ax is None:
            _, ax = plt.subplots()
        ax.scatter(self.const.real, self.const.imag, **kwargs)
        ax.set_xlabel('In-phase')
        ax.set_ylabel('Quadrature')
        ax.set_title(self.name)
        ax.set_aspect('equal')
        ax.grid(True)