compare.py

import sys
import h5py
import numpy as np

def pygame(h5, name, board, moves):
    for idx in range(board.shape[0]):
        for plane in range(board.shape[1]):
            if plane < 44:
                # different board orderings
                f = board[idx, plane, ...].T[:, ::-1]
                m = moves[idx, ...]
                m = (m[1], 18 - m[0])
                yield h5, name, f, m, idx, plane, board[idx, ...]

def aggame(h5, name, board, moves):
    for idx in range(board.shape[0]):
        for plane in range(board.shape[1]):
            if plane < 44:
                yield h5, name, board[idx, plane, ...], moves[idx, ...], idx, plane, board[idx, ...]

def pyfuego_gen(f):
    names = f['gamefiles'][...]
    starts = f['gameoffsets'][...]
    ends = list(f['gameoffsets'][...])[1:] + [f['X'].shape[0]]
    for n, s, e in sorted(zip(names, starts, ends)):
        yield pygame(f, n, f['X'][s:e, ...], f['y'][s:e, ...])

def ag_gen(f):
    names = list(f['file_offsets'])
    for n in sorted(names):
        s, e = f['file_offsets'][n][...]
        e = s + e
        yield aggame(f, n, f['states'][s:e, ...], f['actions'][s:e, ...])

f1 = h5py.File(sys.argv[1], 'r')
f2 = h5py.File(sys.argv[2], 'r')

if 'X' in f1.keys():
    gen1 = pyfuego_gen(f1)
else:
    gen1 = ag_gen(f1)

if 'X' in f2.keys():
    gen2 = pyfuego_gen(f2)
else:
    gen2 = ag_gen(f2)

for game1, game2 in zip(gen1, gen2):
    for (h1, n1, f1, m1, i1, p1, b1), (h2, n2, f2, m2, i2, p2, b2) in zip(game1, game2):
        if not np.all(f1 == f2) and 'komi' not in n1:
            import pdb
            print("difference in plane {} at move {} game {} {}".format(p1, i1, n1, n2))
            pdb.set_trace()