Logistic Regression

"""Logistic regression by gradient descent - the objective, transcribed.

The gradient has the same shape as linear regression's: features weighted by
residuals. Only the link function differs, which is what makes both of these
generalized linear models.
"""

import math


def sigmoid(z):
    # Split by sign to avoid exp() overflowing on large negative z.
    if z >= 0:
        return 1.0 / (1.0 + math.exp(-z))
    exp_z = math.exp(z)
    return exp_z / (1.0 + exp_z)


def fit(X, y, lr=0.1, epochs=1_000):
    n = len(X)
    d = len(X[0])
    theta = [0.0] * d
    bias = 0.0

    for _ in range(epochs):
        grad = [0.0] * d
        grad_bias = 0.0

        for i in range(n):
            z = bias
            for j in range(d):
                z += theta[j] * X[i][j]

            residual = sigmoid(z) - y[i]        # exactly the linear-model form
            for j in range(d):
                grad[j] += residual * X[i][j]
            grad_bias += residual

        for j in range(d):
            theta[j] -= lr * grad[j] / n
        bias -= lr * grad_bias / n

    return theta, bias


def log_loss(X, y, theta, bias):
    """Negative log-likelihood. Clipped, because log(0) is -inf and one
    confident mistake would otherwise make the whole loss infinite."""
    total = 0.0
    for x, target in zip(X, y):
        z = bias + sum(t * v for t, v in zip(theta, x))
        p = min(max(sigmoid(z), 1e-15), 1.0 - 1e-15)
        total -= target * math.log(p) + (1 - target) * math.log(1.0 - p)
    return total / len(y)

"""Logistic regression - typed, validated, calibration-aware."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import StandardScaler

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


@dataclass(frozen=True)
class ClassifierReport:
    """A score is not enough — a threshold decision needs calibration too."""

    model: Pipeline
    log_loss: float
    brier_score: float
    base_rate: float

    @property
    def is_calibrated(self) -> bool:
        """Brier below the base-rate variance means it beats predicting the
        base rate for everyone — a low bar, and one models do fail."""
        return self.brier_score < self.base_rate * (1.0 - self.base_rate)


def fit_classifier(
    X: Matrix,
    y: Vector,
    C: float = 1.0,
    balanced: bool = False,
) -> ClassifierReport:
    if X.ndim != 2:
        raise ValueError(f"X must be 2-D, got shape {X.shape}")
    if X.shape[0] != y.shape[0]:
        raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")
    if not np.isin(y, (0, 1)).all():
        raise ValueError("y must be binary 0/1")
    if len(np.unique(y)) < 2:
        raise ValueError("y contains only one class - nothing to separate")

    # Scaling is inside the pipeline, not applied beforehand: the L2 penalty is
    # scale-dependent, and a scaler fitted separately will silently drift from
    # the model it was fitted with.
    model = Pipeline(
        [
            ("scale", StandardScaler()),
            (
                "clf",
                LogisticRegression(
                    C=C,
                    max_iter=1_000,
                    class_weight="balanced" if balanced else None,
                ),
            ),
        ]
    )
    model.fit(X, y)

    probabilities = model.predict_proba(X)[:, 1]
    clipped = np.clip(probabilities, 1e-15, 1 - 1e-15)

    return ClassifierReport(
        model=model,
        log_loss=float(-np.mean(y * np.log(clipped) + (1 - y) * np.log(1 - clipped))),
        brier_score=float(np.mean((probabilities - y) ** 2)),
        base_rate=float(y.mean()),
    )


def threshold_for_expected_cost(
    probabilities: Vector, cost_false_positive: float, cost_false_negative: float
) -> float:
    """The expected-cost-minimizing cutoff, which is a ratio of costs and has
    nothing to do with 0.5. Defaulting to 0.5 silently assumes the two errors
    are equally expensive, which they almost never are."""
    if cost_false_positive <= 0 or cost_false_negative <= 0:
        raise ValueError("costs must be positive")
    return cost_false_positive / (cost_false_positive + cost_false_negative)

"""Logistic regression - stable log-loss, single-pass gradient, warm starts."""

import numpy as np
from numpy.typing import NDArray
from scipy.special import expit, log_expit

Vector = NDArray[np.float64]
Matrix = NDArray[np.float64]


def log_loss_stable(X: Matrix, y: Vector, theta: Vector) -> float:
    """Log loss computed from logits directly, never from probabilities.

    The naive route (sigmoid, then clip, then log) loses precision exactly
    where it matters - on confident predictions, whose contribution underflows
    to the clip bound and stops being distinguishable. log_expit is the
    numerically stable log-sigmoid, so no clipping is needed at all.
    """
    logits = X @ theta
    # log(sigma(z)) for positives, log(1 - sigma(z)) = log(sigma(-z)) for negatives.
    return float(-np.mean(y * log_expit(logits) + (1 - y) * log_expit(-logits)))


def gradient(X: Matrix, y: Vector, theta: Vector, out: Vector) -> Vector:
    """Gradient written into a caller-owned buffer.

    Called once per optimizer iteration, so a fresh (d,) allocation per call
    is an allocation per iteration.
    """
    residual = expit(X @ theta)
    residual -= y                       # in place; no second array
    np.dot(X.T, residual, out=out)
    out /= X.shape[0]
    return out


def fit_newton(X: Matrix, y: Vector, l2: float = 1.0, iterations: int = 25) -> Vector:
    """IRLS / Newton steps.

    Second-order converges in tens of iterations rather than thousands, and at
    modest d the O(d^3) solve per step is far cheaper than the extra passes
    over the data that first-order methods need.
    """
    X = np.ascontiguousarray(X, dtype=np.float64)
    n, d = X.shape
    theta = np.zeros(d, dtype=np.float64)
    grad_buffer = np.empty(d, dtype=np.float64)

    for _ in range(iterations):
        p = expit(X @ theta)
        # w = p(1-p), the IRLS weights. Computed in place to avoid two temporaries.
        w = p * (1.0 - p)
        np.maximum(w, 1e-10, out=w)     # guard the Hessian against singularity

        gradient(X, y, theta, grad_buffer)
        grad_buffer += l2 * theta / n

        # Hessian = X^T W X + l2 I. einsum fuses the scaling and the product,
        # so the n x d weighted copy of X is never materialized.
        hessian = np.einsum("ij,i,ik->jk", X, w, X) / n
        hessian[np.diag_indices(d)] += l2 / n

        # solve(), not inv(): forming the inverse squares the condition number
        # and costs more.
        theta -= np.linalg.solve(hessian, grad_buffer)

    return theta