analyze_extend.py

import os
import argparse
import sys
import chess
import chess.pgn
import chess.uci
from timeit import default_timer as timer



def parse_args():
    """
    Define an argument parser and return the parsed arguments
    """
    parser = argparse.ArgumentParser(
        prog='analyze',
        description='takes a list of PGN files and analyzes last position'
        'of each PGN')
    parser.add_argument("--file", "-f",
                        help="list of PGN files",
                        required=True,
                        metavar="list_FILE")
    parser.add_argument("--engine", "-e",
                        help="analysis engine (default: %(default)s)",
                        default="/home/mkolar/prog/stockfish/stockfish_10_x64_bmi2")
    parser.add_argument("--Threads", "-t",
                        help="number of threads used by engine",
                        default="1",
                        type=int)
    parser.add_argument("--depthShallow", "-ds",
                        help="depth of analysis shallow (when comming up with candidate moves)",
                        default="26",
                        type=int)
    parser.add_argument("--depthDeep", "-dd",
                        help="depth of analysis (when evaluating position)",
                        default="30",
                        type=int)
    parser.add_argument("--add_plies", "-a",
                        help="extend the pgn by this many plies",
                        default="2",
                        type=int)
    parser.add_argument("--Hash", "-H",
                        help="hash size (1, ..., 4096, 8192, 16384, 32768, 65536, 131072)",
                        default="4096",
                        type=int)
    return parser.parse_args()


def analyze_node(node, plies_left, engine, handler, args, my_side):
    # obtain evaluation
    engine.setoption({
        "MultiPV": 1
    })
    board = node.board()
    engine.position(board)
    engine.go(depth=args.depthDeep)
    board_score = handler.info["score"][1].cp/100.0
    if board.turn:
        node.comment = format(board_score)
    else:
        node.comment = format(-board_score)

    if plies_left == 0:
        return

    # find candidate moves
    engine.setoption({
        "MultiPV": 500
    })
    engine.position(board)
    engine.go(depth=args.depthShallow)
    candidate_moves = []
    for l in range(1, len(handler.info["pv"]) + 1):
        move_score = handler.info["score"][l].cp/100.0
        if my_side == board.turn:
            if board_score - move_score > 0.5:
                break
        else:
            if board_score - move_score > 0.5:
                break

        candidate_moves.append(handler.info["pv"][l][0])

    print(candidate_moves)
    for move in candidate_moves:
        node_next = node.add_variation(move)
        analyze_node(node_next, plies_left-1, engine, handler, args, my_side)

    return

def main():
    args = parse_args()

    try:
        with open(args.file) as pgn_list:
            # setup engine
            try:
                engine = chess.uci.popen_engine(args.engine)
            except FileNotFoundError:
                errormsg = "Engine '{}' was not found. Aborting...".format(args.engine)
                raise
            except PermissionError:
                errormsg = "Engine '{}' could not be executed. Aborting...".format(args.engine)
                raise

            engine.uci()
            handler = chess.uci.InfoHandler()
            engine.info_handlers.append(handler)

            # hash:  4096, 8192, 16384, 32768, 65536, 131072
            engine.setoption({
                "Threads": args.Threads,
                "Hash": args.Hash
            })
            engine.isready()

            engine.ucinewgame()
            engine.setoption({
                "UCI_AnalyseMode": True,
                "Analysis Contempt": "Off"
            })

            # load PGNs
            for fName in pgn_list.readlines():
                print("Processing: %s ..." % fName.rstrip())
                start = timer()
                game = chess.pgn.read_game(open(fName.rstrip()))
                node = game.end()

                analyze_node(node, args.add_plies, engine, handler, args, node.board().turn)

                with open("%s.done" % (fName.rstrip()), "w") as fOut:
                    print(game.root(), file=fOut, end="\n\n")

                print("... done")


            engine.quit()


    except PermissionError:
        errormsg = "Input file not readable. Aborting..."
        raise

    sys.exit()

if __name__ == "__main__":
    main()