core.py

"""
adaptive-eval — core IRT math (pure Python, zero dependencies)
=============================================================

Item Response Theory (2-parameter logistic) is the engine that lets you stop an
evaluation early: it models each benchmark item by a *difficulty* (b) and a
*discrimination* (a), and each model by a single *ability* (theta). Given those,
the Fisher information an item carries about a model's ability is
``a^2 * p * (1 - p)`` — highest for items whose difficulty sits near the model's
current ability estimate. Administer those first and you pin the ability (and
therefore the ranking) with a fraction of the items.

This module is deliberately small and readable — the math is old and clean; the
value is a usable, dependency-free implementation. See README for the reference
research (ATLAS, "Confident Rankings with Fewer Items").
"""

from __future__ import annotations

import math

Vec = list[float]


def prob(theta: float, a: float, b: float) -> float:
    """2PL probability of a correct response."""
    return 1.0 / (1.0 + math.exp(-a * (theta - b)))


def fisher_information(theta: float, a: float, b: float) -> float:
    """Fisher information item (a, b) carries about ability `theta`."""
    p = prob(theta, a, b)
    return a * a * p * (1.0 - p)


def estimate_theta(responses: list[int], items: list[tuple[float, float]],
                   *, iters: int = 40) -> tuple[float, float]:
    """MLE of ability from responses to (a, b) items, via Newton-Raphson.

    Returns (theta, standard_error). Ability is clamped to [-4, 4]; an all-correct
    or all-incorrect response set pins theta to a bound with a large SE (the
    honest signal that you can't yet separate this model).
    """
    if not responses:
        return 0.0, float("inf")
    theta = 0.0
    for _ in range(iters):
        num = 0.0   # first derivative of log-likelihood
        den = 0.0   # negative second derivative == total Fisher information
        for u, (a, b) in zip(responses, items):
            p = prob(theta, a, b)
            num += a * (u - p)
            den += a * a * p * (1.0 - p)
        if den < 1e-9:
            break
        step = num / den
        theta += max(-1.0, min(1.0, step))  # damped step for stability
        theta = max(-4.0, min(4.0, theta))
        if abs(step) < 1e-5:
            break
    info = sum(fisher_information(theta, a, b) for a, b in items)
    se = 1.0 / math.sqrt(info) if info > 1e-9 else float("inf")
    return theta, se


def _logit(p: float) -> float:
    p = min(0.98, max(0.02, p))
    return math.log(p / (1.0 - p))


def calibrate(matrix: list[list[int]], *, model_2pl: bool = True,
              iters: int = 300, lr: float = 0.02) -> list[tuple[float, float]]:
    """Estimate item parameters (a, b) from a models x items response matrix.

    Joint MLE: alternately estimate abilities and item parameters by gradient
    ascent, standardizing abilities each round for identifiability. Rasch (1PL,
    a fixed to 1) when ``model_2pl=False`` — more stable on small/ragged data.

    This is a solid baseline calibration; supply your own (a, b) if you have a
    better fit. Returns a list of (a, b) per item (column).
    """
    n_models = len(matrix)
    n_items = len(matrix[0]) if n_models else 0
    if n_models == 0 or n_items == 0:
        return []
    # init
    theta = [_logit(sum(row) / n_items) for row in matrix]
    col_mean = [sum(matrix[j][i] for j in range(n_models)) / n_models for i in range(n_items)]
    b = [-_logit(m) for m in col_mean]          # easier item (high p) -> lower difficulty
    a = [1.0] * n_items

    for _ in range(iters):
        # --- update abilities (a few Newton steps each) ---
        for j in range(n_models):
            items = list(zip(a, b))
            theta[j], _ = estimate_theta(matrix[j], items, iters=8)
        # standardize abilities (identifiability)
        mean = sum(theta) / n_models
        sd = math.sqrt(sum((t - mean) ** 2 for t in theta) / n_models) or 1.0
        theta = [(t - mean) / sd for t in theta]
        # --- update item params by gradient ascent ---
        for i in range(n_items):
            g_a = 0.0
            g_b = 0.0
            for j in range(n_models):
                p = prob(theta[j], a[i], b[i])
                resid = matrix[j][i] - p
                g_a += resid * (theta[j] - b[i])
                g_b += -a[i] * resid
            if model_2pl:
                a[i] = max(0.2, min(4.0, a[i] + lr * g_a / n_models))
            b[i] = max(-4.0, min(4.0, b[i] + lr * g_b / n_models))
    return list(zip(a, b))