"""
 problem 9.17 from MacKay's book, pg 155 (167 of 640 in .pdf)
"""

from numpy import *

# X = input:  0-4
# Y = output: 0-9
# P(Y|X) = prob y given x = PyGx is a 10x5 matrix.
#
q = 0.25
             # 0  1  2  3  4          = x
PyGx = array([[q, 0, 0, 0, q],   # 0  = y
              [q, 0, 0, 0, q],   # 1
              [q, q, 0, 0, 0],   # 2
              [q, q, 0, 0, 0],   # 3
              [0, q, q, 0, 0],   # 4
              [0, q, q, 0, 0],   # 5
              [0, 0, q, q, 0],   # 6
              [0, 0, q, q, 0],   # 7
              [0, 0, 0, q, q],   # 8
              [0, 0, 0, q, q]])  # 9

              
# Let's start with all P(x) = same; it looks pretty symmetric
Px = ones(5)/5.0     # which are 5 copies of 0.2
# Then probability of y is just found with matrix multiplication.
Py = dot(PyGx, Px)   # which are 10 copies of 0.1
# And P(Y,X) = prob y and x = PyAx = P(y|x)*P(x). But with P(x) constant,
PyAx = 0.2 * PyGx

# And so now we can find the entropies and then the mutual information

def entropy(pYandX, pYgivenX=None):
    """ Return -sum(pYandX*log2(pYgivenX)) for nonzero pYgivenX terms
        This can calculate any of
        * H(x) from p(x) as entropy(pX)
        * H(y,x) from p(y,x) joint distribution as entropy(pXandY)
        * H(y|x) from p(y,x) and p(y|x) as entropy(pYandX, pYgivenX)
    """
    if pYgivenX == None:
        pYgivenX = pYandX
    nz = pYgivenX.nonzero()
    return -sum( pYandX[nz] * log2( pYgivenX[nz] ))

Hy = entropy(Py)
HyGx = entropy(PyAx, PyGx)
Iyx = Hy - HyGx
print " -- channel capacity calculation --"
print "P(x) = "
print Px
print "P(y|x) = "
print PyGx
print "P(y) = "
print Py
print "P(x,y) ="
print PyAx
print "H(y) = ", Hy
print "H(y|x) = ", HyGx
print "I(y;x) = ", Iyx