{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Water Filling in Communications\n",
    "by Robert Gowers, Roger Hill, Sami Al-Izzi, Timothy Pollington and Keith Briggs\n",
    "\n",
    "from Boyd and Vandenberghe, Convex Optimization, example 5.2 page 145\n",
    "\n",
    "Convex optimisation can be used to solve the classic water filling problem.  This problem is where a total amount of power $P$ has to be assigned to $n$ communication channels, with the objective of maximising the total communication rate.  The communication rate of the $i$th channel is given by:\n",
    "\n",
    "$\\log(\\alpha_i + x_i)$\n",
    "\n",
    "where $x_i$ represents the power allocated to channel $i$ and $\\alpha_i$ represents the floor above the baseline at which power can be added to the channel.  Since $-\\log(X)$ is convex, we can write the water-filling problem as a convex optimisation problem:\n",
    "\n",
    "minimise $\\sum_{i=1}^N -\\log(\\alpha_i + x_i)$ \n",
    "\n",
    "subject to $x_i \\succeq 0$ and $\\sum_{i=1}^N x_i = P$\n",
    "\n",
    "This form is also very straightforward to put into DCP format and thus can be simply solved using CVXPY."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "#!/usr/bin/env python3\n",
    "# @author: R. Gowers, S. Al-Izzi, T. Pollington, R. Hill & K. Briggs\n",
    "import numpy as np\n",
    "import cvxpy as cp"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def water_filling(n, a, sum_x=1):\n",
    "    '''\n",
    "    Boyd and Vandenberghe, Convex Optimization, example 5.2 page 145\n",
    "    Water-filling.\n",
    "  \n",
    "    This problem arises in information theory, in allocating power to a set of\n",
    "    n communication channels in order to maximise the total channel capacity.\n",
    "    The variable x_i represents the transmitter power allocated to the ith channel, \n",
    "    and log(α_i+x_i) gives the capacity or maximum communication rate of the channel. \n",
    "    The objective is to minimise -∑log(α_i+x_i) subject to the constraint ∑x_i = 1 \n",
    "    '''\n",
    "    \n",
    "    # Declare variables and parameters\n",
    "    x = cp.Variable(shape=n)\n",
    "    alpha = cp.Parameter(n, nonneg=True)\n",
    "    alpha.value = a\n",
    "\n",
    "    # Choose objective function. Interpret as maximising the total communication rate of all the channels\n",
    "    obj = cp.Maximize(cp.sum(cp.log(alpha + x)))\n",
    "\n",
    "    # Declare constraints\n",
    "    constraints = [x >= 0, cp.sum(x) - sum_x == 0]\n",
    "      \n",
    "    # Solve\n",
    "    prob = cp.Problem(obj, constraints)\n",
    "    prob.solve()\n",
    "    if(prob.status=='optimal'):\n",
    "        return prob.status, prob.value, x.value\n",
    "    else:\n",
    "        return prob.status, np.nan, np.nan"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example\n",
    "As a simple example, we set $N = 3$, $P = 1$ and $\\boldsymbol{\\alpha} = (0.8,1.0,1.2)$.\n",
    "\n",
    "The function outputs whether the problem status, the maximum communication rate and the power allocation required is achieved with this maximum communication rate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# As an example, we will solve the water filling problem with 3 buckets, each with different α\n",
    "np.set_printoptions(precision=3) \n",
    "buckets = 3\n",
    "alpha = np.array([0.8, 1.0, 1.2])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem status: optimal\n",
      "Optimal communication rate = 0.863 \n",
      "Transmitter powers:\n",
      "[0.533 0.333 0.133]\n"
     ]
    }
   ],
   "source": [
    "stat, prob, x = water_filling(buckets, alpha)\n",
    "print('Problem status: {}'.format(stat))\n",
    "print('Optimal communication rate = {:.4g} '.format(prob))\n",
    "print('Transmitter powers:\\n{}'.format(x))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "collapsed": true
   },
   "source": [
    "To illustrate the water filling principle, we will plot $\\alpha_i + x_i$ and check that this level is flat where power has been allocated:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "image/png": "[encoded data, 13864 characters]\n",
      "text/plain": [
       "<Figure size 432x288 with 1 Axes>"
      ]
     },
     "metadata": {
      "needs_background": "light"
     },
     "output_type": "display_data"
    }
   ],
   "source": [
    "import matplotlib\n",
    "import matplotlib.pylab as plt\n",
    "%matplotlib inline\n",
    "\n",
    "matplotlib.rcParams.update({'font.size': 14})\n",
    "\n",
    "axis = np.arange(0.5,buckets+1.5,1)\n",
    "index = axis+0.5\n",
    "X = x.copy()\n",
    "Y = alpha + X\n",
    "\n",
    "# to include the last data point as a step, we need to repeat it\n",
    "A = np.concatenate((alpha,[alpha[-1]]))\n",
    "X = np.concatenate((X,[X[-1]]))\n",
    "Y = np.concatenate((Y,[Y[-1]]))\n",
    "\n",
    "plt.xticks(index)\n",
    "plt.xlim(0.5,buckets+0.5)\n",
    "plt.ylim(0,1.5)\n",
    "plt.step(axis,A,where='post',label =r'$\\alpha$',lw=2)\n",
    "plt.step(axis,Y,where='post',label=r'$\\alpha + x$',lw=2)\n",
    "plt.legend(loc='lower right')\n",
    "plt.xlabel('Bucket Number')\n",
    "plt.ylabel('Power Level')\n",
    "plt.title('Water Filling Solution')\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}
