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

1 files changed, 0 insertions, 885 deletions

diff --git a/src/nn_train.py b/src/nn_train.py
deleted file mode 100755
index ae78e43..0000000
--- a/src/nn_train.py
+++ /dev/null
@@ -1,885 +0,0 @@
-#!/usr/bin/env -S uv run --script
-#
-# /// script
-# requires-python = ">=3.14"
-# dependencies = [
-#   "numpy>=2.5.3",
-#   "torch>=2.14.0",
-#   "tqdm>=4.70.0",
-# ]
-#
-# [tool.uv.sources]
-# torch = { index = "pytorch-cpu" }
-#
-# [[tool.uv.index]]
-# name = "pytorch-cpu"
-# url = "https://download.pytorch.org/whl/cpu"
-# explicit = true
-# ///
-
-#
-# 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 argparse
-import concurrent.futures
-import json
-import math
-import os
-import random
-import subprocess
-from collections.abc import Iterator
-from dataclasses import dataclass
-from itertools import product
-from pathlib import Path
-
-import numpy as np
-import torch
-from torch import Tensor
-from torch.nn import Embedding, Linear, Module, MSELoss
-from tqdm import tqdm
-
-# ===================================================================
-# Net dimensions
-# ===================================================================
-
-# (input dim, hidden dim) per board size; must match include/weights_<size>.h
-NET_DIMS: dict[int, tuple[int, int]] = {
-    5: (27, 64),  # FIVE_INP_NUM, FIVE_DENSE_NUM
-    6: (38, 64),  # SIX_INP_NUM, SIX_DENSE_NUM
-}
-
-
-# ===================================================================
-# C weight formatting
-# ===================================================================
-
-
-def _fmt_flat(a: np.ndarray, per_line: int = 8) -> str:
-    a = np.ravel(np.asarray(a, dtype=np.float32))
-    lines: list[str] = []
-    for i in range(0, len(a), per_line):
-        chunk = ", ".join(f"{v:.9g}" for v in a[i : i + per_line])
-        if i + per_line < len(a):
-            chunk += ","
-        lines.append("    " + chunk)
-    return "\n".join(lines)
-
-
-def _fmt_matrix(mat: np.ndarray, per_line: int = 8) -> str:
-    mat = np.asarray(mat, dtype=np.float32)
-    rows: list[str] = []
-    for r in mat:
-        chunks = [
-            ", ".join(f"{v:.9g}" for v in r[i : i + per_line])
-            for i in range(0, len(r), per_line)
-        ]
-        rows.append("    {" + ",\n     ".join(chunks) + "},")
-    return "\n".join(rows)
-
-
-# ===================================================================
-# Config
-# ===================================================================
-
-
-@dataclass(frozen=True)
-class DataGenConfig:
-    size: int
-    depth: int
-    min_ply: int
-    games: int
-    procs: int
-    replay_iters: int
-    root: Path
-    train_only_self_play: bool
-    k: float
-
-
-@dataclass(frozen=True)
-class TrainingGridSearchConfig:
-    smooth_factor: list[float]
-    epochs: list[int]
-    batch: list[int]
-    lr: list[float]
-    val_frac: list[float]
-
-    def make_configs(self, data_gen_config: DataGenConfig) -> list["Candidate"]:
-        inp_dim, hidden = NET_DIMS[data_gen_config.size]
-        flds = self.__dataclass_fields__.keys()
-        almost_cands = [
-            dict(zip(flds, vals)) for vals in product(*(getattr(self, k) for k in flds))
-        ]
-        for ident, almost_cand in enumerate(almost_cands):
-            net = Net(inp_dim, hidden)
-            net.to_c(
-                size=data_gen_config.size,
-                ident=ident,
-                path=data_gen_config.root
-                / "include"
-                / f"cand_{ident}_weights_{data_gen_config.size}.c",
-                note=f"{ident=}, v000, random init",
-            )
-            cand_copy_for(
-                6 if data_gen_config.size == 5 else 5, ident, data_gen_config.root
-            )
-            almost_cand["net"] = net
-            almost_cand["train_only_self_play"] = data_gen_config.train_only_self_play
-        return [Candidate(**d) for d in almost_cands]
-
-
-@dataclass
-class Candidate:
-    train_only_self_play: bool
-    smooth_factor: float
-    epochs: int
-    batch: int
-    lr: float
-    val_frac: float
-    net: "Net"
-    elo: float = 1500.0
-    rated_games: int = 0
-
-    def __str__(self) -> str:
-        return f"(λ={self.smooth_factor}, E={self.epochs}, BS={self.batch}, η={self.lr:.2e}, V%={self.val_frac * 100.0:.1f})"
-
-    def training_pool(self, pool: Pool) -> Pool:
-        if self.train_only_self_play:
-            return pool.select_chains(pool.chain_self_play)
-        return pool
-
-    def fit(self, pool: Pool) -> tuple[float, float]:
-        # The train/val split is fixed but targets are refreshed
-        split = DataSplit.build(self.training_pool(pool), self.val_frac)
-        if split.n == 0:
-            return math.nan, math.nan
-
-        optimizer = torch.optim.AdamW(self.net.parameters(), lr=self.lr)
-        loss_fn = MSELoss()
-
-        pbar = tqdm(range(self.epochs), desc="train", unit="ep")
-        for _ in pbar:
-            split.refresh_targets(self.net, self.smooth_factor)
-
-            self.net.train()
-            for head, cells, parity, target in split.batches(self.batch):
-                optimizer.zero_grad()
-                predictions = self.net(head, cells) * parity
-                loss = loss_fn(predictions, target)
-                loss.backward()
-                optimizer.step()
-
-            with torch.no_grad():
-                self.net.eval()
-                val = split.val_idx
-                predictions = self.net(split.pool.head[val], split.pool.cells[val])
-                predictions = predictions * split.pool.parity[val]
-                targets = split.targets[val]
-                mse = loss_fn(predictions, targets).item()
-                corr = torch.corrcoef(torch.stack([predictions, targets]))[0, 1].item()
-
-            pbar.set_postfix(val_corr=f"{corr:.4f}", val_mse=f"{mse:.4f}")
-
-        return corr, mse
-
-
-# ===================================================================
-# Self-play records: games (for ratings) and chains (for training)
-# ===================================================================
-
-
-@dataclass(frozen=True)
-class Games:
-    black: np.ndarray  # [G] candidate id playing black
-    white: np.ndarray  # [G] candidate id playing white
-    result: np.ndarray  # [G] outcome from black's perspective, {-1, 0, 1}
-    newest: np.ndarray  # [G] game comes from the last loaded file
-
-    def head_to_head(self) -> Iterator[tuple[int, int, float]]:
-        for i in np.flatnonzero(self.newest & (self.black != self.white)):
-            black, white, result = (
-                int(self.black[i]),
-                int(self.white[i]),
-                float(self.result[i]),
-            )
-            # elo score: 1 win, 0.5 draw, 0 loss, from a's perspective
-            yield black, white, (result + 1.0) / 2.0
-
-
-@dataclass(frozen=True)
-class Pool:
-    head: Tensor  # [n, 2] fixed head inputs of the root state s_t
-    cells: Tensor  # [n, n_cells] cell states 0..6 of s_t
-    leaf_head: Tensor  # [n, 2] head inputs of the PV leaf pv_t
-    leaf_cells: Tensor  # [n, n_cells] cell states of pv_t
-    parity: Tensor  # [n] +1 when black is to move at s_t
-    leaf_parity: Tensor  # [n] +1 when black is to move at pv_t
-    terminal: Tensor  # [n] leaf is an exact terminal
-    terminal_value: Tensor  # [n] exact terminal value, black's perspective
-    chain_lens: Tensor  # [C] rows per chain
-    chain_result: Tensor  # [C] game result, black's perspective
-    chain_self_play: Tensor  # [C] chain comes from a self vs self game
-
-    @classmethod
-    def empty(cls, n_cells: int) -> "Pool":
-        def rows(dtype: torch.dtype, width: int = 0) -> Tensor:
-            return torch.empty((0, width) if width else (0,), dtype=dtype)
-
-        return cls(
-            head=rows(torch.float32, 2),
-            cells=rows(torch.uint8, n_cells),
-            leaf_head=rows(torch.float32, 2),
-            leaf_cells=rows(torch.uint8, n_cells),
-            parity=rows(torch.float32),
-            leaf_parity=rows(torch.float32),
-            terminal=rows(torch.bool),
-            terminal_value=rows(torch.float32),
-            chain_lens=rows(torch.int64),
-            chain_result=rows(torch.float32),
-            chain_self_play=rows(torch.bool),
-        )
-
-    @property
-    def n_rows(self) -> int:
-        return self.head.shape[0]
-
-    def select_chains(self, keep: Tensor) -> "Pool":
-        keep = keep.bool()
-        rows = torch.repeat_interleave(
-            torch.arange(self.chain_lens.shape[0], dtype=torch.int64), self.chain_lens
-        )
-        idx = torch.nonzero(keep[rows]).flatten()
-        return Pool(
-            head=self.head[idx],
-            cells=self.cells[idx],
-            leaf_head=self.leaf_head[idx],
-            leaf_cells=self.leaf_cells[idx],
-            parity=self.parity[idx],
-            leaf_parity=self.leaf_parity[idx],
-            terminal=self.terminal[idx],
-            terminal_value=self.terminal_value[idx],
-            chain_lens=self.chain_lens[keep],
-            chain_result=self.chain_result[keep],
-            chain_self_play=self.chain_self_play[keep],
-        )
-
-
-# ===================================================================
-# Reading game records
-# ===================================================================
-
-_ROW_FIELDS = (
-    "head",
-    "cells",
-    "leaf_head",
-    "leaf_cells",
-    "parity",
-    "leaf_parity",
-    "terminal",
-    "terminal_value",
-)
-
-
-def _chain_arrays(
-    rows: list[dict], n_cells: int, source: Path
-) -> dict[str, np.ndarray]:
-
-    def column(part: str, key: str, dtype: np.dtype) -> np.ndarray:
-        return np.array([r[part][key] for r in rows], dtype=dtype)
-
-    out = {
-        "head": column("root", "head", np.float32),
-        "cells": column("root", "cells", np.uint8),
-        "leaf_head": column("leaf", "head", np.float32),
-        "leaf_cells": column("leaf", "cells", np.uint8),
-        "parity": column("root", "parity", np.float32),
-        "leaf_parity": column("leaf", "parity", np.float32),
-        "terminal": column("leaf", "terminal", np.bool_),
-        "terminal_value": column("leaf", "terminal_value", np.float32),
-    }
-    for name in ("cells", "leaf_cells"):
-        if out[name].shape[1] != n_cells:
-            raise ValueError(
-                f"{source}: {name} width {out[name].shape[1]}, expected {n_cells}"
-            )
-    return out
-
-
-def _game_colours(game: dict) -> tuple[int, int]:
-    a, b = int(game["cand_A"]), int(game["cand_B"])
-    for r in game["rows"]:
-        if r["root"]["parity"] > 0:
-            cand = int(r["cand"])
-            return cand, b if cand == a else a
-    recorded = {int(r["cand"]) for r in game["rows"]}
-    return (a, b) if a not in recorded else (b, a)
-
-
-def load_window(
-    files: list[Path], size: int, n_candidates: int
-) -> tuple[list[Pool], Games]:
-    n_cells = NET_DIMS[size][0] - 2
-    blocks = [{f: [] for f in _ROW_FIELDS} for _ in range(n_candidates)]
-    chain_lens: list[list[int]] = [[] for _ in range(n_candidates)]
-    chain_result: list[list[float]] = [[] for _ in range(n_candidates)]
-    chain_self_play: list[list[bool]] = [[] for _ in range(n_candidates)]
-    black, white, result, newest = [], [], [], []
-
-    for file_index, path in enumerate(files):
-        is_newest = file_index == len(files) - 1
-        with open(path) as f:
-            for line in f:
-                if not line.strip():
-                    continue
-                game = json.loads(line)
-                game_result = float(game["result"])
-                b, w = _game_colours(game)
-                black.append(b)
-                white.append(w)
-                result.append(game_result)
-                newest.append(is_newest)
-
-                by_cand: dict[int, list[dict]] = {}
-                for r in game["rows"]:
-                    by_cand.setdefault(int(r["cand"]), []).append(r)
-
-                for cand, rows in by_cand.items():
-                    parity = [r["root"]["parity"] for r in rows]
-                    for sign in (1, -1):
-                        chain = [
-                            r for r, p in zip(rows, parity) if (p > 0) == (sign > 0)
-                        ]
-                        if not chain:
-                            continue
-                        for name, array in _chain_arrays(chain, n_cells, path).items():
-                            blocks[cand][name].append(array)
-                        chain_lens[cand].append(len(chain))
-                        chain_result[cand].append(game_result)
-                        chain_self_play[cand].append(b == w)
-
-    def pool_of(cand: int) -> Pool:
-        if not chain_lens[cand]:
-            return Pool.empty(n_cells)
-        arrays = {
-            name: torch.from_numpy(np.concatenate(blocks[cand][name]))
-            for name in _ROW_FIELDS
-        }
-        return Pool(
-            **arrays,
-            chain_lens=torch.tensor(chain_lens[cand], dtype=torch.int64),
-            chain_result=torch.tensor(chain_result[cand], dtype=torch.float32),
-            chain_self_play=torch.tensor(chain_self_play[cand], dtype=torch.bool),
-        )
-
-    games = Games(
-        black=np.array(black, dtype=np.int16),
-        white=np.array(white, dtype=np.int16),
-        result=np.array(result, dtype=np.float32),
-        newest=np.array(newest, dtype=np.bool_),
-    )
-    return [pool_of(c) for c in range(n_candidates)], games
-
-
-# ===================================================================
-# Training data: fixed rows, refreshed targets
-# ===================================================================
-
-
-@dataclass
-class DataSplit:
-    pool: Pool
-    train_idx: Tensor
-    val_idx: Tensor
-    targets: Tensor  # [n] aligned with pool rows, refreshed per epoch
-    chain_end: Tensor  # [C] last row index of each chain
-    wave_order: Tensor  # [n] rows ordered by distance from their chain's end
-    wave_bounds: Tensor  # [d] start of distance d in wave_order
-
-    @property
-    def n(self) -> int:
-        return self.pool.n_rows
-
-    @classmethod
-    def build(cls, pool: Pool, val_frac: float) -> "DataSplit":
-        n = pool.n_rows
-        if n == 0:
-            empty = torch.empty(0, dtype=torch.int64)
-            return cls(pool, empty, empty, torch.empty(0), empty, empty, empty)
-
-        # Smoothing is G[t] = (1-λ)V[t] + λG[t+1] with G at the chain end.
-        # Chains are contiguous, so every row that is d plies from its chain's
-        # end can be updated in one vectorised wavefront step per d.
-        lens = pool.chain_lens
-        chain_end = torch.cumsum(lens, dim=0) - 1
-        chain_of_row = torch.repeat_interleave(
-            torch.arange(lens.shape[0], dtype=torch.int64), lens
-        )
-        dist = chain_end[chain_of_row] - torch.arange(n, dtype=torch.int64)
-        wave_order = torch.argsort(dist, stable=True)
-        wave_bounds = torch.searchsorted(
-            dist[wave_order], torch.arange(int(dist.max()) + 2, dtype=torch.int64)
-        )
-
-        n_val = max(1, int(n * val_frac))
-        perm = torch.randperm(n)
-        return cls(
-            pool=pool,
-            train_idx=perm[n_val:],
-            val_idx=perm[:n_val],
-            targets=torch.zeros(n),
-            chain_end=chain_end,
-            wave_order=wave_order,
-            wave_bounds=wave_bounds,
-        )
-
-    def refresh_targets(self, net: Module, lam: float) -> None:
-        pool = self.pool
-        net.eval()
-        with torch.no_grad():
-            # Net(pv_t), black's perspective; exact terminals override the net
-            leaf = net(pool.leaf_head, pool.leaf_cells) * pool.leaf_parity
-            leaf = torch.where(pool.terminal, pool.terminal_value, leaf)
-
-        targets = torch.zeros_like(leaf)
-        targets[self.chain_end] = pool.chain_result
-        for lo, hi in zip(self.wave_bounds[1:-1], self.wave_bounds[2:]):
-            idx = self.wave_order[lo:hi]
-            targets[idx] = (1.0 - lam) * leaf[idx] + lam * targets[idx + 1]
-        self.targets = targets
-
-    def batches(
-        self, batch_size: int
-    ) -> Iterator[tuple[Tensor, Tensor, Tensor, Tensor]]:
-        order = torch.randperm(self.train_idx.shape[0])
-        train_targets = self.targets[self.train_idx]
-        for start in range(0, order.shape[0], batch_size):
-            idx = order[start : start + batch_size]
-            rows = self.train_idx[idx]
-            yield (
-                self.pool.head[rows],
-                self.pool.cells[rows],
-                self.pool.parity[rows],
-                train_targets[idx],
-            )
-
-
-# ===================================================================
-# Network
-# ===================================================================
-
-
-def fast_tanh(x: Tensor) -> Tensor:
-    # Must match fast_tanh in include/nn.c: degree-11 fit of tanh on [-3, 3],
-    # saturating to +-1 outside.
-    x2 = x * x
-    poly = x * (
-        0.997136032
-        + x2
-        * (
-            -0.305215124
-            + x2
-            * (
-                0.0832466371
-                + x2 * (-0.0139947924 + x2 * (0.00123106314 + x2 * -4.29748644e-05))
-            )
-        )
-    )
-    return torch.where(x > 3.0, 1.0, torch.where(x < -3.0, -1.0, poly))
-
-
-def fast_silu(x: Tensor) -> Tensor:
-    # Must match fast_silu in include/nn.c: [4,2] Pade approximant with cutoff.
-    # The cutoff below is evaluated at x = 0 (instead of x) because torch
-    # evaluates both branches of where: 1 + 0.1 * x is 0 for x = -10.
-    xs = torch.where(x < -4.6, 0.0, x)
-    poly = xs * (0.5 + xs * (0.25 + xs * (0.05 + 0.004167 * xs))) / (1 + 0.1 * xs)
-    return torch.where(x < -4.6, 0.0, torch.where(x > 4.6, x, poly))
-
-
-class Net(Module):
-    def __init__(self, inp_dim: int, hidden: int) -> None:
-        super().__init__()
-        # learned cell-state embedding, see nn.c
-        self.embedding = Embedding(7, 1)
-        self.hidden = Linear(inp_dim, hidden, bias=True)
-        self.output = Linear(hidden, 1, bias=True)
-
-    def forward(self, head: Tensor, cells: Tensor) -> Tensor:
-        # head: [n, 2] fixed head inputs; cells: [n, n_cells] raw cell states
-        # (0..6) - the only place they become anything other than integers
-        inp = torch.cat([head, self.embedding(cells.long()).flatten(1)], dim=1)
-        hidden = fast_silu(self.hidden(inp))
-        return fast_tanh(self.output(hidden)).squeeze(-1)
-
-    def to_c(
-        self,
-        *,
-        size: int,
-        ident: int | None,
-        path: Path,
-        note: str,
-    ) -> None:
-        upper = "FIVE" if size == 5 else "SIX"
-        prefix = (f"cand_{ident}_" if ident is not None else "") + upper
-        guard = f"WEIGHTS_{upper}" + (f"_{ident}" if ident is not None else "")
-
-        E = self.embedding.weight.detach().cpu().numpy()
-        W = self.hidden.weight.detach().cpu().numpy()
-        B = self.hidden.bias.detach().cpu().numpy()
-        LW = self.output.weight.detach().cpu().numpy().ravel()
-        LB = float(self.output.bias.detach().cpu().item())
-
-        with open(path, "w") as f:
-            lines = [
-                "/* generated by nn_train.py, do not edit by hand */\n",
-                f"/* {note} */\n",
-                f"#ifndef {guard}\n#define {guard}\n",
-                f'#include "weights_{size}.h"\n\n',
-                f"const float {prefix}_embedding[{len(E)}] =\n{{\n{_fmt_flat(E, per_line=7)}\n}};\n\n",
-                f"const float {prefix}_dense_weights[{upper}_INP_NUM][{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_matrix(W.T)}\n}};\n\n",
-                f"const float {prefix}_dense_biases[{upper}_DENSE_NUM] =\n{{\n{_fmt_flat(B)}\n}};\n\n",
-                f"const float {prefix}_logit_weights[{upper}_DENSE_NUM] NN_ALIGNED =\n{{\n{_fmt_flat(LW)}\n}};\n\n",
-                f"const float {prefix}_logit_bias = {LB:.9g};\n",
-                "#endif",
-            ]
-            for line in lines:
-                _ = f.write(line)
-
-
-# ===================================================================
-# Process helpers
-# ===================================================================
-
-
-def fan_out(
-    cmds: list[list[str]], n: int, procs: int, tag: str, root: Path
-) -> Iterator[str]:
-    with concurrent.futures.ThreadPoolExecutor(max_workers=procs) as pool:
-        futs = [
-            pool.submit(
-                subprocess.run,
-                cmd,
-                cwd=root,
-                check=True,
-                capture_output=True,
-                text=True,
-            )
-            for _ in range(n)
-            for cmd in cmds
-        ]
-        for fut in tqdm(
-            concurrent.futures.as_completed(futs),
-            total=n * len(cmds),
-            desc=tag,
-            unit="game",
-        ):
-            yield fut.result().stdout
-
-
-def cand_copy_for(size: int, id: int, root: Path) -> None:
-    src = root / "include" / f"weights_{size}.c"
-    if not src.is_file():
-        raise SystemExit(f"{src} not found; build a size-{size} champion first")
-    upper = "FIVE" if size == 5 else "SIX"
-    txt = src.read_text()
-    for sym in (
-        "embedding",
-        "dense_weights",
-        "dense_biases",
-        "logit_weights",
-        "logit_bias",
-    ):
-        txt = txt.replace(f"{upper}_{sym}", f"cand_{id}_{upper}_{sym}")
-    txt = txt.replace(f"WEIGHTS_{upper}", f"WEIGHTS_{upper}_{id}")
-    _ = (root / "include" / f"cand_{id}_weights_{size}.c").write_text(txt)
-
-
-# ===================================================================
-# The bootstrap loop
-# ===================================================================
-
-
-def update_elo(a: Candidate, b: Candidate, score: float, k: float) -> None:
-    # score is a's result as an elo score: 1.0 win, 0.5 draw, 0.0 loss
-    ea = 1.0 / (1.0 + 10.0 ** ((b.elo - a.elo) / 400.0))
-    delta = k * (score - ea)
-    a.elo += delta
-    b.elo -= delta
-
-
-def run_step(
-    *,
-    candidates: list[Candidate],
-    data_gen_config: DataGenConfig,
-    iter: int,
-):
-    def data_path(i: int) -> Path:
-        return (
-            data_gen_config.root
-            / "data"
-            / f"selfplay-{data_gen_config.size}-v{i:03d}.json"
-        )
-
-    print(f"\n== iteration {iter} ==")
-
-    num_candidates = len(candidates)
-    cand_pairs = [
-        (i, j) for i in range(num_candidates) for j in range(i, num_candidates)
-    ]
-
-    # 1. rebuild
-    for i, j in cand_pairs:
-        _ = subprocess.run(
-            ["make", f"selfplay-{i}-{j}"],
-            cwd=data_gen_config.root,
-            check=True,
-            capture_output=True,
-        )
-
-    # 2. generate self-play data
-    data = data_path(iter)
-    print(
-        f"selfplay: {num_candidates=}, {data_gen_config.games=}, "
-        f"{data_gen_config.size=}, {data_gen_config.depth=}, {data_gen_config.min_ply=}, "
-        f"{data_gen_config.train_only_self_play=}"
-    )
-    out_games: list[str] = []
-    for out in fan_out(
-        [
-            [
-                f"./selfplay-{i}-{j}",
-                str(data_gen_config.size),
-                str(data_gen_config.depth),
-                str(data_gen_config.min_ply),
-            ]
-            for (i, j) in cand_pairs
-        ],
-        data_gen_config.games,
-        data_gen_config.procs,
-        "selfplay",
-        data_gen_config.root,
-    ):
-        out_games.extend(out.splitlines())
-    print(f"writing {len(out_games):,} games to {data}")
-    with open(data, "w") as f:
-        _ = f.write("\n".join(out_games))
-        if out_games:
-            _ = f.write("\n")
-
-    # 3. rate every head-to-head game played in this iteration
-    data_files = [
-        data_path(i)
-        for i in range(max(iter - data_gen_config.replay_iters, 1), iter + 1)
-    ]
-    print("Loading data from", [str(c) for c in data_files], "...", end=" ")
-    pools, games = load_window(data_files, data_gen_config.size, num_candidates)
-    print("candidates have", [f"{p.n_rows:,}" for p in pools], "rows.")
-
-    fresh = list(games.head_to_head())
-    random.shuffle(fresh)
-    for a, b, score in fresh:
-        update_elo(candidates[a], candidates[b], score, data_gen_config.k)
-        candidates[a].rated_games += 1
-        candidates[b].rated_games += 1
-
-    total_rated_games = sum(c.rated_games for c in candidates)
-    print(
-        f"Tournament standings ({total_rated_games:,} rated games, {total_rated_games // len(candidates):,} games/candidate):"
-    )
-    for idx, cand in sorted(
-        enumerate(candidates), key=lambda t: t[1].elo, reverse=True
-    ):
-        print(f"\tcandidate {idx:03d}: Elo {cand.elo:4.1f} Config {cand}")
-
-    # 4. write champ, unconditionally
-    idx, champ = max(enumerate(candidates), key=lambda t: t[1].elo)
-    path = data_gen_config.root / "include" / f"weights_{data_gen_config.size}.c"
-    champ.net.to_c(
-        size=data_gen_config.size,
-        ident=None,
-        path=path,
-        note=f"{iter=}, elo={champ.elo:.1f}, rated_games={champ.rated_games:,}",
-    )
-    _ = subprocess.run(
-        ["make", "clean"], cwd=data_gen_config.root, check=True, capture_output=True
-    )
-    _ = subprocess.run(
-        ["make", "native"], cwd=data_gen_config.root, check=True, capture_output=True
-    )
-    print(f"Wrote candidate {idx:03d} to {path!s}")
-
-    # 5. fine-tune on the hybrid targets -> candidate
-    for idx, cand in enumerate(candidates):
-        corr, mse = cand.fit(pools[idx])
-        cand.net.to_c(
-            size=data_gen_config.size,
-            ident=idx,
-            path=data_gen_config.root
-            / "include"
-            / f"cand_{idx}_weights_{data_gen_config.size}.c",
-            note=f"id={idx}, v{iter:03d}, val_corr={corr:.4f}, val_mse={mse:.4f}",
-        )
-
-
-# ===================================================================
-# Main
-# ===================================================================
-
-
-def main() -> None:
-    root = Path(__file__).resolve().parent.parent
-    p = argparse.ArgumentParser(
-        description="Train the neural network evaluators of ctak",
-        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
-    )
-    _ = p.add_argument(
-        "--size",
-        type=int,
-        choices=sorted(NET_DIMS),
-        default=5,
-        help="board size",
-    )
-    _ = p.add_argument(
-        "--iters",
-        type=int,
-        default=100,
-        help="max loop iterations",
-    )
-    _ = p.add_argument(
-        "--games",
-        type=int,
-        default=200,
-        help="games per iteration",
-    )
-    _ = p.add_argument(
-        "--replay-iters",
-        type=int,
-        default=10,
-        help="grid points for number of previous data sources to retain",
-    )
-    _ = p.add_argument(
-        "--smooth-factor",
-        type=float,
-        nargs="+",
-        default=[0.3, 0.5, 0.7, 0.9],
-        help="grid points for the EMA smoothing factor",
-    )
-    _ = p.add_argument(
-        "--epochs",
-        type=int,
-        nargs="+",
-        default=[5],
-        help="grid points for candidate training epochs",
-    )
-    _ = p.add_argument(
-        "--lr",
-        type=float,
-        nargs="+",
-        default=[1e-4],
-        help="grid points for learning rate",
-    )
-    _ = p.add_argument(
-        "--batch",
-        type=int,
-        nargs="+",
-        default=[128],
-        help="grid points for training batch size",
-    )
-    _ = p.add_argument(
-        "--val-frac",
-        type=float,
-        nargs="+",
-        default=[0.05],
-        help="grid points for fraction of rows held out for validation",
-    )
-    _ = p.add_argument(
-        "--procs",
-        type=int,
-        default=None,
-        help="parallel game processes (number of CPUs if unset)",
-    )
-    _ = p.add_argument(
-        "--depth",
-        type=int,
-        default=4,
-        help="search depth for selfplay and the matches",
-    )
-    _ = p.add_argument(
-        "--min-ply",
-        type=int,
-        default=8,
-        help="length of random opening in games",
-    )
-    _ = p.add_argument(
-        "--train-all-games",
-        action="store_true",
-        help="train on all selfplay games, including matches between "
-        "different candidates (default: self vs self games only)",
-    )
-    _ = p.add_argument(
-        "--k-factor",
-        type=float,
-        default=32.0,
-        help="elo k-factor",
-    )
-
-    args = p.parse_args()
-
-    # avoid counter-play start
-    if args.min_ply < 2:
-        _ = p.error("--min-ply must be >= 2")
-
-    if args.procs is None:
-        args.procs = max(1, (os.cpu_count() or 1))
-
-    data_gen = DataGenConfig(
-        size=args.size,
-        depth=args.depth,
-        min_ply=args.min_ply,
-        games=args.games,
-        procs=args.procs,
-        replay_iters=args.replay_iters,
-        root=root,
-        train_only_self_play=not args.train_all_games,
-        k=args.k_factor,
-    )
-    tgsc = TrainingGridSearchConfig(
-        smooth_factor=args.smooth_factor,
-        epochs=args.epochs,
-        batch=args.batch,
-        lr=args.lr,
-        val_frac=args.val_frac,
-    )
-
-    candidates = tgsc.make_configs(data_gen)
-
-    (root / "data").mkdir(exist_ok=True)
-
-    num_params = sum(
-        p.numel() for p in candidates[0].net.parameters() if p.requires_grad
-    )
-    print(f"Number of trainable parameters: {num_params:,}")
-
-    for iter in range(1, args.iters + 1):
-        run_step(
-            candidates=candidates,
-            data_gen_config=data_gen,
-            iter=iter,
-        )
-    print(f"\n== done, champion in include/weights_{args.size}.c ==")
-
-
-if __name__ == "__main__":
-    main()