Naive Bayes

"""Multinomial Naive Bayes - counting, then Bayes' rule, transcribed.

Training is the two count tables in the update rule above: how much total token
mass each class holds, and how often each token appears within it. Prediction
adds one log term per token, which is the naive assumption made literal.
"""

import math


def fit(documents, labels, vocabulary_size, alpha=1.0):
    """documents[i] is a list of token ids; labels[i] is a class id."""
    classes = sorted(set(labels))
    class_counts = {c: 0 for c in classes}
    token_totals = {c: 0 for c in classes}
    token_counts = {c: [0] * vocabulary_size for c in classes}

    for i in range(len(documents)):
        label = labels[i]
        class_counts[label] += 1
        for token in documents[i]:
            token_counts[label][token] += 1
            token_totals[label] += 1

    n = len(documents)
    log_prior = {}
    log_likelihood = {}

    for c in classes:
        log_prior[c] = math.log(class_counts[c] / n)

        # Smoothing is not optional: with alpha = 0 an unseen token gives
        # log(0), and one -inf term annihilates every other token's evidence.
        denominator = token_totals[c] + alpha * vocabulary_size
        row = [0.0] * vocabulary_size
        for token in range(vocabulary_size):
            row[token] = math.log((token_counts[c][token] + alpha) / denominator)
        log_likelihood[c] = row

    return log_prior, log_likelihood


def predict(log_prior, log_likelihood, document):
    """argmax_c  log P(c) + sum_j log P(token_j | c)"""
    best_class = None
    best_score = None

    for c in log_prior:
        # Everything is in log space because a product of thousands of
        # probabilities underflows float64 to exactly zero.
        score = log_prior[c]
        for token in document:
            score += log_likelihood[c][token]

        if best_score is None or score > best_score:
            best_class = c
            best_score = score

    return best_class

"""Multinomial Naive Bayes - typed, validated, sparse-aware."""

from dataclasses import dataclass

import numpy as np
from numpy.typing import NDArray

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


@dataclass(frozen=True)
class MultinomialNB:
    """A fitted model: one log-probability row per class, plus the log priors.

    Everything is stored in log space. Converting back is never necessary and
    is actively harmful - the exponentiated posterior saturates at 0 or 1 and
    should not be exposed as a confidence.
    """

    classes: NDArray
    log_prior: Vector           # (C,)
    log_likelihood: Matrix      # (C, V)
    alpha: float

    def joint_log_likelihood(self, X: Matrix) -> Matrix:
        """(n, C) matrix of scores. The whole prediction rule, as one product."""
        if X.ndim != 2 or X.shape[1] != self.log_likelihood.shape[1]:
            raise ValueError(
                f"expected (n, {self.log_likelihood.shape[1]}) counts, got {X.shape}"
            )
        return X @ self.log_likelihood.T + self.log_prior

    def predict(self, X: Matrix) -> NDArray:
        return self.classes[self.joint_log_likelihood(X).argmax(axis=1)]

    def log_proba(self, X: Matrix) -> Matrix:
        """Normalized log posteriors. Correctly computed, and still not calibrated."""
        joint = self.joint_log_likelihood(X)
        # log-sum-exp with the max factored out; the naive version overflows.
        shift = joint.max(axis=1, keepdims=True)
        return joint - (shift + np.log(np.exp(joint - shift).sum(axis=1, keepdims=True)))


def fit(X: Matrix, y: NDArray, alpha: float = 1.0) -> MultinomialNB:
    """Fit by counting. Raises ValueError on malformed input."""
    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 alpha <= 0.0:
        raise ValueError(f"alpha must be positive; alpha=0 makes unseen tokens fatal")
    if (X < 0).any():
        raise ValueError("multinomial NB expects non-negative counts")

    classes, encoded = np.unique(y, return_inverse=True)
    n_classes = classes.size
    vocabulary = X.shape[1]

    # One-hot the labels and the entire count table is a single matrix product:
    # indicator.T @ X sums each class's rows in one pass.
    indicator = np.zeros((X.shape[0], n_classes), dtype=np.float64)
    indicator[np.arange(X.shape[0]), encoded] = 1.0

    token_counts = indicator.T @ X                     # (C, V)
    class_counts = indicator.sum(axis=0)               # (C,)

    smoothed = token_counts + alpha
    log_likelihood = np.log(smoothed) - np.log(smoothed.sum(axis=1, keepdims=True))
    log_prior = np.log(class_counts) - np.log(class_counts.sum())

    return MultinomialNB(
        classes=classes,
        log_prior=log_prior,
        log_likelihood=log_likelihood,
        alpha=alpha,
    )

"""Multinomial Naive Bayes - sparse counts, one SpMM, batched scoring."""

import numpy as np
from numpy.typing import NDArray
from scipy import sparse

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


class SparseMultinomialNB:
    """The same model, taking sparsity seriously.

    On text, X is a document-term matrix that is over 99% zeros. The dense
    formulation touches every vocabulary entry for every document; the sparse
    one touches only the tokens actually present, which is the difference
    between O(n * V) and O(nnz). Nothing about the model changes - only which
    zeros are visited.
    """

    def __init__(self, alpha: float = 1.0) -> None:
        if alpha <= 0.0:
            raise ValueError("alpha must be positive")
        self._alpha = alpha
        self._classes: NDArray | None = None
        self._log_prior: Vector | None = None
        self._log_likelihood: Matrix | None = None

    def fit(self, X: sparse.csr_matrix, y: NDArray) -> "SparseMultinomialNB":
        if X.shape[0] != y.shape[0]:
            raise ValueError(f"X has {X.shape[0]} rows but y has {y.shape[0]}")

        classes, encoded = np.unique(y, return_inverse=True)
        n, vocabulary = X.shape

        # The one-hot indicator is itself sparse, so the count table is one
        # sparse-times-sparse product rather than a dense C x n intermediate.
        indicator = sparse.csr_matrix(
            (np.ones(n), (encoded, np.arange(n))), shape=(classes.size, n)
        )
        token_counts = np.asarray((indicator @ X).todense(), dtype=np.float64)

        smoothed = token_counts + self._alpha
        # In place: the (C, V) table is the largest array in the model, and on
        # a large vocabulary each temporary copy is real memory pressure.
        np.log(smoothed, out=smoothed)
        totals = np.log(np.exp(smoothed).sum(axis=1, keepdims=True))
        smoothed -= totals

        class_counts = np.bincount(encoded, minlength=classes.size).astype(np.float64)

        self._classes = classes
        self._log_prior = np.log(class_counts) - np.log(n)
        # C-contiguous so the SpMM below reads each class row sequentially.
        self._log_likelihood = np.ascontiguousarray(smoothed)
        return self

    def predict(self, X: sparse.csr_matrix, batch: int = 8192) -> NDArray:
        if self._log_likelihood is None or self._classes is None:
            raise RuntimeError("predict called before fit")

        out = np.empty(X.shape[0], dtype=self._classes.dtype)   # allocated once

        # Batching bounds the (batch, C) score block and keeps the class table
        # resident in cache across many documents.
        for start in range(0, X.shape[0], batch):
            stop = min(start + batch, X.shape[0])
            joint = X[start:stop] @ self._log_likelihood.T      # SpMM
            joint += self._log_prior
            out[start:stop] = self._classes[joint.argmax(axis=1)]

        return out