""" engine.py -- a small, readable chess engine on top of python-chess. Two parts: 1. Evaluation. The score of a position is a *linear* function of a short feature vector: eval(s) = w . f(s). Every feature is computed as (White's count) - (Black's count), so the evaluation is from White's point of view and is exactly antisymmetric under colour flip. There are 5 material features (pawn, knight, bishop, rook, queen counts) 12 positional features (centralisation, pawn structure, mobility, ...) -- 17 weights in total, small enough to learn from a few hundred self-play games (see learn.py). 2. Search. Negamax alpha-beta with * quiescence search over captures/promotions (and check evasions), * MVV-LVA move ordering (Most Valuable Victim, Least Valuable Attacker), * an optional transposition table (TT), * iterative deepening with node and time limits, * draw detection by repetition / 50-move rule inside the tree. Scores inside the search are from the side-to-move's point of view ("negamax" convention). `SearchResult.score` is therefore also relative to the side to move at the root; `white_score()` converts it. The module also renders the piece-square "heatmaps" implied by a weight vector (what a single piece on each square is worth to the evaluation). """ from __future__ import annotations import json import math import time from dataclasses import dataclass, field from typing import Dict, List, Optional, Sequence, Tuple import chess # -------------------------------------------------------------------------- # Feature definitions # -------------------------------------------------------------------------- MATERIAL_FEATURES = ["pawn", "knight", "bishop", "rook", "queen"] # Each positional feature is scaled so that a typical |value| is O(1); this # keeps a single learning rate sensible for all weights. POSITIONAL_FEATURES = [ "pawn_advance", # sum over pawns of (ranks advanced) / 5 (0 .. 1 per pawn) "passed_pawn", # number of passed pawns "doubled_pawn", # number of extra pawns on a file (2 on a file -> 1) "isolated_pawn", # pawns with no friendly pawn on an adjacent file "knight_center", # sum over knights of centrality / 3 (0 .. 1 per knight) "bishop_center", # same for bishops "queen_center", # same for queens "king_center", # centrality / 3 of the king (sign should depend on phase!) "mobility", # pseudo-legal moves of N,B,R,Q / 10 (attacked squares not own-occupied) "king_shield", # own pawns on the 3 files around the king, 1-2 ranks ahead "bishop_pair", # 1 if the side has >= 2 bishops "rook_open_file", # rooks on a file without pawns (1) or without own pawns (0.5) ] ALL_FEATURES = MATERIAL_FEATURES + POSITIONAL_FEATURES N_MATERIAL = len(MATERIAL_FEATURES) N_FEATURES = len(ALL_FEATURES) FEATURE_INDEX = {name: i for i, name in enumerate(ALL_FEATURES)} # Conventional values, used ONLY for move ordering (MVV-LVA) and for the # optional ladder adjudication. The learner never sees them. STANDARD_VALUES = {chess.PAWN: 1, chess.KNIGHT: 3, chess.BISHOP: 3, chess.ROOK: 5, chess.QUEEN: 9, chess.KING: 0} # ---- precomputed bitboard masks ------------------------------------------- def _centrality(sq: int) -> int: """0 on the rim, 3 on the four centre squares (d4, e4, d5, e5).""" f, r = chess.square_file(sq), chess.square_rank(sq) return min(f, 7 - f, r, 7 - r) # RING[k] = bitboard of squares with centrality k RING = [0, 0, 0, 0] for _sq in chess.SQUARES: RING[_centrality(_sq)] |= chess.BB_SQUARES[_sq] FILE_MASKS = list(chess.BB_FILES) ADJ_FILE_MASKS = [ (FILE_MASKS[f - 1] if f > 0 else 0) | (FILE_MASKS[f + 1] if f < 7 else 0) for f in range(8) ] RANK_MASKS = list(chess.BB_RANKS) def _passed_mask(color: bool, sq: int) -> int: """Squares in front of `sq` (same + adjacent files) that must be free of enemy pawns for a pawn of `color` on `sq` to be passed.""" f, r = chess.square_file(sq), chess.square_rank(sq) mask = 0 ranks = range(r + 1, 8) if color == chess.WHITE else range(0, r) for rr in ranks: for ff in (f - 1, f, f + 1): if 0 <= ff < 8: mask |= chess.BB_SQUARES[chess.square(ff, rr)] return mask def _shield_mask(color: bool, sq: int) -> int: """The (up to) 6 squares 1-2 ranks in front of a king on `sq`.""" f, r = chess.square_file(sq), chess.square_rank(sq) step = 1 if color == chess.WHITE else -1 mask = 0 for dr in (1, 2): rr = r + step * dr if 0 <= rr < 8: for ff in (f - 1, f, f + 1): if 0 <= ff < 8: mask |= chess.BB_SQUARES[chess.square(ff, rr)] return mask PASSED_MASK = {c: [_passed_mask(c, s) for s in chess.SQUARES] for c in chess.COLORS} SHIELD_MASK = {c: [_shield_mask(c, s) for s in chess.SQUARES] for c in chess.COLORS} def _popcount(x: int) -> int: return x.bit_count() def _center_sum(mask: int) -> int: """Sum of centrality over all pieces in `mask` (uses the ring bitboards).""" return (_popcount(mask & RING[1]) + 2 * _popcount(mask & RING[2]) + 3 * _popcount(mask & RING[3])) def _pawn_advance_sum(pawns: int, color: bool) -> int: """Sum over pawns of the number of ranks advanced from the start rank.""" total = 0 for r in range(8): n = _popcount(pawns & RANK_MASKS[r]) if n: adv = (r - 1) if color == chess.WHITE else (6 - r) total += n * adv return total def _side_positional(board: chess.Board, color: bool) -> List[float]: """Positional features for one side (non-negative counts).""" own = board.occupied_co[color] pawns = board.pieces_mask(chess.PAWN, color) opp_pawns = board.pieces_mask(chess.PAWN, not color) knights = board.pieces_mask(chess.KNIGHT, color) bishops = board.pieces_mask(chess.BISHOP, color) rooks = board.pieces_mask(chess.ROOK, color) queens = board.pieces_mask(chess.QUEEN, color) # --- pawn structure --- passed = 0 for sq in chess.scan_forward(pawns): if not (PASSED_MASK[color][sq] & opp_pawns): passed += 1 doubled = 0 isolated = 0 for f in range(8): n = _popcount(pawns & FILE_MASKS[f]) if n > 1: doubled += n - 1 if n and not (pawns & ADJ_FILE_MASKS[f]): isolated += n # --- mobility of pieces (pseudo-legal, cheap) --- mob = 0 for sq in chess.scan_forward(knights | bishops | rooks | queens): mob += _popcount(board.attacks_mask(sq) & ~own) # --- king safety / centrality --- ksq = board.king(color) shield = _popcount(SHIELD_MASK[color][ksq] & pawns) if ksq is not None else 0 kcent = _centrality(ksq) if ksq is not None else 0 # --- rooks on open / semi-open files --- all_pawns = pawns | opp_pawns rook_open = 0.0 for sq in chess.scan_forward(rooks): fm = FILE_MASKS[chess.square_file(sq)] if not (fm & all_pawns): rook_open += 1.0 elif not (fm & pawns): rook_open += 0.5 return [ _pawn_advance_sum(pawns, color) / 5.0, float(passed), float(doubled), float(isolated), _center_sum(knights) / 3.0, _center_sum(bishops) / 3.0, _center_sum(queens) / 3.0, kcent / 3.0, mob / 10.0, float(shield), 1.0 if _popcount(bishops) >= 2 else 0.0, rook_open, ] def material_features(board: chess.Board) -> List[float]: """[dP, dN, dB, dR, dQ] (White count minus Black count).""" w, b = board.occupied_co[chess.WHITE], board.occupied_co[chess.BLACK] return [ float(_popcount(board.pawns & w) - _popcount(board.pawns & b)), float(_popcount(board.knights & w) - _popcount(board.knights & b)), float(_popcount(board.bishops & w) - _popcount(board.bishops & b)), float(_popcount(board.rooks & w) - _popcount(board.rooks & b)), float(_popcount(board.queens & w) - _popcount(board.queens & b)), ] def features(board: chess.Board, positional: bool = True) -> List[float]: """Full feature vector (length N_FEATURES), White minus Black. If positional=False the positional block is zeros (and not computed).""" f = material_features(board) if positional: pw = _side_positional(board, chess.WHITE) pb = _side_positional(board, chess.BLACK) f.extend(a - b for a, b in zip(pw, pb)) else: f.extend([0.0] * len(POSITIONAL_FEATURES)) return f # -------------------------------------------------------------------------- # Weights # -------------------------------------------------------------------------- def equal_material_weights(value: float = 1.0) -> List[float]: """Starting point for learning: every piece worth `value`, no positional terms.""" return [value] * N_MATERIAL + [0.0] * len(POSITIONAL_FEATURES) def standard_weights() -> List[float]: """Textbook 1/3/3/5/9 values (for tests and as a sanity reference).""" return [1.0, 3.0, 3.0, 5.0, 9.0] + [0.0] * len(POSITIONAL_FEATURES) def weights_to_dict(w: Sequence[float]) -> Dict[str, float]: return {name: float(v) for name, v in zip(ALL_FEATURES, w)} def weights_from_dict(d: Dict[str, float]) -> List[float]: return [float(d.get(name, 0.0)) for name in ALL_FEATURES] def save_weights(path: str, w: Sequence[float], meta: Optional[dict] = None) -> None: with open(path, "w") as fh: json.dump({"weights": weights_to_dict(w), "meta": meta or {}}, fh, indent=2) def load_weights(path: str) -> List[float]: with open(path) as fh: return weights_from_dict(json.load(fh)["weights"]) def normalized_piece_values(w: Sequence[float]) -> Dict[str, float]: """Piece values divided by the pawn weight (pawn = 1).""" p = w[0] return {name: (w[i] / p if p != 0 else float("nan")) for i, name in enumerate(MATERIAL_FEATURES)} class Evaluator: """Linear evaluation w . f(board). Positional features are only computed if at least one positional weight is non-zero (material-only eval is ~10x faster, which matters a lot in pure Python).""" def __init__(self, weights: Sequence[float], positional: Optional[bool] = None): self.w = [float(x) for x in weights] assert len(self.w) == N_FEATURES, f"expected {N_FEATURES} weights" # positional=None -> decide automatically from the weights; the # learner passes True so positional features are computed even while # their weights are still exactly zero. if positional is None: positional = any(abs(x) > 0 for x in self.w[N_MATERIAL:]) self.use_positional = positional def features(self, board: chess.Board) -> List[float]: return features(board, positional=self.use_positional) def white_eval(self, board: chess.Board) -> float: f = self.features(board) return sum(a * b for a, b in zip(self.w, f)) def stm_eval(self, board: chess.Board) -> float: """Evaluation from the side-to-move's point of view (negamax).""" v = self.white_eval(board) return v if board.turn == chess.WHITE else -v @property def pawn_unit(self) -> float: """Value of one pawn, used to print evaluations in 'pawns'.""" return self.w[0] if self.w[0] > 1e-9 else 1.0 # -------------------------------------------------------------------------- # Search # -------------------------------------------------------------------------- MATE = 100000.0 # mate score (far above any material eval) MATE_THRESHOLD = MATE - 1000.0 INF = float("inf") TT_EXACT, TT_LOWER, TT_UPPER = 0, 1, 2 class SearchAborted(Exception): """Raised internally when the node or time budget runs out.""" @dataclass class SearchResult: move: Optional[chess.Move] score: float # side-to-move-at-root perspective pv: List[chess.Move] = field(default_factory=list) depth: int = 0 # deepest *completed* iteration nodes: int = 0 seconds: float = 0.0 def white_score(self, board: chess.Board) -> float: return self.score if board.turn == chess.WHITE else -self.score def is_mate_score(self) -> bool: return abs(self.score) >= MATE_THRESHOLD def is_mate_score(score: float) -> bool: return abs(score) >= MATE_THRESHOLD def _mvv_lva(board: chess.Board, move: chess.Move) -> int: """Ordering key for captures: big victim first, then small attacker.""" victim = board.piece_type_at(move.to_square) if victim is None and board.is_en_passant(move): victim = chess.PAWN attacker = board.piece_type_at(move.from_square) score = 10 * STANDARD_VALUES.get(victim, 0) - STANDARD_VALUES.get(attacker, 0) if move.promotion: score += 10 * STANDARD_VALUES[move.promotion] return score class Searcher: """Negamax alpha-beta + quiescence. Parameters ---------- evaluator : Evaluator use_tt : enable the transposition table (cleared every `search` call). qdepth : maximum quiescence depth (captures only, plus check evasions). """ def __init__(self, evaluator: Evaluator, use_tt: bool = True, qdepth: int = 6): self.ev = evaluator self.use_tt = use_tt self.qdepth = qdepth self.tt: Dict[tuple, Tuple[int, float, int, Optional[chess.Move]]] = {} self.nodes = 0 self._node_limit: Optional[int] = None self._deadline: Optional[float] = None self._path: List[tuple] = [] # position keys: game history + search path # ---------------- public API ---------------- def search(self, board: chess.Board, depth: int, node_limit: Optional[int] = None, time_limit: Optional[float] = None) -> SearchResult: """Iterative deepening to `depth` plies (depth >= 1). If the node or time budget runs out, returns the result of the last finished iteration (or the best move found so far in the first iteration).""" t0 = time.time() self.nodes = 0 self._node_limit = node_limit self._deadline = (t0 + time_limit) if time_limit else None self.tt.clear() self._path = self._history_keys(board) root_moves = list(board.legal_moves) if not root_moves: return SearchResult(None, self._terminal_score(board, 0), [], 0, 0, 0.0) root_moves.sort(key=lambda m: self._order_key(board, m, None), reverse=True) best = SearchResult(root_moves[0], -INF, [root_moves[0]], 0, 0, 0.0) for d in range(1, max(1, depth) + 1): try: score, pv = self._root(board, d, root_moves) except SearchAborted: if best.depth == 0 and self._partial is not None: best = SearchResult(self._partial[1][0], self._partial[0], self._partial[1], 0) break best = SearchResult(pv[0], score, pv, d) # search the previous best move first in the next iteration root_moves.remove(pv[0]) root_moves.insert(0, pv[0]) if is_mate_score(score): break # a forced mate was found; deeper search can't improve it best.nodes = self.nodes best.seconds = time.time() - t0 if best.score == -INF: # aborted before any move was scored best.score = self.ev.stm_eval(board) return best def score_root_moves(self, board: chess.Board, depth: int ) -> List[Tuple[chess.Move, float, List[chess.Move]]]: """Exact (full-window) score and PV for *every* legal root move. Used by the learner for softmax/epsilon exploration. Scores are from the root side-to-move's perspective. `depth` counts the root move, so depth=1 means: play the move, then quiescence search.""" self.nodes = 0 self._node_limit = None self._deadline = None self.tt.clear() self._path = self._history_keys(board) out = [] for m in board.legal_moves: board.push(m) self._path.append(board._transposition_key()) s, pv = self._negamax(board, depth - 1, -INF, INF, 1) self._path.pop() board.pop() out.append((m, -s, [m] + pv)) return out # ---------------- internals ---------------- _partial: Optional[Tuple[float, List[chess.Move]]] = None def _history_keys(self, board: chess.Board) -> List[tuple]: """Keys of the positions since the last irreversible move, so the search can see repetitions of game positions.""" keys = [] b = board.copy(stack=True) keys.append(b._transposition_key()) for _ in range(min(board.halfmove_clock, len(board.move_stack))): b.pop() keys.append(b._transposition_key()) keys.reverse() return keys def _check_limits(self) -> None: if self._node_limit is not None and self.nodes >= self._node_limit: raise SearchAborted if self._deadline is not None and (self.nodes & 255) == 0 \ and time.time() > self._deadline: raise SearchAborted def _terminal_score(self, board: chess.Board, ply: int) -> float: """Score of a position with no legal moves (side to move's view).""" return -(MATE - ply) if board.is_check() else 0.0 def _order_key(self, board: chess.Board, move: chess.Move, tt_move: Optional[chess.Move]) -> int: if move == tt_move: return 1_000_000 if board.is_capture(move) or move.promotion: return 10_000 + _mvv_lva(board, move) return 0 def _root(self, board: chess.Board, depth: int, root_moves: List[chess.Move]): alpha, beta = -INF, INF best_score, best_pv = -INF, [root_moves[0]] self._partial = None for m in root_moves: board.push(m) self._path.append(board._transposition_key()) try: s, pv = self._negamax(board, depth - 1, -beta, -alpha, 1) finally: self._path.pop() board.pop() s = -s if s > best_score: best_score, best_pv = s, [m] + pv self._partial = (best_score, best_pv) alpha = max(alpha, s) return best_score, best_pv def _negamax(self, board: chess.Board, depth: int, alpha: float, beta: float, ply: int) -> Tuple[float, List[chess.Move]]: self.nodes += 1 self._check_limits() # Draws inside the tree: 50-move rule and repetition of any position # on the current game/search path (treated as a draw immediately). if board.halfmove_clock >= 100: return 0.0, [] key = self._path[-1] if key in self._path[:-1]: return 0.0, [] if depth <= 0: return self._quiesce(board, alpha, beta, ply, self.qdepth) alpha0 = alpha tt_move = None if self.use_tt: entry = self.tt.get(key) if entry is not None: e_depth, e_score, e_flag, tt_move = entry if e_depth >= depth: if e_flag == TT_EXACT: return e_score, ([tt_move] if tt_move else []) if e_flag == TT_LOWER and e_score >= beta: return e_score, ([tt_move] if tt_move else []) if e_flag == TT_UPPER and e_score <= alpha: return e_score, ([tt_move] if tt_move else []) moves = list(board.legal_moves) if not moves: return self._terminal_score(board, ply), [] moves.sort(key=lambda m: self._order_key(board, m, tt_move), reverse=True) best, best_pv = -INF, [] for m in moves: board.push(m) self._path.append(board._transposition_key()) try: s, pv = self._negamax(board, depth - 1, -beta, -alpha, ply + 1) finally: self._path.pop() board.pop() s = -s if s > best: best, best_pv = s, [m] + pv if s > alpha: alpha = s if alpha >= beta: break if self.use_tt: flag = TT_EXACT if best <= alpha0: flag = TT_UPPER elif best >= beta: flag = TT_LOWER self.tt[key] = (depth, best, flag, best_pv[0] if best_pv else None) return best, best_pv def _quiesce(self, board: chess.Board, alpha: float, beta: float, ply: int, qdepth: int) -> Tuple[float, List[chess.Move]]: """Search captures/promotions until the position is quiet. When in check, all evasions are searched (no stand-pat), so mates delivered at the horizon are detected.""" self.nodes += 1 self._check_limits() in_check = board.is_check() if in_check: moves = list(board.legal_moves) if not moves: return -(MATE - ply), [] if qdepth <= 0: return self.ev.stm_eval(board), [] best, best_pv = -INF, [] else: stand = self.ev.stm_eval(board) if qdepth <= 0 or stand >= beta: return stand, [] best, best_pv = stand, [] if stand > alpha: alpha = stand moves = [m for m in board.generate_legal_captures()] moves += [m for m in board.generate_legal_moves(board.occupied_co[board.turn] & board.pawns) if m.promotion and not board.is_capture(m)] moves.sort(key=lambda m: _mvv_lva(board, m) if (board.is_capture(m) or m.promotion) else -100, reverse=True) for m in moves: board.push(m) try: s, pv = self._quiesce(board, -beta, -alpha, ply + 1, qdepth - 1) finally: board.pop() s = -s if s > best: best, best_pv = s, [m] + pv if s > alpha: alpha = s if alpha >= beta: break return best, best_pv def leaf_board(board: chess.Board, pv: Sequence[chess.Move]) -> chess.Board: """The position at the end of a principal variation (for TD-Leaf).""" b = board.copy(stack=False) for m in pv: if not b.is_legal(m): break b.push(m) return b def pv_to_san(board: chess.Board, pv: Sequence[chess.Move], max_len: int = 10) -> str: """Render a PV in SAN, e.g. '12. Nf3 Nc6 13. Bb5'.""" b = board.copy(stack=False) parts = [] for m in list(pv)[:max_len]: if not b.is_legal(m): break prefix = "" if b.turn == chess.WHITE: prefix = f"{b.fullmove_number}. " elif not parts: prefix = f"{b.fullmove_number}... " parts.append(prefix + b.san(m)) b.push(m) return " ".join(parts) # -------------------------------------------------------------------------- # Implied piece-square heatmaps # -------------------------------------------------------------------------- def implied_square_values(w: Sequence[float], piece_type: int, color: bool = chess.WHITE) -> List[List[float]]: """Value (in pawns, i.e. divided by the pawn weight) that the *linear* evaluation assigns to a lone piece of `piece_type` on each square, counting only the terms that depend on that square: material + centrality term (+ pawn advancement for pawns) + mobility term for N/B/R/Q, measured on an otherwise EMPTY board (so it is an upper bound on real mobility; on a crowded board it is lower). Pawn-structure and king-shield terms depend on other pieces and are omitted. Returns an 8x8 grid indexed [rank][file].""" wd = weights_to_dict(w) unit = wd["pawn"] if abs(wd["pawn"]) > 1e-9 else 1.0 base = {chess.PAWN: wd["pawn"], chess.KNIGHT: wd["knight"], chess.BISHOP: wd["bishop"], chess.ROOK: wd["rook"], chess.QUEEN: wd["queen"], chess.KING: 0.0}[piece_type] grid = [[0.0] * 8 for _ in range(8)] for sq in chess.SQUARES: f, r = chess.square_file(sq), chess.square_rank(sq) v = base c = _centrality(sq) / 3.0 if piece_type in (chess.KNIGHT, chess.BISHOP, chess.ROOK, chess.QUEEN): lone = chess.Board(None) lone.set_piece_at(sq, chess.Piece(piece_type, color)) v += wd["mobility"] * _popcount(lone.attacks_mask(sq)) / 10.0 if piece_type == chess.KNIGHT: v += wd["knight_center"] * c elif piece_type == chess.BISHOP: v += wd["bishop_center"] * c elif piece_type == chess.QUEEN: v += wd["queen_center"] * c elif piece_type == chess.KING: v += wd["king_center"] * c elif piece_type == chess.PAWN: adv = (r - 1) if color == chess.WHITE else (6 - r) if not 0 <= adv <= 5: # pawns never stand on the first rank grid[r][f] = float("nan") continue v += wd["pawn_advance"] * adv / 5.0 grid[r][f] = v / unit return grid def plot_heatmaps(w: Sequence[float], path: str, title: str = "") -> None: """Save a 2x3 figure of implied piece-square values (White's view).""" import matplotlib matplotlib.use("Agg") import matplotlib.pyplot as plt import numpy as np pieces = [(chess.PAWN, "Pawn"), (chess.KNIGHT, "Knight"), (chess.BISHOP, "Bishop"), (chess.ROOK, "Rook"), (chess.QUEEN, "Queen"), (chess.KING, "King")] fig, axes = plt.subplots(2, 3, figsize=(11, 7.5)) for ax, (pt, name) in zip(axes.flat, pieces): grid = np.array(implied_square_values(w, pt)) im = ax.imshow(grid, origin="lower", cmap="viridis") ax.set_xticks(range(8), list("abcdefgh")) ax.set_yticks(range(8), [str(i) for i in range(1, 9)]) ax.set_title(name) for r in range(8): for f in range(8): if not math.isnan(grid[r, f]): ax.text(f, r, f"{grid[r, f]:.2f}", ha="center", va="center", fontsize=6, color="white") fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04) fig.suptitle((title + "\n" if title else "") + "value of a lone piece on each square, in pawns " "(material + centrality/advance + empty-board mobility)", fontsize=10) fig.tight_layout() fig.savefig(path, dpi=110) plt.close(fig)