Mutual Information

Compute the mutual information I(X; Y) from a joint probability table.

Mutual information is how many bits knowing one variable tells you about the other. It is zero exactly when X and Y are independent. Unlike correlation, it also catches dependence that is not linear, which is why it is used for feature selection.

Given joint[x, y] = P(X = x, Y = y):

I(X; Y) = sum_{x, y} P(x, y) * log2(P(x, y) / (P(x) * P(y)))

Skip terms where P(x, y) = 0.

Function signature

def mutual_information(joint):

Example

Input

joint = np.array([[0.4, 0.1],
                  [0.1, 0.4]])
mutual_information(joint)

Output

0.27807191

Reasoning

Both marginals are [0.5, 0.5], so P(x) * P(y) = 0.25 in every cell. The two diagonal cells give 0.4 * log2(0.4 / 0.25) = 0.2712 each, and the two off-diagonal cells give 0.1 * log2(0.1 / 0.25) = -0.1322 each. The sum is 0.2781 bits. Single cells can be negative, but the total never is.

Constraints and notes

• joint has shape (n_x, n_y), non-negative, and sums to 1.
• Use log base 2.
• I(X; Y) = I(Y; X), so transposing joint must not change the result.