Linear Algebra: Least Squares and Regression

import numpy as np
import matplotlib.pyplot as plt

%matplotlib inline

def least_squares_normal(A, b):
    """Solve least squares using normal equations."""
    ATA = A.T @ A
    ATb = A.T @ b
    return np.linalg.solve(ATA, ATb)

# Generate noisy data
np.random.seed(42)
x_data = np.linspace(0, 10, 20)
y_data = 2.5 * x_data + 1.5 + np.random.randn(20) * 2

# Set up the design matrix A = [x, 1]
A = np.column_stack([x_data, np.ones_like(x_data)])
b = y_data

print(f"Design matrix A shape: {A.shape}")
print(f"First 5 rows of A:\n{A[:5]}")

# Solve using normal equations
coeffs = least_squares_normal(A, b)
m, c = coeffs

print(f"Fitted line: y = {m:.4f}x + {c:.4f}")
print(f"True parameters: y = 2.5x + 1.5")

# Solve using NumPy's lstsq
coeffs_numpy, residuals, rank, s = np.linalg.lstsq(A, b, rcond=None)
print(f"NumPy lstsq: y = {coeffs_numpy[0]:.4f}x + {coeffs_numpy[1]:.4f}")

# Visualize the fit
plt.figure(figsize=(10, 6))
plt.scatter(x_data, y_data, label='Data', s=50)
x_line = np.linspace(0, 10, 100)
y_line = m * x_line + c
plt.plot(x_line, y_line, 'r-', linewidth=2, label=f'Fit: y = {m:.2f}x + {c:.2f}')
plt.xlabel('x')
plt.ylabel('y')
plt.legend()
plt.title('Simple Linear Regression via Least Squares')
plt.grid(True, alpha=0.3)
plt.show()

# Generate nonlinear data
x_data = np.linspace(-3, 3, 30)
y_data = 0.5*x_data**3 - 2*x_data + np.random.randn(30) * 2

def fit_polynomial(x, y, degree):
    """Fit a polynomial of given degree using least squares."""
    # Design matrix: [1, x, x^2, ..., x^n]
    A = np.column_stack([x**i for i in range(degree + 1)])
    coeffs, _, _, _ = np.linalg.lstsq(A, y, rcond=None)
    return coeffs

# Fit polynomials of different degrees
fig, axes = plt.subplots(1, 3, figsize=(15, 4))
degrees = [1, 3, 10]

for ax, deg in zip(axes, degrees):
    coeffs = fit_polynomial(x_data, y_data, deg)
    x_plot = np.linspace(-3, 3, 100)
    y_plot = sum(c * x_plot**i for i, c in enumerate(coeffs))

    ax.scatter(x_data, y_data, s=30, alpha=0.7)
    ax.plot(x_plot, y_plot, 'r-', linewidth=2)
    ax.set_title(f'Degree {deg} polynomial')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.grid(True, alpha=0.3)

plt.tight_layout()
plt.show()

# Generate data with two features
np.random.seed(42)
n_samples = 100

x1 = np.random.randn(n_samples)
x2 = np.random.randn(n_samples)
y = 3 + 2*x1 - 1.5*x2 + np.random.randn(n_samples) * 0.5

# Design matrix
A = np.column_stack([np.ones(n_samples), x1, x2])

# Fit model
beta, _, _, _ = np.linalg.lstsq(A, y, rcond=None)

print(f"True coefficients: [3, 2, -1.5]")
print(f"Fitted coefficients: {beta}")

# Compute residuals
x_ls = least_squares_normal(A, y)
y_pred = A @ x_ls
residuals = y - y_pred

# Verify orthogonality: A^T @ r should be ~0
ATr = A.T @ residuals
print(f"A^T @ residuals (should be ~0):\n{ATr}")

# Compute R-squared
ss_res = np.sum(residuals**2)
ss_tot = np.sum((y - np.mean(y))**2)
r_squared = 1 - ss_res / ss_tot
print(f"\nR-squared: {r_squared:.6f}")

# Plot residuals
fig, axes = plt.subplots(1, 2, figsize=(12, 4))

# Residual plot
axes[0].scatter(y_pred, residuals, alpha=0.6)
axes[0].axhline(y=0, color='r', linestyle='--')
axes[0].set_xlabel('Predicted values')
axes[0].set_ylabel('Residuals')
axes[0].set_title('Residual Plot')
axes[0].grid(True, alpha=0.3)

# QQ plot (residuals should be normally distributed)
from scipy import stats
stats.probplot(residuals, dist="norm", plot=axes[1])
axes[1].set_title('Q-Q Plot')

plt.tight_layout()
plt.show()

def weighted_least_squares(A, b, weights):
    """Solve weighted least squares."""
    W = np.diag(weights)
    AW = A.T @ W
    return np.linalg.solve(AW @ A, AW @ b)

# Example with heteroscedastic data
x = np.linspace(0, 10, 50)
y_true = 2 * x + 1
noise = np.random.randn(50) * (0.5 + 0.3 * x)  # Noise increases with x
y = y_true + noise

# Equal weights (standard least squares)
A = np.column_stack([x, np.ones_like(x)])
coeffs_unweighted = least_squares_normal(A, y)

# Weights inversely proportional to variance
weights = 1 / (0.5 + 0.3 * x)**2
coeffs_weighted = weighted_least_squares(A, y, weights)

print(f"True: y = 2x + 1")
print(f"Unweighted: y = {coeffs_unweighted[0]:.4f}x + {coeffs_unweighted[1]:.4f}")
print(f"Weighted: y = {coeffs_weighted[0]:.4f}x + {coeffs_weighted[1]:.4f}")

def ridge_regression(A, b, alpha):
    """Solve ridge regression."""
    n_features = A.shape[1]
    ATA = A.T @ A + alpha * np.eye(n_features)
    ATb = A.T @ b
    return np.linalg.solve(ATA, ATb)

# Compare standard vs ridge regression
# Create ill-conditioned problem
np.random.seed(42)
X = np.random.randn(100, 10)
X[:, 9] = X[:, 0] + 0.001 * np.random.randn(100)  # Near-duplicate column
y = X @ np.ones(10) + np.random.randn(100) * 0.1

A = np.column_stack([np.ones(100), X])

# Standard least squares
try:
    coeffs_std = least_squares_normal(A, y)
    print(f"Standard LS coefficients: {coeffs_std[:5]}...")
except np.linalg.LinAlgError:
    print("Standard LS failed")

# Ridge regression
coeffs_ridge = ridge_regression(A, y, alpha=0.1)
print(f"Ridge coefficients: {coeffs_ridge[:5]}...")