Browse Source

I/O now with chunks of data, not full MS. Bug fix

master
Benito Marcote 1 year ago
parent
commit
76ad2ed630
1 changed files with 28 additions and 36 deletions
  1. +28
    -36
      polswap.py

+ 28
- 36
polswap.py View File

@@ -11,9 +11,17 @@ Options:
use either string 'Ef, Mc' or a non-spaced str:
Ef,Mc,Ys.

Version: 1.2
Date: July 2018
Version: 2.0
Date: March 2020
Written by Benito Marcote (marcote@jive.eu)

version 2.0 changes (March 2020)
- MS read and modified in chunks of data.
- Simplified code.
- Bug Fix? Issues in a particular MS with the wrong station only available part of the time.
version 1.2 changes (July 2018)
- Several bug fixes.

"""

import sys
@@ -48,7 +56,7 @@ parser.add_argument('antenna', type=str, help=help_antenna)
parser.add_argument('-t1', '--starttime', default=None, type=str, help=help_t1)
parser.add_argument('-t2', '--endtime', default=None, type=str, help=help_t2)
parser.add_argument('-v', '--version', action='version', version='%(prog)s 1.0')
parser.add_argument('--verbose', default=False, action='store_true')
# parser.add_argument('--verbose', default=False, action='store_true')
# parser.add_argument('--timing', default=False, action='store_true')

arguments = parser.parse_args()
@@ -95,6 +103,13 @@ class Stokes(IntEnum):



def chunkert(f, l, cs, verbose=True):
while f<l:
n = min(cs, l-f)
yield (f, n)
f = f + n


def atime2datetime(atime):
"""Converts a string with the form YYYY/MM/DD/hh:mm:ss or YYYY/DOY/hh:mm:ss to datetime"""
if atime.count('/') == 3:
@@ -147,18 +162,12 @@ with pt.table(msdata, readonly=False, ack=False) as ms:
pols_order = [Stokes(i) for i in ms_pol.getcol('CORR_TYPE')[0]]
# Check that the stokes are the correct ones to do a cross pol.
# Only change it if circular or linear pols.
if arguments.verbose:
print('DEBUG: Stokes in the data: {}'.format(pols_order))

for a_pol_order in pols_order:
if (a_pol_order not in (Stokes.RR, Stokes.RL, Stokes.LR, Stokes.LL)) and \
(a_pol_order not in (Stokes.XX, Stokes.XY, Stokes.YX, Stokes.YY)) and \
(a_pol_order not in (Stokes.RX, Stokes.RY, Stokes.LX, Stokes.LY)) and \
(a_pol_order not in (Stokes.XR, Stokes.XL, Stokes.YR, Stokes.YL)):

if arguments.verbose:
print('DEBUG: checking {}'.format(a_pol_order))

print('Polswap only works for circular or linear pols (or both combined).')
print('These data contain the following stokes: {}'.format(pols_order))
ms.close()
@@ -170,56 +179,39 @@ with pt.table(msdata, readonly=False, ack=False) as ms:
changes = [get_nedded_move(pols_prod, i) for i in (0, 1)]

# transpose data for columns DATA, WEIGHT_SPECTRUM (if exists)
if arguments.verbose:
print('DEBUG: getting ANTENNA columns...')

ants = [ms.getcol('ANTENNA1'), ms.getcol('ANTENNA2')]

if arguments.verbose:
print('DEBUG: getting TIME column...')

datetimes = dt.datetime(1858, 11, 17, 0, 0, 2) + ms.getcol('TIME')*dt.timedelta(seconds=1)
# Get the timerange to apply to polswap
if arguments.starttime is not None:
datetimes_start = atime2datetime(arguments.starttime)
if arguments.verbose:
print('DEBUG: restricted to times after {}'.format(datetimes_start))
else:
datetimes_start = datetimes[0] - dt.timedelta(seconds=1)
if arguments.verbose:
print('DEBUG: No starting time selected. Including everything after {}'.format(datetimes_start))

if arguments.endtime is not None:
datetimes_end = atime2datetime(arguments.endtime)
if arguments.verbose:
print('DEBUG: restricted to times before {}'.format(datetimes_end))
else:
datetimes_end = datetimes[-1] + dt.timedelta(seconds=1)
if arguments.verbose:
print('DEBUG: No ending time selected. Including everything before {}'.format(datetimes_end))

for anti, changei, antpos in zip(ants, changes, ('ANTENNA1','ANTENNA2')):
cond = np.where((anti == antenna_number) & (datetimes > datetimes_start) & (datetimes < datetimes_end))
for a_col in ('DATA', 'FLOAT_DATA', 'FLAG', 'SIGMA_SPECTRUM', 'WEIGHT_SPECTRUM'):
if a_col in ms.colnames():
print('Modifying {} column when {} is {}'.format(a_col, arguments.antenna, antpos))
ms_col = ms.getcol(a_col)
ms_col[cond,] = ms_col[(cond),][:,:,:,changei,]
ms.putcol(a_col, ms_col)
else:
if arguments.verbose:
print('DEBUG: {} column not present in the dataset'.format(a_col))
for (start, nrow) in chunkert(0, len(ms), 5000):
# shape: (nrow, npol, nfreq)
ms_col = ms.getcol(a_col, startrow=start, nrow=nrow)
ms_col[cond,] = ms_col[cond,][:,:,:,changei,]
ms.putcol(a_col, ms_col, startrow=start, nrow=nrow)

for a_col in ('WEIGHT', 'SIGMA'):
if a_col in ms.colnames():
print('Modifying {} column'.format(a_col))
ms_col = ms.getcol(a_col)
ms_col[cond,] = ms_col[cond,][:,:,(changei),]
ms.putcol(a_col, ms_col)
else:
if arguments.verbose:
print('DEBUG: {} column not present in the dataset'.format(a_col))
# Not including FLAG_CATEGORY because it is not recognized when openned in normal MSs
for (start, nrow) in chunkert(0, len(ms), 5000):
# shape: (nrow, npol)
ms_col = ms.getcol(a_col, startrow=start, nrow=nrow)
ms_col[cond,] = ms_col[cond,][:,:,changei,]
ms.putcol(a_col, ms_col, startrow=start, nrow=nrow)


print('\n{} modified correctly.'.format(msdata))


Loading…
Cancel
Save