comparison-script

import matplotlib
matplotlib.use('Agg')
import yt

def scan_blocklist(ds):
    leaf_blocks = {}
    other_blocks = {}
    for g in ds.index.grids:
        tmp = (g.block_name, g.filename)
        if (g.Level < 0) or (len(g.Children) > 1):
            other_blocks[g.block_name] = (g.filename, g.id)
        else:
            leaf_blocks[g.block_name] = (g.filename, g.id)

    return leaf_blocks, other_blocks


def walk_block_pairs(ds1, ds2):
    # this actually seems like a fundamental datatype

    leaf_blocks_1, _ = scan_blocklist(ds1)
    leaf_blocks_2, _ = scan_blocklist(ds2)

    num = len(leaf_blocks_1)
    assert num == len(leaf_blocks_2)

    for i, (key, val_1) in enumerate(leaf_blocks_1.items()):
        print(f"{i+1}/{num}: {key}")
        assert key in leaf_blocks_2
        val_2 = leaf_blocks_2[key]
        yield key, val_1, val_2

def get_fields_to_compare(blocklist_path):
    ds = yt.load(blocklist_path)
    particle_section = ds.parameters['Particle']
    if 'list' in ds.parameters['Particle']:
        particle_types = ds.parameters['Particle']['list']
    else:
        particle_types = []

    fields = []
    for field in ds.field_list:
        if field[0] in ['enzoe', 'enzop']:
            fields.append(field)
        elif field[0] in particle_types:
            fields.append(field)
    return fields

def compare_leafs(blocklist_path_ref, blocklist_path_other):

    fields_to_compare = get_fields_to_compare(blocklist_path_ref)

    ds1 = yt.load(blocklist_path_ref)
    ds2 = yt.load(blocklist_path_other)

    for name, (fname1,id1), (fname2,id2) in walk_block_pairs(ds1, ds2):
        g1 = ds1.index.grids[id1]
        g2 = ds2.index.grids[id2]

        for field in fields_to_compare:
            if not (g1[field] == g2[field]).all():
                raise RuntimeError(f'mismatch of {field}')
        g1.clear_data()
        g2.clear_data()

if __name__ == '__main__':
    blocklist_path_ref = #path/to/reference/output
    blocklist_path_other = #path/to/alt/output
    compare_leafs(blocklist_path_ref, blocklist_path_other)