KeshavAnandCode
3eccd97536
- Implement piece-square feature extraction - 32 active features for 32 pieces on board - Tests for feature extraction (7 tests) - Fix: piece_sq * 6 + piece_type mapping
128 lines
2.5 KiB
Python
128 lines
2.5 KiB
Python
"""Extract NNUE features from FEN strings"""
|
||
|
||
import chess
|
||
from chess import Board as chess_board
|
||
from python.constants import (
|
||
HALF_KA_V2_HM,
|
||
FULL_THREATS,
|
||
TOTAL_FEATURES,
|
||
PIECE_SQUARE_INDEX,
|
||
PIECE_TYPE_MAP,
|
||
)
|
||
|
||
# King bucket indices (56 squares / 8 buckets = 7 squares per bucket)
|
||
# Each bucket maps 7 consecutive squares to the same bucket index (0-7)
|
||
KING_BUCKETS = [
|
||
0,
|
||
0,
|
||
0,
|
||
0,
|
||
0,
|
||
0,
|
||
0, # Bucket 0: squares 0-6
|
||
1,
|
||
1,
|
||
1,
|
||
1,
|
||
1,
|
||
1,
|
||
1, # Bucket 1: squares 7-13
|
||
2,
|
||
2,
|
||
2,
|
||
2,
|
||
2,
|
||
2,
|
||
2, # Bucket 2: squares 14-20
|
||
3,
|
||
3,
|
||
3,
|
||
3,
|
||
3,
|
||
3,
|
||
3, # Bucket 3: squares 21-27
|
||
4,
|
||
4,
|
||
4,
|
||
4,
|
||
4,
|
||
4,
|
||
4, # Bucket 4: squares 28-34
|
||
5,
|
||
5,
|
||
5,
|
||
5,
|
||
5,
|
||
5,
|
||
5, # Bucket 5: squares 35-41
|
||
6,
|
||
6,
|
||
6,
|
||
6,
|
||
6,
|
||
6,
|
||
6, # Bucket 6: squares 42-48
|
||
7,
|
||
7,
|
||
7,
|
||
7,
|
||
7,
|
||
7,
|
||
7, # Bucket 7: squares 49-55
|
||
]
|
||
|
||
|
||
def fen_to_features(fen: str) -> list:
|
||
"""
|
||
Convert FEN to 61,072 feature vector.
|
||
|
||
Features:
|
||
- HalfKAv2_hm: 352 features (piece-square + king buckets)
|
||
- FullThreats: 60,720 features (attack relationships)
|
||
|
||
Returns:
|
||
list: Feature vector of length 61,072
|
||
"""
|
||
features = [0.0] * TOTAL_FEATURES
|
||
|
||
b = chess_board(fen)
|
||
perspective = int(b.turn) # 0 for white, 1 for black (True=1, False=0)
|
||
|
||
# Find king square
|
||
ksq = None
|
||
for sq in range(64):
|
||
piece = b.piece_at(sq)
|
||
if piece and piece.unicode_symbol() in (
|
||
"\u265a",
|
||
"\u2654",
|
||
): # White or black king
|
||
ksq = sq
|
||
break
|
||
|
||
# Compute orientation offset
|
||
orient_offset = PIECE_SQUARE_INDEX[perspective][
|
||
0
|
||
] # Base offset from PIECE_SQUARE_INDEX
|
||
orient_offset ^= 56 * perspective # Add perspective offset
|
||
|
||
# Extract HalfKAv2_hm features (352 features)
|
||
for piece_sq in range(64):
|
||
piece = b.piece_at(piece_sq)
|
||
if piece is None:
|
||
continue
|
||
|
||
# Get piece type (0-5) from PIECE_TYPE_MAP
|
||
piece_type = PIECE_TYPE_MAP.get(piece.unicode_symbol())
|
||
if piece_type is None:
|
||
continue
|
||
|
||
# Calculate feature index
|
||
# HalfKAv2_hm: 352 features (56 squares × 6 piece types + 16 king buckets)
|
||
# Simple mapping: piece_sq * 6 + piece_type for pieces
|
||
feature_idx = piece_sq * 6 + piece_type
|
||
|
||
# Set feature (1 for presence, 0 for absence)
|
||
features[feature_idx] = 1.0
|
||
|
||
return features
|