Mutual Information from Marginals

Problem

Given a marginal px and a conditional distribution cond[i][j] = P(Y = j | X = i), compute I(X; Y) in bits:

I = sum_i px[i] * sum_j P(j|i) * log2(P(j|i) / p(y_j))

where p(y_j) = sum_i px[i] * P(j|i). Skip zero terms.

Hint

First compute the output marginal p(y), then average the log-ratio over the input marginal.

case 1

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

expected: 1

case 2

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

expected: 0

case 3

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

expected: 0.204434002924965

case 4

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

expected: 0.5310044064107189

Starter codePython

import math
def mi_from_marginals(px, cond):
    # Your code here
    pass

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