adaptive.py

"""
adaptive-eval — adaptive test loop
==================================

Given calibrated items and a way to *administer* an item to a model
(``answer_fn(item_index) -> 0 | 1``), run a Computerized-Adaptive-Test:
repeatedly pick the item that carries the most Fisher information about the
model's current ability estimate, administer it, re-estimate, and stop as soon
as the ability is pinned tightly enough (or a budget is hit).

The model never sees the whole bank — that's the 10x saving.
"""

from __future__ import annotations

import random
from dataclasses import dataclass, field
from collections.abc import Callable

from .core import estimate_theta, fisher_information

Item = tuple[float, float]           # (discrimination a, difficulty b)
AnswerFn = Callable[[int], int]      # item_index -> 1 (correct) | 0 (incorrect)


@dataclass
class TestResult:
    theta: float
    se: float
    administered: list[int] = field(default_factory=list)   # item indices, in order
    responses: list[int] = field(default_factory=list)

    @property
    def n_items(self) -> int:
        return len(self.administered)


def adaptive_test(
    items: list[Item],
    answer_fn: AnswerFn,
    *,
    se_threshold: float = 0.30,
    max_items: int | None = None,
    min_items: int = 3,
) -> TestResult:
    """Run an adaptive test. Stops when SE(theta) <= se_threshold (ability pinned)
    or `max_items` is reached or the bank is exhausted."""
    remaining = set(range(len(items)))
    max_items = max_items or len(items)
    theta = 0.0
    order: list[int] = []
    responses: list[int] = []
    se = float("inf")

    while remaining and len(order) < max_items:
        # pick the most informative remaining item at the current ability estimate
        idx = max(remaining, key=lambda i: fisher_information(theta, *items[i]))
        remaining.discard(idx)
        u = int(answer_fn(idx))
        order.append(idx)
        responses.append(u)
        theta, se = estimate_theta(responses, [items[i] for i in order])
        if len(order) >= min_items and se <= se_threshold:
            break
    return TestResult(theta=theta, se=se, administered=order, responses=responses)


def fixed_test(
    items: list[Item],
    answer_fn: AnswerFn,
    *,
    n_items: int,
    seed: int = 0,
) -> TestResult:
    """Baseline: administer `n_items` chosen at random (what non-adaptive
    subsampling does). Same budget as an adaptive run, for a fair comparison."""
    rng = random.Random(seed)
    idxs = rng.sample(range(len(items)), min(n_items, len(items)))
    responses = [int(answer_fn(i)) for i in idxs]
    theta, se = estimate_theta(responses, [items[i] for i in idxs])
    return TestResult(theta=theta, se=se, administered=idxs, responses=responses)