{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Capacity of a Communication Channel\n",
    "by Robert Gowers, Roger Hill, Sami Al-Izzi, Timothy Pollington and Keith Briggs\n",
    "\n",
    "from Boyd and Vandenberghe, Convex Optimization, exercise 4.57 pages 207-8\n",
    "\n",
    "Convex optimization can be used to find the channel capacity $C$ of a discrete memoryless channel. Consider a communication channel with input $X(t) \\in \\{1,2,...,n\\}$ and output $Y(t) \\in \\{1,2,...m\\}$.   This means that the random variables $X$ and $Y$ can take $n$ and $m$ different values, respectively.  \n",
    "\n",
    "In a discrete memoryless channel, the relation between the input and the output is given by the transition probability:\n",
    "\n",
    "$p_{ij} = \\mathbb{P}(Y(t)=i | X(t)=j)$\n",
    "\n",
    "These transition probabilities form the channel transition matrix $P$, with $P \\in \\mathbb{R}^{m\\times n}$.\n",
    "\n",
    "Assume that $X$ has a probability distribution denoted by $x \\in \\mathbb{R}^n$, meaning that: \n",
    "\n",
    "$x_j = \\mathbb{P}(X(t) = j) \\quad j \\in \\{1,...,n\\}$.\n",
    "\n",
    "From Shannon, the channel capacity is given by the maximum possible mutual information $I$ between $X$ and $Y$:\n",
    "\n",
    "$C = \\sup_x I(X;Y)$\n",
    "\n",
    "where,\n",
    "\n",
    "$I(X;Y) = -\\sum_{i=1}^{m} y_i \\log_2y_i + \\sum_{j=1}^{n}\\sum_{i=1}^{m}x_j p_{ij}\\log_2p_{ij}$\n",
    "\n",
    "Given that $x\\log x$ is convex for $x \\geq 0$, we can formulate this as a convex optimization problem:\n",
    "\n",
    "minimise $-I(X;Y)$\n",
    "\n",
    "subject to $\\sum_{i=1}^{n}x_i = 1 \\quad x \\succeq 0 \\quad$ since $x$ describes a probability \n",
    "\n",
    "Due to the entropy function in CVXPY, this can be written quite easily in DCP."
   ]
  },
  {
   "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",
    "\n",
    "import cvxpy as cp\n",
    "import numpy as np\n",
    "import math\n",
    "from scipy.special import xlogy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "def channel_capacity(n, m, P, sum_x=1):\n",
    "    '''\n",
    "    Boyd and Vandenberghe, Convex Optimization, exercise 4.57 page 207\n",
    "    Capacity of a communication channel.\n",
    "    \n",
    "    We consider a communication channel, with input X(t)∈{1,..,n} and\n",
    "    output Y(t)∈{1,...,m}, for t=1,2,... .The relation between the\n",
    "    input and output is given statistically:\n",
    "    p_(i,j) = ℙ(Y(t)=i|X(t)=j), i=1,..,m  j=1,...,n\n",
    "    \n",
    "    The matrix P ∈ ℝ^(m*n) is called the channel transition matrix, and\n",
    "    the channel is called a discrete memoryless channel. Assuming X has a\n",
    "    probability distribution denoted x ∈ ℝ^n, i.e.,\n",
    "    x_j = ℙ(X=j), j=1,...,n\n",
    "    \n",
    "    The mutual information between X and Y is given by\n",
    "    ∑(∑(x_j p_(i,j)log_2(p_(i,j)/∑(x_k p_(i,k)))))\n",
    "    Then channel capacity C is given by\n",
    "    C = sup I(X;Y).\n",
    "    With a variable change of y = Px this becomes\n",
    "    I(X;Y)=  c^T x - ∑(y_i log_2 y_i)\n",
    "    where c_j = ∑(p_(i,j)log_2(p_(i,j)))\n",
    "    '''\n",
    "    \n",
    "    # n is the number of different input values\n",
    "    # m is the number of different output values\n",
    "    if n*m == 0:\n",
    "        print('The range of both input and output values must be greater than zero')\n",
    "        return 'failed', np.nan, np.nan\n",
    "\n",
    "    # x is probability distribution of the input signal X(t)\n",
    "    x = cp.Variable(shape=n)\n",
    "    \n",
    "    # y is the probability distribution of the output signal Y(t)\n",
    "    # P is the channel transition matrix\n",
    "    y = P@x\n",
    "    \n",
    "    # I is the mutual information between x and y\n",
    "    c = np.sum(np.array((xlogy(P, P) / math.log(2))), axis=0)\n",
    "    I = c@x + cp.sum(cp.entr(y) / math.log(2))\n",
    "\n",
    "    # Channel capacity maximised by maximising the mutual information\n",
    "    obj = cp.Maximize(I)\n",
    "    constraints = [cp.sum(x) == sum_x,x >= 0]\n",
    "    \n",
    "    # Form and solve problem\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\n",
    "    "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Example\n",
    "\n",
    "In this example we consider a communication channel with two possible inputs and outputs, so $n = m = 2$.  The channel transition matrix we use in this case is:\n",
    "\n",
    "$P = \\pmatrix{0.75,0.25\\\\0.25,0.75}$\n",
    "\n",
    "Note that the columns of $P$ must sum to 1 and all elements of $P$ must be positive."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Problem status:  optimal\n",
      "Optimal value of C = 0.1887\n",
      "Optimal variable x = \n",
      " [0.5 0.5]\n"
     ]
    }
   ],
   "source": [
    "np.set_printoptions(precision=3)\n",
    "n = 2\n",
    "m = 2\n",
    "P = np.array([[0.75,0.25],\n",
    "             [0.25,0.75]])\n",
    "stat, C, x = channel_capacity(n, m, P)\n",
    "print('Problem status: ',stat)\n",
    "print('Optimal value of C = {:.4g}'.format(C))\n",
    "print('Optimal variable x = \\n', x)"
   ]
  }
 ],
 "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
}