Stepwise Regression

"""Forward stepwise selection by AIC - the objective, transcribed.

Each step refits one least-squares model per remaining candidate and keeps the
feature whose addition lowers AIC the most. Features and target are assumed
centred, which is the same as including an intercept in every model: the
intercept and the noise variance add the same constant to every candidate's
AIC, so they drop out of the comparison.
"""

import math

SINGULAR_TOLERANCE = 1e-12


def solve(matrix, rhs):
    """Gaussian elimination with partial pivoting. None if singular."""
    size = len(rhs)
    a = [row[:] for row in matrix]
    b = rhs[:]
    for col in range(size):
        pivot = col
        for row in range(col + 1, size):
            if abs(a[row][col]) > abs(a[pivot][col]):
                pivot = row
        if abs(a[pivot][col]) < SINGULAR_TOLERANCE:
            return None
        a[col], a[pivot] = a[pivot], a[col]
        b[col], b[pivot] = b[pivot], b[col]
        for row in range(col + 1, size):
            factor = a[row][col] / a[col][col]
            for inner in range(col, size):
                a[row][inner] -= factor * a[col][inner]
            b[row] -= factor * b[col]

    solution = [0.0] * size
    for row in range(size - 1, -1, -1):
        total = b[row]
        for col in range(row + 1, size):
            total -= a[row][col] * solution[col]
        solution[row] = total / a[row][row]
    return solution


def aic(n, rss, k):
    # AIC = n ln(RSS / n) + 2k for Gaussian errors, up to a constant.
    return n * math.log(max(rss, 1e-300) / n) + 2.0 * k


def rss_of(X, y, subset):
    """Residual sum of squares of the least-squares fit on the given columns."""
    n = len(y)
    k = len(subset)
    if k == 0:
        return sum(value * value for value in y)

    # Normal equations (X_S^T X_S) beta = X_S^T y, built entry by entry.
    gram = [[0.0] * k for _ in range(k)]
    rhs = [0.0] * k
    for a in range(k):
        for b in range(k):
            total = 0.0
            for i in range(n):
                total += X[i][subset[a]] * X[i][subset[b]]
            gram[a][b] = total
        total = 0.0
        for i in range(n):
            total += X[i][subset[a]] * y[i]
        rhs[a] = total

    beta = solve(gram, rhs)
    if beta is None:
        return math.inf            # collinear with the current model: never chosen

    rss = 0.0
    for i in range(n):
        prediction = 0.0
        for a in range(k):
            prediction += X[i][subset[a]] * beta[a]
        rss += (y[i] - prediction) ** 2
    return rss


def forward_stepwise(X, y):
    n = len(y)
    p = len(X[0])
    selected = []
    best_aic = aic(n, rss_of(X, y, selected), 0)

    for _ in range(p):                       # at most p additions
        best_feature = None
        best_candidate = best_aic
        for j in range(p):
            if j in selected:
                continue
            trial = selected + [j]
            score = aic(n, rss_of(X, y, trial), len(trial))
            if score < best_candidate:
                best_candidate = score
                best_feature = j

        if best_feature is None:             # no addition lowers AIC: stop
            break
        selected.append(best_feature)
        best_aic = best_candidate

    return selected

"""Stepwise selection by AIC or BIC - typed, validated, direction-aware."""

from dataclasses import dataclass
from typing import Literal

import numpy as np
from numpy.typing import NDArray

Matrix = NDArray[np.float64]
Vector = NDArray[np.float64]
Direction = Literal["forward", "backward", "both"]
Criterion = Literal["aic", "bic"]

MIN_RSS = float(np.finfo(np.float64).tiny)


@dataclass(frozen=True)
class Step:
    """One accepted move on the path, kept so the search can be audited."""

    action: Literal["add", "drop"]
    feature: int
    criterion: float


@dataclass(frozen=True)
class StepwiseResult:
    selected: tuple[int, ...]
    criterion: float
    path: tuple[Step, ...]


def _criterion_value(X: Matrix, y: Vector, subset: list[int], penalty: float) -> float:
    n = y.size
    if subset:
        design = X[:, subset]
        coefficients, *_ = np.linalg.lstsq(design, y, rcond=None)
        residual = y - design @ coefficients
    else:
        residual = y
    rss = max(float(residual @ residual), MIN_RSS)
    return n * float(np.log(rss / n)) + penalty * len(subset)


def stepwise(
    X: Matrix,
    y: Vector,
    direction: Direction = "both",
    criterion: Criterion = "aic",
    max_steps: int | None = None,
) -> StepwiseResult:
    """Greedy subset search. 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 criterion not in ("aic", "bic"):
        raise ValueError(f"criterion must be 'aic' or 'bic', got {criterion!r}")

    n, p = X.shape
    if direction == "backward" and n <= p:
        raise ValueError(f"backward elimination needs n > p, got n={n}, p={p}")

    penalty = 2.0 if criterion == "aic" else float(np.log(n))

    # Centring stands in for an intercept that is always in the model.
    design = X - X.mean(axis=0)
    target = y - y.mean()

    selected: list[int] = list(range(p)) if direction == "backward" else []
    current = _criterion_value(design, target, selected, penalty)
    path: list[Step] = []

    # Bounded loop: each accepted move strictly lowers the criterion, so the
    # walk cannot cycle, but the explicit cap makes termination checkable.
    for _ in range(max_steps if max_steps is not None else 2 * p):
        moves: list[tuple[str, int, float]] = []
        if direction in ("forward", "both"):
            for feature in range(p):
                if feature in selected:
                    continue
                trial = sorted([*selected, feature])
                moves.append(("add", feature, _criterion_value(design, target, trial, penalty)))
        if direction in ("backward", "both"):
            for feature in selected:
                trial = [kept for kept in selected if kept != feature]
                moves.append(("drop", feature, _criterion_value(design, target, trial, penalty)))

        if not moves:
            break
        action, feature, score = min(moves, key=lambda move: move[2])
        if score >= current:
            break

        selected = sorted([*selected, feature]) if action == "add" else [
            kept for kept in selected if kept != feature
        ]
        path.append(Step(action=action, feature=feature, criterion=score))  # type: ignore[arg-type]
        current = score

    return StepwiseResult(selected=tuple(selected), criterion=current, path=tuple(path))

"""Forward selection by incremental orthogonalization - one GEMM per step."""

import numpy as np
from numpy.typing import NDArray

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

COLLINEARITY_TOLERANCE = 1e-10
MIN_RSS = float(np.finfo(np.float64).tiny)


def forward_select(
    X: Matrix,
    y: Vector,
    penalty: float = 2.0,
    max_terms: int | None = None,
) -> list[int]:
    """Forward selection without a single refit.

    Each step keeps an orthonormal basis Q of the features chosen so far.
    Residualizing every candidate against Q turns "how much would RSS fall if I
    added column j" into a closed form: (q_j . r)^2, where q_j is that candidate
    after orthogonalization and r is the current residual. All p candidates are
    scored together, so the p refits per step become two matrix products.
    """
    if X.ndim != 2 or X.shape[0] != y.shape[0]:
        raise ValueError(f"shape mismatch: X={X.shape}, y={y.shape}")

    n, p = X.shape
    limit = min(p, n - 2) if max_terms is None else min(max_terms, p, n - 2)
    if limit <= 0:
        return []

    design = np.ascontiguousarray(X, dtype=np.float64)
    design = design - design.mean(axis=0)          # intercept via centring
    residual = np.ascontiguousarray(y, dtype=np.float64) - y.mean()
    column_energy = np.einsum("ij,ij->j", design, design)

    basis = np.empty((n, limit), dtype=np.float64)  # allocated once, filled per step
    available = np.ones(p, dtype=bool)
    selected: list[int] = []
    rss = max(float(residual @ residual), MIN_RSS)

    for step in range(limit):
        q = basis[:, :step]
        # Residualize every candidate against the current basis: two GEMMs.
        projected = design - q @ (q.T @ design)
        energy = np.einsum("ij,ij->j", projected, projected)
        usable = available & (energy > COLLINEARITY_TOLERANCE * column_energy)

        # The residual is already orthogonal to q, so projected.T @ residual
        # equals design.T @ residual and the extra product is skipped.
        covariance = design.T @ residual
        gain = np.full(p, -np.inf)
        gain[usable] = covariance[usable] ** 2 / energy[usable]

        best = int(np.argmax(gain))
        if not usable[best]:
            break

        new_rss = max(rss - gain[best], MIN_RSS)
        if n * np.log(new_rss / rss) + penalty >= 0.0:    # delta-AIC: no longer pays
            break

        direction = projected[:, best] / np.sqrt(energy[best])
        basis[:, step] = direction
        residual -= direction * (direction @ residual)    # in place
        rss = new_rss
        available[best] = False
        selected.append(best)

    return selected