Blahut-Arimoto Iteration

Problem

Perform one Blahut-Arimoto iteration for a discrete memoryless channel and return the updated input distribution. With q(y) = sum_x px[x] * P(y|x), set px_new[x] proportional to exp(sum_y P(y|x) * ln(P(y|x) / q(y))), skipping terms where P(y|x) = 0. Subtract the largest exponent before exponentiating for numerical stability.

Hint

For a noiseless channel the update pushes mass toward previously unlikely inputs.

Test cases

case 1

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

expected: [0.5,0.5]

case 2

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

expected: [0.22915688898295755,0.7708431110170424]

case 3

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

expected: [0.5,0.5]

case 4

in: [[0.3,0.7],[[1,0],[0,1]]]

expected: [0.7,0.3]

Starter codePython

import math
def ba_iteration(px, channel):
    # Your code here
    pass

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