Super_Nesterov_Accelerated_Deep_Unfolded_WMMSE.ipynb

{
  "nbformat": 4,
  "nbformat_minor": 0,
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "name": "python3",
      "display_name": "Python 3"
    },
    "accelerator": "TPU"
  },
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "X6NBSyvW9e5Q"
      },
      "source": [
        "\n",
        "# Nesterov accelerated deep unfolded WMMSE versus WMMSE\n",
        "The WMMSE algorithm is implemented as described in the paper \"An Iteratively Weighted MMSE Approach to Distributed Sum-Utility Maximization for \n",
        "a MIMO Interfering Broadcast Channel\", by Q. Shi et al. \n",
        "\n",
        "Note:\n",
        "*   Functions and variables with \"_nn\" (i.e. neural network) in the name refer to the deep unfolded WMMSE algorithm\n",
        "*   In the code we use the term precoder to indicate the beamformer \n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "xF5eK6Qt-J7g"
      },
      "source": [
        "### Import libraries and set variables"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "cRHBPbJT9_dN"
      },
      "source": [
        "# Import libraries\n",
        "import tensorflow as tf # tensorflow_version 1.x is needed\n",
        "import numpy as np\n",
        "import copy\n",
        "from copy import deepcopy\n",
        "import time\n",
        "import matplotlib.pyplot as plt\n",
        "\n",
        "\n",
        "# Set variables \n",
        "nr_of_users = 4\n",
        "nr_of_BS_antennas = 4\n",
        "scheduled_users = [0,1,2,3] # array of scheduled users. Note that we schedule all the users.\n",
        "epsilon = 0.0001 # used to end the iterations of the WMMSE algorithm in Shi et al. when the number of iterations is not fixed (note that the stopping criterion has precendence over the fixed number of iterations)\n",
        "power_tolerance = 0.0001 # used to end the bisection search in the WMMSE algorithm in Shi et al.\n",
        "total_power = 10 # power constraint in the weighted sum rate maximization problem\n",
        "noise_power = 1\n",
        "path_loss_option = False # used to add a random path loss (drawn from a uniform distribution) to the channel of each user\n",
        "path_loss_range = [-5,5] # interval of the uniform distribution from which the path loss if drawn (in dB)\n",
        "nr_of_batches_training = 20000 # used for training\n",
        "nr_of_batches_test = 1000 # used for testing\n",
        "nr_of_samples_per_batch = 100 \n",
        "nr_of_iterations = 1 # for WMMSE algorithm in Shi et al. \n",
        "nr_of_iterations_nn = 1 # for the deep unfolded WMMSE in our paper\n",
        "\n",
        "# User weights in the weighted sum rate (denoted by alpha in our paper)\n",
        "user_weights = np.reshape(np.ones(nr_of_users*nr_of_samples_per_batch),(nr_of_samples_per_batch,nr_of_users,1))\n",
        "user_weights_for_regular_WMMSE = np.ones(nr_of_users)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "2w-i4Wq5-FAV"
      },
      "source": [
        "### Define utility functions\n"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "vyYStVlV9f-w"
      },
      "source": [
        "# Compute power for bisection search in the optimization of the transmitter precoder \n",
        "# - eq. (18) in the paper by Shi et al.\n",
        "def compute_P(Phi_diag_elements, Sigma_diag_elements, mu):\n",
        "  nr_of_BS_antennas = Phi_diag_elements.size\n",
        "  mu_array = mu*np.ones(Phi_diag_elements.size)\n",
        "  result = np.divide(Phi_diag_elements,(Sigma_diag_elements + mu_array)**2)\n",
        "  result = np.sum(result)\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_norm_of_complex_array(x):\n",
        "  result = np.sqrt(np.sum((np.absolute(x))**2))\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_sinr(channel, precoder, noise_power, user_id, selected_users):\n",
        "    nr_of_users = np.size(channel,0)\n",
        "    numerator = (np.absolute(np.matmul(np.conj(channel[user_id,:]),precoder[user_id,:])))**2\n",
        "\n",
        "    inter_user_interference = 0\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index != user_id and user_index in selected_users:\n",
        "        inter_user_interference = inter_user_interference + (np.absolute(np.matmul(np.conj(channel[user_id,:]),precoder[user_index,:])))**2\n",
        "    denominator = noise_power + inter_user_interference\n",
        "\n",
        "    result = numerator/denominator\n",
        "    return result\n",
        "\n",
        "\n",
        "def compute_user_weights(nr_of_users, selected_users):\n",
        "  result = np.ones(nr_of_users)\n",
        "  for user_index in range(nr_of_users):\n",
        "    if not (user_index in selected_users):\n",
        "      result[user_index] = 0\n",
        "  return result\n",
        "\n",
        "\n",
        "def compute_weighted_sum_rate(user_weights, channel, precoder, noise_power, selected_users):\n",
        "   result = 0\n",
        "   nr_of_users = np.size(channel,0)\n",
        "   \n",
        "   for user_index in range(nr_of_users):\n",
        "     if user_index in selected_users:\n",
        "       user_sinr = compute_sinr(channel, precoder, noise_power, user_index, selected_users)\n",
        "       result = result + user_weights[user_index]*np.log2(1 + user_sinr)\n",
        "    \n",
        "   return result\n",
        "\n",
        "\n",
        "def compute_sinr_nn(channel, precoder, noise_power, user_id, nr_of_users):\n",
        "\n",
        "    numerator = tf.reduce_sum((tf.matmul(tf.transpose(channel[user_id]),precoder[user_id]))**2)\n",
        "    inter_user_interference = 0\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index != user_id:\n",
        "        inter_user_interference = inter_user_interference +  tf.reduce_sum((tf.matmul(tf.transpose(channel[user_id]),precoder[user_index]))**2)\n",
        "    denominator = noise_power + inter_user_interference\n",
        "\n",
        "    result = numerator/denominator\n",
        "    return result\n",
        "\n",
        "\n",
        "def compute_WSR_nn(user_weights, channel, precoder, noise_power, nr_of_users):\n",
        "\n",
        "   result = 0\n",
        "   for batch_index in range(nr_of_samples_per_batch):\n",
        "    for user_index in range(nr_of_users):\n",
        "        user_sinr = compute_sinr_nn(channel[batch_index], precoder[batch_index], noise_power, user_index,nr_of_users)\n",
        "        result = result + user_weights[batch_index][user_index]*(tf.log(1 + user_sinr)/tf.log(tf.cast(2.0,tf.float64)))\n",
        "\n",
        "   return result\n",
        "\n",
        "\n",
        "# Computes a channel realization and returns it in two formats, one for the WMMSE and one for the deep unfolded WMMSE.\n",
        "# It also returns the initialization value of the transmitter precoder, which is used as input in the computation graph of the deep unfolded WMMSE.\n",
        "def compute_channel(nr_of_BS_antennas, nr_of_users, total_power, path_loss_option = False, path_loss_range = [-5,5] ):\n",
        "  channel_nn = []\n",
        "  initial_transmitter_precoder = []\n",
        "  channel_WMMSE = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas))\n",
        "\n",
        "  \n",
        "  for i in range(nr_of_users):\n",
        "\n",
        "      regularization_parameter_for_RZF_solution = 0\n",
        "      path_loss = 0 # path loss is 0 dB by default, otherwise it is drawn randomly from a uniform distribution (N.B. it is different for each user)\n",
        "      if path_loss_option == True:\n",
        "        path_loss = np.random.uniform(path_loss_range[0],path_loss_range[-1])\n",
        "        regularization_parameter_for_RZF_solution = regularization_parameter_for_RZF_solution + 1/((10**(path_loss/10))*total_power) # computed as in \"MMSE precoding for multiuser MISO downlink transmission with non-homogeneous user SNR conditions\" by D.H. Nguyen and T. Le-Ngoc\n",
        "\n",
        "      result_real = np.sqrt(10**(path_loss/10))*np.sqrt(0.5)*np.random.normal(size = (nr_of_BS_antennas,1))\n",
        "      result_imag  =  np.sqrt(10**(path_loss/10))*np.sqrt(0.5)*np.random.normal(size = (nr_of_BS_antennas,1))\n",
        "      channel_WMMSE[i,:] = np.reshape(result_real,(1,nr_of_BS_antennas)) + 1j*np.reshape(result_imag, (1,nr_of_BS_antennas))\n",
        "      result_col_1 = np.vstack((result_real,result_imag))\n",
        "      result_col_2 = np.vstack((-result_imag,result_real))\n",
        "      result =  np.hstack((result_col_1, result_col_2))\n",
        "      initial_transmitter_precoder.append(result_col_1)\n",
        "      channel_nn.append(result)\n",
        "\n",
        "  initial_transmitter_precoder_array = np.array(initial_transmitter_precoder)\n",
        "  initial_transmitter_precoder_array = np.sqrt(total_power)*initial_transmitter_precoder_array/np.linalg.norm(initial_transmitter_precoder_array)\n",
        "  initial_transmitter_precoder = []\n",
        "\n",
        "  for i in range(nr_of_users):\n",
        "    initial_transmitter_precoder.append(initial_transmitter_precoder_array[i])\n",
        "\n",
        "  return channel_nn, initial_transmitter_precoder, channel_WMMSE, regularization_parameter_for_RZF_solution\n",
        "\n",
        "\n",
        "# Computes the zero-forcing solution as in \"MMSE precoding for multiuser MISO downlink transmission with non-homogeneous user SNR conditions\" by D.H. Nguyen and T. Le-Ngoc\n",
        "def zero_forcing(channel_realization, total_power):\n",
        "  \n",
        "  ZF_solution = np.matmul((np.transpose(channel_realization)),np.linalg.inv(np.matmul(np.conj(channel_realization),(np.transpose(channel_realization)))))\n",
        "  ZF_solution = ZF_solution*np.sqrt(total_power)/np.linalg.norm(ZF_solution) # scaled according to the power constraint\n",
        "\n",
        "  return np.transpose(ZF_solution)\n",
        "\n",
        "\n",
        "# Computes the regularized zero-forcing solution as in \"MMSE precoding for multiuser MISO downlink transmission with non-homogeneous user SNR conditions\" by D.H. Nguyen and T. Le-Ngoc\n",
        "def regularized_zero_forcing(channel_realization, total_power, regularization_parameter = 0, path_loss_option = False):\n",
        "  \n",
        "  if path_loss_option == False:\n",
        "    RZF_solution = np.matmul((np.transpose(channel_realization)),np.linalg.inv(np.matmul(np.conj(channel_realization),(np.transpose(channel_realization))) + nr_of_users/total_power*np.eye(nr_of_users, nr_of_users)))\n",
        "  else:\n",
        "    RZF_solution = np.matmul((np.transpose(channel_realization)),np.linalg.inv(np.matmul(np.conj(channel_realization),(np.transpose(channel_realization))) + regularization_parameter*np.eye(nr_of_users, nr_of_users)))\n",
        "\n",
        "  RZF_solution = RZF_solution*np.sqrt(total_power)/np.linalg.norm(RZF_solution) # scaled according to the power constraint\n",
        "\n",
        "  return np.transpose(RZF_solution)\n",
        "\n",
        "\n",
        "# Builds one Super Nesterov accelerated PGD iteration in the deep unfolded WMMSE network\n",
        "def SupNesterov_acc_PGD_step( init_step_size, init_momentum1, init_momentum2, name, mse_weights, user_weights, receiver_precoder,channel, initial_transmitter_precoder_present,initial_transmitter_precoder_past, total_power):\n",
        "\n",
        "  with tf.variable_scope(name): \n",
        "\n",
        "    step_size =  tf.Variable(tf.constant(init_step_size, dtype = tf.float64), name = \"step_size\", dtype = tf.float64)\n",
        "    momentum1 =  tf.Variable(tf.constant(init_momentum1, dtype = tf.float64), name = \"momentum_1\", dtype = tf.float64)\n",
        "    momentum2 =  tf.Variable(tf.constant(init_momentum2, dtype = tf.float64), name = \"momentum_2\", dtype = tf.float64)\n",
        "\n",
        "    # First iteration\n",
        "    a1_exp = tf.tile(tf.expand_dims(mse_weights[:,0,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "    a2_exp = tf.tile(tf.expand_dims(user_weights[:,0,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "    a3_exp = tf.tile(tf.expand_dims(tf.reduce_sum((receiver_precoder[:,0,:,:])**2,axis =-2),-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])    \n",
        "    temp = a1_exp*a2_exp*a3_exp*tf.matmul(channel[:,0,:,:],tf.transpose(channel[:,0,:,:],perm = [0,2,1]))\n",
        "    \n",
        "    # Next iterations\n",
        "    for i in range(1, nr_of_users):\n",
        "      a1_exp = tf.tile(tf.expand_dims(mse_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "      a2_exp = tf.tile(tf.expand_dims(user_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "      a3_exp = tf.tile(tf.expand_dims(tf.reduce_sum((receiver_precoder[:,i,:,:])**2,axis =-2),-1),[1,2*nr_of_BS_antennas,2*nr_of_BS_antennas])\n",
        "      temp = temp + a1_exp*a2_exp*a3_exp*tf.matmul(channel[:,i,:,:],tf.transpose(channel[:,i,:,:],perm = [0,2,1]))\n",
        "\n",
        "    sum_gradient = temp \n",
        "    \n",
        "    # Super Nesterov acceleration (part 1)\n",
        "    initial_transmitter_precoder = initial_transmitter_precoder_present + momentum2*(initial_transmitter_precoder_present - initial_transmitter_precoder_past)\n",
        "\n",
        "    gradient = []\n",
        "\n",
        "    # Gradient computation\n",
        "    for i in range(nr_of_users):\n",
        "      a1_exp = tf.tile(tf.expand_dims(mse_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,1])\n",
        "      a2_exp = tf.tile(tf.expand_dims(user_weights[:,i,:],-1),[1,2*nr_of_BS_antennas,1])\n",
        "      gradient.append(step_size*(-2.0*a1_exp*a2_exp*tf.matmul(channel[:,i,:,:],receiver_precoder[:,i,:,:])+ 2*tf.matmul(sum_gradient,initial_transmitter_precoder[:,i,:,:]))) \n",
        "      \n",
        "    tf.stack(gradient)\n",
        "    gradient = tf.transpose( tf.stack(gradient), perm=[1,0,2,3])\n",
        "\n",
        "    # Super Nesterov acceleration (part 2)\n",
        "    output_temp = initial_transmitter_precoder_present + momentum1*(initial_transmitter_precoder_present - initial_transmitter_precoder_past) - gradient\n",
        "\n",
        "    output = []\n",
        "    for i in range(nr_of_samples_per_batch):\n",
        "      output.append(tf.cond((tf.linalg.norm(output_temp[i]))**2 < total_power, lambda: output_temp[i] , lambda: tf.sqrt(tf.cast(total_power, tf.float64))*output_temp[i]/tf.linalg.norm(output_temp[i]))) \n",
        "\n",
        "    return tf.stack(output), initial_transmitter_precoder_present, step_size, momentum1, momentum2\n",
        "\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "q37r-dBd-0K1"
      },
      "source": [
        "### Function to run the WMMSE algorithm described in the paper by Shi et al."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "F8me_JRP-0bw"
      },
      "source": [
        "def run_WMMSE(epsilon, channel, selected_users, total_power, noise_power, user_weights, max_nr_of_iterations, log = False):\n",
        "\n",
        "  nr_of_users = np.size(channel,0)\n",
        "  nr_of_BS_antennas = np.size(channel,1)\n",
        "  WSR=[] # to check if the WSR (our cost function) increases at each iteration of the WMMSE\n",
        "  break_condition = epsilon + 1 # break condition to stop the WMMSE iterations and exit the while\n",
        "  receiver_precoder = np.zeros(nr_of_users) + 1j*np.zeros(nr_of_users) # receiver_precoder is \"u\" in the paper of Shi et al. (it's a an array of complex scalars)\n",
        "  mse_weights = np.ones(nr_of_users) # mse_weights is \"w\" in the paper of Shi et al. (it's a an array of real scalars)\n",
        "  transmitter_precoder = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas))# transmitter_precoder is \"v\" in the paper of Shi et al. (it's a complex matrix)\n",
        "  \n",
        "  new_receiver_precoder = np.zeros(nr_of_users) + 1j*np.zeros(nr_of_users) # for the first iteration \n",
        "  new_mse_weights = np.zeros(nr_of_users) # for the first iteration\n",
        "  new_transmitter_precoder = np.zeros((nr_of_users, nr_of_BS_antennas)) + 1j*np.zeros((nr_of_users, nr_of_BS_antennas)) # for the first iteration\n",
        "\n",
        "  \n",
        "  # Initialization of transmitter precoder\n",
        "  for user_index in range(nr_of_users):\n",
        "    if user_index in selected_users:\n",
        "      transmitter_precoder[user_index,:] = channel[user_index,:]\n",
        "  transmitter_precoder = transmitter_precoder/np.linalg.norm(transmitter_precoder)*np.sqrt(total_power)\n",
        "  \n",
        "  # Store the WSR obtained with the initialized trasmitter precoder    \n",
        "  WSR.append(compute_weighted_sum_rate(user_weights, channel, transmitter_precoder, noise_power, selected_users))\n",
        "\n",
        "  # Compute the initial power of the transmitter precoder\n",
        "  initial_power = 0\n",
        "  for user_index in range(nr_of_users):\n",
        "    if user_index in selected_users:\n",
        "      initial_power = initial_power + (compute_norm_of_complex_array(transmitter_precoder[user_index,:]))**2 \n",
        "  if log == True:\n",
        "    print(\"Power of the initialized transmitter precoder:\", initial_power)\n",
        "\n",
        "  nr_of_iteration_counter = 0 # to keep track of the number of iteration of the WMMSE\n",
        "\n",
        "  while break_condition >= epsilon and nr_of_iteration_counter<=max_nr_of_iterations:\n",
        "    \n",
        "    nr_of_iteration_counter = nr_of_iteration_counter + 1\n",
        "    if log == True:\n",
        "      print(\"WMMSE ITERATION: \", nr_of_iteration_counter)\n",
        "\n",
        "    # Optimize receiver precoder - eq. (5) in the paper of Shi et al.\n",
        "    for user_index_1 in range(nr_of_users):\n",
        "      if user_index_1 in selected_users:\n",
        "        user_interference = 0.0\n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 in selected_users:\n",
        "            user_interference = user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "\n",
        "        new_receiver_precoder[user_index_1] = np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_1,:]) / (noise_power + user_interference)\n",
        "\n",
        "    # Optimize mse_weights - eq. (13) in the paper of Shi et al.\n",
        "    for user_index_1 in range(nr_of_users):\n",
        "      if user_index_1 in selected_users:\n",
        "\n",
        "        user_interference = 0 # it includes the channel of all selected users\n",
        "        inter_user_interference = 0 # it includes the channel of all selected users apart from the current one\n",
        "        \n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 in selected_users:\n",
        "            user_interference = user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "        for user_index_2 in range(nr_of_users):\n",
        "          if user_index_2 != user_index_1 and user_index_2 in selected_users:\n",
        "            inter_user_interference = inter_user_interference + (np.absolute(np.matmul(np.conj(channel[user_index_1,:]),transmitter_precoder[user_index_2,:])))**2\n",
        "        \n",
        "        new_mse_weights[user_index_1] = (noise_power + user_interference)/(noise_power + inter_user_interference)\n",
        "\n",
        "    A = np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))+1j*np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        # hh should be an hermitian matrix of size (nr_of_BS_antennas X nr_of_BS_antennas)\n",
        "        hh = np.matmul(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)),np.conj(np.transpose(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))))\n",
        "        A = A + (new_mse_weights[user_index]*user_weights[user_index]*(np.absolute(new_receiver_precoder[user_index]))**2)*hh\n",
        "\n",
        "    Sigma_diag_elements_true, U = np.linalg.eigh(A)\n",
        "    Sigma_diag_elements = copy.deepcopy(np.real(Sigma_diag_elements_true))\n",
        "    Lambda = np.zeros((nr_of_BS_antennas,nr_of_BS_antennas)) + 1j*np.zeros((nr_of_BS_antennas,nr_of_BS_antennas))\n",
        "    \n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:     \n",
        "        hh = np.matmul(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)),np.conj(np.transpose(np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))))\n",
        "        Lambda = Lambda + ((user_weights[user_index])**2)*((new_mse_weights[user_index])**2)*((np.absolute(new_receiver_precoder[user_index]))**2)*hh\n",
        "\n",
        "    Phi = np.matmul(np.matmul(np.conj(np.transpose(U)),Lambda),U)\n",
        "    Phi_diag_elements_true = np.diag(Phi)\n",
        "    Phi_diag_elements = copy.deepcopy(Phi_diag_elements_true)\n",
        "    Phi_diag_elements = np.real(Phi_diag_elements)\n",
        "\n",
        "    for i in range(len(Phi_diag_elements)):\n",
        "      if Phi_diag_elements[i]<np.finfo(float).eps:\n",
        "        Phi_diag_elements[i] = np.finfo(float).eps\n",
        "      if (Sigma_diag_elements[i])<np.finfo(float).eps:\n",
        "        Sigma_diag_elements[i] = 0\n",
        "\n",
        "    # Check if mu = 0 is a solution (eq.s (15) and (16) of in the paper of Shi et al.)\n",
        "    power = 0 # the power of transmitter precoder (i.e. sum of the squared norm)\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        if np.linalg.det(A) != 0:\n",
        "          w = np.matmul(np.linalg.inv(A),np.reshape(channel[user_index,:],(nr_of_BS_antennas,1)))*user_weights[user_index]*new_mse_weights[user_index]*(new_receiver_precoder[user_index])\n",
        "          power = power + (compute_norm_of_complex_array(w))**2\n",
        "\n",
        "    # If mu = 0 is a solution, then mu_star = 0\n",
        "    if np.linalg.det(A) != 0 and power <= total_power:\n",
        "      mu_star = 0\n",
        "    # If mu = 0 is not a solution then we search for the \"optimal\" mu by bisection\n",
        "    else:\n",
        "      power_distance = [] # list to store the distance from total_power in the bisection algorithm \n",
        "      mu_low = np.sqrt(1/total_power*np.sum(Phi_diag_elements))\n",
        "      mu_high = 0\n",
        "      low_point = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_low)\n",
        "      high_point = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_high)\n",
        "\n",
        "      obtained_power = total_power + 2*power_tolerance # initialization of the obtained power such that we enter the while \n",
        "\n",
        "      # Bisection search\n",
        "      while np.absolute(total_power - obtained_power) > power_tolerance:\n",
        "        mu_new = (mu_high + mu_low)/2\n",
        "        obtained_power = compute_P(Phi_diag_elements, Sigma_diag_elements, mu_new) # eq. (18) in the paper of Shi et al.\n",
        "        power_distance.append(np.absolute(total_power - obtained_power))\n",
        "        if obtained_power > total_power:\n",
        "          mu_high = mu_new\n",
        "        if obtained_power < total_power:\n",
        "          mu_low = mu_new\n",
        "      mu_star = mu_new\n",
        "      if log == True:\n",
        "        print(\"first value:\", power_distance[0])\n",
        "        plt.title(\"Distance from the target value in bisection (it should decrease)\")\n",
        "        plt.plot(power_distance)\n",
        "        plt.show()\n",
        "\n",
        "    for user_index in range(nr_of_users):\n",
        "      if user_index in selected_users:\n",
        "        new_transmitter_precoder[user_index,:] = np.matmul(np.linalg.inv(A + mu_star*np.eye(nr_of_BS_antennas)),channel[user_index,:])*user_weights[user_index]*new_mse_weights[user_index]*(new_receiver_precoder[user_index]) \n",
        "\n",
        "    # To select only the weights of the selected users to check the break condition\n",
        "    mse_weights_selected_users = []\n",
        "    new_mse_weights_selected_users = []\n",
        "    for user_index in range(nr_of_users): \n",
        "      if user_index in selected_users:\n",
        "        mse_weights_selected_users.append(mse_weights[user_index])\n",
        "        new_mse_weights_selected_users.append(new_mse_weights[user_index])\n",
        "\n",
        "    mse_weights = deepcopy(new_mse_weights)\n",
        "    transmitter_precoder = deepcopy(new_transmitter_precoder)\n",
        "    receiver_precoder = deepcopy(new_receiver_precoder)\n",
        "\n",
        "    WSR.append(compute_weighted_sum_rate(user_weights, channel, transmitter_precoder, noise_power, selected_users))\n",
        "    break_condition = np.absolute(np.log2(np.prod(new_mse_weights_selected_users))-np.log2(np.prod(mse_weights_selected_users)))\n",
        "\n",
        "  if log == True:\n",
        "    plt.title(\"Change of the WSR at each iteration of the WMMSE (it should increase)\")\n",
        "    plt.plot(WSR,'bo')\n",
        "    plt.show()\n",
        "\n",
        "  return transmitter_precoder, receiver_precoder, mse_weights, WSR[-1]\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "MwsPve2s-1f_"
      },
      "source": [
        "### Tensorflow computation graph to run the accelerated deep unfolded WMMSE algorithm described in our paper"
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "LSbUVrNg-1r6"
      },
      "source": [
        "tf.reset_default_graph()\n",
        "\n",
        "channel_input = tf.placeholder(tf.float64, shape = None, name = 'channel_input')\n",
        "initial_tp = tf.placeholder(tf.float64, shape = None, name = 'initial_transmitter_precoder')\n",
        "initial_tp_past = tf.placeholder(tf.float64, shape = None, name = 'initial_transmitter_precoder_past') # needed for the acceleration scheme\n",
        "initial_transmitter_precoder = initial_tp\n",
        "initial_transmitter_precoder_past = initial_tp_past\n",
        "\n",
        "# Arrays that contain the initialization values of the step sizes.\n",
        "# The number of step sizes depends on the selected number of PGD layers, the number of elements for each step size initializer depends on the selected number of deep unfolded iterations\n",
        "step_size1_init = [1.0]\n",
        "step_size2_init = [1.0]\n",
        "step_size3_init = [1.0]\n",
        "step_size4_init = [1.0]\n",
        "\n",
        "# Arrays that contain the initialization values of the momentum parameters for Super Nesterov acceleration. N.B. there are two sets of momentum parammeters\n",
        "# The number of momentum parameters (in each set) depends on the selected number of PGD layers, the number of elements for each momentum initializer depends on the selected number of deep unfolded iterations\n",
        "momentum1_1_init = [0.0]\n",
        "momentum1_2_init = [0.0]\n",
        "momentum1_3_init = [0.0]\n",
        "momentum1_4_init = [0.0]\n",
        "\n",
        "momentum2_1_init = [0.0]\n",
        "momentum2_2_init = [0.0]\n",
        "momentum2_3_init = [0.0]\n",
        "momentum2_4_init = [0.0]\n",
        "\n",
        "# Used to collect the step sizes at each iteration\n",
        "all_step_size1_temp = []\n",
        "all_step_size2_temp = []\n",
        "all_step_size3_temp = []\n",
        "all_step_size4_temp = []\n",
        "\n",
        "# Used to collect the momentum parameters at each iteration\n",
        "all_momentum1_1_temp = []\n",
        "all_momentum1_2_temp = []\n",
        "all_momentum1_3_temp = []\n",
        "all_momentum1_4_temp = []\n",
        "\n",
        "all_momentum2_1_temp = []\n",
        "all_momentum2_2_temp = []\n",
        "all_momentum2_3_temp = []\n",
        "all_momentum2_4_temp = []\n",
        "\n",
        "profit = [] # stores the WSR obtained at each iteration\n",
        "\n",
        "for loop in range(0,nr_of_iterations_nn):\n",
        "\n",
        "  user_interference2 = []\n",
        "\n",
        "  for batch_index in range(nr_of_samples_per_batch): \n",
        "    user_interference_single = []\n",
        "    for i in range(nr_of_users):\n",
        "      temp = 0.0\n",
        "      for j in range(nr_of_users):\n",
        "        temp = temp + tf.reduce_sum((tf.matmul(tf.transpose(channel_input[batch_index, i,:,:]),initial_transmitter_precoder[batch_index,j,:,:]))**2)\n",
        "      user_interference_single.append(temp + noise_power) \n",
        "    user_interference2.append(user_interference_single)\n",
        "    \n",
        "  tf.stack(user_interference2)\n",
        "\n",
        "  user_interference_exp2 = tf.tile(tf.expand_dims(tf.tile(tf.expand_dims(user_interference2,-1),[1,1,2]),-1),[1,1,1,1])\n",
        "\n",
        "  receiver_precoder_temp = (tf.matmul(tf.transpose(channel_input,perm = [0,1,3,2]),initial_transmitter_precoder))\n",
        "  # Optimize the receiver precoder \n",
        "  receiver_precoder = tf.divide(receiver_precoder_temp,user_interference_exp2)\n",
        "\n",
        "  # Optimize the mmse weights \n",
        "  self_interference = tf.reduce_sum((tf.matmul(tf.transpose(channel_input,perm = [0,1,3,2]),initial_transmitter_precoder))**2, axis = 2)\n",
        "\n",
        "  inter_user_interference_total = []\n",
        "\n",
        "  for batch_index in range(nr_of_samples_per_batch):\n",
        "    inter_user_interference_temp = []\n",
        "    for i in range(nr_of_users):\n",
        "      temp = 0.0\n",
        "      for j in range(nr_of_users):\n",
        "        if j != i:\n",
        "          temp = temp + tf.reduce_sum((tf.matmul(tf.transpose(channel_input[batch_index, i,:,:]),initial_transmitter_precoder[batch_index,j,:,:]))**2)\n",
        "      inter_user_interference_temp.append(temp + noise_power) # $sum{|(h_i)*H,v_i}|**2 + noise_power$\n",
        "    inter_user_interference = tf.reshape(tf.stack(inter_user_interference_temp),(nr_of_users,1)) # Nx1 $sum{|(h_i)*H,v_i}|**2 + noise_power$\n",
        "    inter_user_interference_total.append(inter_user_interference)\n",
        "      \n",
        "  mse_weights = (tf.divide(self_interference,inter_user_interference_total)) + 1.0\n",
        "\n",
        "  # Optimize the transmitter precoder through Super Nesterov accelerated PGD\n",
        "  transmitter_precoder1, transmitter_precoder1_past, step_size1, momentum1_1, momentum2_1 = SupNesterov_acc_PGD_step( step_size1_init[loop], momentum1_1_init[loop], momentum2_1_init[loop], 'PGD_step1', mse_weights, user_weights, receiver_precoder,channel_input, initial_transmitter_precoder, initial_transmitter_precoder_past, total_power)\n",
        "  transmitter_precoder2, transmitter_precoder2_past, step_size2, momentum1_2, momentum2_2 = SupNesterov_acc_PGD_step( step_size2_init[loop], momentum1_2_init[loop], momentum2_2_init[loop],'PGD_step2', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder1, transmitter_precoder1_past, total_power)\n",
        "  transmitter_precoder3, transmitter_precoder3_past, step_size3, momentum1_3, momentum2_3 = SupNesterov_acc_PGD_step( step_size3_init[loop], momentum1_3_init[loop], momentum2_3_init[loop], 'PGD_step3', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder2, transmitter_precoder2_past, total_power)\n",
        "  transmitter_precoder, transmitter_precoder3_past, step_size4, momentum1_4, momentum2_4 = SupNesterov_acc_PGD_step( step_size4_init[loop], momentum1_4_init[loop], momentum2_4_init[loop], 'PGD_step4', mse_weights, user_weights, receiver_precoder,channel_input, transmitter_precoder3, transmitter_precoder3_past, total_power)\n",
        "\n",
        "  initial_transmitter_precoder = transmitter_precoder\n",
        "\n",
        "  all_step_size1_temp.append(step_size1)\n",
        "  all_step_size2_temp.append(step_size2)\n",
        "  all_step_size3_temp.append(step_size3)\n",
        "  all_step_size4_temp.append(step_size4)\n",
        "\n",
        "  all_momentum1_1_temp.append(momentum1_1)\n",
        "  all_momentum1_2_temp.append(momentum1_2)\n",
        "  all_momentum1_3_temp.append(momentum1_3)\n",
        "  all_momentum1_4_temp.append(momentum1_4)\n",
        "\n",
        "  all_momentum2_1_temp.append(momentum2_1)\n",
        "  all_momentum2_2_temp.append(momentum2_2)\n",
        "  all_momentum2_3_temp.append(momentum2_3)\n",
        "  all_momentum2_4_temp.append(momentum2_4)\n",
        "\n",
        "  # The WSR achieved with the transmitter precoder obtained at the current iteration is appended\n",
        "  profit.append(compute_WSR_nn(user_weights, channel_input, initial_transmitter_precoder, noise_power,nr_of_users))\n",
        "\n",
        "all_step_size1 = tf.stack(all_step_size1_temp)\n",
        "all_step_size2 = tf.stack(all_step_size2_temp)\n",
        "all_step_size3 = tf.stack(all_step_size3_temp)\n",
        "all_step_size4 = tf.stack(all_step_size4_temp)\n",
        "\n",
        "all_momentum1_1 = tf.stack(all_momentum1_1_temp)\n",
        "all_momentum1_2 = tf.stack(all_momentum1_2_temp)\n",
        "all_momentum1_3 = tf.stack(all_momentum1_3_temp)\n",
        "all_momentum1_4 = tf.stack(all_momentum1_4_temp)\n",
        "\n",
        "all_momentum2_1 = tf.stack(all_momentum2_1_temp)\n",
        "all_momentum2_2 = tf.stack(all_momentum2_2_temp)\n",
        "all_momentum2_3 = tf.stack(all_momentum2_3_temp)\n",
        "all_momentum2_4 = tf.stack(all_momentum2_4_temp)\n",
        "\n",
        "\n",
        "final_precoder = initial_transmitter_precoder # this is the last transmitter precoder, i.e. the one that will be actually used for transmission\n",
        "\n",
        "WSR = tf.reduce_sum(tf.stack(profit)) # this is the cost function to maximize, i.e. the WSR obtained if we use the transmitter precoder that we have at each round of the loop \n",
        "WSR_final = compute_WSR_nn(user_weights, channel_input, final_precoder, noise_power,nr_of_users)/nr_of_samples_per_batch # this is the WSR computed using the \"final_precoder\"\n",
        "\n",
        "optimizer = tf.train.AdamOptimizer(learning_rate=0.001).minimize(-WSR)\n"
      ],
      "execution_count": null,
      "outputs": []
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "IbxrxLHm_FYp"
      },
      "source": [
        "### Running the accelerated deep unfolded WMMSE and the WMMSE\n",
        "In this code we train the neural network and we compare its performance with the WMMSE, the zero-forcing, and the regularized zero-forcing solutions, by testing all algorithms on the same dataset of channel realizations."
      ]
    },
    {
      "cell_type": "code",
      "metadata": {
        "id": "dtZIqgqM_FkX"
      },
      "source": [
        "WSR_WMMSE = [] # to store the WSR attained by the WMMSE\n",
        "WSR_ZF = [] # to store the WSR attained by the zero-forcing \n",
        "WSR_RZF = [] # to store the WSR attained by the regularized zero-forcing\n",
        "WSR_nn = [] # to store the WSR attained by the deep unfolded WMMSE\n",
        "training_loss = []\n",
        "\n",
        "with tf.Session() as sess:\n",
        "    print(\"start of session\")\n",
        "    start_of_time = time.time()\n",
        "    sess.run(tf.global_variables_initializer())\n",
        "    \n",
        "    for i in range(nr_of_batches_training):\n",
        "      batch_for_training = []\n",
        "      initial_transmitter_precoder_batch = []\n",
        "\n",
        "      # Building a batch for training\n",
        "      for ii in range(nr_of_samples_per_batch):\n",
        "        channel_realization_nn,init_transmitter_precoder,_,_ = compute_channel(nr_of_BS_antennas, nr_of_users, total_power, path_loss_option, path_loss_range)\n",
        "        batch_for_training.append(channel_realization_nn)\n",
        "        initial_transmitter_precoder_batch.append(init_transmitter_precoder)\n",
        "      initial_transmitter_precoder_batch_past = np.zeros((nr_of_samples_per_batch,nr_of_users,nr_of_BS_antennas*2,1)) # for the acceleration scheme\n",
        "      # Training\n",
        "      sess.run(optimizer, feed_dict={channel_input:batch_for_training,initial_tp:initial_transmitter_precoder_batch,initial_tp_past:initial_transmitter_precoder_batch_past })\n",
        "      training_loss.append(-1*(sess.run(WSR, feed_dict={channel_input:batch_for_training,initial_tp:initial_transmitter_precoder_batch,initial_tp_past:initial_transmitter_precoder_batch_past })))\n",
        "    \n",
        "    print(\"step size1\", sess.run(all_step_size1))\n",
        "    print(\"step size2\", sess.run(all_step_size2))\n",
        "    print(\"step size3\", sess.run(all_step_size3))\n",
        "    print(\"step size4\", sess.run(all_step_size4))\n",
        "    print(\"momentum 1_1\", sess.run(all_momentum1_1))\n",
        "    print(\"momentum 1_2\", sess.run(all_momentum1_2))\n",
        "    print(\"momentum 1_3\", sess.run(all_momentum1_3))\n",
        "    print(\"momentum 1_4\", sess.run(all_momentum1_4))\n",
        "    print(\"momentum 2_1\", sess.run(all_momentum2_1))\n",
        "    print(\"momentum 2_2\", sess.run(all_momentum2_2))\n",
        "    print(\"momentum 2_3\", sess.run(all_momentum2_3))\n",
        "    print(\"momentum 2_4\", sess.run(all_momentum2_4))\n",
        "\n",
        "    print(\"Training took:\", time.time()-start_of_time)\n",
        "\n",
        "    # For repeatability\n",
        "    np.random.seed(1234)\n",
        "\n",
        "    for i in range(nr_of_batches_test):    \n",
        "      batch_for_testing = []\n",
        "      initial_transmitter_precoder_batch = []\n",
        "      WSR_WMMSE_batch = 0.0\n",
        "      WSR_ZF_batch = 0.0\n",
        "      WSR_RZF_batch = 0.0\n",
        "\n",
        "      # Building a batch for testing\n",
        "      for ii in range(nr_of_samples_per_batch):       \n",
        "        channel_realization_nn,init_transmitter_precoder,channel_realization_regular,regularization_parameter_for_RZF_solution = compute_channel(nr_of_BS_antennas, nr_of_users, total_power, path_loss_option, path_loss_range)    \n",
        "        \n",
        "        # Compute the WMMSE solution\n",
        "        _,_,_,WSR_WMMSE_one_sample = run_WMMSE(epsilon, channel_realization_regular, scheduled_users, total_power, noise_power, user_weights_for_regular_WMMSE, nr_of_iterations-1, log = False)      \n",
        "        WSR_WMMSE_batch =  WSR_WMMSE_batch + WSR_WMMSE_one_sample\n",
        "        \n",
        "        # Compute the zero-forcing solution\n",
        "        ZF_solution = zero_forcing(channel_realization_regular, total_power)\n",
        "        WSR_ZF_one_sample = compute_weighted_sum_rate(user_weights_for_regular_WMMSE , channel_realization_regular, ZF_solution, noise_power, scheduled_users)\n",
        "        WSR_ZF_batch =  WSR_ZF_batch + WSR_ZF_one_sample\n",
        "        \n",
        "        # Compute the regilarized zero-forcing solution\n",
        "        RZF_solution = regularized_zero_forcing(channel_realization_regular, total_power, regularization_parameter_for_RZF_solution, path_loss_option)\n",
        "        WSR_RZF_one_sample = compute_weighted_sum_rate(user_weights_for_regular_WMMSE , channel_realization_regular, RZF_solution, noise_power, scheduled_users)\n",
        "        WSR_RZF_batch =  WSR_RZF_batch + WSR_RZF_one_sample\n",
        "\n",
        "        batch_for_testing.append(channel_realization_nn)\n",
        "        initial_transmitter_precoder_batch.append(init_transmitter_precoder)\n",
        "      initial_transmitter_precoder_batch_past = np.zeros((nr_of_samples_per_batch,nr_of_users,nr_of_BS_antennas*2,1)) # for the acceleration scheme\n",
        "\n",
        "      #Testing\n",
        "      WSR_nn.append(sess.run(WSR_final, feed_dict={channel_input:batch_for_testing,initial_tp:initial_transmitter_precoder_batch,initial_tp:initial_transmitter_precoder_batch,initial_tp_past:initial_transmitter_precoder_batch_past }))\n",
        "      WSR_WMMSE.append(WSR_WMMSE_batch/nr_of_samples_per_batch)\n",
        "      WSR_ZF.append(WSR_ZF_batch/nr_of_samples_per_batch)\n",
        "      WSR_RZF.append(WSR_RZF_batch/nr_of_samples_per_batch)   \n",
        "      \n",
        "print(\"Training and testing took:\", time.time()-start_of_time)\n",
        "print(\"The WSR acheived with the deep unfolded WMMSE algorithm is: \",np.mean(WSR_nn))\n",
        "print(\"The WSR acheived with the WMMSE algorithm is: \",np.mean(WSR_WMMSE))\n",
        "print(\"The WSR acheived with the zero forcing is: \",np.mean(WSR_ZF))\n",
        "print(\"The WSR acheived with the regularized zero forcing is: \",np.mean(WSR_RZF))\n",
        "\n",
        "plt.figure()\n",
        "plt.plot(training_loss)\n",
        "plt.ylabel(\"Training loss\")\n",
        "plt.xlabel(\"Sample index\")\n"
      ],
      "execution_count": null,
      "outputs": []
    }
  ]
}