factor out nn_train.py into a separate python project to make it more manageable, add plotting and offline strength fit

1 files changed, 118 insertions, 0 deletions

diff --git a/train/src/ctak_train/strength.py b/train/src/ctak_train/strength.py
new file mode 100644
index 0000000..2302606
--- /dev/null
+++ b/train/src/ctak_train/strength.py
@@ -0,0 +1,118 @@
+#
+# This file is part of ctak.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with ctak. If not, see <https://www.gnu.org/licenses/>.
+
+import math
+from pathlib import Path
+
+import matplotlib
+import matplotlib.pyplot as plt
+import numpy as np
+from ctak_train.candidate import Candidate
+from ctak_train.data import GameLog
+from scipy.optimize import minimize
+
+matplotlib.use("Agg")
+
+# Elo points per unit of logit strength
+ELO_PER_LOGIT = 400.0 / math.log(10.0)
+
+# Spread of the zero-mean prior on the ratings, in logit units. Head-to-head
+# results alone leave the likelihood unbounded while the data is thin (a pair
+# that has never split draws runs off to +-inf).
+PRIOR_SD_LOGIT = 1.0
+
+
+def fit_strengths(log: GameLog, n_candidates: int) -> tuple[np.ndarray, np.ndarray]:
+    """Bradley-Terry logistic model (draws scored 0.5) fit in natural-logit
+    units, mean anchored to 0, reported as Elo points around 1500. Returns
+    (strengths, +-1 sd).
+    """
+    rated = log.black != log.white
+    black = log.black[rated].astype(np.int64)
+    white = log.white[rated].astype(np.int64)
+    if n_candidates < 2 or black.size == 0:
+        return np.full(n_candidates, 1500.0), np.zeros(n_candidates)
+    score = (log.result[rated].astype(np.float64) + 1.0) / 2.0
+
+    # n_candidates - 1 free parameters; the last candidate anchors the mean to 0
+    design = np.eye(n_candidates, n_candidates - 1)
+    design[n_candidates - 1] = -1.0
+
+    def nll_grad(theta: np.ndarray) -> tuple[float, np.ndarray]:
+        rating = design @ theta
+        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
+        p = 1.0 / (1.0 + np.exp(-x))
+        sp = np.maximum(x, 0.0) + np.log1p(np.exp(-np.abs(x)))  # softplus(x)
+        sp_neg = np.maximum(-x, 0.0) + np.log1p(np.exp(-np.abs(x)))
+        # nll = score*softplus(-x) + (1-score)*softplus(x): finite for all x
+        nll = float(np.sum(score * sp_neg + (1.0 - score) * sp))
+        nll += 0.5 * float(rating @ rating) / PRIOR_SD_LOGIT**2
+
+        grad_rating = rating / PRIOR_SD_LOGIT**2
+        np.add.at(grad_rating, black, p - score)
+        np.add.at(grad_rating, white, score - p)
+        return nll, design.T @ grad_rating
+
+    def hessian(theta: np.ndarray) -> np.ndarray:
+        rating = design @ theta
+        x = np.clip(rating[black] - rating[white], -350.0, 350.0)
+        p = 1.0 / (1.0 + np.exp(-x))
+        w = p * (1.0 - p)
+
+        diag = np.bincount(black, weights=w, minlength=n_candidates) + np.bincount(
+            white, weights=w, minlength=n_candidates
+        )
+        off = np.bincount(
+            black * n_candidates + white, weights=-w, minlength=n_candidates**2
+        ).reshape(n_candidates, n_candidates)
+        hess_rating = (
+            off
+            + np.diag(diag - np.diag(off))
+            + np.eye(n_candidates) / PRIOR_SD_LOGIT**2
+        )
+        return design.T @ hess_rating @ design
+
+    res = minimize(nll_grad, np.zeros(n_candidates - 1), jac=True, method="L-BFGS-B")
+    if not res.success:
+        print(f"strength fit did not converge: {res.message}")
+    cov = design @ np.linalg.pinv(hessian(res.x)) @ design.T
+    se = np.sqrt(np.maximum(np.diag(cov), 0.0)) * ELO_PER_LOGIT
+    return 1500.0 + (design @ res.x) * ELO_PER_LOGIT, se
+
+
+def plot_strengths(
+    history: list[np.ndarray], candidates: list[Candidate], path: Path
+) -> None:
+    fig, ax = plt.subplots(figsize=(10, 6))
+    iterations = range(1, len(history) + 1)
+    for idx, cand in enumerate(candidates):
+        ax.plot(
+            iterations,
+            [row[idx] for row in history],
+            marker=".",
+            ms=3,
+            lw=1.2,
+            label=f"cand {idx:03d} {cand}",
+        )
+    ax.axhline(1500.0, color="gray", lw=0.5)
+    ax.set_xscale("log")
+    ax.set_xlabel("Iteration")
+    ax.set_ylabel("Strength (fit on all games so far, mean 1500)")
+    ax.set_title("Candidate strength, refitted on the full history each iteration")
+    ax.legend(fontsize=8)
+    fig.tight_layout()
+    _ = fig.savefig(path, dpi=150)
+    plt.close(fig)