Upload encoding_simple.py with huggingface_hub

encoding_simple.py

@@ -0,0 +1,317 @@
+import numpy as np
+import torch
+import bulletchess
+from typing import List, Tuple, Optional
+from .vocab import policy_index
+
+# Precompute policy string to index mapping for O(1) lookups
+policy_to_idx = {u: i for i, u in enumerate(policy_index)}
+
+
+def _board_to_12_piece_planes(board: bulletchess.Board) -> np.ndarray:
+    piece_types = [bulletchess.PAWN, bulletchess.KNIGHT, bulletchess.BISHOP, bulletchess.ROOK, bulletchess.QUEEN, bulletchess.KING]
+    piece_colors = [bulletchess.WHITE, bulletchess.BLACK]
+
+    planes = []
+    for color in piece_colors:
+        for piece_type in piece_types:
+            mask = np.zeros((8, 8), dtype=np.float32)
+            # Use board[color, piece_type] to get Bitboard, then iterate over squares
+            bitboard = board[color, piece_type]
+            for square in bitboard:
+                # In bulletchess, squares have an index() method that returns 0-63
+                square_idx = square.index()
+                rank = square_idx // 8
+                file = square_idx % 8
+                mask[rank][file] = 1.0
+            planes.append(mask)
+    # Shape (8,8,12)
+    return np.transpose(np.array(planes, dtype=np.float32), (1, 2, 0))
+
+
+def _castling_planes(board: bulletchess.Board) -> np.ndarray:
+    # Order must match existing model expectation via ustotheirs:
+    # [WQ, WK, BQ, BK]
+    wq = 1.0 if bulletchess.WHITE_QUEENSIDE in board.castling_rights else 0.0
+    wk = 1.0 if bulletchess.WHITE_KINGSIDE in board.castling_rights else 0.0
+    bq = 1.0 if bulletchess.BLACK_QUEENSIDE in board.castling_rights else 0.0
+    bk = 1.0 if bulletchess.BLACK_KINGSIDE in board.castling_rights else 0.0
+    planes = [
+        np.full((8, 8), wq, dtype=np.float32),
+        np.full((8, 8), wk, dtype=np.float32),
+        np.full((8, 8), bq, dtype=np.float32),
+        np.full((8, 8), bk, dtype=np.float32),
+    ]
+    return np.stack(planes, axis=0)  # (4,8,8)
+
+
+def _mirror_board(board: bulletchess.Board) -> bulletchess.Board:
+    """
+    Fast mirror implementation for bulletchess.Board.
+    Mirrors the board (flips ranks 1<->8, 2<->7, etc.) and flips colors.
+    """
+    # Create empty board
+    mirrored = bulletchess.Board.empty()
+    
+    # Mirror all pieces
+    for square in bulletchess.SQUARES:
+        piece = board[square]
+        if piece is not None:
+            # Calculate mirrored square: flip rank (0-7 -> 7-0), keep file
+            square_idx = square.index()
+            rank = square_idx // 8
+            file = square_idx % 8
+            mirrored_rank = 7 - rank
+            mirrored_idx = mirrored_rank * 8 + file
+            mirrored_square = bulletchess.SQUARES[mirrored_idx]
+            
+            # Flip piece color
+            mirrored_color = piece.color.opposite
+            mirrored[mirrored_square] = bulletchess.Piece(mirrored_color, piece.piece_type)
+    
+    # Mirror castling rights: swap white<->black
+    # Build castling rights by checking each type and creating CastlingRights
+    new_castling_types = []
+    if bulletchess.WHITE_KINGSIDE in board.castling_rights:
+        new_castling_types.append(bulletchess.BLACK_KINGSIDE)
+    if bulletchess.WHITE_QUEENSIDE in board.castling_rights:
+        new_castling_types.append(bulletchess.BLACK_QUEENSIDE)
+    if bulletchess.BLACK_KINGSIDE in board.castling_rights:
+        new_castling_types.append(bulletchess.WHITE_KINGSIDE)
+    if bulletchess.BLACK_QUEENSIDE in board.castling_rights:
+        new_castling_types.append(bulletchess.WHITE_QUEENSIDE)
+    
+    # Build CastlingRights from list of types
+    if new_castling_types:
+        mirrored.castling_rights = bulletchess.CastlingRights(new_castling_types)
+    else:
+        mirrored.castling_rights = bulletchess.NO_CASTLING
+    
+    # Flip turn
+    mirrored.turn = board.turn.opposite
+    
+    # Mirror en passant square if exists
+    if board.en_passant_square is not None:
+        ep_idx = board.en_passant_square.index()
+        ep_rank = ep_idx // 8
+        ep_file = ep_idx % 8
+        mirrored_ep_rank = 7 - ep_rank
+        mirrored_ep_idx = mirrored_ep_rank * 8 + ep_file
+        mirrored.en_passant_square = bulletchess.SQUARES[mirrored_ep_idx]
+    
+    # Copy move counters
+    mirrored.halfmove_clock = board.halfmove_clock
+    mirrored.fullmove_number = board.fullmove_number
+    
+    return mirrored
+
+
+def _build_snapshots(board: bulletchess.Board) -> List[bulletchess.Board]:
+    # snapshots[0] is current, snapshots[1] one ply ago, ... up to 7 plies ago
+    temp = board.copy()
+    snaps: List[bulletchess.Board] = [temp.copy()]
+    for _ in range(7):
+        # Check if there are moves to undo by checking if undo() returns None
+        try:
+            temp.undo()
+            snaps.append(temp.copy())
+        except (IndexError, AttributeError):
+            # No more moves to undo
+            snaps.append(None)  # type: ignore
+    return snaps
+
+
+def encode_moves_to_tensor(uci_moves: List[str], starting_fen: Optional[str] = None) -> Tuple[torch.Tensor, np.ndarray]:
+    board = bulletchess.Board.from_fen(starting_fen) if starting_fen is not None else bulletchess.Board()
+    for mv in uci_moves:
+        move = bulletchess.Move.from_uci(mv)
+        board.apply(move)
+
+    # Build history snapshots (current first)
+    snapshots = _build_snapshots(board)
+
+    # Always encode from white's perspective; mirror all snapshots if black to move
+    mirror = (board.turn == bulletchess.BLACK)
+    if mirror:
+        snapshots = [_mirror_board(s) if s is not None else None for s in snapshots]
+
+    # Assemble 112-channel tensor
+    # 8 groups: each 12 piece planes + 1 blank = 13; total 104
+    channels: List[np.ndarray] = []
+    for i in range(8):
+        if snapshots[i] is not None:
+            planes12 = _board_to_12_piece_planes(snapshots[i])  # (8,8,12)
+            channels.append(planes12)
+        else:
+            channels.append(np.zeros((8, 8, 12), dtype=np.float32))
+        # blank plane
+        channels.append(np.zeros((8, 8, 1), dtype=np.float32))
+
+    # Special planes: WQ, WK, BQ, BK, is_black_to_move, blank, blank, ones
+    current_for_flags = snapshots[0]
+    assert current_for_flags is not None
+    castling = _castling_planes(current_for_flags)  # (4,8,8)
+    is_black_to_move = 1.0 if (board.turn == bulletchess.BLACK) else 0.0
+    specials = [
+        castling[0:1, :, :],  # WQ
+        castling[1:2, :, :],  # WK
+        castling[2:3, :, :],  # BQ
+        castling[3:4, :, :],  # BK
+        np.full((1, 8, 8), is_black_to_move, dtype=np.float32),
+        np.zeros((1, 8, 8), dtype=np.float32),
+        np.zeros((1, 8, 8), dtype=np.float32),
+        np.ones((1, 8, 8), dtype=np.float32),
+    ]
+
+    # Concatenate to (8,8,112)
+    stacked = np.concatenate(channels, axis=2)  # (8,8,104)
+    specials_hwk = np.transpose(np.concatenate(specials, axis=0), (1, 2, 0))  # (8,8,8)
+    final_hwk = np.concatenate([stacked, specials_hwk], axis=2)  # (8,8,112)
+
+    # Convert to tensor (1,112,8,8)
+    final_tensor = torch.from_numpy(final_hwk).permute(2, 0, 1).unsqueeze(0).float()
+
+    # Legal moves mask built from mirrored board to match policy_index perspective
+    board_for_mask = _mirror_board(board) if (board.turn == bulletchess.BLACK) else board.copy()
+    lm = np.ones(1858, dtype=np.float32) * (-1000)
+    
+    # Collect all legal moves as UCI strings
+    legal_moves_uci = set()
+    for possible in board_for_mask.legal_moves():
+        u = possible.uci()
+        if u[-1] != 'n':
+            legal_moves_uci.add(u)
+        else:
+            legal_moves_uci.add(u[:-1])
+    
+    # Mark all legal moves
+    for u in legal_moves_uci:
+        idx = policy_to_idx.get(u)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Add castling moves as king-to-rook-square moves ONLY if the corresponding 
+    # standard castling move is actually legal (to verify castling is possible)
+    # White kingside: e1h1 (king to rook square) if e1g1 is legal
+    if "e1g1" in legal_moves_uci:
+        castling_move = "e1h1"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # White queenside: e1a1 (king to rook square) if e1c1 is legal
+    if "e1c1" in legal_moves_uci:
+        castling_move = "e1a1"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Black kingside: e8h8 (king to rook square) if e8g8 is legal
+    if "e8g8" in legal_moves_uci:
+        castling_move = "e8h8"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Black queenside: e8a8 (king to rook square) if e8c8 is legal
+    if "e8c8" in legal_moves_uci:
+        castling_move = "e8a8"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+
+    return final_tensor, lm
+
+
+def encode_fen_to_tensor(fen: str) -> Tuple[torch.Tensor, np.ndarray]:
+    board = bulletchess.Board.from_fen(fen)
+
+    # History: only current snapshot, others are zeros
+    snapshots = [board.copy()] + [None] * 7
+
+    # Mirror snapshots if black to move so encoding is from white's perspective
+    if board.turn == bulletchess.BLACK:
+        snapshots = [_mirror_board(s) if s is not None else None for s in snapshots]
+
+    # Assemble 112-channel tensor
+    channels: List[np.ndarray] = []
+    for i in range(8):
+        if snapshots[i] is not None:
+            planes12 = _board_to_12_piece_planes(snapshots[i])
+            channels.append(planes12)
+        else:
+            channels.append(np.zeros((8, 8, 12), dtype=np.float32))
+        channels.append(np.zeros((8, 8, 1), dtype=np.float32))
+
+    current_for_flags = snapshots[0]
+    assert current_for_flags is not None
+    castling = _castling_planes(current_for_flags)
+    is_black_to_move = 1.0 if (board.turn == bulletchess.BLACK) else 0.0
+    specials = [
+        castling[0:1, :, :],
+        castling[1:2, :, :],
+        castling[2:3, :, :],
+        castling[3:4, :, :],
+        np.full((1, 8, 8), is_black_to_move, dtype=np.float32),
+        np.zeros((1, 8, 8), dtype=np.float32),
+        np.zeros((1, 8, 8), dtype=np.float32),
+        np.ones((1, 8, 8), dtype=np.float32),
+    ]
+
+    stacked = np.concatenate(channels, axis=2)
+    specials_hwk = np.transpose(np.concatenate(specials, axis=0), (1, 2, 0))
+    final_hwk = np.concatenate([stacked, specials_hwk], axis=2)
+
+    final_tensor = torch.from_numpy(final_hwk).permute(2, 0, 1).unsqueeze(0).float()
+
+    # Legal moves mask from mirrored perspective when black to move
+    board_for_mask = _mirror_board(board) if (board.turn == bulletchess.BLACK) else board.copy()
+    lm = np.ones(1858, dtype=np.float32) * (-1000)
+    
+    # Collect all legal moves as UCI strings
+    legal_moves_uci = set()
+    for possible in board_for_mask.legal_moves():
+        u = possible.uci()
+        if u[-1] != 'n':
+            legal_moves_uci.add(u)
+        else:
+            legal_moves_uci.add(u[:-1])
+    
+    # Mark all legal moves
+    for u in legal_moves_uci:
+        idx = policy_to_idx.get(u)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Add castling moves as king-to-rook-square moves ONLY if the corresponding 
+    # standard castling move is actually legal (to verify castling is possible)
+    # White kingside: e1h1 (king to rook square) if e1g1 is legal
+    if "e1g1" in legal_moves_uci:
+        castling_move = "e1h1"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # White queenside: e1a1 (king to rook square) if e1c1 is legal
+    if "e1c1" in legal_moves_uci:
+        castling_move = "e1a1"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Black kingside: e8h8 (king to rook square) if e8g8 is legal
+    if "e8g8" in legal_moves_uci:
+        castling_move = "e8h8"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+    
+    # Black queenside: e8a8 (king to rook square) if e8c8 is legal
+    if "e8c8" in legal_moves_uci:
+        castling_move = "e8a8"
+        idx = policy_to_idx.get(castling_move)
+        if idx is not None:
+            lm[idx] = 0
+
+    return final_tensor, lm
+