check_workflow_results.py

#! /usr/bin/env python3
"""
Checks the outputs of a workflow run
"""
import argparse
from argparse import RawTextHelpFormatter
from losoto.h5parm import h5parm
import os
import sys
from filecmp import dircmp
import numpy as np


def main(results_path, control_path):
    """
    Checks the outputs of an end-to-end workflow job

    Parameters
    ----------
    results_path : str
        Path to output results directory of job
    control_path : str
        Path to control results directory
    """
    error = False
    # Check that all expected output files are present
    dcmp = dircmp(results_path, control_path)
    if len(dcmp.left_only) > 0 or len(dcmp.right_only) > 0:
        if len(dcmp.left_only) > 0:
            print('ERROR: The following files are present in the output but not in the '
                  'control: {}'.format(dcmp.left_only))
        if len(dcmp.right_only) > 0:
            print('ERROR: The following files are present in the control but not in the '
                  'output: {}'.format(dcmp.right_only))
        error = True

    # Check that the calibration solutions match the control ones
    check_h5parm = os.path.join(results_path, 'cal_values', 'cal_solutions.h5')
    control_h5parm = os.path.join(control_path, 'cal_values', 'cal_solutions.h5')
    check_h5 = h5parm(check_h5parm, readonly=True)
    control_h5 = h5parm(control_h5parm, readonly=True)
    for check_solset in check_h5.getSolsets():
        control_solset = control_h5.getSolset(check_solset.name)
        for check_soltab in check_solset.getSoltabs():
            control_soltab = control_solset.getSoltab(check_soltab.name)
            if not np.allclose(check_soltab.val, control_soltab.val, rtol=1e-03,
                               atol=1e-03, equal_nan=True):
                error = True
                print('ERROR: Val array of soltab {} of solset {} does not match '
                      'the control'.format(check_soltab.name, check_solset.name))
                with open("check_soltab.{}.val".format(check_soltab.name), "w") as f:
                    f.write(str(check_soltab.val))
                with open("control_soltab.{}.val".format(control_soltab.name), "w") as f:
                    f.write(str(control_soltab.val))
            else:
                print('INFO: Val array of soltab {} of solset {} matches '
                      'the control'.format(check_soltab.name, check_solset.name))
    control_h5.close()
    check_h5.close()
    sys.exit(error)

if __name__ == '__main__':
    descriptiontext = "Checks the ouput of a workflow run.\n"

    parser = argparse.ArgumentParser(description=descriptiontext, formatter_class=RawTextHelpFormatter)
    parser.add_argument('results_path', help='Path to output results')
    parser.add_argument('control_path', help='Path to control results')
    args = parser.parse_args()

    main(args.results_path, args.control_path)