sgd_classifier.py

"""A NumPy-only linear classifier trained with mini-batch SGD."""

from __future__ import annotations

from numbers import Integral, Real
from typing import Literal

import numpy as np


SUPPORTED_LOSSES = (
    "hinge",
    "log_loss",
    "modified_huber",
    "squared_hinge",
    "perceptron",
    "squared_error",
    "huber",
    "epsilon_insensitive",
    "squared_epsilon_insensitive",
)


class SGDClassifier:
    """Linear one-vs-rest classifier optimized by mini-batch SGD.

    ``random`` batch selection independently samples each batch without
    replacement. ``permutation`` shuffles the dataset and consumes
    non-overlapping batches, then reshuffles at the start of the next epoch.

    The implementation depends only on NumPy. For multiclass problems, each
    class is fitted against all other classes using a one-vs-rest target.
    """

    def __init__(
        self,
        loss: str = "hinge",
        *,
        batch_size: int = 32,
        batch_selection: Literal["random", "permutation"] = "permutation",
        max_epochs: int = 10,
        learning_rate: float = 0.01,
        penalty: Literal["l2", "l1", "elasticnet"] | None = "l2",
        alpha: float = 0.000, # no regularization by default
        l1_ratio: float = 0.15,
        fit_intercept: bool = True,
        epsilon: float = 0.1,
        random_state: int | None = None,
    ) -> None:
        self.loss = loss
        self.batch_size = batch_size
        self.batch_selection = batch_selection
        self.max_epochs = max_epochs
        self.learning_rate = learning_rate
        self.penalty = penalty
        self.alpha = alpha
        self.l1_ratio = l1_ratio
        self.fit_intercept = fit_intercept
        self.epsilon = epsilon
        self.random_state = random_state

    def _validate_params(self) -> None:
        if self.loss not in SUPPORTED_LOSSES:
            raise ValueError(
                f"Unsupported loss {self.loss!r}; expected one of {SUPPORTED_LOSSES}."
            )
        if not isinstance(self.batch_size, Integral) or self.batch_size <= 0:
            raise ValueError("batch_size must be a positive integer.")
        if self.batch_selection not in {"random", "permutation"}:
            raise ValueError("batch_selection must be 'random' or 'permutation'.")
        if not isinstance(self.max_epochs, Integral) or self.max_epochs <= 0:
            raise ValueError("max_epochs must be a positive integer.")
        if not isinstance(self.learning_rate, Real) or self.learning_rate <= 0:
            raise ValueError("learning_rate must be positive.")
        if self.penalty not in {None, "l1", "l2", "elasticnet"}:
            raise ValueError("penalty must be None, 'l1', 'l2', or 'elasticnet'.")
        if not isinstance(self.alpha, Real) or self.alpha < 0:
            raise ValueError("alpha must be non-negative.")
        if not 0 <= self.l1_ratio <= 1:
            raise ValueError("l1_ratio must be between 0 and 1.")
        if self.epsilon < 0:
            raise ValueError("epsilon must be non-negative.")

    @staticmethod
    def _check_X_y(X, y):
        X = np.asarray(X, dtype=float)
        y = np.asarray(y)
        if X.ndim != 2:
            raise ValueError("X must be a two-dimensional array.")
        if y.ndim != 1:
            y = np.ravel(y)
        if len(X) != len(y) or len(y) == 0:
            raise ValueError("X and y must contain the same non-zero sample count.")
        if not np.all(np.isfinite(X)):
            raise ValueError("X must contain only finite values.")
        return X, y

    def _initialize(self, n_features: int, classes) -> None:
        self.classes_ = np.asarray(classes)
        if self.classes_.ndim != 1 or len(self.classes_) < 2:
            raise ValueError("At least two classes are required.")
        self.coef_ = np.zeros((len(self.classes_), n_features), dtype=float)
        self.intercept_ = np.zeros(len(self.classes_), dtype=float)
        self.n_features_in_ = n_features
        self.n_updates_ = 0

    def _loss_gradient(self, scores, targets):
        """Derivative of each requested loss with respect to the scores."""
        margin = targets * scores
        if self.loss == "hinge": # (1-margin)_{+}
            return np.where(margin < 1, -targets, 0.0)
        if self.loss == "log_loss": # ln()
            # Stable form of -y / (1 + exp(y * score)).
            z = np.clip(margin, -50, 50)
            return -targets / (1.0 + np.exp(z))
        if self.loss == "modified_huber":
            return np.where(
                margin >= 1,
                0.0,
                np.where(margin >= -1, -2 * targets * (1 - margin), -4 * targets),
            )
        if self.loss == "squared_hinge": # (1-margin)^2
            return np.where(margin < 1, -2 * targets * (1 - margin), 0.0)
        if self.loss == "perceptron": # (-margin)_{+}
            return np.where(margin <= 0, -targets, 0.0)

        residual = scores - targets
        if self.loss == "squared_error":
            return residual
        if self.loss == "huber":
            return np.clip(residual, -self.epsilon, self.epsilon)
        outside = np.abs(residual) > self.epsilon
        if self.loss == "epsilon_insensitive":
            return np.where(outside, np.sign(residual), 0.0)
        # squared_epsilon_insensitive
        return np.where(
            outside, 2 * np.sign(residual) * (np.abs(residual) - self.epsilon), 0.0
        )

    def _loss_values(self, scores, targets):
        """Element-wise loss corresponding to :meth:`_loss_gradient`."""
        margin = targets * scores
        if self.loss == "hinge":
            return np.maximum(0.0, 1 - margin)
        if self.loss == "log_loss":
            return np.logaddexp(0.0, -margin)
        if self.loss == "modified_huber":
            return np.where(
                margin >= 1,
                0.0,
                np.where(margin >= -1, (1 - margin) ** 2, -4 * margin),
            )
        if self.loss == "squared_hinge":
            return np.maximum(0.0, 1 - margin) ** 2
        if self.loss == "perceptron":
            return np.maximum(0.0, -margin)

        residual = scores - targets
        absolute = np.abs(residual)
        if self.loss == "squared_error":
            return 0.5 * residual**2
        if self.loss == "huber":
            return np.where(
                absolute <= self.epsilon,
                0.5 * residual**2,
                self.epsilon * (absolute - 0.5 * self.epsilon),
            )
        if self.loss == "epsilon_insensitive":
            return np.maximum(0.0, absolute - self.epsilon)
        return np.maximum(0.0, absolute - self.epsilon) ** 2

    def evaluate(self, X, y):
        """Return mean loss and accuracy over an entire labelled dataset."""
        X, y = self._check_X_y(X, y)
        if not hasattr(self, "coef_"):
            raise RuntimeError("The classifier has not been initialized.")
        if X.shape[1] != self.n_features_in_:
            raise ValueError(f"X must have {self.n_features_in_} features.")
        scores = X @ self.coef_.T + self.intercept_
        targets = np.where(y[:, None] == self.classes_[None, :], 1.0, -1.0)
        loss = float(np.mean(self._loss_values(scores, targets)))
        predictions = self.classes_[np.argmax(scores, axis=1)]
        return {"loss": loss, "accuracy": float(np.mean(predictions == y))}

    def _regularization_gradient(self):
        if self.penalty is None or self.alpha == 0:
            return np.zeros_like(self.coef_)
        if self.penalty == "l2":
            return self.alpha * self.coef_
        if self.penalty == "l1":
            return self.alpha * np.sign(self.coef_)
        return self.alpha * (
            self.l1_ratio * np.sign(self.coef_)
            + (1 - self.l1_ratio) * self.coef_
        )

    def get_update(self, X, y, *, classes=None, sample_weight=None):
        """Return the parameter update for a batch without applying it.

        On the first call, ``classes`` may be supplied to declare labels that
        are absent from this particular batch. Otherwise classes are inferred
        from ``y``. Initializing the parameter arrays is the only state change
        this method can make.
        """
        self._validate_params()
        X, y = self._check_X_y(X, y)
        if not hasattr(self, "coef_"):
            initial_classes = np.unique(y) if classes is None else np.asarray(classes)
            self._initialize(X.shape[1], initial_classes)
        elif X.shape[1] != self.n_features_in_:
            raise ValueError(f"X must have {self.n_features_in_} features.")

        known = np.isin(y, self.classes_)
        if not np.all(known):
            raise ValueError(f"Unknown class labels: {np.unique(y[~known])!r}.")

        if sample_weight is None:
            weights = np.ones(len(y), dtype=float)
        else:
            weights = np.asarray(sample_weight, dtype=float)
            if weights.ndim != 1 or len(weights) != len(y):
                raise ValueError("sample_weight must have one value per sample.")
            if np.any(weights < 0) or not np.all(np.isfinite(weights)):
                raise ValueError("sample_weight must be finite and non-negative.")
        weight_sum = weights.sum()
        if weight_sum <= 0:
            raise ValueError("sample_weight must sum to a positive value.")

        scores = X @ self.coef_.T + self.intercept_
        targets = np.where(y[:, None] == self.classes_[None, :], 1.0, -1.0)
        derivatives = self._loss_gradient(scores, targets) * weights[:, None]
        coef_gradient = derivatives.T @ X / weight_sum
        coef_gradient += self._regularization_gradient()
        coef_update = -self.learning_rate * coef_gradient
        if self.fit_intercept:
            intercept_update = -self.learning_rate * derivatives.sum(axis=0) / weight_sum
        else:
            intercept_update = np.zeros_like(self.intercept_)
        return coef_update, intercept_update

    def train_step(self, X, y, *, classes=None, sample_weight=None):
        """Perform one mini-batch SGD update and return this model."""
        coef_update, intercept_update = self.get_update(
            X, y, classes=classes, sample_weight=sample_weight
        )
        self.coef_ += coef_update
        self.intercept_ += intercept_update
        self.n_updates_ += 1
        return self

    def _batches(self, n_samples: int, rng):
        size = min(int(self.batch_size), n_samples)
        updates = (n_samples + size - 1) // size
        if self.batch_selection == "permutation":
            order = rng.permutation(n_samples)
            for start in range(0, n_samples, size):
                yield order[start : start + size]
        else:
            for _ in range(updates):
                yield rng.choice(n_samples, size=size, replace=False)

    def fit(self, X, y, sample_weight=None):
        """Reset and train the model, returning ``self``."""
        self._validate_params()
        X, y = self._check_X_y(X, y)
        classes = np.unique(y)
        self._initialize(X.shape[1], classes)
        if sample_weight is not None:
            sample_weight = np.asarray(sample_weight, dtype=float)
            if sample_weight.ndim != 1 or len(sample_weight) != len(y):
                raise ValueError("sample_weight must have one value per sample.")

        rng = np.random.default_rng(self.random_state)
        for _ in range(int(self.max_epochs)):
            for indices in self._batches(len(y), rng):
                batch_weights = None if sample_weight is None else sample_weight[indices]
                self.train_step(X[indices], y[indices], sample_weight=batch_weights)
        return self

    def _check_fitted_X(self, X):
        if not hasattr(self, "coef_"):
            raise RuntimeError("The classifier has not been trained.")
        X = np.asarray(X, dtype=float)
        if X.ndim != 2 or X.shape[1] != self.n_features_in_:
            raise ValueError(f"X must be two-dimensional with {self.n_features_in_} features.")
        return X

    def decision_function(self, X):
        X = self._check_fitted_X(X)
        scores = X @ self.coef_.T + self.intercept_
        return scores[:, 1] - scores[:, 0] if len(self.classes_) == 2 else scores

    def predict(self, X):
        X = self._check_fitted_X(X)
        scores = X @ self.coef_.T + self.intercept_
        return self.classes_[np.argmax(scores, axis=1)]

    def predict_proba(self, X):
        """Return softmax-normalized class scores."""
        X = self._check_fitted_X(X)
        scores = X @ self.coef_.T + self.intercept_
        shifted = scores - scores.max(axis=1, keepdims=True)
        probabilities = np.exp(np.clip(shifted, -700, 0))
        return probabilities / probabilities.sum(axis=1, keepdims=True)

    def predict_log_proba(self, X):
        return np.log(np.clip(self.predict_proba(X), np.finfo(float).tiny, 1.0))


__all__ = ["SUPPORTED_LOSSES", "SGDClassifier"]