Collection of scripts and small programs used by the EVN Support Scientists at JIVE during the regular data processing of EVN observations.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 

146 line
6.2 KiB

  1. #!/usr/bin/env python3
  2. """
  3. Flag visibilities with weights below the provided threshold.
  4. Usage: flag_weights.py msdata threshold
  5. Options:
  6. msdata : str MS data set containing the data to be flagged.
  7. threshold : float Visibilities with a weight below the specified
  8. value will be flagged. Must be positive.
  9. Version: 3.2
  10. Date: Mar 2021
  11. Written by Benito Marcote (marcote@jive.eu)
  12. version 3.2 changes (Mar 2021)
  13. - w > 0 condiction moved to > 0.001 to include possible weights at the ~1e-5 level.
  14. version 3.1 changes (Mar 2020)
  15. - Progress bar added.
  16. version 3.0 changes (Apr 2019)
  17. - Refactoring code (thanks to Harro).
  18. version 2.0 changes
  19. - Major revision. Now it does not modify the weights anymore. Instead, it
  20. flags those data with weights below the given threshold by modifying the
  21. FLAG table.
  22. - Small change in print messages to show '100%' instead of '1e+02%' in certain
  23. cases.
  24. version 1.4 changes
  25. - Now it also reports the percentage or data that were different from
  26. zero and will be flagged (not only the total data as before).
  27. version 1.3 changes
  28. - Minor fixes (prog name in optparse info).
  29. version 1.2 changes
  30. - Minor fixes.
  31. version 1.1 changes
  32. - Added option -v that allows you to just get how many visibilities will
  33. be flagged (but without actually flagging the data).
  34. """
  35. from pyrap import tables as pt
  36. import numpy as np
  37. import sys
  38. __version__ = 3.0
  39. help_msdata = 'Measurement set containing the data to be corrected.'
  40. help_threshold = 'Visibilities with a weight below this value will be flagged. Must be positive.'
  41. help_v = 'Only checks the visibilities to flag (do not flag the data).'
  42. try:
  43. usage = "%(prog)s [-h] <measurement set> <weight threshold>"
  44. description="""Flag visibilities with weights below the provided threshold.
  45. """
  46. import argparse
  47. parser = argparse.ArgumentParser(description=description, prog='flag_weights.py', usage=usage)
  48. parser.add_argument('msdata', type=str, help=help_msdata)
  49. parser.add_argument('threshold', type=float, help=help_threshold)
  50. parser.add_argument('--version', action='version', version='%(prog)s {}'.format(__version__))
  51. parser.add_argument("-v", "--verbose", default=True, action="store_false" , help=help_v)
  52. arguments = parser.parse_args()
  53. #print('The arguments ', arguments)
  54. verbose = arguments.verbose
  55. msdata = arguments.msdata[:-1] if arguments.msdata[-1]=='/' else arguments.msdata
  56. threshold = arguments.threshold
  57. except ImportError:
  58. usage = "%prog [-h] [-v] <measurement set> <weight threshold>"
  59. description="""Flag visibilities with weights below the provided threshold.
  60. """
  61. # Compatibility with Python 2.7 in eee
  62. import optparse
  63. parser = optparse.OptionParser(usage=usage, description=description, prog='flag_weights.py', version='%prog 1.3')
  64. parser.add_option("-v", action="store_false", dest="verbose", default=True, help=help_v)
  65. theparser = parser.parse_args()
  66. verbose = theparser[0].verbose
  67. arguments = theparser[1]
  68. #arguments = parser.parse_args()[1]
  69. if len(arguments) != 2:
  70. print('Two arguments must be provided: flag_weights.py [-h] [-v] <measurement set> <weight threshold>')
  71. print('Use -h to get help.')
  72. sys.exit(1)
  73. msdata = arguments[0][:-1] if arguments[0][-1]=='/' else arguments[0]
  74. threshold = float(arguments[1])
  75. assert threshold > 0.0
  76. def chunkert(f, l, cs, verbose=True):
  77. while f<l:
  78. n = min(cs, l-f)
  79. yield (f, n)
  80. f = f + n
  81. def cli_progress_bar(current_val, end_val, bar_length=40):
  82. percent = current_val/end_val
  83. hashes = '#'*int(round(percent*bar_length))
  84. spaces = ' '*(bar_length-len(hashes))
  85. sys.stdout.write("\rProgress: [{0}] {1}%".format(hashes+spaces, int(round(percent*100))))
  86. sys.stdout.flush()
  87. percent = lambda x, y: (float(x)/float(y))*100.0
  88. with pt.table(msdata, readonly=False, ack=False) as ms:
  89. total_number = 0
  90. flagged_before, flagged_after = (0, 0)
  91. flagged_nonzero, flagged_nonzero_before, flagged_nonzero_after = (0, 0, 0)
  92. # WEIGHT: (nrow, npol)
  93. # WEIGHT_SPECTRUM: (nrow, npol, nfreq)
  94. # flags[weight < threshold] = True
  95. weightcol = 'WEIGHT_SPECTRUM' if 'WEIGHT_SPECTRUM' in ms.colnames() else 'WEIGHT'
  96. transpose = (lambda x:x) if weightcol == 'WEIGHT_SPECTRUM' else (lambda x: x.transpose((1, 0, 2)))
  97. for (start, nrow) in chunkert(0, len(ms), 5000):
  98. cli_progress_bar(start, len(ms), bar_length=40)
  99. # shape: (nrow, npol, nfreq)
  100. flags = transpose(ms.getcol("FLAG", startrow=start, nrow=nrow))
  101. total_number += np.product(flags.shape)
  102. # count how much data is already flagged
  103. flagged_before += np.sum(flags)
  104. # extract weights and compute new flags based on threshold
  105. weights = ms.getcol(weightcol, startrow=start, nrow=nrow)
  106. # how many non-zero did we flag
  107. flagged_nonzero_before = np.logical_and(flags, weights > 0.001)
  108. # join with existing flags and count again
  109. flags = np.logical_or(flags, weights < threshold)
  110. flagged_after += np.sum(flags)
  111. flagged_nonzero_after = np.logical_and(flags, weights > 0.001)
  112. # Saving the total of nonzero flags (in this and previous runs)
  113. # flagged_nonzero += np.sum(np.logical_xor(flagged_nonzero_before, flagged_nonzero_after))
  114. flagged_nonzero += np.sum(flagged_nonzero_after)
  115. # one thing left to do: write the updated flags to disk
  116. #flags = ms.putcol("FLAG", flags.transpose((1, 0 , 2)), startrow=start, nrow=nrow)
  117. if verbose:
  118. flags = ms.putcol("FLAG", transpose(flags), startrow=start, nrow=nrow)
  119. print("\nGot {0:11} visibilities".format(total_number))
  120. print("Got {0:11} visibilities to flag using threshold {1}\n".format(flagged_after-flagged_before,
  121. threshold))
  122. print("{0:.2f}% total vis. flagged ({2:.2f}% to flag in this execution).\n{1:.2f}% data with non-zero weights flagged.\n".format(percent(flagged_after, total_number), percent(flagged_nonzero, total_number), percent(flagged_after-flagged_before, total_number)))
  123. ms.close()
  124. if verbose:
  125. print('Done.')
  126. else:
  127. print('Flags have not been applied.')