spb/llmindex Public

1# llmindex.io — custom 2PL IRT fit (numpy MAP via Adam, missing-aware)2# Author:  Simon-Pierre Boucher3# Contact: contact@spboucher.ai4# License: Proprietary — © Simon-Pierre Boucher, all rights reserved5#6# 2PL: P(correct) = sigmoid(a_i * (theta_m - b_i))7# MAP estimation with priors:8#   theta ~ N(0, 1)      b ~ N(0, 1.5)      log a ~ N(0, 0.5)9# Missing responses are masked. Theta SE comes from the Fisher information10# (plus prior precision). Numpy-based (no torch) so the fit runs anywhere.1112from __future__ import annotations1314from dataclasses import dataclass, field1516import numpy as np171819def _sigmoid(x: np.ndarray) -> np.ndarray:20    return 1.0 / (1.0 + np.exp(-np.clip(x, -30, 30)))212223@dataclass24class Fit2PLResult:25    theta: np.ndarray  # (M,)26    theta_se: np.ndarray  # (M,)27    a: np.ndarray  # (I,) discrimination28    b: np.ndarray  # (I,) difficulty (logits)29    iterations: int30    converged: bool31    final_loglik: float32    diagnostics: dict = field(default_factory=dict)333435def fit_2pl(36    responses: np.ndarray,37    max_iterations: int = 500,38    tolerance: float = 1e-6,39    lr: float = 0.05,40    prior_theta_sd: float = 1.0,41    prior_b_sd: float = 1.5,42    prior_log_a_sd: float = 0.5,43    seed: int = 0,44) -> Fit2PLResult:45    """Fit a 2PL model on an (M models × I items) matrix of {0,1,nan}."""46    X = np.asarray(responses, dtype=float)47    if X.ndim != 2:48        raise ValueError("responses must be a 2D matrix (models × items)")49    M, I = X.shape50    mask = ~np.isnan(X)51    if mask.sum() == 0:52        raise ValueError("no observed responses")53    Xf = np.nan_to_num(X, nan=0.0)5455    rng = np.random.default_rng(seed)56    # Warm start: theta from row accuracy, b from item difficulty (logit of failure rate).57    row_acc = np.where(mask.sum(1) > 0, Xf.sum(1) / np.maximum(mask.sum(1), 1), 0.5)58    col_acc = np.where(mask.sum(0) > 0, Xf.sum(0) / np.maximum(mask.sum(0), 1), 0.5)59    theta = np.clip(np.log(row_acc + 1e-3) - np.log(1 - row_acc + 1e-3), -2, 2)60    theta = theta - theta.mean()61    b = np.clip(-(np.log(col_acc + 1e-3) - np.log(1 - col_acc + 1e-3)), -2.5, 2.5)62    log_a = rng.normal(0.0, 0.01, size=I)6364    # Adam state65    params = [theta, b, log_a]66    m_state = [np.zeros_like(p) for p in params]67    v_state = [np.zeros_like(p) for p in params]68    beta1, beta2, eps = 0.9, 0.999, 1e-86970    prev_obj = -np.inf71    converged = False72    it = 073    for it in range(1, max_iterations + 1):74        a = np.exp(log_a)75        Z = a[None, :] * (theta[:, None] - b[None, :])76        P = _sigmoid(Z)77        R = np.where(mask, Xf - P, 0.0)  # residuals on observed cells7879        g_theta = (R * a[None, :]).sum(1) - theta / prior_theta_sd**280        g_b = (-(R * a[None, :])).sum(0) - b / prior_b_sd**281        g_a = (R * (theta[:, None] - b[None, :])).sum(0)82        g_log_a = g_a * a - log_a / prior_log_a_sd**28384        grads = [g_theta, g_b, g_log_a]85        for j, (p, g) in enumerate(zip(params, grads)):86            m_state[j] = beta1 * m_state[j] + (1 - beta1) * g87            v_state[j] = beta2 * v_state[j] + (1 - beta2) * g * g88            mhat = m_state[j] / (1 - beta1**it)89            vhat = v_state[j] / (1 - beta2**it)90            p += lr * mhat / (np.sqrt(vhat) + eps)9192        # Identification: center abilities each step (scale is pinned by priors).93        shift = theta.mean()94        theta -= shift95        b -= shift9697        with np.errstate(divide="ignore", invalid="ignore"):98            ll = np.where(mask, Xf * np.log(P + 1e-12) + (1 - Xf) * np.log(1 - P + 1e-12), 0.0).sum()99        obj = (100            ll101            - 0.5 * (theta**2).sum() / prior_theta_sd**2102            - 0.5 * (b**2).sum() / prior_b_sd**2103            - 0.5 * (log_a**2).sum() / prior_log_a_sd**2104        )105        if abs(obj - prev_obj) < tolerance * (1 + abs(prev_obj)):106            converged = True107            break108        prev_obj = obj109110    a = np.exp(log_a)111    Z = a[None, :] * (theta[:, None] - b[None, :])112    P = _sigmoid(Z)113    info = (mask * (a[None, :] ** 2) * P * (1 - P)).sum(1) + 1.0 / prior_theta_sd**2114    theta_se = 1.0 / np.sqrt(info)115    final_ll = float(116        np.where(mask, Xf * np.log(P + 1e-12) + (1 - Xf) * np.log(1 - P + 1e-12), 0.0).sum()117    )118119    return Fit2PLResult(120        theta=theta,121        theta_se=theta_se,122        a=a,123        b=b,124        iterations=it,125        converged=converged,126        final_loglik=final_ll,127        diagnostics={128            "observed_cells": int(mask.sum()),129            "models": int(M),130            "items": int(I),131        },132    )133