Expand function documentation and add algorithm explanation

Author: tianyizheng02

machine_learning/local_weighted_learning/local_weighted_learning.py

@@ -1,11 +1,54 @@
+"""
+Locally weighted linear regression, also called local regression, is a type of
+non-parametric linear regression that prioritizes data closest to a given
+prediction point. The algorithm estimates the vector of model coefficients β
+using weighted least squares regression:
+
+β = (XᵀWX)⁻¹(XᵀWy),
+
+where X is the design matrix, y is the response vector, and W is the diagonal
+weight matrix.
+
+This implementation calculates wᵢ, the weight of the ith training sample, using
+the Gaussian weight:
+
+wᵢ = exp(-‖xᵢ - x‖²/(2τ²)),
+
+where xᵢ is the ith training sample, x is the prediction point, τ is the
+"bandwidth", and ‖x‖ is the Euclidean norm (also called the 2-norm or the L²
+norm). The bandwidth τ controls how quickly the weight of a training sample
+decreases as its distance from the prediction point increases. One can think of
+the Gaussian weight as a bell curve centered around the prediction point: a
+training sample is weighted lower if it's farther from the center, and τ
+controls the spread of the bell curve.
+
+Other types of locally weighted regression such as locally estimated scatterplot
+smoothing (LOESS) typically use different weight functions.
+
+References:
+    - https://en.wikipedia.org/wiki/Local_regression
+    - https://en.wikipedia.org/wiki/Weighted_least_squares
+    - https://cs229.stanford.edu/notes2022fall/main_notes.pdf
+"""
+
 import matplotlib.pyplot as plt
 import numpy as np
 
 
 def weight_matrix(point: np.ndarray, x_train: np.ndarray, tau: float) -> np.ndarray:
     """
-    Calculate the weight for every point in the data set.
-    point --> the x value at which we want to make predictions
+    Calculate the weight of every point in the training data around a given
+    prediction point
+
+    Args:
+        point: x-value at which the prediction is being made
+        x_train: ndarray of x-values for training
+        tau: bandwidth value, controls how quickly the weight of training values
+            decreases as the distance from the prediction point increases
+
+    Returns:
+        n x n weight matrix around the prediction point, where n is the size of
+        the training set
     >>> weight_matrix(
     ...     np.array([1., 1.]),
     ...     np.array([[16.99, 10.34], [21.01,23.68], [24.59,25.69]]),
@@ -15,22 +58,30 @@ def weight_matrix(point: np.ndarray, x_train: np.ndarray, tau: float) -> np.ndar
            [0.00000000e+000, 0.00000000e+000, 0.00000000e+000],
            [0.00000000e+000, 0.00000000e+000, 0.00000000e+000]])
     """
-    m, _ = np.shape(x_train)  # m is the number of training samples
-    weights = np.eye(m)  # Initializing weights as identity matrix
-
-    # calculating weights for all training examples [x(i)'s]
-    for j in range(m):
+    n = len(x_train)  # Number of training samples
+    weights = np.eye(n)  # Initialize weights as identity matrix
+    for j in range(n):
         diff = point - x_train[j]
         weights[j, j] = np.exp(diff @ diff.T / (-2.0 * tau**2))
+
     return weights
 
 
 def local_weight(
     point: np.ndarray, x_train: np.ndarray, y_train: np.ndarray, tau: float
 ) -> np.ndarray:
     """
-    Calculate the local weights using the weight_matrix function on training data.
-    Return the weighted matrix.
+    Calculate the local weights at a given prediction point using the weight
+    matrix for that point
+
+    Args:
+        point: x-value at which the prediction is being made
+        x_train: ndarray of x-values for training
+        y_train: ndarray of y-values for training
+        tau: bandwidth value, controls how quickly the weight of training values
+            decreases as the distance from the prediction point increases
+    Returns:
+        ndarray of local weights
     >>> local_weight(
     ...     np.array([1., 1.]),
     ...     np.array([[16.99, 10.34], [21.01,23.68], [24.59,25.69]]),
@@ -52,17 +103,24 @@ def local_weight_regression(
     x_train: np.ndarray, y_train: np.ndarray, tau: float
 ) -> np.ndarray:
     """
-    Calculate predictions for each data point on axis
+    Calculate predictions for each point in the training data
+
+    Args:
+        x_train: ndarray of x-values for training
+        y_train: ndarray of y-values for training
+        tau: bandwidth value, controls how quickly the weight of training values
+            decreases as the distance from the prediction point increases
+
+    Returns:
+        ndarray of predictions
     >>> local_weight_regression(
     ...     np.array([[16.99, 10.34], [21.01, 23.68], [24.59, 25.69]]),
     ...     np.array([[1.01, 1.66, 3.5]]),
     ...     0.6
     ... )
     array([1.07173261, 1.65970737, 3.50160179])
     """
-    m, _ = np.shape(x_train)
-    y_pred = np.zeros(m)
-
+    y_pred = np.zeros(len(x_train))  # Initialize array of predictions
     for i, item in enumerate(x_train):
         y_pred[i] = item @ local_weight(item, x_train, y_train, tau)
 
@@ -74,14 +132,15 @@ def load_data(
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
     """
     Load data from seaborn and split it into x and y points
+    >>> pass    # No doctests, function is for demo purposes only
     """
     import seaborn as sns
 
     data = sns.load_dataset(dataset_name)
-    x_data = np.array(data[x_name])  # total_bill
-    y_data = np.array(data[y_name])  # tip
+    x_data = np.array(data[x_name])
+    y_data = np.array(data[y_name])
 
-    one = np.ones(np.shape(y_data)[0], dtype=int)
+    one = np.ones(len(y_data))
 
     # pairing elements of one and x_data
     x_train = np.column_stack((one, x_data))
@@ -99,6 +158,7 @@ def plot_preds(
 ) -> plt.plot:
     """
     Plot predictions and display the graph
+    >>> pass    # No doctests, function is for demo purposes only
     """
     x_train_sorted = np.sort(x_train, axis=0)
     plt.scatter(x_data, y_data, color="blue")
@@ -119,6 +179,7 @@ def plot_preds(
 
     doctest.testmod()
 
+    # Demo with a dataset from the seaborn module
     training_data_x, total_bill, tip = load_data("tips", "total_bill", "tip")
     predictions = local_weight_regression(training_data_x, tip, 5)
     plot_preds(training_data_x, predictions, total_bill, tip, "total_bill", "tip")