Conditional Mutual Information

Problem

Compute the conditional mutual information I(X; Y | Z) in bits from a 3D joint distribution joint[x][y][z]:

I = sum p(x, y, z) * log2(p(x, y, z) * p(z) / (p(x, z) * p(y, z)))

Skip zero-probability terms.

Hint

Marginalize to p(x, z), p(y, z), and p(z) before forming the log-ratio.

case 1

in: [[[[0.125,0.125],[0.125,0.125]],[[0.125,0.125],[0.125,0.125]]]]

expected: 0

case 2

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

expected: 0

case 3

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

expected: 1

case 4

in: [[[[0.2,0.1],[0.05,0.05]],[[0.05,0.05],[0.1,0.2]]]]

expected: 0.12709440467994398

Starter codePython

import math
def conditional_mutual_information(joint):
    # Your code here
    pass

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