Covariance and Mahalanobis Distance Exercise¶
This notebook turns the theory note 02-mahalanobis-distance.md into a hands-on exercise.
You will implement the important pieces yourself:
- feature-wise mean
- mean centering
- covariance matrix
- Mahalanobis distance
Conventions used here:
Xhas shape(n_samples, n_features)- this is the transpose of the
d x nconvention used in the theory note - all notebook text is in English to match the project language rules for artifacts
Linked theory: research/02-mahalanobis-distance.md · Chạy: uv run jupyter lab từ thư mục latent-anything-theory/ (cần các package trong pyproject.toml).
import numpy as np
import matplotlib.pyplot as plt
plt.style.use("seaborn-v0_8-whitegrid")
np.set_printoptions(precision=3, suppress=True)
1. Build a correlated dataset¶
We will use a 2D dataset with clear correlation between the two features. That makes the difference between Euclidean distance and Mahalanobis distance easy to see.
rng = np.random.default_rng(7)
true_mean = np.array([1.0, -0.5])
true_cov = np.array([
[3.2, 2.4],
[2.4, 2.5],
])
X = rng.multivariate_normal(mean=true_mean, cov=true_cov, size=250)
X.shape
(250, 2)
fig, ax = plt.subplots(figsize=(6, 6))
ax.scatter(X[:, 0], X[:, 1], s=20, alpha=0.65)
ax.set_title("Synthetic correlated dataset")
ax.set_xlabel("feature_1")
ax.set_ylabel("feature_2")
ax.axis("equal")
plt.show()
What you see¶
What the plot shows: The ~250 points form a tilted, cigar-shaped cloud stretched along a diagonal — when
feature_1is large,feature_2tends to be large too (positive correlation). The spread is clearly larger along one diagonal direction (a major axis) than the perpendicular one (a minor axis).Why it looks this way: The data was sampled from a Gaussian with covariance
[[3.2, 2.4], [2.4, 2.5]]. The off-diagonal entry2.4couples the two features, which tilts the cloud; the unequal eigenvalues of that matrix are what make it longer in one direction than the other.Key takeaway: Because the spread is direction-dependent, plain Euclidean distance is misleading here — a point two units away along the thin (minor) axis is far more surprising than a point two units away along the long (major) axis. Correcting for exactly this is what the Mahalanobis distance, built next, does.
2. Exercise 1: mean vector and centered data¶
Implement the mean vector and the centered matrix without calling np.mean inside your function.
Use broadcasting for the centering step.
def compute_mean(X: np.ndarray) -> np.ndarray:
"""Return the feature-wise mean vector with shape (n_features,)."""
# # TODO:
# # 1. Count the number of samples.
# # 2. Sum over the sample axis.
# # 3. Divide by the number of samples.
# raise NotImplementedError("TODO: implement compute_mean")
return np.mean(X, axis=0)
def center_data(X: np.ndarray, mean: np.ndarray) -> np.ndarray:
"""Center every sample by subtracting the mean vector."""
# # TODO:
# # Use NumPy broadcasting to subtract the mean from every row.
# raise NotImplementedError("TODO: implement center_data")
return X - mean
expected_mean = np.mean(X, axis=0)
mean_vector = compute_mean(X)
assert mean_vector.shape == (X.shape[1],)
assert np.allclose(mean_vector, expected_mean)
centered_X = center_data(X, mean_vector)
assert centered_X.shape == X.shape
assert np.allclose(centered_X, X - expected_mean)
print("Exercise 1 looks good.")
print("Mean vector:", mean_vector)
Exercise 1 looks good. Mean vector: [ 1.211 -0.472]
3. Exercise 2: covariance matrix¶
Now compute the covariance matrix from the centered data.
Target formula for row-wise samples: $$\Sigma = \frac{1}{n-1} X_{centered}^T X_{centered}$$
Use n - 1 when unbiased=True, otherwise use n.
def compute_covariance(centered_X: np.ndarray, unbiased: bool = True) -> np.ndarray:
"""Return the covariance matrix with shape (n_features, n_features)."""
# # TODO:
# # 1. Get n_samples from centered_X.
# # 2. Pick the denominator: n - 1 if unbiased else n.
# # 3. Compute centered_X.T @ centered_X / denominator.
# raise NotImplementedError("TODO: implement compute_covariance")
n_samples = centered_X.shape[0]
denominator = n_samples - 1 if unbiased else n_samples
return centered_X.T @ centered_X / denominator
covariance = compute_covariance(centered_X, unbiased=True)
reference_covariance = np.cov(X, rowvar=False, ddof=1)
assert covariance.shape == (X.shape[1], X.shape[1])
assert np.allclose(covariance, covariance.T)
assert np.allclose(covariance, reference_covariance)
print("Exercise 2 looks good.")
print(covariance)
Exercise 2 looks good. [[2.884 2.149] [2.149 2.201]]
eigenvalues, eigenvectors = np.linalg.eigh(covariance)
order = np.argsort(eigenvalues)[::-1]
eigenvalues = eigenvalues[order]
eigenvectors = eigenvectors[:, order]
print("Eigenvalues:", eigenvalues)
print("Eigenvectors:\n", eigenvectors)
Eigenvalues: [4.719 0.366] Eigenvectors: [[-0.761 0.649] [-0.649 -0.761]]
order
array([1, 0])
4. Exercise 3: Mahalanobis distance¶
Implement Mahalanobis distance from a point x to the distribution center mean.
Formula: $$d_M(x, \mu) = \sqrt{(x-\mu)^T \Sigma^{-1} (x-\mu)}$$
Tip: use np.linalg.solve(covariance, delta) instead of forming the inverse directly.
def mahalanobis_distance(x: np.ndarray, mean: np.ndarray, covariance: np.ndarray) -> float:
"""Return the Mahalanobis distance from x to the distribution center."""
# # TODO:
# # 1. Compute delta = x - mean.
# # 2. Solve covariance @ z = delta.
# # 3. Return sqrt(delta @ z) as a Python float.
# raise NotImplementedError("TODO: implement mahalanobis_distance")
delta = x - mean
z = np.linalg.solve(covariance, delta)
return float(np.sqrt(delta @ z))
sample_point = X[0]
delta = sample_point - mean_vector
reference_distance = float(np.sqrt(delta @ np.linalg.inv(covariance) @ delta))
distance = mahalanobis_distance(sample_point, mean_vector, covariance)
assert np.isclose(distance, reference_distance)
print("Exercise 3 looks good.")
print("Reference distance:", reference_distance)
print("Your distance:", distance)
Exercise 3 looks good. Reference distance: 0.5196010671267661 Your distance: 0.5196010671267661
5. Compare Euclidean vs Mahalanobis distance¶
The first two points below are both exactly 2.5 sigma away, but along different principal directions.
Mahalanobis distance should treat them similarly, while Euclidean distance reacts to raw scale.
major_axis = eigenvectors[:, 0]
minor_axis = eigenvectors[:, 1]
major_sigma = np.sqrt(eigenvalues[0])
minor_sigma = np.sqrt(eigenvalues[1])
probe_points = {
"2.5 sigma on major axis": mean_vector + 2.5 * major_sigma * major_axis,
"2.5 sigma on minor axis": mean_vector + 2.5 * minor_sigma * minor_axis,
"mixed direction": mean_vector + 1.5 * major_sigma * major_axis + 2.0 * minor_sigma * minor_axis,
}
for name, point in probe_points.items():
euclidean = float(np.linalg.norm(point - mean_vector))
mahalanobis = mahalanobis_distance(point, mean_vector, covariance)
print(f"{name:24s} | Euclidean = {euclidean:6.3f} | Mahalanobis = {mahalanobis:6.3f}")
2.5 sigma on major axis | Euclidean = 5.431 | Mahalanobis = 2.500 2.5 sigma on minor axis | Euclidean = 1.513 | Mahalanobis = 2.500 mixed direction | Euclidean = 3.476 | Mahalanobis = 2.500
padding = 1.5
x_min, y_min = X.min(axis=0) - padding
x_max, y_max = X.max(axis=0) + padding
grid_x = np.linspace(x_min, x_max, 140)
grid_y = np.linspace(y_min, y_max, 140)
distance_grid = np.empty((grid_y.size, grid_x.size))
for row_index, y_value in enumerate(grid_y):
for col_index, x_value in enumerate(grid_x):
point = np.array([x_value, y_value])
distance_grid[row_index, col_index] = mahalanobis_distance(point, mean_vector, covariance)
fig, ax = plt.subplots(figsize=(7, 7))
ax.scatter(X[:, 0], X[:, 1], s=16, alpha=0.35, label="dataset")
contours = ax.contour(grid_x, grid_y, distance_grid, levels=[1.0, 2.0, 3.0], colors=["#f28e2b", "#e15759", "#4e79a7"])
ax.clabel(contours, inline=True, fontsize=9)
ax.scatter(mean_vector[0], mean_vector[1], color="black", s=80, marker="x", label="mean")
for name, point in probe_points.items():
ax.scatter(point[0], point[1], s=70, label=name)
ax.set_title("Mahalanobis distance contours")
ax.set_xlabel("feature_1")
ax.set_ylabel("feature_2")
ax.axis("equal")
ax.legend(loc="upper left", bbox_to_anchor=(1.02, 1.0))
plt.show()
What you see¶
What the plot shows: The printed table compares two probe points placed at 2.5σ along the major axis versus 2.5σ along the minor axis. Their Euclidean distances differ a lot (the major-axis point is physically much farther), yet their Mahalanobis distances are both ≈ 2.5. The contour plot draws the Mahalanobis level sets as ellipses aligned with the data cloud, and the probe points land near the same ellipse despite sitting at very different Euclidean radii.
Why it looks this way: Mahalanobis distance $d_M = \sqrt{(x-\mu)^T \Sigma^{-1} (x-\mu)}$ rescales every principal direction by its standard deviation. The inverse covariance $\Sigma^{-1}$ stretches the thin direction and shrinks the long one, so sets of equal distance become ellipses matching $\Sigma$ rather than circles.
Key takeaway: Mahalanobis distance measures "how many standard deviations away, in the data's own geometry," which is the right notion of closeness for correlated, anisotropic features. That is exactly the property the bonus-challenge classifier relies on when it assigns a point to the nearest class distribution.
Bonus challenge¶
Extend this notebook into a class-wise classifier:
- Generate two classes with different covariance matrices.
- Compute one
(mean, covariance)pair per class. - For a new point, assign the class with the smallest Mahalanobis distance.
That directly connects this exercise back to the discussion in the theory note.