check_workflow_results.py

#!/usr/bin/env python3
"""
Checks the outputs of a workflow run
"""
import argparse
from argparse import RawTextHelpFormatter
import h5py
import sys
import numpy as np

def main(check, control):
    """
    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
    """
    check_path = check + '/cal_values/cal_solutions.h5'
    control_path = control + '/cal_values/cal_solutions.h5'
    check_h5 = h5py.File(check_path, 'r')
    control_h5 = h5py.File(control_path, 'r')

    for solsetname in check_h5:
        soltabnames = [name for name in check_h5[solsetname] 
                            if name not in ('source', 'antenna')]
        try:
            control_h5[solsetname]
        except:
            print("Error: solset {} not present in control.".format(solsetname))
            sys.exit(1)
        for soltabname in soltabnames:
            try:
                control_h5[solsetname][soltabname]
            except:
                print("Error: soltab {} not present in control.".format(soltabname))
                sys.exit(1)
            check_soltabval = check_h5[solsetname][soltabname]['val']
            control_soltabval = control_h5[solsetname][soltabname]['val']
            matching_vals = control_soltabval
            if 'freq' in check_h5[solsetname][soltabname].keys():
                check_axes = check_soltabval.attrs['AXES'].decode('utf-8').split(',')
                freq_axis_index = check_axes.index('freq')
                check_soltabfreq = check_h5[solsetname][soltabname]['freq'][:]
                control_soltabfreq = control_h5[solsetname][soltabname]['freq'][:]
                matches = np.isclose(control_soltabfreq[:, np.newaxis], check_soltabfreq)
                matching_freq_indices = np.where(matches)[0]
                matching_vals = np.take(control_soltabval, 
                                        matching_freq_indices, 
                                        axis=freq_axis_index)
            if not np.allclose(check_soltabval, matching_vals, 
                               rtol=1e-03, atol=1e-03, equal_nan=True):
                print("Val array of soltab {} does not match the control".format(soltabname))
                with open("check_soltab.{}.val".format(soltabname), "w") as f:
                    f.write(str(check_soltabval[:]))
                with open("control_soltab.{}.val".format(soltabname), "w") as f:
                    f.write(str(check_soltabval[:]))
            else:
                print("Val array of soltab {} matches the control".format(soltabname))
    check_h5.close()
    control_h5.close()
    sys.exit(0)

if __name__ == '__main__':
    descriptiontext = "Checks the output 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)