WSR maximization using WMMSE for MIMO-BC beamforming design using Python

import io
import scipy.io as sio
import numpy as np
from tqdm import tqdm
import time

import math
import matplotlib.pyplot as plt

K = 4  # num of users
P = 8  # Tx AnT.
Q = 2  # Rx AnT.
iter = 100
SNRdB = 10
SNR = 10.0 ** (SNRdB/10.0)
E_tx = 10.0

u = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]

def conj_tran(input) :
    return np.transpose(input.conj())
#np.random.seed(12334)

# Hk 행렬 생성
#H = []
B = []
B_bar = []
A = []
E = []
W = []
Rvkvk = []

for i in range(K) :
    #H.append(np.random.randn(Q, P) + 1j * np.random.randn(Q, P))
    B.append(np.random.randn(P, Q) + 1j * np.random.randn(P, Q))
    A.append(np.random.randn(Q, Q) + 1j * np.random.randn(Q, Q))
    E.append(np.random.randn(Q, Q) + 1j * np.random.randn(Q, Q))
    Rvkvk.append(np.random.randn(Q, Q) + 1j * np.random.randn(Q, Q))
    W.append(np.random.randn(Q, Q) + 1j * np.random.randn(Q, Q))
    #B_bar.append(np.random.randn(P, P) + 1j * np.random.randn(P, P))
    

B_bar = np.random.randn(P, Q*K) + 1j * np.random.randn(P, Q*K)

#for i in range(K) :
 #   H[i] = np.random.normal(0, 1, size = (Q, P)) * math.sqrt(SNR/2) + 1j* np.random.normal(0, 1, size = (Q, P)) * math.sqrt(SNR/2)

#H = np.array(H)
B = np.array(B)
A = np.array(A)
E = np.array(E)
W = np.array(W)
Rvkvk = np.array(Rvkvk)

b = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
R = []
WSR = []

mat_file = sio.loadmat("./H_test.mat")
H_mat = mat_file.get("H")

H = np.array(np.split(H_mat, K))


for iteration in tqdm(range(iter)) :
    time.sleep(0.01)

    R = 0.0 + 0.0 * 1j

    for k in range(K) :

        sum_in_Rvkvk = 0

        for j in range(K) :
            if k != j :
                sum_in_Rvkvk += H[k] @ B[j] @ conj_tran(B[j]) @ conj_tran(H[k])
            else :
                sum_in_Rvkvk += 0
        Rvkvk[k] = np.identity(Q) + sum_in_Rvkvk

        A[k] = conj_tran(B[k]) @ conj_tran(H[k]) @ np.linalg.inv( H[k] @ B[k] @ conj_tran(B[k]) @ conj_tran(H[k]) + Rvkvk[k] )

########################################################################################################################
        E[k] = np.linalg.inv(np.identity(Q) + conj_tran(B[k]) @ conj_tran(H[k]) @ np.linalg.inv(Rvkvk[k]) @ H[k] @ B[k])

        W[k] = u[k] * np.linalg.inv(E[k])

########################################################################################################################

        W_final = np.zeros((Q * K, Q * K), dtype=complex)
        A_final = np.zeros((Q * K, Q * K), dtype=complex)

        for q in range(K) :
        
            W_final[q*Q :(q+1)*Q, q*Q :(q+1)*Q ] = W[q]

            A_final[q*Q :(q+1)*Q, q*Q :(q+1)*Q ] = A[q]


        mid = np.transpose(H[0])
        for q in range(K-1) :
            mid = np.concatenate((mid,np.transpose(H[q+1])), axis = 1)
        H_final = np.transpose(mid)

        B_bar = (np.linalg.inv(conj_tran(H_final) @ conj_tran(A_final) @ W_final @ A_final @ H_final + (((W_final @ A_final @ conj_tran(A_final)).trace()) / E_tx) * np.identity(P))) @ conj_tran(H_final) @ conj_tran(A_final) @ W_final

        b = np.sqrt(E_tx /  (B_bar @ conj_tran(B_bar)).trace()  )

        mid_B = b * B_bar


        for f in range(K) :
            B[f] = mid_B[ 0:P   , Q*f:Q*f+Q]


########################################################################################################################

        R += np.log2(np.linalg.det(      np.identity(Q) + conj_tran(B[k]) @ conj_tran(H[k]) @ np.linalg.inv(Rvkvk[k]) @ H[k] @ B[k]                                                  ))

    WSR.append(R)


plt.plot(list(range(1, iter+1)), WSR, linestyle = 'dashed', color = 'blue')