utils/information.py

import numpy as np

def mutual_information(x, y, bins=30):
    c_xy = np.histogram2d(x, y, bins)[0]
    p_xy = c_xy / np.sum(c_xy)
    p_x = np.sum(p_xy, axis=1)
    p_y = np.sum(p_xy, axis=0)

    mi = 0.0
    for i in range(len(p_x)):
        for j in range(len(p_y)):
            if p_xy[i, j] > 0:
                mi += p_xy[i, j] * np.log(
                    p_xy[i, j] / (p_x[i] * p_y[j] + 1e-12)
                )
    return mi