Replace check_workflow_results.py

Former-commit-id: b4841233 [formerly 7026d589]
Former-commit-id: 46758d32
Former-commit-id: d86b916d

test_jobs/check_workflow_results.py

-#! /usr/bin/env python3
+#!/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 h5py
 import sys
-from filecmp import dircmp
 import numpy as np
 
-
-def main(results_path, control_path):
+def main(check, control):
     """
     Checks the outputs of an end-to-end workflow job
 
@@ -22,45 +19,53 @@ def main(results_path, control_path):
     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_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')
 
-    # 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))
+    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('INFO: Val array of soltab {} of solset {} matches '
-                      'the control'.format(check_soltab.name, check_solset.name))
-    control_h5.close()
+                print("Val array of soltab {} matches the control".format(soltabname))
     check_h5.close()
-    sys.exit(error)
+    control_h5.close()
+    sys.exit(0)
 
 if __name__ == '__main__':
-    descriptiontext = "Checks the ouput of a workflow run.\n"
+    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')