Channel Capacity via Blahut-Arimoto

Problem

Compute the capacity in bits of a discrete memoryless channel given by a row-stochastic matrix channel[x][y] = P(y | x). Run 500 iterations of the Blahut-Arimoto algorithm from a uniform input distribution:

px_new[x] is proportional to exp(sum_y P(y|x) * ln(P(y|x) / q(y)))

where q(y) = sum_x px[x] * P(y|x). Then return I(X; Y) in bits at the final distribution.

Hint

Subtract the max before exponentiating so the update is numerically stable.

Test cases

case 1

in: [[[1,0],[0,1]]]

expected: 1

case 2

in: [[[0.9,0.1],[0.1,0.9]]]

expected: 0.5310044064107189

case 3

in: [[[0.75,0.25,0],[0,0.25,0.75]]]

expected: 0.75

case 4

in: [[[0.5,0.5],[0.5,0.5]]]

expected: 0

Starter codePython

import math
def channel_capacity(channel):
    # Your code here
    pass

4 test cases · run your code against them in the editor