reformat

box_helpers/move_boxes.py

@@ -6,6 +6,8 @@
 
 from glob import glob
 import os
+from argparse import ArgumentParser
+
 
 def move_boxes(file, folder):
     try:
@@ -30,8 +32,9 @@ def move_boxes(file, folder):
         print("Failing to open ds9 to verify box selection, check if installed and try to run on the commandline"
               "\nds9 {FILE} -regions load all '{DATALOC}/boxes/*.reg'".format(FILE=file, DATALOC=folder))
 
-if __name__ == '__main__':
-    from argparse import ArgumentParser
+def parse_args():
+    """Command line parser"""
+
     parser = ArgumentParser()
     parser.add_argument('-l', '--location', type=str, help='data location folder name', default='.')
     parser.add_argument('-f', '--file', type=str, help='fitsfile name',
@@ -42,4 +45,14 @@ def move_boxes(file, folder):
     if folder[-1] == '/':
         folder = folder[0:-1]
 
+    return args, folder
+
+def main():
+    """Main function"""
+
+    args, folder = parse_args()
+
     move_boxes(args.file, folder)
+
+if __name__ == '__main__':
+    main()

h5_helpers/add_h5_dirs.py

@@ -3,17 +3,30 @@
 import os
 import numpy as np
 
-if __name__ == '__main__':
+
+def parse_args():
+    """
+    Command line argument parser
+
+    :return: parsed arguments
+    """
 
     parser = ArgumentParser()
-    parser.add_argument('--h5_in', type=str, help='input h5 (from which to extract the frequencies and antennas)', required=True)
+    parser.add_argument('--h5_in', type=str, help='input h5 (from which to extract the frequencies and antennas)',
+                        required=True)
     parser.add_argument('--h5_dirs', type=str, help='h5 with directions to copy from', required=True)
     parser.add_argument('--h5_out', type=str, help='output h5')
-    args = parser.parse_args()
+    return parser.parse_args()
+
+
+def main():
+    """Main function"""
+
+    args = parse_args()
 
     os.system(f'cp {args.h5_in} {args.h5_out}')
 
-    print("Make "+args.h5_out)
+    print("Make " + args.h5_out)
 
     H = tables.open_file(args.h5_out, "r+")
     T = tables.open_file(args.h5_dirs)
@@ -41,4 +54,8 @@
             H.root.sol000._f_get_child(table)._f_get_child(axes).attrs['AXES'] = AXES
 
         H.root.sol000._f_get_child(table).dir._f_remove()
-        H.create_array(H.root.sol000._f_get_child(table), 'dir', T.root.sol000._f_get_child(table).dir[:])
+        H.create_array(H.root.sol000._f_get_child(table), 'dir', T.root.sol000._f_get_child(table).dir[:])
+
+
+if __name__ == '__main__':
+    main()

h5_helpers/find_closest_h5.py

@@ -5,10 +5,12 @@
 import sys
 from argparse import ArgumentParser
 
+
 class FindClosestDir:
     """
     Make template h5
     """
+
     def __init__(self, h5_in, template_name):
         os.system(' '.join(['cp', h5_in, template_name]))
         print(f'Created {template_name}')
@@ -20,7 +22,6 @@ def __init__(self, h5_in, template_name):
         self.dirs = T.root.sol000.source[:]['dir']
         T.close()
 
-
     def make_template(self):
         """
         Make template h5 with 1 direction
@@ -32,7 +33,7 @@ def make_template(self):
         for solset in self.h5.root._v_groups.keys():
             ss = self.h5.root._f_get_child(solset)
             ss._f_get_child('source')._f_remove()
-            values = np.array([(b'Dir00', [0. ,  0. ])], dtype=[('name', 'S128'), ('dir', '<f4', (2,))])
+            values = np.array([(b'Dir00', [0., 0.])], dtype=[('name', 'S128'), ('dir', '<f4', (2,))])
             title = 'Source names and directions'
             self.h5.create_table(ss, 'source', values, title=title)
 
@@ -46,9 +47,9 @@ def make_template(self):
                     shape[dir_idx] = 1
 
                     # only phases are zeros, others are ones
-                    if 'phase' in soltab and axes!='weight':
+                    if 'phase' in soltab and axes != 'weight':
                         newvals = np.zeros(shape)
-                    elif 'amplitude' in soltab or axes=='weight':
+                    elif 'amplitude' in soltab or axes == 'weight':
                         newvals = np.ones(shape)
                     else:
                         newvals = np.zeros(shape)
@@ -87,11 +88,11 @@ def closest_dir(self, coor):
             c1 = SkyCoord(dir[0], dir[1], unit='arcsecond', frame='icrs')  # your coords
             c2 = SkyCoord(coor[0], coor[1], unit='arcsecond', frame='icrs')
             sep = c1.separation(c2).value
-            if sep<min_sep:
+            if sep < min_sep:
                 min_sep_idx = n
 
         # Make sure this function selects a closest direction
-        assert min_sep_idx!=999, "ERROR: coordinate comparison went wrong?! (most likely a bug in h5 or code)"
+        assert min_sep_idx != 999, "ERROR: coordinate comparison went wrong?! (most likely a bug in h5 or code)"
         print(f"Closest direction for {coor} is {self.dirs[min_sep_idx]}")
 
         return min_sep_idx
@@ -128,7 +129,7 @@ def add_closest_values(self, coor):
                     # select values corresponding to closest direction index (idx variable)
                     allvals = self.h5_in.root._f_get_child(solset)._f_get_child(soltab)._f_get_child(axes)[:]
                     newvals = np.take(allvals, indices=[idx], axis=dir_idx)
-                    del allvals # save memory
+                    del allvals  # save memory
 
                     # get correct value type
                     valtype = str(st._f_get_child(axes).dtype)
@@ -149,22 +150,35 @@ def add_closest_values(self, coor):
         self.h5_out.close()
         return self
 
+
 def make_list(arglist):
     try:
-        return [[float(d) for d in dir.replace('[','').replace(']','').replace('(','').replace(')','').replace(' ','').split(',')] for dir in arglist]
+        return [[float(d) for d in
+                 dir.replace('[', '').replace(']', '').replace('(', '').replace(')', '').replace(' ', '').split(',')]
+                for dir in arglist]
     except ValueError:
         sys.exit("ERROR: --directions input invalid\nDo not use any spaces, example input: [0.1,0.2] [1.2,4.1]")
 
 
+def parse_args():
+    """
+    Command line argument parser
 
-
-if __name__ == '__main__':
+    :return: parsed arguments
+    """
 
     parser = ArgumentParser()
     parser.add_argument('--h5_in', help='Input h5parm', required=True)
-    parser.add_argument('--directions', nargs='+', help='directions to find the closest h5_in direction to (Example: (0.1, 0.2) (1.2, 3.1) (3.5, 1.2)', default=None)
-    args = parser.parse_args()
+    parser.add_argument('--directions', nargs='+',
+                        help='directions to find the closest h5_in direction to (Example: (0.1, 0.2) (1.2, 3.1) (3.5, 1.2)',
+                        default=None)
+    return parser.parse_args()
+
+
+def main():
+    """Main function"""
 
+    args = parse_args()
     inputh5 = args.h5_in
     dirs = make_list(args.directions)
 
@@ -179,3 +193,7 @@ def make_list(arglist):
     print('See output in --> output_h5s\n'
           'You can optionally use h5_merger.py to merge all the directions into 1 big h5 file\n'
           'Example command --> python h5_merger.py -in output_h5s/source_*.h5 -out merged.h5 --propagate_flags')
+
+
+if __name__ == '__main__':
+    main()

h5_helpers/h5_filter.py

@@ -16,6 +16,7 @@
 from glob import glob
 import re
 
+
 def str2bool(v):
     v = str(v)
     if v.lower() in ('yes', 'true', 't', 'y', '1'):
@@ -25,24 +26,29 @@ def str2bool(v):
     else:
         raise ArgumentTypeError('Boolean value expected.')
 
+
 def degree_to_radian(inp):
     """degree to radian"""
     return float(inp) / 360 * pi * 2
 
+
 def radian_to_degree(inp):
     """radion to degree"""
     return float(inp) * 360 / (pi * 2)
 
+
 def angular_distance(p1, p2):
     """angular distance for points in ra and dec in degrees"""
-    if p1[0]>2*pi:
+    if p1[0] > 2 * pi:
         p1 = [degree_to_radian(p) for p in p1]
         p2 = [degree_to_radian(p) for p in p2]
-    return radian_to_degree(acos(sin(p1[1])*sin(p2[1])+cos(p1[1])*cos(p2[1])*cos(p1[0]-p2[0])))
+    return radian_to_degree(acos(sin(p1[1]) * sin(p2[1]) + cos(p1[1]) * cos(p2[1]) * cos(p1[0] - p2[0])))
+
 
 def remove_numbers(inp):
     return "".join(re.findall("[a-zA-z]+", inp))
 
+
 def create_new_soltab(h5_in_name, h5_out_name, directions, sources):
     """
     Create a new dataset in the h5 table
@@ -68,7 +74,9 @@ def create_new_soltab(h5_in_name, h5_out_name, directions, sources):
             solsetout.obj.source.append(new_sources)
 
         for st in h5_in.getSolset(ss).getSoltabNames():
-            print('Filter {solset}/{soltab} from {h5_in} into {h5_out}'.format(solset=ss, soltab=st, h5_in=h5_in_name.split('/')[-1], h5_out=h5_out_name.split('/')[-1]))
+            print('Filter {solset}/{soltab} from {h5_in} into {h5_out}'.format(solset=ss, soltab=st,
+                                                                               h5_in=h5_in_name.split('/')[-1],
+                                                                               h5_out=h5_out_name.split('/')[-1]))
 
             solutiontable = h5_in.getSolset(ss).getSoltab(st)
             axes = solutiontable.getValues()[1]
@@ -77,65 +85,86 @@ def create_new_soltab(h5_in_name, h5_out_name, directions, sources):
             axes['dir'] = [ns[0] for ns in new_sources]
             dir_index = solutiontable.getAxesNames().index('dir')
             shape = list(values_in.shape)
-            shape[dir_index]=len(directions)
+            shape[dir_index] = len(directions)
             values_new = zeros(shape)
 
             for idx_new, idx_old in enumerate(indexes):
                 if dir_index == 0:
-                    values_new[idx_new,...] += values_in[idx_old, ...]
+                    values_new[idx_new, ...] += values_in[idx_old, ...]
                 elif dir_index == 1:
-                    values_new[:, idx_new,...] += values_in[:, idx_old, ...]
+                    values_new[:, idx_new, ...] += values_in[:, idx_old, ...]
                 elif dir_index == 2:
-                    values_new[:, :, idx_new,...] += values_in[:, :, idx_old, ...]
+                    values_new[:, :, idx_new, ...] += values_in[:, :, idx_old, ...]
                 elif dir_index == 3:
-                    values_new[:, :, :, idx_new,...] += values_in[:, :, :, idx_old, ...]
+                    values_new[:, :, :, idx_new, ...] += values_in[:, :, :, idx_old, ...]
                 elif dir_index == 4:
-                    values_new[:, :, :, :, idx_new,...] += values_in[:, :, :, :, idx_old,...]
+                    values_new[:, :, :, :, idx_new, ...] += values_in[:, :, :, :, idx_old, ...]
 
             print('New number of sources {num}'.format(num=len(sources)))
             print('Filtered output shape {shape}'.format(shape=values_new.shape))
 
             weights = ones(values_new.shape)
-            solsetout.makeSoltab(remove_numbers(st), axesNames=list(axes.keys()), axesVals=list(axes.values()), vals=values_new,
-                             weights=weights)
+            solsetout.makeSoltab(remove_numbers(st), axesNames=list(axes.keys()), axesVals=list(axes.values()),
+                                 vals=values_new,
+                                 weights=weights)
 
     h5_in.close()
     h5_out.close()
 
-parser = ArgumentParser()
-parser.add_argument('-f', '--fits', type=str, help='fitsfile name')
-parser.add_argument('-ac', '--angular_cutoff', type=float, default=None, help='angular distances higher than this value from the center will be excluded from the box selection')
-parser.add_argument('-in', '--inside', type=str2bool, default=False, help='keep directions inside the angular cutoff')
-parser.add_argument('-h5out', '--h5_file_out', type=str, help='h5 output name')
-parser.add_argument('-h5in', '--h5_file_in', type=str, help='h5 input name (to filter)')
-args = parser.parse_args()
-
-# clean up files
-if args.h5_file_out.split('/')[-1] in [f.split('/')[-1] for f in glob(args.h5_file_out)]:
-    os.system('rm -rf {file}'.format(file=args.h5_file_out))
-
-# get header from fits file
-hdu = fits.open(args.fits)[0]
-header = WCS(hdu.header, naxis=2).to_header()
-# get center of field
-center = (degree_to_radian(header['CRVAL1']), degree_to_radian(header['CRVAL2']))
-
-# return list of directions that have to be included
-H = tables.open_file(args.h5_file_in)
-sources = []
-directions = []
-for dir in H.root.sol000.source[:]:
-    print(angular_distance(center, dir[1]))
-    if args.inside and angular_distance(center, dir[1])<args.angular_cutoff:
-        print('Keep {dir}'.format(dir=dir))
-        directions.append(dir[0])
-        sources.append(dir)
-    elif not args.inside and angular_distance(center, dir[1])>=args.angular_cutoff:
-        print('Keep {dir}'.format(dir=dir))
-        directions.append(dir[0])
-        sources.append(dir)
-    else:
-        print('Remove {dir}'.format(dir=dir))
-H.close()
 
-create_new_soltab(args.h5_file_in, args.h5_file_out, directions, sources)
+def parse_args():
+    """
+    Command line argument parser
+
+    :return: parsed arguments
+    """
+
+    parser = ArgumentParser()
+    parser.add_argument('-f', '--fits', type=str, help='fitsfile name')
+    parser.add_argument('-ac', '--angular_cutoff', type=float, default=None,
+                        help='angular distances higher than this value from the center will be excluded from the box selection')
+    parser.add_argument('-in', '--inside', type=str2bool, default=False,
+                        help='keep directions inside the angular cutoff')
+    parser.add_argument('-h5out', '--h5_file_out', type=str, help='h5 output name')
+    parser.add_argument('-h5in', '--h5_file_in', type=str, help='h5 input name (to filter)')
+    return parser.parse_args()
+
+
+def main():
+    """Main function"""
+
+    args = parse_args()
+
+    # clean up files
+    if args.h5_file_out.split('/')[-1] in [f.split('/')[-1] for f in glob(args.h5_file_out)]:
+        os.system('rm -rf {file}'.format(file=args.h5_file_out))
+
+    # get header from fits file
+    hdu = fits.open(args.fits)[0]
+    header = WCS(hdu.header, naxis=2).to_header()
+    # get center of field
+    center = (degree_to_radian(header['CRVAL1']), degree_to_radian(header['CRVAL2']))
+
+    # return list of directions that have to be included
+    H = tables.open_file(args.h5_file_in)
+    sources = []
+    directions = []
+    for dir in H.root.sol000.source[:]:
+        print(angular_distance(center, dir[1]))
+        if args.inside and angular_distance(center, dir[1]) < args.angular_cutoff:
+            print('Keep {dir}'.format(dir=dir))
+            directions.append(dir[0])
+            sources.append(dir)
+        elif not args.inside and angular_distance(center, dir[1]) >= args.angular_cutoff:
+            print('Keep {dir}'.format(dir=dir))
+            directions.append(dir[0])
+            sources.append(dir)
+        else:
+            print('Remove {dir}'.format(dir=dir))
+    H.close()
+
+    create_new_soltab(args.h5_file_in, args.h5_file_out, directions, sources)
+
+
+if __name__ == '__main__':
+    main()