PrecodeNet_CustomEnv.py

#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Created on Thu May 19 20:43:14 2022

@author: jayanths
"""

import numpy as np
import gym
import torch 

from gym import spaces

# %%
"""
To Do :
    Create custom environment for PrecoderNet Paper :
        1. Create the environment with the details given in the paper
        2. Check if the gym.Env can be the parent class for this environment
        3. If 2 is satisfied then use either SB3 or rllib codes
"""


class HybridBeamForming():
    """
    Description :
        In the Hybrid beam forming the entire transmission system at BS is the
        environment. The analog beamformer V_rf of size (N_t x N_rf) is assumed
        to be given. Objective is to find V_bb and W_rf such that the upper
        bound on the rate is maximized.
    Observation Space :
        V_bb : (N_t x N_rf) complex matrix
        W_rf : (N_r x N_rf) complex matrix
        The real and imaginary parts of V_bb and W_rf are stacked. Hence each
        observation will be of size (N_t * N_rf * 2 + N_r + N_rf * 2, 1)
    Action Space :
        Predict the next V_bb, W_rf.
        The real and imaginary parts of V_bb and W_rf are stacked. Hence each
        action will be of size (N_t * N_rf * 2 + N_r + N_rf * 2, 1)
    Reward :
        1. Calculate W_bb using W_rf, V_rf, V_bb, H, beta, sigma
        2. Then the corresponding reward is calculated using equation 6.
    """

    def __init__(self, Sysparams):
        
        """
        Arguments:
            N_s    : Number of data streams
            Nt_rf  : Number of RF chains at BS (Tx)
            N_t    : Number of trasmit antennas 
            N_r    : Number of receive antennas
            Nr_rf  : Number of RF chains at Rx
            N_cl   : No. of scattering clusters
            N_ray  : No. of scattering rays in each cluster
            Cl_var : Variance of complex path gains of all rays
            lambd  : Carrier wavelength
            d      : Antenna spacing
            beta2  : Level of channel imperfection (between 0 and 1)
            sigma2 : Noise variance
        """
        
        # Tx parameters
        self.N_s   = Sysparams['N_s']
        self.Nt_rf = Sysparams['Nt_rf']
        self.N_t   = Sysparams['N_t']
        
        # Rx parameters
        self.N_r   = Sysparams['N_r']
        self.Nr_rf = Sysparams['Nr_rf']
        
        # Saleh-Valenzuela channel model parameters
        self.N_cl   = Sysparams['N_cl']
        self.N_ray  = Sysparams['N_ray']
        self.Cl_var = Sysparams['Cl_var']
        self.lambd  = Sysparams['lambda']
        self.d      = Sysparams['d']
        
        # Other parameters
        self.beta2  = Sysparams['beta2']
        self.SNRdB  = Sysparams['SNRdB']
        self.SNR    = Sysparams['SNR']
        self.sigma2 = Sysparams['sigma2']

        self.Vrf_real = np.zeros((self.N_t, self.Nt_rf), dtype=np.float32)
        self.Vrf_img = np.zeros((self.N_t, self.Nt_rf), dtype=np.float32)

        # Default data type for complex numbers is complex128
        self.V_rf = np.zeros((self.N_t, self.Nt_rf), dtype=np.complex64)

        self.V_rf = self.Vrf_real + self.Vrf_img * 1j

        self.Vbb_real = np.zeros((self.Nt_rf, self.N_s), dtype=np.float32)
        self.Vbb_img = np.zeros((self.Nt_rf, self.N_s), dtype=np.float32)

        # self.V_bb = torch.complex(self.Vbb_real, self.Vrf_img)
        self.V_bb = self.Vbb_real + self.Vbb_img * 1j

        self.Wrf_real = np.zeros((self.N_r, self.Nr_rf), dtype=np.float32)
        self.Wrf_img = np.zeros((self.N_r, self.Nr_rf), dtype=np.float32)

        self.W_rf = self.Wrf_real + self.Wrf_img * 1j
        # self.W_rf = torch.complex(self.Wrf_real, self.Vrf_img)

        self.Wbb_real = np.zeros((self.Nr_rf, self.N_s), dtype=np.float32)
        self.Wbb_img = np.zeros((self.Nr_rf, self.N_s), dtype=np.float32)

        # self.W_bb = torch.complex(self.Vrf_real, self.Vrf_img)
        self.W_bb = self.Wbb_real + self.Wbb_img * 1j
    
    def Sample_SVChannel(self):
        
        """
         Sample Geometric Saleh-Valenzuela channel
        """
        
        self.H = np.zeros((self.N_r, self.N_t))
        
        # Complex path gains of all rays associated with all clusters
        alpha  = np.sqrt(self.Cl_var/2)* (np.random.randn(self.N_cl, self.N_ray) + np.random.randn(self.N_cl, self.N_ray)* 1j)
        
        # Angle of arrival(AoA) and angle of departure(AoD) for all rays associated with all clusters
        # Angle is uniformly randomly sampled in the range [60 120]
        phi_t  = np.random.uniform(60, 120, (self.N_cl, self.N_ray))
        phi_r  = np.random.uniform(60, 120, (self.N_cl, self.N_ray))
        
        # Compute the channel matrix according the SV channel model 
        for i in range(self.N_cl):
            
            for j in range(self.N_ray):
                
                phase_t = (2*np.pi*self.d/self.lambd) * np.sin(phi_t[i,j])
                phase_r = (2*np.pi*self.d/self.lambd) * np.sin(phi_r[i,j])
                
                # Computing normalized Tx and Rx antenna array response
                f_t = (1/self.N_t) * np.exp(1j*np.reshape(np.arange(0,self.N_t,1),(self.N_t,1))*phase_t)
                f_r = (1/self.N_r) * np.exp(1j*np.reshape(np.arange(0,self.N_r,1),(self.N_r,1))*phase_r)
                
                self.H = self.H + alpha[i,j]* f_r @ np.transpose(np.conjugate(f_t))
        
        self.H = np.sqrt((self.N_r*self.N_t)/(self.N_cl*self.N_ray)) * self.H
        return self.H
        
    def Reset(self):
        
        # retrun initial state
        
        self.W_rf = 1
        self.V_bb = 1
        pass

    def Step(self, state, action):
        
        """
        Assuming inputs are numpy arrays
        Input : (state, action)
            state: (Vbb_prev, Wrf_prev)
            action : (V_bb, W_rf)
            V_bb : 1D stacked vector with entries (Real(V_bb), Imag(V_bb))
            W_rf : 1D stacked vector with entries (Real(W_rf), Imag(W_rf))
        Returns :  (next_state, reward)
            next_state : same as the action taken
        """
        
        Vbb = action[0: 2 * self.Nt_rf * self.N_s]
        Wrf = action[2 * self.Nt_rf * self.N_s + 1:]

        # Obtain the complex Vbb matix

        Vbb_real = Vbb[0: self.Nt_rf * self.N_s]
        self.Vbb_real = np.reshape(Vbb_real, newshape=(self.Nt_rf, self.N_s))

        Vbb_imag = Vbb[self.Nt_rf * self.N_s + 1:]
        self.Vbb_imag = np.reshape(Vbb_imag, newshape=(self.Nt_rf, self.N_s))

        self.V_bb = self.Vbb_real + self.Vbb_image * 1j

        # Obatin the complex Wrf matix

        Wrf_real = Wrf[0: self.N_r * self.Nr_rf]
        self.Wrf_real = np.reshape(Wrf_real, newshape=(self.N_r, self.Nr_rf))

        Wrf_imag = Wrf[self.N_r * self.Nr_rf + 1:]
        self.Wrf_imag = np.reshape(Wrf_imag, newshape=(self.N_r, self.Nr_rf))

        self.Wrf = self.Wrf_real + self.Wrf_image * 1j

        # We need to condition that absolute value of each component of Wrf
        # should be equal to 1.

        # Need to calculate Wbb
        # 1. Calculate psi
        # 2. Then use psi to calculate Wbb as given in eqn 8

        # Understand the MMSE criterion on how formula for Wbb is obtained

        P = self.SNR * self.sigma
        Psi = (((1 - self.beta) * self.H @ (self.V_rf @ self.V_bb)
               * np.transpose(np.conjugate(self.V_rf @ self.V_bb))
               * np.transpose(np.conjugate(self.H)))
               + (self.beta * P + self.sigma) * np.eye(self.N_r))

        self.W_bb = (np.sqrt(1-self.beta) * np.linalg.inv(np.transpose(np.conjugate(self.W_rf)) @ Psi @ self.W_rf)
                     @ (np.transpose(np.conjugate(self.W_rf))) @ self.H @ self.V_rf @ self.V_bb)

        reward = self.cal_reward()

        return (action, reward)

    def Call_reward(self):

        # Calculate C and take inverse of it

        C = self.sigma * np.transpose(np.conjugate(self.W_bb)) @ np.transpose(np.conjugate(self.W_rf)) @ self.W_rf @ self.W_bb

        C_inv = np.linalg.inv(C)

        # Split the rate equation into two terms and finally add them and take logdet

        t1 = (1 + self.beta * self.SNR) * torch.eye(self.N_s)
        t2 = ((1 - self.beta) * C_inv
              @ np.transpose(np.conjugate(self.W_bb)) @ np.transpose(np.conjugate(self.W_rf))
              @ self.H @ self.V_rf @ self.V_bb
              @ np.transpose(np.conjugate(self.V_bb))
              @ np.transpose(np.conjugate(self.V_rf))
              @ np.transpose(np.conjugate(self.H))
              @ self.W_rf @ self.W_bb
              )

        Rbar = np.log2(np.linalg.det(t1 + t2))

        return Rbar
# %%

import numpy as np
from numpy.random import rand
# Randomly choose real and imaginary parts.
# Treat last axis as the real and imaginary parts.
A = rand(100, 2)
# Cast the array as a complex array
# Note that this will now be a 100x1 array
A_comp = A.view(dtype=np.complex128)
# To get the original array A back from the complex version
# A = A.view(dtype=np.float64)