Python
Mutual Information
Data Science
Machine Learning
Information Theory

Python's implementation of Mutual Information

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

Introduction

Mutual information measures how much knowing one variable reduces uncertainty about another. In Python, the right implementation depends on the type of data you have: discrete labels, continuous features, or a feature-selection workflow inside machine learning.

Mutual Information for Discrete Labels

If both variables are discrete categories, scikit-learn provides a direct implementation through mutual_info_score.

python
1from sklearn.metrics import mutual_info_score
2
3x = ["red", "red", "blue", "blue", "green"]
4y = ["A", "A", "B", "B", "A"]
5
6score = mutual_info_score(x, y)
7print(score)

This function compares two discrete label arrays and returns their mutual information. It is symmetrical, so swapping x and y gives the same result.

This is a good fit for clustering evaluation, label agreement, or small discrete experiments where you want a direct information-theoretic dependency measure.

Feature Selection With Continuous Inputs

For machine learning, you often want mutual information between features and a target. scikit-learn exposes specialized estimators for that use case.

Use mutual_info_classif for classification targets:

python
1from sklearn.feature_selection import mutual_info_classif
2from sklearn.datasets import load_iris
3
4X, y = load_iris(return_X_y=True)
5scores = mutual_info_classif(X, y, random_state=42)
6
7print(scores)

Use mutual_info_regression for continuous targets:

python
1from sklearn.feature_selection import mutual_info_regression
2import numpy as np
3
4X = np.array([[0.0], [1.0], [2.0], [3.0], [4.0]])
5y = np.array([0.1, 1.1, 1.9, 3.2, 4.1])
6
7scores = mutual_info_regression(X, y, random_state=42)
8print(scores)

These estimators use neighbor-based methods rather than a simple table of counts, which makes them appropriate for continuous variables.

A Manual Discrete Implementation

If you want to understand what the library is doing, you can compute mutual information manually for discrete data using counts and probabilities.

python
1from collections import Counter
2from math import log
3
4
5def mutual_information(x, y):
6    total = len(x)
7    joint = Counter(zip(x, y))
8    x_counts = Counter(x)
9    y_counts = Counter(y)
10
11    mi = 0.0
12    for (xv, yv), joint_count in joint.items():
13        p_xy = joint_count / total
14        p_x = x_counts[xv] / total
15        p_y = y_counts[yv] / total
16        mi += p_xy * log(p_xy / (p_x * p_y))
17
18    return mi
19
20
21print(mutual_information(["red", "red", "blue"], ["A", "A", "B"]))

This version is useful for learning, testing, and situations where you want exact control over the calculation on categorical data.

Interpreting the Result

A higher mutual information score means stronger dependency, but the value is not normalized by default. That means scores are useful for ranking features or comparing variables within the same problem, but not always for comparing unrelated datasets directly.

Mutual information also detects nonlinear dependence. That is one reason it is attractive for feature selection. Correlation can miss nonlinear structure that mutual information still captures.

In practice, many teams use mutual information as a screening metric rather than a final decision rule. You might rank features with mutual information first, then validate the best candidates with cross-validation, model-specific importance measures, or domain checks.

Common Pitfalls

One common mistake is feeding raw continuous values into mutual_info_score. That function is intended for discrete labels, not arbitrary floating-point measurements.

Another pitfall is over-interpreting the absolute value. Mutual information is often more useful for relative comparison than for saying that one score is universally "good" in isolation.

It is also easy to forget that neighbor-based estimators involve randomness and finite-sample behavior. For reproducible experiments, set random_state when the API supports it.

Summary

  • Use mutual_info_score for discrete label arrays.
  • Use mutual_info_classif or mutual_info_regression for feature selection with continuous inputs.
  • A manual count-based implementation is helpful for understanding discrete mutual information.
  • Mutual information captures nonlinear dependency, not just linear correlation.
  • Choose the estimator that matches your data type, or the score can be misleading.

Related reading
Free course
Beginner
7 lessons
2 hours
Tackling System Design Interview Problems

A short course that equips you with the skills to approach system design interviews methodically.

Start the free course
Track what you have practised

A free account saves your progress, solutions and study plan across every problem on Codemia.

ML System Design practice on Codemia

Design recommenders, ranking systems and training pipelines the way ML interviews actually ask for them, with worked solutions.

Practice ML system design

All Rights Reserved.