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.
 
 

244 lines
10 KiB

  1. #! /usr/bin/env python3
  2. """
  3. Swap polarizations for specified antennas and for a specific timerange.
  4. Usage: polswap.py msdata antennas [-sb SUBBANDS]
  5. Options:
  6. msdata : str MS data set containing the data to be swapged.
  7. antennas : str Antenna or list of antennas that require to be
  8. swapped. Use the two-letter name (e.g. Ef, Mc,
  9. or Jb2). In case of more than one station, please
  10. use either string 'Ef, Mc' or a non-spaced str:
  11. Ef,Mc,Ys.
  12. Version: 3.0
  13. Date: March 2020
  14. Written by Benito Marcote (marcote@jive.eu)
  15. version 3.0 changes (March 2020)
  16. - Columns with different dimensions grouped together in code.
  17. - Fix memory bug.
  18. - Not it modifies all expected columns for each chunk of (time) data.
  19. Before it was the opposite: modify one column at the time.
  20. - Progress bar implemented.
  21. version 2.0 changes (March 2020)
  22. - MS read and modified in chunks of data.
  23. - Simplified code.
  24. - Bug Fix? Issues in a particular MS with the wrong station only available part of the time.
  25. version 1.2 changes (July 2018)
  26. - Several bug fixes.
  27. """
  28. import sys
  29. import copy
  30. import time
  31. import argparse
  32. from enum import IntEnum
  33. import datetime as dt
  34. import numpy as np
  35. from pyrap import tables as pt
  36. usage = "%(prog)s [-h] [-v] [-t1 STARTTIME] [-t2 ENDTIME] <measurement set> <antenna>"
  37. description="""Swap polarizations for specified antennas.
  38. Fixes the polarizations of an antenna that have been labeled incorrectly (R or X corresponds to L or Y,
  39. respectively; and L or Y to R or X). It also changes accordingly the cross-pols per each baseline
  40. containing the mentioned antenna (RL,LR, or XY,YX).
  41. polswap.py works for both types of polarizations: circular and linear pols.
  42. """
  43. help_msdata = 'Measurement Set containing the data to be corrected.'
  44. help_antenna = 'Name of the antenna to be corrected as it appears in the MS (case insensitive).'
  45. help_t1 = 'Start time of the data that need to be corrected. By default the beginning of the observation.'\
  46. +'In Aips format: YYYY/MM/DD/hh:mm:ss or YYYY/DOY/hh:mm:ss'
  47. help_t2 = 'Ending time of the data that need to be corrected. By default the ending of the observation.'\
  48. +'In Aips format: YYYY/MM/DD/hh:mm:ss or YYYY/DOY/hh:mm:ss'
  49. parser = argparse.ArgumentParser(description=description, prog='polswap.py', usage=usage)
  50. parser.add_argument('msdata', type=str, help=help_msdata)
  51. parser.add_argument('antenna', type=str, help=help_antenna)
  52. parser.add_argument('-t1', '--starttime', default=None, type=str, help=help_t1)
  53. parser.add_argument('-t2', '--endtime', default=None, type=str, help=help_t2)
  54. parser.add_argument('-v', '--version', action='version', version='%(prog)s 1.0')
  55. # parser.add_argument('--verbose', default=False, action='store_true')
  56. # parser.add_argument('--timing', default=False, action='store_true')
  57. arguments = parser.parse_args()
  58. msdata = arguments.msdata[:-1] if arguments.msdata[-1]=='/' else arguments.msdata
  59. def cli_progress_bar(current_val, end_val, bar_length=40):
  60. percent = current_val/end_val
  61. hashes = '#'*int(round(percent*bar_length))
  62. spaces = ' '*(bar_length-len(hashes))
  63. sys.stdout.write("\rProgress: [{0}] {1}%".format(hashes+spaces, int(round(percent*100))))
  64. sys.stdout.flush()
  65. class Stokes(IntEnum):
  66. """The Stokes types defined as in the enum class from casacore code.
  67. """
  68. Undefined = 0 # Undefined value
  69. I = 1 # standard stokes parameters
  70. Q = 2
  71. U = 3
  72. V = 4
  73. RR = 5 # circular correlation products
  74. RL = 6
  75. LR = 7
  76. LL = 8
  77. XX = 9 # linear correlation products
  78. XY = 10
  79. YX = 11
  80. YY = 12
  81. RX = 13 # mixed correlation products
  82. RY = 14
  83. LX = 15
  84. LY = 16
  85. XR = 17
  86. XL = 18
  87. YR = 19
  88. YL = 20
  89. PP = 21 # general quasi-orthogonal correlation products
  90. PQ = 22
  91. QP = 23
  92. QQ = 24
  93. RCircular = 25 # single dish polarization types
  94. LCircular = 26
  95. Linear = 27
  96. Ptotal = 28 # Polarized intensity ((Q^2+U^2+V^2)^(1/2))
  97. Plinear = 29 # Linearly Polarized intensity ((Q^2+U^2)^(1/2))
  98. PFtotal = 30 # Polarization Fraction (Ptotal/I)
  99. PFlinear = 31 # linear Polarization Fraction (Plinear/I)
  100. Pangle = 32 # linear polarization angle (0.5 arctan(U/Q)) (in radians)
  101. def chunkert(f, l, cs, verbose=True):
  102. while f<l:
  103. n = min(cs, l-f)
  104. yield (f, n)
  105. f = f + n
  106. def atime2datetime(atime):
  107. """Converts a string with the form YYYY/MM/DD/hh:mm:ss or YYYY/DOY/hh:mm:ss to datetime"""
  108. if atime.count('/') == 3:
  109. # Format: YYYY/MM/DD/hh:mm:ss
  110. return dt.datetime.strptime(atime, '%Y/%m/%d/%H:%M:%S')
  111. if atime.count('/') == 2:
  112. # Format: YYYY/DOY/hh:mm:ss
  113. return dt.datetime.strptime(atime, '%Y/%j/%H:%M:%S')
  114. else:
  115. raise ValueError('Date format must be YYYY/MM/DD/hh:mm:ss or YYYY/DOY/hh:mm:ss')
  116. def get_nedded_move(products, ant_order):
  117. """Returns the transposing necessary to do a polswap in one of the stations.
  118. Inputs
  119. ------
  120. products : 2-D array-like
  121. is the CORR_PRODUCT of the MS data. Sets the mapping of the stokes given
  122. two stations. e.g. [[0, 0], [1, 1], [0, 1], [1, 0]] represents that there
  123. are four stokes products, where the first two rows are the direct hands
  124. between antenna 1 (first column) and antenna 2 (second column).
  125. ant_order : int
  126. The position of the antenna in the CORR_PRODUCT (if the antenna to change
  127. is the ANT1 or ANT2, i.e. 0 or 1.
  128. Outputs
  129. -------
  130. changes : 1-D array-like
  131. The transposition of the columns necessary to make a swap pol for the antenna
  132. specified in ant_order.
  133. e.g. for the case mentioned before, if ANT1 is the one that needs to be converted,
  134. then the output 'changes' is [3, 2, 1, 0], as the stokes wanted at the end are:
  135. [[1, 0], [0, 1], [1, 1], [0, 0]]
  136. """
  137. pols_prod = list([list(j) for j in products])
  138. pols_prod_mod = np.copy(products)
  139. pols_prod_mod[:,ant_order] = products[:,ant_order] ^ 1
  140. pols_prod_mod = list([list(j) for j in pols_prod_mod])
  141. return np.array([pols_prod.index(i) for i in pols_prod_mod])
  142. with pt.table(msdata, readonly=False, ack=False) as ms:
  143. # Get CORR_TYPE, the polarization entries in the data
  144. changes = None
  145. with pt.table(ms.getkeyword('ANTENNA'), readonly=True, ack=False) as ms_ant:
  146. antenna_number = [i.upper() for i in ms_ant.getcol('NAME')].index(arguments.antenna.upper())
  147. with pt.table(ms.getkeyword('POLARIZATION'), readonly=True, ack=False) as ms_pol:
  148. pols_order = [Stokes(i) for i in ms_pol.getcol('CORR_TYPE')[0]]
  149. # Check that the stokes are the correct ones to do a cross pol.
  150. # Only change it if circular or linear pols.
  151. for a_pol_order in pols_order:
  152. if (a_pol_order not in (Stokes.RR, Stokes.RL, Stokes.LR, Stokes.LL)) and \
  153. (a_pol_order not in (Stokes.XX, Stokes.XY, Stokes.YX, Stokes.YY)) and \
  154. (a_pol_order not in (Stokes.RX, Stokes.RY, Stokes.LX, Stokes.LY)) and \
  155. (a_pol_order not in (Stokes.XR, Stokes.XL, Stokes.YR, Stokes.YL)):
  156. print('Polswap only works for circular or linear pols (or both combined).')
  157. print('These data contain the following stokes: {}'.format(pols_order))
  158. ms.close()
  159. raise ValueError('Wrong stokes type.')
  160. # Get the column changes that are necessary
  161. pols_prod = ms_pol.getcol('CORR_PRODUCT')[0]
  162. # print(pols_prod)
  163. changes = [get_nedded_move(pols_prod, i) for i in (0, 1)]
  164. # transpose data for columns DATA, WEIGHT_SPECTRUM (if exists)
  165. # ants = [ms.getcol('ANTENNA1'), ms.getcol('ANTENNA2')]
  166. with pt.table(ms.getkeyword('OBSERVATION'), readonly=True, ack=False) as ms_obs:
  167. time_range = (dt.datetime(1858, 11, 17, 0, 0, 2) + \
  168. ms_obs.getcol('TIME_RANGE')*dt.timedelta(seconds=1))[0]
  169. # Get the timerange to apply to polswap
  170. if arguments.starttime is not None:
  171. datetimes_start = atime2datetime(arguments.starttime)
  172. else:
  173. datetimes_start = time_range[0] - dt.timedelta(seconds=1)
  174. if arguments.endtime is not None:
  175. datetimes_end = atime2datetime(arguments.endtime)
  176. else:
  177. datetimes_end = time_range[1] + dt.timedelta(seconds=1)
  178. # shapes of DATA, FLOAT_DATA, FLAG, SIGMA_SPECTRUM, WEIGHT_SPECTRUM: (nrow, npol, nfreq)
  179. # shapes of WEIGHT, SIGMA: (nrow, npol)
  180. columns = ('DATA', 'FLOAT_DATA', 'FLAG', 'SIGMA_SPECTRUM', 'WEIGHT_SPECTRUM',
  181. 'WEIGHT', 'SIGMA')
  182. # Only leave the ones that are in the MS. Not all of them are always present.
  183. columns = [a_col for a_col in columns if a_col in ms.colnames()]
  184. print('\nThe following columns will be modified: {}.\n'.format(', '.join(columns)))
  185. for (start, nrow) in chunkert(0, len(ms), 5000):
  186. cli_progress_bar(start, len(ms), bar_length=40)
  187. for changei, antpos in zip(changes, ('ANTENNA1','ANTENNA2')):
  188. ants = ms.getcol(antpos, startrow=start, nrow=nrow)
  189. datetimes = dt.datetime(1858, 11, 17, 0, 0, 2) + \
  190. ms.getcol('TIME', startrow=start, nrow=nrow)*dt.timedelta(seconds=1)
  191. cond = np.where((ants == antenna_number) & (datetimes > datetimes_start) & (datetimes < datetimes_end))
  192. if len(cond[0]) > 0:
  193. for a_col in columns:
  194. ms_col = ms.getcol(a_col, startrow=start, nrow=nrow)
  195. if len(ms_col.shape) == 3:
  196. ms_col[cond,] = ms_col[cond,][:,:,:,changei,]
  197. elif len(ms_col.shape) == 2:
  198. ms_col[cond,] = ms_col[cond,][:,:,changei,]
  199. elif len(ms_col.shape) == 1:
  200. ms_col[cond,] = ms_col[cond,][:,changei,]
  201. else:
  202. raise ValueError('Unexpected dimensions for {} column.'.format(a_col))
  203. ms.putcol(a_col, ms_col, startrow=start, nrow=nrow)
  204. print('\n{} modified correctly.'.format(msdata))