path: root/annotate_pgns/extract_selfplay_positions.py

#!/usr/bin/env python3
"""Extract game-shaped records (game_id, result, positions[]) directly
from a self-play PGN written by eval_tune/match_play.py's
build_pgn_game -- a different source format from extract_positions.py
(which mines annotate.py's raw shard_NN.script/.out logs), so this is
a sibling script, not a variant of that one: the per-move score here
already lives as a PGN comment on each move
("{+-}NN.NN/DEPTH PV..." or "{+-}M/DEPTH PV..." for mate), not a
tellothers log line, and needs a PGN parser (python-chess) rather than
a log-line regex. Output schema matches build_games.py's games.jsonl
exactly (that script's own docstring anticipated this: "a future
loader for self-play games can produce the same games.jsonl schema
and both sources mix freely at training-set-assembly time") -- drops
straight into fit_k_and_label.py -> quiet_filter_games.py ->
build_training_tensors.py unchanged, no separate split_holdout.py/
build_games.py step needed since python-chess already gives real
per-game grouping directly.

2026-10-16, see docs/NNUE.md and the same day's qsearch-futility/NNUE-
calibration investigation: this is the tool for pulling trainable
positions out of the self-play corpus (pgn/self_play_games.pgn)
specifically to fill the "decisive material, not yet mate" gap the
TWIC/human-game corpus structurally can't supply (strong humans
resign long before that point; these self-play games, generated with
ResignThreshold disabled -- see match_play.py -- are deliberately
played to real checkmate/stalemate/50-move/EGTB conclusions instead).

Score sign convention matches extract_positions.py's output: each
row's "score" is relative to the side to move AT "fen" (not a fixed
White-POV), matching the training pipeline's existing convention
elsewhere. match_play.py's own comments are stored White-relative
(_format_eval_comment's score_white) -- converted back to mover-
relative here.

Per-move "mover" field records exactly which PGN White/Black header
name played that move (e.g. "candidate", "baseline",
"candidate_nnue_loose") -- lets a caller filter out any evaluator
whose own self-reported scores shouldn't be trusted as training
labels (see relabel_selfplay_pgn.py and the --exclude-movers flag
below) without needing to know which games/dates those came from.

Usage:
    python3 extract_selfplay_positions.py \\
        ../pgn/self_play_games.pgn \\
        --out runs/selfplay_1013_1014/games.jsonl \\
        --dates 2026-10-13 2026-10-14 \\
        --event "typhoon self-play match" \\
        --exclude-movers candidate_nnue_loose
"""
import argparse
import json
import re
import sys

import chess
import chess.pgn

COMMENT_RE = re.compile(
    r"^([+-]M|[+-]?\d+\.\d+)(?:/(\d+))?(?:\s+(.*))?$"
)

# match_play.py's own mate-token path ("+M"/"-M", _format_eval_comment)
# only fires when |score_white| >= NMATE_THRESHOLD=32000 (centipawns),
# a threshold keyed to typhoon's own NMATE sentinel. It is never reached
# when the mover is a *foreign* engine driven through a UCI<->xboard
# adapter (see src/stockfish_xboard -- Stockfish via PolyGlot), because
# PolyGlot's own UCI mate-score-to-xboard-cp conversion caps out far
# below that: empirically, mate-in-N on this corpus (35,295
# Stockfish-vs-Stockfish games, pgn/stockfish_self_play_training_data.pgn)
# always lands as a plain float in [+90.00, +99.99] (or the negated
# range), one per decisive game at depth 1 -- 5528 "+99.99/1" rows vs
# exactly 5528 "1-0" results, 2728 "-99.99/1" vs 2729 "0-1". A direct
# histogram of |score| confirms a clean gap: every value is either
# below 70.00 or in [99.00, 99.99] (51,488 rows total, matching the
# count of |score| >= 90 exactly) -- no ambiguous values near the
# boundary. So any foreign-engine corpus needs this magnitude check in
# addition to the literal "+M"/"-M" token case; MATE_CLAMP_THRESHOLD is
# set well below the observed 99.00 floor (and well above any real
# positional/material evaluation) to catch this with margin.
MATE_CLAMP_THRESHOLD = 90.0


def parse_comment(comment):
    """Returns (score_white_pawns_or_None, depth_or_None, pv_or_None,
    exact, exact_reason, outcome_white_or_None). score_white is None
    for mate ("+M"/"-M", distance not recoverable from this comment
    format -- see match_play.py's _format_eval_comment) -- exact=True
    covers that case the same way extract_positions.py treats MATE
    tokens, just without a distance number to carry. Also None for a
    foreign-engine mate score that was clamped into an ordinary-looking
    float instead of using match_play.py's own "+M"/"-M" convention
    (see MATE_CLAMP_THRESHOLD above) -- same treatment, different
    detection path.

    outcome_white is the exact, White-relative game-theoretic result
    (1.0/0.0/0.5) for exact rows where it can be read straight off this
    comment, independent of score_white (which is deliberately nulled
    out above for the mate/clamped cases -- the sign is still known
    even though the magnitude isn't, and that sign is exactly the
    label this field exists to carry). None when the comment doesn't
    determine an outcome on its own (ordinary non-exact rows, and the
    "unparseable" case) -- see NNUE.md's EGTB-ground-truth-mixing
    decision: this is deliberately NOT derived from the game's final
    result, since an exact-at-this-position verdict (TB/mate-forced)
    is a fact about the position, not about whether the game later
    actually reached that forced outcome."""
    m = COMMENT_RE.match(comment.strip())
    if not m:
        return None, None, None, True, f"unparseable:{comment!r}", None
    score_tok, depth_tok, pv = m.groups()
    depth = int(depth_tok) if depth_tok else None
    pv = pv.strip() if pv else None

    exact, reason = False, None
    if pv:
        last_move = pv.rsplit(" ", 1)[-1]
        if last_move in ("<EGTB>", "<REC>", "<=>", "<#>"):
            exact = True
            reason = {"<=>": "draw_recognizer",
                      "<#>": "mate_recognizer"}.get(
                          last_move, last_move.strip("<>").lower())

    if reason == "draw_recognizer":
        return None, depth, pv, True, reason, 0.5

    if score_tok in ("+M", "-M"):
        outcome_white = 1.0 if score_tok == "+M" else 0.0
        return None, depth, pv, True, (reason or "mate"), outcome_white

    score_white = float(score_tok)
    if abs(score_white) >= MATE_CLAMP_THRESHOLD:
        outcome_white = 1.0 if score_white > 0 else 0.0
        return None, depth, pv, True, (reason or "mate_clamped"), outcome_white

    if reason == "mate_recognizer":
        outcome_white = 1.0 if score_white > 0 else 0.0
        return score_white, depth, pv, True, reason, outcome_white

    return score_white, depth, pv, exact, reason, None


def main():
    ap = argparse.ArgumentParser(description=__doc__,
                                  formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("pgn_file")
    ap.add_argument("--out", required=True)
    ap.add_argument("--event", default=None,
                     help="only extract games with this Event header "
                          "(default: any)")
    ap.add_argument("--dates", nargs="*", default=None,
                     help="only extract games with Date header in this "
                          "list (default: any)")
    ap.add_argument("--exclude-movers", nargs="*", default=[],
                     help="skip any ply whose mover's PGN name is in "
                          "this list (e.g. candidate_nnue_loose) -- the "
                          "move/FEN sequence for the rest of that same "
                          "game is still usable, only that one side's "
                          "own score annotations are dropped")
    args = ap.parse_args()

    target_dates = set(args.dates) if args.dates else None
    exclude_movers = set(args.exclude_movers)

    total_games = 0
    games_in_scope = 0
    total_rows = 0
    total_exact = 0
    total_excluded_mover = 0

    with open(args.pgn_file, encoding="utf-8", errors="surrogateescape") as fin, \
         open(args.out, "w") as fout:
        while True:
            game = chess.pgn.read_game(fin)
            if game is None:
                break
            total_games += 1
            if args.event is not None and game.headers.get("Event") != args.event:
                continue
            if target_dates is not None and game.headers.get("Date") not in target_dates:
                continue
            games_in_scope += 1

            white_name = game.headers.get("White", "")
            black_name = game.headers.get("Black", "")
            game_result = game.headers.get("Result", "*")
            board = game.board()
            node = game
            game_positions = []
            ply_num = 0
            while node.variations:
                next_node = node.variations[0]
                move = next_node.move
                mover_is_white = board.turn == chess.WHITE
                mover_name = white_name if mover_is_white else black_name
                fen_before = board.fen()
                comment = next_node.comment
                board.push(move)
                node = next_node
                ply_num += 1
                if not comment:
                    continue
                if mover_name in exclude_movers:
                    total_excluded_mover += 1
                    continue
                (score_white, depth, pv, exact, reason,
                 outcome_white) = parse_comment(comment)
                score = None
                if score_white is not None:
                    score = score_white if mover_is_white else -score_white
                rec = {
                    "fen": fen_before,
                    "depth": depth,
                    "score": score,
                    "score_raw": comment,
                    "pv": pv,
                    "exact": exact,
                    "mover": mover_name,
                }
                if reason:
                    rec["exact_reason"] = reason
                if outcome_white is not None:
                    rec["r_outcome"] = (outcome_white if mover_is_white
                                         else 1.0 - outcome_white)
                game_positions.append(rec)
                total_rows += 1
                if exact:
                    total_exact += 1

            if not game_positions:
                continue
            game_rec = {
                "game_id": f"selfplay_{games_in_scope}_{game.headers.get('Date', '?')}_{ply_num}ply",
                "result": game_result,
                "positions": game_positions,
            }
            fout.write(json.dumps(game_rec) + "\n")

    print(f"{total_games} games scanned, {games_in_scope} in scope "
          f"(event={args.event!r}, dates={sorted(target_dates) if target_dates else None})",
          file=sys.stderr)
    print(f"{total_rows} positions written ({total_exact} exact), "
          f"{total_excluded_mover} plies skipped for excluded mover",
          file=sys.stderr)


if __name__ == "__main__":
    sys.exit(main())