#!/usr/bin/env python3 """Build a larger cross-sectional position pool (successor to tests/cross_sectional_200k.ep_) out of the full twic.pgn + lichess_elite*.pgn corpus. Unlike sample_eval_calibration_positions.py's random-byte-offset sampling (one position per game, uniform-random ply), this does an exhaustive streaming pass over every game and pulls 5 positions per qualifying game at fixed fractional points in the game (early/early-mid/mid/mid-late/late) -- Scott's own framing. Fractional- ply sampling structurally balances game-stage representation (every kept game contributes exactly one position to each of the 5 stage buckets), which random-ply sampling does not. Deliberately does NOT filter for quiescence -- these positions are `go` roots for the engine's own search (which resolves tactics via its own qsearch at every node), not Texel-tuning eval-only positions. See sample_eval_calibration_positions.py's own docstring for the same reasoning: a cross-sectional pool should look like "all positions a real search tree touches", not just the quiet subset. Ply fractions: 5 samples come from 6 equal segments (1/6..5/6 of the game's ply count), not literal 1/5..5/5 -- 5/5 is the final move of the game, frequently checkmate/stalemate/resignation-adjacent, which `go` can't search meaningfully. Interior boundaries only. """ import argparse import collections import io import sys from pathlib import Path import chess import chess.pgn STAGE_NAMES = ["early", "early_mid", "mid", "mid_late", "late"] FRACTIONS = [1 / 6, 2 / 6, 3 / 6, 4 / 6, 5 / 6] PIECE_VALUES = {chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9} def material_bucket(board): """Same bucketing as eval.c's _MaterialBucket / the calibration script's own copy: combined non-pawn army value (pawn units) / 8, capped at 7 (0 = bare kings, 7 = full army).""" total = 0 for pt, v in PIECE_VALUES.items(): total += v * len(board.pieces(pt, chess.WHITE)) total += v * len(board.pieces(pt, chess.BLACK)) return min(7, total // 8) def iter_game_texts(pgn_path): """Stream raw PGN text blocks, one game at a time, without loading the whole file -- these corpora run into the GB range (twic.pgn alone is 3.5GB).""" with open(pgn_path, "r", encoding="utf-8", errors="replace") as f: header_lines = [] movetext_lines = [] in_headers = True for line in f: if line.startswith("["): if not in_headers and (header_lines or movetext_lines): yield "".join(header_lines) + "\n" + "".join(movetext_lines) header_lines, movetext_lines = [], [] header_lines.append(line) in_headers = True elif line.strip() == "": if header_lines and not movetext_lines: in_headers = False else: in_headers = False movetext_lines.append(line) if header_lines or movetext_lines: yield "".join(header_lines) + "\n" + "".join(movetext_lines) def normalized_fen_key(fen): """Board+side-to-move+castling+ep, dropping halfmove/fullmove counters -- used only for dedup accounting, never written out.""" parts = fen.split() return " ".join(parts[:4]) def process_pgn(path, min_elo, min_ply, stats, seen_keys, out_f): kept_games = 0 scanned_games = 0 for raw in iter_game_texts(path): scanned_games += 1 game = chess.pgn.read_game(io.StringIO(raw)) if game is None: continue headers = game.headers try: we = int(headers.get("WhiteElo", "0")) be = int(headers.get("BlackElo", "0")) except ValueError: we = be = 0 if min_elo and (we < min_elo or be < min_elo): stats["rejected_elo"] += 1 continue moves = list(game.mainline_moves()) if len(moves) < min_ply: stats["rejected_short"] += 1 continue result = headers.get("Result", "*") stats["games_by_result"][result] += 1 stats["source_games"][path.name] += 1 kept_games += 1 board = game.board() move_idx = 0 boundaries = [round(f * len(moves)) for f in FRACTIONS] # de-dup consecutive identical boundaries (very short-but-kept # games can round two fractions to the same ply) seen_boundaries = set() bi = 0 for target_ply in boundaries: target_ply = max(1, min(target_ply, len(moves) - 1)) while move_idx < target_ply: board.push(moves[move_idx]) move_idx += 1 stage = STAGE_NAMES[bi] bi += 1 if target_ply in seen_boundaries: stats["skipped_duplicate_boundary"] += 1 continue seen_boundaries.add(target_ply) if board.is_game_over(): stats["skipped_game_over"] += 1 continue fen = board.fen() key = normalized_fen_key(fen) stats["stage_total"][stage] += 1 stats["material_bucket"][material_bucket(board)] += 1 stats["side_to_move"]["white" if board.turn else "black"] += 1 if key in seen_keys: stats["duplicate_fens"] += 1 stats["stage_duplicates"][stage] += 1 continue seen_keys.add(key) out_f.write(f"setboard {fen}\n") out_f.write("go\n") stats["written"] += 1 stats["stage_written"][stage] += 1 stats["scanned_games"] += scanned_games stats["kept_games"] += kept_games return kept_games def main(): ap = argparse.ArgumentParser() ap.add_argument("pgn_paths", nargs="+") ap.add_argument("--min-elo", type=int, default=2000) ap.add_argument("--min-ply", type=int, default=30, help="minimum game length in plies to be sampled " "at all -- filters short draws/miniatures") ap.add_argument("--out", required=True) ap.add_argument("--progress-every", type=int, default=50000) args = ap.parse_args() stats = { "scanned_games": 0, "kept_games": 0, "rejected_elo": 0, "rejected_short": 0, "skipped_game_over": 0, "skipped_duplicate_boundary": 0, "duplicate_fens": 0, "written": 0, "games_by_result": collections.Counter(), "source_games": collections.Counter(), "stage_total": collections.Counter(), "stage_written": collections.Counter(), "stage_duplicates": collections.Counter(), "material_bucket": collections.Counter(), "side_to_move": collections.Counter(), } seen_keys = set() with open(args.out, "w") as out_f: for p in args.pgn_paths: path = Path(p) print(f"processing {path.name}...", file=sys.stderr) before = stats["kept_games"] process_pgn(path, args.min_elo, args.min_ply, stats, seen_keys, out_f) print(f" kept {stats['kept_games'] - before} games " f"(running total written: {stats['written']})", file=sys.stderr) print("\n=== summary ===") print(f"games scanned: {stats['scanned_games']:,}") print(f"games kept: {stats['kept_games']:,}") print(f" rejected (elo): {stats['rejected_elo']:,}") print(f" rejected (too short): {stats['rejected_short']:,}") print(f"positions written: {stats['written']:,}") print(f"duplicate fens skipped: {stats['duplicate_fens']:,} " f"({100.0*stats['duplicate_fens']/max(1,stats['written']+stats['duplicate_fens']):.1f}%)") print(f"game-over boundaries skipped: {stats['skipped_game_over']:,}") print() print("by result:") for r, c in stats["games_by_result"].most_common(): print(f" {r:>8}: {c:,}") print() print("by source file:") for s, c in stats["source_games"].most_common(): print(f" {s}: {c:,}") print() print("by game stage (written / total-before-dedup):") for stage in STAGE_NAMES: w = stats["stage_written"][stage] t = stats["stage_total"][stage] d = stats["stage_duplicates"][stage] print(f" {stage:>10}: written={w:,} dup_rate={100.0*d/max(1,t):.1f}%") print() print("by material bucket (0=bare kings .. 7=full army):") for b in range(8): print(f" bucket {b}: {stats['material_bucket'][b]:,}") print() print("by side to move:") for k, c in stats["side_to_move"].items(): print(f" {k}: {c:,}") if __name__ == "__main__": main()