Chapter 176 — Matrix Calculus Introduction

# --- Visualization: Gradient of ||Ax - b||^2 on a 2D grid ---
import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-v0_8-whitegrid')

# Simple 2x2 case: f(x) = ||Ax - b||^2
A = np.array([[2.0, 1.0], [1.0, 3.0]])
b = np.array([4.0, 5.0])
x_star = np.linalg.solve(A.T @ A, A.T @ b)  # optimal x

# Evaluate f on a grid
x1 = np.linspace(-1, 4, 80)
x2 = np.linspace(-1, 4, 80)
X1, X2 = np.meshgrid(x1, x2)
F = np.zeros_like(X1)
for i in range(X1.shape[0]):
    for j in range(X1.shape[1]):
        x = np.array([X1[i,j], X2[i,j]])
        r = A @ x - b
        F[i,j] = r @ r

fig, axes = plt.subplots(1, 2, figsize=(13, 5))

# Contour plot
ax = axes[0]
cs = ax.contourf(X1, X2, F, levels=30, cmap='Blues')
plt.colorbar(cs, ax=ax)
ax.plot(*x_star, 'r*', markersize=15, label='Minimum x*')
ax.set_xlabel('x₁'); ax.set_ylabel('x₂')
ax.set_title('f(x) = ||Ax - b||² — loss landscape')
ax.legend()

# Gradient field
ax2 = axes[1]
x1g = np.linspace(-1, 4, 12)
x2g = np.linspace(-1, 4, 12)
X1g, X2g = np.meshgrid(x1g, x2g)
# Gradient = 2 A^T (Ax - b)
ATA2 = 2 * A.T @ A
ATb2 = 2 * A.T @ b
G1 = ATA2[0,0]*X1g + ATA2[0,1]*X2g - ATb2[0]
G2 = ATA2[1,0]*X1g + ATA2[1,1]*X2g - ATb2[1]
ax2.quiver(X1g, X2g, -G1, -G2, alpha=0.5)  # negative gradient = descent direction
cs2 = ax2.contour(X1, X2, F, levels=15, cmap='Blues', alpha=0.5)
ax2.plot(*x_star, 'r*', markersize=15)
ax2.set_xlabel('x₁'); ax2.set_ylabel('x₂')
ax2.set_title('Negative gradient field (descent directions)')

plt.tight_layout()
plt.show()

# Worked numeric example: verify gradient analytically vs numerically
import numpy as np

def numerical_gradient(f, x, eps=1e-5):
    """
    Compute gradient of scalar f at x using central finite differences.

    Args:
        f: function R^n -> R
        x: (n,) array

    Returns:
        grad: (n,) numerical gradient
    """
    grad = np.zeros_like(x, dtype=float)
    for i in range(len(x)):
        x_plus = x.copy(); x_plus[i] += eps
        x_minus = x.copy(); x_minus[i] -= eps
        grad[i] = (f(x_plus) - f(x_minus)) / (2 * eps)
    return grad


A = np.array([[2.0, 1.0], [1.0, 3.0], [0.5, 2.0]])
b = np.array([4.0, 5.0, 2.0])
x0 = np.array([1.0, 2.0])

# f(x) = ||Ax - b||^2
f = lambda x: np.sum((A @ x - b)**2)

# Analytic gradient: 2 A^T (Ax - b)
analytic_grad = 2 * A.T @ (A @ x0 - b)
numerical_grad = numerical_gradient(f, x0)

print("Analytic gradient: ", analytic_grad)
print("Numerical gradient:", numerical_grad)
print(f"Max error: {np.max(np.abs(analytic_grad - numerical_grad)):.2e}")

# --- Implementation: Gradient checker (validate any analytic gradient) ---
import numpy as np

def gradient_check(f, grad_f, x, eps=1e-5, atol=1e-6):
    """
    Check analytic gradient against numerical gradient.

    Args:
        f: scalar function R^n -> R
        grad_f: gradient function R^n -> R^n
        x: evaluation point (flattened)
        eps: finite difference step
        atol: absolute tolerance

    Returns:
        passed: bool
        max_error: float
    """
    x = x.flatten()
    analytic = grad_f(x).flatten()
    numeric = numerical_gradient(f, x, eps)
    max_error = np.max(np.abs(analytic - numeric))
    return max_error < atol, max_error


# Test 1: f(x) = x^T A x (quadratic form)
A = np.array([[3.0, 1.0], [1.0, 2.0]])
f1 = lambda x: x @ A @ x
grad_f1 = lambda x: (A + A.T) @ x  # 2Ax for symmetric A

x_test = np.array([2.0, -1.0])
passed, err = gradient_check(f1, grad_f1, x_test)
print(f"Quadratic form gradient check: {'PASS' if passed else 'FAIL'} (max error {err:.2e})")

# Test 2: f(x) = ||Wx - b||^2, gradient w.r.t. W
W = np.random.randn(3, 4)
x_vec = np.random.randn(4)
b_vec = np.random.randn(3)

f2 = lambda w_flat: np.sum((w_flat.reshape(3,4) @ x_vec - b_vec)**2)
grad_f2 = lambda w_flat: (2 * np.outer(W.reshape(3,4) @ x_vec - b_vec, x_vec)).flatten()

passed2, err2 = gradient_check(f2, grad_f2, W)
print(f"Matrix gradient check:        {'PASS' if passed2 else 'FAIL'} (max error {err2:.2e})")

# --- Experiment 1: Gradient descent on ||Ax - b||^2 ---
# Hypothesis: Following the negative gradient converges to the least-squares solution
# Try changing: LEARNING_RATE

import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-v0_8-whitegrid')

LEARNING_RATE = 0.05  # try: 0.001, 0.01, 0.05, 0.2 (instability!)
N_STEPS = 200

np.random.seed(0)
A = np.random.randn(10, 3)
b = np.random.randn(10)
x = np.zeros(3)

x_star = np.linalg.lstsq(A, b, rcond=None)[0]

losses = []
for _ in range(N_STEPS):
    residual = A @ x - b
    loss = residual @ residual
    grad = 2 * A.T @ residual
    x = x - LEARNING_RATE * grad
    losses.append(loss)

plt.figure(figsize=(8, 4))
plt.semilogy(losses)
plt.axhline(np.sum((A @ x_star - b)**2), color='red', ls='--', label='Optimal loss')
plt.xlabel('Iteration')
plt.ylabel('Loss (log scale)')
plt.title(f'Gradient descent on ||Ax-b||² (lr={LEARNING_RATE})')
plt.legend()
plt.tight_layout()
plt.show()
print(f"Final loss: {losses[-1]:.6f}")
print(f"Optimal loss: {np.sum((A @ x_star - b)**2):.6f}")

# --- Mini Project: Matrix calculus in action — linear regression by gradient descent ---
# Problem: Fit a linear model y = Xw to data using gradient descent,
#          and compare to the closed-form solution from ch182.

import numpy as np
import matplotlib.pyplot as plt
plt.style.use('seaborn-v0_8-whitegrid')

np.random.seed(42)

# Generate synthetic regression data
n, p = 100, 5
X = np.random.randn(n, p)
w_true = np.array([2.0, -1.5, 0.8, 3.0, -0.5])
y = X @ w_true + 0.5 * np.random.randn(n)

# Closed-form solution: w* = (X^T X)^{-1} X^T y
w_closed = np.linalg.solve(X.T @ X, X.T @ y)

# Gradient descent: grad_w = 2/n X^T (Xw - y)
LR = 0.1
N_ITER = 500
w_gd = np.zeros(p)
losses_gd = []

for _ in range(N_ITER):
    residual = X @ w_gd - y
    loss = np.mean(residual**2)
    grad = (2/n) * X.T @ residual
    w_gd = w_gd - LR * grad
    losses_gd.append(loss)

fig, axes = plt.subplots(1, 2, figsize=(13, 4))

axes[0].semilogy(losses_gd, color='steelblue')
axes[0].axhline(np.mean((X @ w_closed - y)**2), color='red', ls='--', label='Closed-form MSE')
axes[0].set_xlabel('Iteration')
axes[0].set_ylabel('MSE (log)')
axes[0].set_title('Gradient descent on MSE loss')
axes[0].legend()

axes[1].bar(range(p), w_true, alpha=0.7, label='True w', color='steelblue')
axes[1].bar(range(p), w_closed, alpha=0.5, label='Closed-form', color='green')
axes[1].bar(range(p), w_gd, alpha=0.4, label='GD', color='tomato')
axes[1].set_xlabel('Feature')
axes[1].set_ylabel('Weight value')
axes[1].set_title('Weight comparison')
axes[1].legend()

plt.tight_layout()
plt.show()
print(f"||w_closed - w_true||: {np.linalg.norm(w_closed - w_true):.4f}")
print(f"||w_gd - w_true||:     {np.linalg.norm(w_gd - w_true):.4f}")