Analyze chess960 starting positions for chess or chess variants.

atomic960.py

import argparse
import csv
import chess, chess.variant, chess.uci


def parse_args():
    parser = argparse.ArgumentParser(description='Evaluate atomic960 positions.')
    parser.add_argument('-e', '--engine', required=True, help='Path to engine.')
    parser.add_argument('-o', '--output', required=True, help='Output csv file.')
    parser.add_argument('-d', '--depth', type=int, default=10, help='Search depth.')
    parser.add_argument('-m', '--memory', type=int, default=16, help='Hash size in MB.')
    parser.add_argument('-a', '--start', type=int, default=0, help='Start from this position.')
    parser.add_argument('-b', '--stop', type=int, default=959, help='Stop at this position.')
    parser.add_argument('-v', '--variant', type=str, default='atomic', help='Variant name.')
    
    return parser.parse_args()


def main(args):
    board = chess.variant.find_variant(args.variant)()

    engine = chess.uci.popen_engine(args.engine)
    info_handler = chess.uci.InfoHandler()
    engine.info_handlers.append(info_handler)

    engine.uci()
    engine.setoption({'UCI_Variant': board.uci_variant, 'UCI_Chess960': True, 'Hash': args.memory})

    with open(args.output, 'w') as csv_file:
        columns = ['index', 'FEN', 'depth', 'score']
        csv_writer = csv.DictWriter(csv_file, fieldnames=columns, dialect='excel-tab')
        csv_writer.writeheader()

        print('Starting to analyze positions %d to %d.\nSettings: depth=%d, hash=%dMB'
              % (args.start, args.stop, args.depth, args.memory))
        for i in range(args.start, args.stop + 1):
            print('Analyzing position %d.' % i)
            board.set_chess960_pos(i)
            engine.position(board)
            engine.go(depth=args.depth)
            score = (info_handler.info["score"][1].cp if info_handler.info["score"][1].cp is not None
                     else ('#%s' % info_handler.info["score"][1].mate))
            csv_writer.writerow({'index': i, 'FEN': board.fen(), 'depth': args.depth, 'score': score})
        print('Done.\nOutput written to file %s.' % args.output)


if __name__ == "__main__":
    main(parse_args())