Logistic_softmax_regression.py

import os
from dataclasses import dataclass
import numpy as np
import matplotlib.pyplot as plt


################################################################################
#                              Helper utilities                                #
################################################################################

def _ensure_exists(path: str) -> None:
    if not os.path.exists(path):
        raise FileNotFoundError(
            f"Could not find data file: {path}. "
            f"Make sure it is in the same folder as this script."
        )


def _sigmoid_stable(z: np.ndarray) -> np.ndarray:
    """Numerically-stable sigmoid."""
    # Same trick you used: avoid overflow for negative values.
    return np.where(
        z >= 0,
        1.0 / (1.0 + np.exp(-z)),
        np.exp(z) / (1.0 + np.exp(z)),
    )


def _softmax_stable(z: np.ndarray) -> np.ndarray:
    """Numerically-stable softmax over axis=1."""
    z = z - np.max(z, axis=1, keepdims=True)
    e = np.exp(z)
    return e / np.sum(e, axis=1, keepdims=True)


################################################################################
#                                   PART A                                     #
#                          Data Visualization                                  #
################################################################################

def part_a_visualize_data(data_path: str) -> None:
    """Visualize the first 3 examples from the training data."""
    print("\n" + "=" * 80)
    print("PART A: Data Visualization")
    print("=" * 80)

    _ensure_exists(data_path)
    data = np.load(data_path, allow_pickle=True)
    X_train = data["X_train"]
    y_train = data["y_train"]  # In the PDF it is called y0_train

    for i in range(3):
        image_matrix = X_train[i].reshape(28, 28)
        label = int(y_train[i])

        plt.figure(figsize=(3, 3))
        plt.imshow(image_matrix, cmap="gray")
        plt.title(f"Row #{i} - Label: {label}")
        plt.axis("off")
        plt.show()

        print(f"Image {i}: label in y_train = {label}")


################################################################################
#                                   PART B                                     #
#                        Logistic Regression (Binary)                           #
################################################################################

class LogisticRegression:
    """Binary logistic regression with L2 regularization (weights only)."""

    def __init__(self, n_features: int, seed: int | None = None):
        if seed is not None:
            np.random.seed(seed)
        self.n_features = n_features
        self.w = np.random.randn(n_features)
        self.b = 0.0

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        z = X @ self.w + self.b
        return _sigmoid_stable(z)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return (self.predict_proba(X) > 0.5).astype(int)

    def compute_loss(self, X: np.ndarray, y: np.ndarray, lambda_val: float) -> float:
        y_pred = self.predict_proba(X)
        eps = 1e-15
        y_pred = np.clip(y_pred, eps, 1 - eps)

        ce = -np.mean(y * np.log(y_pred) + (1 - y) * np.log(1 - y_pred))
        reg = (lambda_val / 2.0) * np.sum(self.w ** 2)  # regularize weights only
        return float(ce + reg)

    def compute_gradients(
        self, X: np.ndarray, y: np.ndarray, lambda_val: float
    ) -> tuple[np.ndarray, float]:
        y_pred = self.predict_proba(X)
        error = y_pred - y
        n = X.shape[0]

        dw = (X.T @ error) / n + (lambda_val * self.w)
        db = float(np.mean(error))
        return dw, db

    def calc_accuracy(self, X: np.ndarray, y: np.ndarray) -> float:
        return float(np.mean(self.predict(X) == y))

    def fit(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
        learning_rate: float,
        lambda_val: float,
        epochs: int,
    ) -> tuple[list[float], list[float], list[float], list[float]]:
        train_loss: list[float] = []
        valid_loss: list[float] = []
        train_acc: list[float] = []
        valid_acc: list[float] = []

        for _ in range(epochs):
            dw, db = self.compute_gradients(X_train, y_train, lambda_val)
            self.w -= learning_rate * dw
            self.b -= learning_rate * db

            train_loss.append(self.compute_loss(X_train, y_train, lambda_val))
            valid_loss.append(self.compute_loss(X_val, y_val, lambda_val))
            train_acc.append(self.calc_accuracy(X_train, y_train))
            valid_acc.append(self.calc_accuracy(X_val, y_val))

        return train_loss, valid_loss, train_acc, valid_acc


def part_b_train_and_plot(data_path: str) -> LogisticRegression:
    """Train a binary logistic regression model and plot loss/accuracy."""
    print("\n" + "=" * 80)
    print("PART B: Logistic Regression Training")
    print("=" * 80)

    _ensure_exists(data_path)
    data = np.load(data_path, allow_pickle=True)
    X_train = data["X_train"]
    y_train = data["y_train"]
    X_val = data["X_validation"]
    y_val = data["y_validation"]

    model = LogisticRegression(n_features=784)
    t_loss, v_loss, t_acc, v_acc = model.fit(
        X_train,
        y_train,
        X_val,
        y_val,
        learning_rate=0.1,
        lambda_val=0.0,
        epochs=1000,
    )

    plt.figure(figsize=(10, 5))
    plt.plot(t_loss, label="Train Loss")
    plt.plot(v_loss, label="Validation Loss")
    plt.xlabel("Epochs")
    plt.ylabel("Cross Entropy Loss")
    plt.title("Loss over Epochs (Learning Rate=0.1, Lambda=0)")
    plt.legend()
    plt.grid(True)
    plt.savefig('results.png', dpi=300, bbox_inches='tight')
    plt.show()

    plt.figure(figsize=(10, 5))
    plt.plot(t_acc, label="Train Accuracy")
    plt.plot(v_acc, label="Validation Accuracy")
    plt.axhline(y=0.5, linestyle="--", label="Random Guess (0.5)")
    plt.xlabel("Epochs")
    plt.ylabel("Accuracy")
    plt.title("Accuracy over Epochs (Learning Rate=0.1, Lambda=0)")
    plt.legend()
    plt.grid(True)
    plt.show()

    return model


################################################################################
#                                   PART C                                     #
#                     Hyperparameter Tuning & Visualization                    #
################################################################################

def find_best_hyperparameters(
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    epochs: int = 1000,
    seed: int = 42,
) -> tuple[dict, list[dict]]:
    learning_rates = np.logspace(-5, 2, num=8)
    lambdas = np.logspace(-5, 2, num=8)

    results: list[dict] = []
    for lr in learning_rates:
        for lam in lambdas:
            model = LogisticRegression(n_features=784, seed=seed)
            t_loss, v_loss, t_acc, v_acc = model.fit(
                X_train,
                y_train,
                X_val,
                y_val,
                learning_rate=float(lr),
                lambda_val=float(lam),
                epochs=epochs,
            )
            results.append(
                {
                    "lr": float(lr),
                    "lambda": float(lam),
                    "final_train_loss": t_loss[-1],
                    "final_val_loss": v_loss[-1],
                    "final_train_acc": t_acc[-1],
                    "final_val_acc": v_acc[-1],
                }
            )

    best = max(results, key=lambda r: r["final_val_acc"])
    print("Best Result Found:")
    print(best)
    return best, results


def plot_metric(
    data: list[dict],
    x_key: str,
    train_key: str,
    val_key: str,
    xlabel: str,
    ylabel: str,
    title: str,
) -> None:
    data = sorted(data, key=lambda r: r[x_key])
    x = [r[x_key] for r in data]

    plt.figure(figsize=(8, 5))
    plt.plot(x, [r[train_key] for r in data], "o-", label="Train")
    plt.plot(x, [r[val_key] for r in data], "o-", label="Validation")
    plt.xscale("log")
    plt.xlabel(xlabel)
    plt.ylabel(ylabel)
    plt.title(title)
    plt.legend()
    plt.grid(True)
    plt.show()


def show_weights_binary(model: LogisticRegression, title: str) -> None:
    w_img = model.w.reshape(28, 28)
    plt.figure(figsize=(6, 5))
    plt.imshow(w_img, cmap="bwr")
    plt.colorbar(label="Weight Value")
    plt.title(title)
    plt.axis("off")
    plt.show()


def plot_results(
    results: list[dict],
    best: dict,
    X_train: np.ndarray,
    y_train: np.ndarray,
    X_val: np.ndarray,
    y_val: np.ndarray,
    epochs: int = 1000,
) -> None:
    best_lr = best["lr"]
    best_lambda = best["lambda"]
    print(f"Plotting graphs for Best LR={best_lr}, Best Lambda={best_lambda}")

    lr_data = [r for r in results if r["lambda"] == best_lambda]
    lam_data = [r for r in results if r["lr"] == best_lr]

    plot_metric(
        lr_data,
        "lr",
        "final_train_loss",
        "final_val_loss",
        "Learning Rate",
        "Cross Entropy Loss",
        f"Loss vs Learning Rate (Fixed Lambda={best_lambda})",
    )
    plot_metric(
        lr_data,
        "lr",
        "final_train_acc",
        "final_val_acc",
        "Learning Rate",
        "Accuracy",
        f"Accuracy vs Learning Rate (Fixed Lambda={best_lambda})",
    )
    plot_metric(
        lam_data,
        "lambda",
        "final_train_loss",
        "final_val_loss",
        "Lambda (Regularization)",
        "Cross Entropy Loss",
        f"Loss vs Lambda (Fixed LR={best_lr})",
    )
    plot_metric(
        lam_data,
        "lambda",
        "final_train_acc",
        "final_val_acc",
        "Lambda (Regularization)",
        "Accuracy",
        f"Accuracy vs Lambda (Fixed LR={best_lr})",
    )

    # Weight visualizations
    for lam, title in [
        (best_lambda, f"Optimal Weights (Lambda={best_lambda})"),
        (0.0, "Weights without Regularization (Lambda=0)"),
        (1.0, "Weights with Strong Regularization (Lambda=1)"),
    ]:
        print(f"Retraining model for visualization: LR={best_lr}, Lambda={lam}...")
        model = LogisticRegression(n_features=784, seed=42)
        model.fit(
            X_train,
            y_train,
            X_val,
            y_val,
            learning_rate=best_lr,
            lambda_val=float(lam),
            epochs=epochs,
        )
        show_weights_binary(model, title)


def part_c_tuning_and_evaluation(data_path: str) -> tuple[LogisticRegression, dict]:
    """Hyperparameter tuning + final evaluation on test."""
    print("\n" + "=" * 80)
    print("PART C: Hyperparameter Tuning & Evaluation")
    print("=" * 80)

    _ensure_exists(data_path)
    data = np.load(data_path, allow_pickle=True)
    X_train, y_train = data["X_train"], data["y_train"]
    X_val, y_val = data["X_validation"], data["y_validation"]
    X_test, y_test = data["X_test"], data["y_test"]

    best, results = find_best_hyperparameters(X_train, y_train, X_val, y_val)
    plot_results(results, best, X_train, y_train, X_val, y_val)

    final_model = LogisticRegression(n_features=784, seed=42)
    final_model.fit(
        X_train,
        y_train,
        X_val,
        y_val,
        learning_rate=best["lr"],
        lambda_val=best["lambda"],
        epochs=1000,
    )
    test_acc = final_model.calc_accuracy(X_test, y_test)
    print("\n" + "=" * 30)
    print("FINAL TEST EVALUATION")
    print("=" * 30)
    print(f"Final Test Accuracy: {test_acc:.4f}")
    print("=" * 30)

    return final_model, best


################################################################################
#                                   BONUS A                                    #
#                    Model response to random noise images                      #
################################################################################

def bonus_a_noise_response(
    model: LogisticRegression,
    num_samples: int = 10_000,
    seed: int = 0,
) -> None:
    """
    Bonus A:
    Sample random images ~ U[0,1], feed them to the binary model and
    analyze the output distribution.
    """
    print("\n" + "=" * 80)
    print("BONUS A: Model response to random noise")
    print("=" * 80)

    rng = np.random.default_rng(seed)

    # Generate random noise images (same shape as flattened MNIST)
    X_noise = rng.uniform(0.0, 1.0, size=(num_samples, 784)).astype(np.float32)

    probs = model.predict_proba(X_noise)

    mean_p = float(np.mean(probs))
    std_p = float(np.std(probs))
    frac_class1 = float(np.mean(probs > 0.5))

    print(f"Mean predicted probability: {mean_p:.4f}")
    print(f"Std  predicted probability: {std_p:.4f}")
    print(f"Fraction classified as class-1 (p>0.5): {frac_class1:.4f}")

    plt.figure(figsize=(9, 5))
    plt.hist(probs, bins=60, alpha=0.75)
    plt.axvline(mean_p, color="red", linestyle="--", label=f"Mean = {mean_p:.3f}")
    plt.xlabel("Predicted probability p(y=1 | x)")
    plt.ylabel("Count")
    plt.title("Binary model output on random noise images U[0,1]")
    plt.legend()
    plt.grid(True)
    plt.show()

################################################################################
#                                   BONUS B                                    #
#            Multinomial (softmax) regression for 10 classes                    #
################################################################################

@dataclass
class SoftmaxConfig:
    learning_rate: float = 0.3
    lambda_val: float = 5e-4
    epochs: int = 20
    batch_size: int = 512
    seed: int = 42


class SoftmaxRegression:
    """
    Multinomial logistic regression (Softmax),
    trained with mini-batch gradient descent.
    """

    def __init__(self, n_features: int, n_classes: int, seed: int = 42):
        rng = np.random.default_rng(seed)
        self.W = (rng.standard_normal((n_features, n_classes)).astype(np.float32)) * 0.01
        self.b = np.zeros((n_classes,), dtype=np.float32)

    def predict_proba(self, X: np.ndarray) -> np.ndarray:
        logits = X @ self.W + self.b
        return _softmax_stable(logits)

    def predict(self, X: np.ndarray) -> np.ndarray:
        return np.argmax(X @ self.W + self.b, axis=1)

    def loss(self, X: np.ndarray, y: np.ndarray, lambda_val: float) -> float:
        P = self.predict_proba(X)
        eps = 1e-15
        p_true = np.clip(P[np.arange(len(y)), y], eps, 1.0)
        ce = -np.mean(np.log(p_true))
        reg = (lambda_val / 2.0) * float(np.sum(self.W ** 2))
        return float(ce + reg)

    def accuracy(self, X: np.ndarray, y: np.ndarray) -> float:
        return float(np.mean(self.predict(X) == y))

    def fit_minibatch(
        self,
        X_train: np.ndarray,
        y_train: np.ndarray,
        X_val: np.ndarray,
        y_val: np.ndarray,
        cfg: SoftmaxConfig,
    ) -> dict[str, list[float]]:

        rng = np.random.default_rng(cfg.seed)
        n = X_train.shape[0]

        history = {
            "train_loss": [],
            "val_loss": [],
            "train_acc": [],
            "val_acc": [],
        }

        for ep in range(cfg.epochs):
            perm = rng.permutation(n)

            for start in range(0, n, cfg.batch_size):
                idx = perm[start:start + cfg.batch_size]
                Xb = X_train[idx]
                yb = y_train[idx]

                P = self.predict_proba(Xb)
                P[np.arange(len(yb)), yb] -= 1.0
                P /= len(yb)

                dW = Xb.T @ P + cfg.lambda_val * self.W
                db = np.sum(P, axis=0)

                self.W -= cfg.learning_rate * dW
                self.b -= cfg.learning_rate * db

            history["train_loss"].append(self.loss(X_train, y_train, cfg.lambda_val))
            history["val_loss"].append(self.loss(X_val, y_val, cfg.lambda_val))
            history["train_acc"].append(self.accuracy(X_train, y_train))
            history["val_acc"].append(self.accuracy(X_val, y_val))

            print(
                f"[epoch {ep+1:02d}/{cfg.epochs}] "
                f"train_acc={history['train_acc'][-1]:.4f} | "
                f"val_acc={history['val_acc'][-1]:.4f}"
            )

        return history


def bonus_b_train_softmax(bonus_data_path: str, cfg: SoftmaxConfig | None = None) -> None:
    """
    Bonus B: 10-class classification with Softmax Regression
    """
    print("\n" + "=" * 80)
    print("BONUS B: 10-class softmax regression")
    print("=" * 80)

    if cfg is None:
        cfg = SoftmaxConfig()

    _ensure_exists(bonus_data_path)
    data = np.load(bonus_data_path, allow_pickle=True)

    # Flatten + normalize
    X_train = data["X_train"].reshape(-1, 784).astype(np.float32) / 255.0
    y_train = data["y_train"].astype(int)
    X_val = data["X_validation"].reshape(-1, 784).astype(np.float32) / 255.0
    y_val = data["y_validation"].astype(int)
    X_test = data["X_test"].reshape(-1, 784).astype(np.float32) / 255.0
    y_test = data["y_test"].astype(int)

    # -------- Feature Standardization  --------
    mean = X_train.mean(axis=0)
    std = X_train.std(axis=0) + 1e-6

    X_train = (X_train - mean) / std
    X_val = (X_val - mean) / std
    X_test = (X_test - mean) / std
    # -------------------------------------------------------------

    model = SoftmaxRegression(n_features=784, n_classes=10, seed=cfg.seed)
    hist = model.fit_minibatch(X_train, y_train, X_val, y_val, cfg)

    test_acc = model.accuracy(X_test, y_test)
    print(f"\nFinal test accuracy (BONUS B): {test_acc:.4f}")

    # Plots
    plt.figure(figsize=(10, 5))
    plt.plot(hist["train_loss"], label="Train Loss")
    plt.plot(hist["val_loss"], label="Validation Loss")
    plt.xlabel("Epoch")
    plt.ylabel("Cross Entropy Loss")
    plt.title("Softmax Loss (BONUS B)")
    plt.legend()
    plt.grid(True)
    plt.show()

    plt.figure(figsize=(10, 5))
    plt.plot(hist["train_acc"], label="Train Accuracy")
    plt.plot(hist["val_acc"], label="Validation Accuracy")
    plt.xlabel("Epoch")
    plt.ylabel("Accuracy")
    plt.title("Softmax Accuracy (BONUS B)")
    plt.legend()
    plt.grid(True)
    plt.show()


################################################################################
#                                   MAIN                                       #
################################################################################

if __name__ == "__main__":
    # Part I / II: binary_class.npz
    binary_data_path = "binary_class.npz"

    # Run Part A
    part_a_visualize_data(binary_data_path)

    # Run Part B
    _ = part_b_train_and_plot(binary_data_path)

    # Run Part C + get final model (used for Bonus A)
    final_model, _best = part_c_tuning_and_evaluation(binary_data_path)

    # Bonus A (needs a trained binary model)
    bonus_a_noise_response(final_model)

    # Bonus B (runs only if bonus.npz exists next to this file)
    if os.path.exists("bonus.npz"):
        bonus_b_train_softmax("bonus.npz")
    else:
        print("bonus.npz not found - skipping Bonus B.")