Mutual Information

import numpy as np
from collections import Counter
from sklearn.feature_selection import mutual_info_regression, mutual_info_classif

def mutual_information(x, y, base=2):
    """Compute MI between two discrete arrays."""
    n = len(x)
    x, y = np.array(x), np.array(y)

    # Joint distribution
    xy = np.stack([x, y], axis=1)
    xy_counter = Counter(map(tuple, xy))
    pxy = np.array([count / n for count in xy_counter.values()])

    # Marginals
    x_counter = Counter(x)
    y_counter = Counter(y)
    px = np.array([x_counter[k] / n for k in xy_counter.keys().__iter__()])
    # Rebuild properly
    px = np.array([sum(pxy[i] for i in range(len(pxy))
                       if list(xy_counter.keys())[i][0] == k)
                   for k in x_counter.keys()])

    # Simpler approach: compute from joint
    joint = np.zeros((len(set(x)), len(set(y))))
    x_vals, y_vals = sorted(set(x)), sorted(set(y))
    x_map = {v: i for i, v in enumerate(x_vals)}
    y_map = {v: i for i, v in enumerate(y_vals)}
    for xi, yi in zip(x, y):
        joint[x_map[xi], y_map[yi]] += 1
    joint /= n

    # MI from joint
    mi = 0.0
    for i in range(joint.shape[0]):
        for j in range(joint.shape[1]):
            if joint[i, j] > 0:
                px = joint[i].sum()
                py = joint[:, j].sum()
                mi += joint[i, j] * np.log2(joint[i, j] / (px * py))
    return mi

def conditional_mi(x, y, z):
    """Estimate I(X; Y | Z) by averaging I(X; Y) within each Z value."""
    z_vals = np.unique(z)
    total = 0.0
    for z_val in z_vals:
        mask = z == z_val
        if mask.sum() > 1:
            total += mutual_information(x[mask], y[mask]) * mask.mean()
    return total

# --- Feature Selection Example ---
np.random.seed(42)
X = np.random.randn(1000, 5)
y = X[:, 0] * 2 + X[:, 1] * 0.5 + np.random.randn(1000) * 0.1

mi_scores = mutual_info_regression(X, y, random_state=42)
print("MI scores per feature:")
for i, score in enumerate(mi_scores):
    print(f"  Feature {i}: {score:.4f}")

top_features = np.argsort(mi_scores)[::-1][:3]
print(f"Top 3 features: {top_features}")
# Feature 0 and 1 should have highest MI