utils.py

import sys
import torch
import numpy as np
import matplotlib.pyplot as plt


def db2power(dB_value):
    return 10 ** (dB_value / 10)

def sum_rate(H, V, sigma2, R, I, K, alpha):
    rate = np.zeros(shape=(I, K))
    for i in range(I):
        for k in range(K):
            temp = np.zeros(shape=(R, R)) + 1j*np.zeros(shape=(R, R))
            for l in range(I):
                for j in range(K):
                    if l != i or j != k:
                        temp = temp + H[:,:,i,k,j]@V[:,:,l,j]@np.conjugate(V[:,:,l,j]).T@np.conjugate(H[:,:,i,k,j]).T
            rate[i,k] = np.log2(np.linalg.det(np.eye(R)+H[:,:,i,k,k]@V[:,:,i,k]@np.conjugate(V[:,:,i,k]).T@np.conjugate(H[:,:,i,k,k]).T@np.linalg.inv(temp+sigma2*np.eye(R))))
    system_rate = np.real(np.sum(rate*alpha))

    return system_rate



def compute_circular_gaussian_channel(R, T, I, K):
    H = np.zeros(shape=(R, T, I, K, K)) + 1j*np.zeros(shape=(R, T, I, K, K))
    for i in range(I):
        for k in range(K):
            for j in range(K):
                H[:,:,i,k,j] = np.sqrt(0.5)*(np.random.normal(loc=0.0, scale=1.0, size=(R,T))+1j*np.random.normal(loc=0.0, scale=1.0, size=(R,T)))
    return H

def compute_channel(K, R, T, I, total_power, path_loss_option=True, path_loss_min=-5, path_loss_max=5):
    H = np.zeros(shape=(R, T, I, K,K)) + 1j*np.zeros(shape=(R, T, I, K, K))

    for i in range(I):
        for k in range(K):
            for j in range(K):
                path_loss = 0
                if path_loss_option == True:
                    path_loss = np.random.uniform(path_loss_min, path_loss_max)

                result_real = np.sqrt(10**(path_loss/10))*np.sqrt(0.5)*np.random.normal(size=(R,T))
                result_imag = np.sqrt(10**(path_loss/10))*np.sqrt(0.5)*np.random.normal(size=(R,T))
                H[:,:,i,k,j] = result_real+1j*result_imag

    return H



def plot_ites(rates, model_name, K, T, R, d, snr, epsilon):
    x = np.arange(len(rates))
    fig, ax = plt.subplots(figsize=(8, 5))

    line, = ax.plot(
        x, rates,
        color="#8B0000",  # deep red
        marker="^",  # filled triangles
        markersize=8,  # markers size
        markerfacecolor="#8B0000",
        markeredgecolor="k",
        linestyle="-",
        linewidth=2,
        label=model_name
    )

    ax.grid(
        True,
        linestyle="--",
        linewidth=0.5,
        color="gray",
        alpha=0.7
    )


    ax.legend(
        loc="best",
        frameon=True,
        shadow=True
    )

    ax.set_xlabel("Iterations", fontsize=12)
    ax.set_ylabel("Sum rate (bits per channel use)", fontsize=12)
    ax.set_title(r"{}, K={}, T={}, R={}, d={}, SNR={}dB, $\epsilon$={}".format(model_name, K, T, R, d, snr, epsilon), fontsize=14)

    plt.tight_layout()
    plt.show()


def get_drop_pct(base, lower):
    return round((lower - base) / base * 100, 1)


def batch_compute_J_U(H, V):
    """
    H: (K, K, I, N, M) 
    V: (K, I, M, d) 
    Returns: (K, I, N, N) 
    """
    S = torch.einsum('j l m d, j l n d -> j m n', V, V.conj())  # 对 l 和 d 求和
    H_ct = H.conj().transpose(-1, -2)  # (K, K, I, M, N)
    output = torch.einsum('j k i n m, j m p, j k i p q -> k i n q', H, S, H_ct)
    return output

def batch_compute_HUW(H, U, W):
    """
    H: (K, K, I, N, M)
    U: (K, I, N, d)
    W: (K, I, d, d)
    Returns: (K, M, M)
    """
    K, _, I, N, M = H.shape
    _, _, d, _ = W.shape
    # C = U @ W @ U^H
    U_flat = U.reshape(K * I, N, -1)  # (K*I, N, d)
    W_flat = W.reshape(K * I, d, d)  # (K*I, d, d)
    temp = torch.bmm(U_flat, W_flat)  # (K*I, N, d)
    C_flat = torch.bmm(temp, U_flat.conj().transpose(1, 2))  # (K*I, N, N)
    C = C_flat.view(K, I, N, N)  # (K, I, N, N)

    # output[k] = sum_{j,l} H[k,j,l]^H @ C[j,l] @ H[k,j,l]
    H_H = H.conj().transpose(-1, -2)  # (K, K, I, M, N)
    # H_H @ C -> (K,K,I,M,N) @ (K,I,N,N) -> (K,K,I,M,N)
    # 注意：C 的索引为 (j,l,N,N)，H_H 索引 (k,j,l,M,N)，对 N 求和，结果保留 j,l 维度
    temp1 = torch.einsum('kjlmn, jlnp -> kjlmp', H_H, C)  # (K, K, I, M, N)
    # temp1 @ H -> (K,K,I,M,N) @ (K,K,I,N,M) -> (K,K,I,M,M) 对 N 求和
    output = torch.einsum('kjlmp, kjlpq -> kjlmq', temp1, H)  # (K, K, I, M, M)
    output = output.sum(dim=(1, 2))  # (K, M, M)
    return output


def batch_compute_HU(H, U):
    """
    H: (K, K, I, N, M)
    U: (K, I, N, d)
    W: (K, I, d, d)
    Returns: (K, M, M)
    """
    K, _, I, N, M = H.shape
    # C = U @ W @ U^H
    U_flat = U.reshape(K * I, N, -1)  # (K*I, N, d)
    C_flat = torch.bmm(U_flat, U_flat.conj().transpose(1, 2))  # (K*I, N, N)
    C = C_flat.view(K, I, N, N)  # (K, I, N, N)

    # output[k] = sum_{j,l} H[k,j,l]^H @ C[j,l] @ H[k,j,l]
    H_H = H.conj().transpose(-1, -2)  # (K, K, I, M, N)
    temp1 = torch.einsum('kjlmn, jlnp -> kjlmp', H_H, C)  # (K, K, I, M, N)
    output = torch.einsum('kjlmp, kjlpq -> kjlmq', temp1, H)  # (K, K, I, M, M)
    output = output.sum(dim=(1, 2))  # (K, M, M)
    return output