130 lines
4.7 KiB
Python
130 lines
4.7 KiB
Python
import multiprocessing as mp
|
|
import random
|
|
import time
|
|
import chess
|
|
import chess.pgn
|
|
from typing import Tuple, List
|
|
from enum import Enum
|
|
from dataclasses import dataclass
|
|
|
|
from chesspp.engine_factory import StrategyEnum, EngineFactory, EngineEnum
|
|
from chesspp.engine import Engine, Limit
|
|
|
|
|
|
class Winner(Enum):
|
|
Engine_A = 0
|
|
Engine_B = 1
|
|
Draw = 2
|
|
|
|
|
|
@dataclass
|
|
class GameStatistics:
|
|
white: str
|
|
black: str
|
|
average_time_white: float
|
|
average_time_black: float
|
|
nodes_white: int
|
|
nodes_black: int
|
|
length: int
|
|
|
|
@dataclass
|
|
class EvaluationResult:
|
|
winner: Winner
|
|
game: str
|
|
statistics: GameStatistics
|
|
|
|
|
|
def simulate_game(white: Engine, black: Engine, limit: Limit, board: chess.Board) -> (chess.pgn.Game, GameStatistics):
|
|
is_white_playing = True
|
|
times_white = []
|
|
times_black = []
|
|
game_length = 0
|
|
while not board.is_game_over():
|
|
start = time.time()
|
|
play_result = white.play(board, limit) if is_white_playing else black.play(board, limit)
|
|
end = time.time()
|
|
times_white.append(end - start) if is_white_playing else times_black.append(end - start)
|
|
board.push(play_result.move)
|
|
is_white_playing = not is_white_playing
|
|
game_length += 1
|
|
|
|
game = chess.pgn.Game.from_board(board)
|
|
game.headers['White'] = white.get_name()
|
|
game.headers['Black'] = black.get_name()
|
|
|
|
if hasattr(white, "node_counts"):
|
|
white_nodes = sum(white.node_counts) // len(white.node_counts)
|
|
else:
|
|
white_nodes = 0
|
|
|
|
if hasattr(black, "node_counts"):
|
|
black_nodes = sum(black.node_counts) // len(black.node_counts)
|
|
else:
|
|
black_nodes = 0
|
|
|
|
statistics = GameStatistics(white=white.get_name(),
|
|
black=black.get_name(),
|
|
average_time_white=(sum(times_white)/len(times_white)),
|
|
average_time_black=(sum(times_black)/len(times_black)),
|
|
nodes_white=white_nodes,
|
|
nodes_black=black_nodes,
|
|
length=game_length
|
|
)
|
|
|
|
return game, statistics
|
|
|
|
|
|
class Evaluation:
|
|
def __init__(self, engine_a: EngineEnum, strategy_a, engine_b: EngineEnum, strategy_b, limit: Limit,
|
|
stockfish_path: str, lc0_path: str, stockfish_elo: int):
|
|
self.engine_a = engine_a
|
|
self.strategy_a = strategy_a
|
|
self.engine_b = engine_b
|
|
self.strategy_b = strategy_b
|
|
self.stockfish_path = stockfish_path
|
|
self.lc0_path = lc0_path
|
|
self.limit = limit
|
|
self.stockfish_elo = stockfish_elo
|
|
|
|
def run(self, n_games=100, proc=mp.cpu_count()) -> List[EvaluationResult]:
|
|
proc = min(proc, mp.cpu_count())
|
|
arg = (self.engine_a, self.strategy_a, self.engine_b, self.strategy_b, self.limit, self.stockfish_path, self.lc0_path, self.stockfish_elo)
|
|
if proc > 1:
|
|
with mp.Pool(proc) as pool:
|
|
args = [arg for i in range(n_games)]
|
|
return pool.map(Evaluation._test_simulate, args)
|
|
return [
|
|
Evaluation._test_simulate(arg)
|
|
for _ in range(n_games)
|
|
]
|
|
|
|
@staticmethod
|
|
def _test_simulate(arg: Tuple[EngineEnum, StrategyEnum, EngineEnum, StrategyEnum, Limit, str, str, int]) -> EvaluationResult:
|
|
engine_a, strategy_a, engine_b, strategy_b, limit, stockfish_path, lc0_path, stockfish_elo = arg
|
|
flip_engines = bool(random.getrandbits(1))
|
|
|
|
if flip_engines:
|
|
black, white = EngineFactory.create_engine(engine_a, strategy_a, chess.BLACK,
|
|
stockfish_path, lc0_path, stockfish_elo), EngineFactory.create_engine(
|
|
engine_b, strategy_b, chess.WHITE, stockfish_path, lc0_path, stockfish_elo)
|
|
else:
|
|
white, black = EngineFactory.create_engine(engine_a, strategy_a, chess.WHITE,
|
|
stockfish_path, lc0_path, stockfish_elo), EngineFactory.create_engine(
|
|
engine_b, strategy_b, chess.BLACK, stockfish_path, lc0_path, stockfish_elo)
|
|
|
|
game, statistics = simulate_game(white, black, limit, chess.Board())
|
|
winner = game.end().board().outcome().winner
|
|
|
|
result = Winner.Draw
|
|
match (winner, flip_engines):
|
|
case (chess.WHITE, True):
|
|
result = Winner.Engine_B
|
|
case (chess.BLACK, True):
|
|
result = Winner.Engine_A
|
|
case (chess.WHITE, False):
|
|
result = Winner.Engine_A
|
|
case (chess.BLACK, False):
|
|
result = Winner.Engine_B
|
|
|
|
return EvaluationResult(result, str(game), statistics)
|