https://github.com/RadioAstronomySoftwareGroup/pyuvdata
Revision 5984fbc3d518153d644728d8b014ce50c9f4b40c authored by Bryna Hazelton on 22 May 2019, 20:52:45 UTC, committed by Bryna Hazelton on 23 May 2019, 18:17:56 UTC
1 parent 3eafff5
Tip revision: 5984fbc3d518153d644728d8b014ce50c9f4b40c authored by Bryna Hazelton on 22 May 2019, 20:52:45 UTC
Add get_fhd_history to the developer docs, update the docstring
Add get_fhd_history to the developer docs, update the docstring
Tip revision: 5984fbc
aipy_extracts.py
# -*- coding: utf-8 -*-
# Copyright 2018 the HERA Project
# Licensed under the 2-clause BSD License
"""This module extracts some Python code from AIPY used in our MIRIAD I/O
routines. It was copied from AIPY commit
6cb5a70876f33dccdd68d4063b076f8d42d9edae, then reformatted. The only items
used by pyuvdata are ``uv_selector`` and ``UV``.
"""
from __future__ import absolute_import, division, print_function
__all__ = [
'uv_selector',
'UV',
]
import numpy as np
import re
try:
from . import _miriad
except ImportError:
# See setup.py for an explanation. The workaround here is kind of annoying
# since the code below needs the name `_miriad.UV` to be a class in order
# to import successfully.
import os
if os.environ.get('PYUVDATA_IGNORE_EXTMOD_IMPORT_FAIL', '') != '1':
raise
class UV(object):
pass
_miriad = UV() # eww gross but it works
_miriad.UV = UV
str2pol = {
'I': 1, # Stokes Paremeters
'Q': 2,
'U': 3,
'V': 4,
'rr': -1, # Circular Polarizations
'll': -2,
'rl': -3,
'lr': -4,
'xx': -5, # Linear Polarizations
'yy': -6,
'xy': -7,
'yx': -8,
}
def bl2ij(bl):
bl = int(bl)
if bl > 65536:
bl -= 65536
mant = 2048
else:
mant = 256
return (bl // mant - 1, bl % mant - 1)
def ij2bl(i, j):
if i > j:
i, j = j, i
if j + 1 < 256:
return 256 * (i + 1) + (j + 1)
return 2048 * (i + 1) + (j + 1) + 65536
ant_re = r'(\(((-?\d+[xy]?,?)+)\)|-?\d+[xy]?)'
bl_re = '(^(%s_%s|%s),?)' % (ant_re, ant_re, ant_re)
def parse_ants(ant_str, nants):
"""Generate list of (baseline, include, pol) tuples based on parsing of the
string associated with the 'ants' command-line option.
"""
rv, cnt = [], 0
while cnt < len(ant_str):
m = re.search(bl_re, ant_str[cnt:])
if m is None:
if ant_str[cnt:].startswith('all'):
rv = []
elif ant_str[cnt:].startswith('auto'):
rv.append(('auto', 1, -1))
elif ant_str[cnt:].startswith('cross'):
rv.append(('auto', 0, -1))
else:
raise ValueError('Unparsable ant argument "%s"' % ant_str)
c = ant_str[cnt:].find(',')
if c >= 0:
cnt += c + 1
else:
cnt = len(ant_str)
else:
m = m.groups()
cnt += len(m[0])
if m[2] is None:
ais = [m[8]]
ajs = list(range(nants))
else:
if m[3] is None:
ais = [m[2]]
else:
ais = m[3].split(',')
if m[6] is None:
ajs = [m[5]]
else:
ajs = m[6].split(',')
for i in ais:
if type(i) == str and i.startswith('-'):
i = i[1:] # nibble the - off the string
include_i = 0
else:
include_i = 1
for j in ajs:
include = None
if type(j) == str and j.startswith('-'):
j = j[1:]
include_j = 0
else:
include_j = 1
include = int(include_i and include_j)
pol = None
i, j = str(i), str(j)
if not i.isdigit():
ai = re.search(r'(\d+)([x,y])', i).groups()
if not j.isdigit():
aj = re.search(r'(\d+)([x,y])', j).groups()
if i.isdigit() and not j.isdigit():
pol = ['x' + aj[1], 'y' + aj[1]]
ai = [i, '']
elif not i.isdigit() and j.isdigit():
pol = [ai[1] + 'x', ai[1] + 'y']
aj = [j, '']
elif not i.isdigit() and not j.isdigit():
pol = [ai[1] + aj[1]]
if pol is not None:
bl = ij2bl(abs(int(ai[0])), abs(int(aj[0])))
for p in pol:
rv.append((bl, include, p))
else:
bl = ij2bl(abs(int(i)), abs(int(j)))
rv.append((bl, include, -1))
return rv
def uv_selector(uv, ants=-1, pol_str=-1):
"""Call uv.select with appropriate options based on string argument for
antennas (can be 'all', 'auto', 'cross', '0,1,2', or '0_1,0_2') and string
for polarization ('xx','yy','xy','yx').
"""
if ants != -1:
if type(ants) == str:
ants = parse_ants(ants, uv['nants'])
for cnt, (bl, include, pol) in enumerate(ants):
if cnt > 0:
if include:
uv.select('or', -1, -1)
else:
uv.select('and', -1, -1)
if pol == -1:
pol = pol_str # default to explicit pol parameter
if bl == 'auto':
uv.select('auto', 0, 0, include=include)
else:
i, j = bl2ij(bl)
uv.select('antennae', i, j, include=include)
if pol != -1:
for p in pol.split(','):
polopt = str2pol[p]
uv.select('polarization', polopt, 0)
elif pol_str != -1:
for p in pol_str.split(','):
polopt = str2pol[p]
uv.select('polarization', polopt, 0)
itemtable = {
'obstype': 'a',
'history': 'a',
'vartable': 'a',
'ngains': 'i',
'nfeeds': 'i',
'ntau': 'i',
'nsols': 'i',
'interval': 'd',
'leakage': 'c',
'freq0': 'd',
'freqs': '?',
'bandpass': 'c',
'nspect0': 'i',
'nchan0': 'i',
'stopt': 'd',
'duration': 'd',
}
def _uv_pipe_default_action(uv, p, d):
return p, d
class UV(_miriad.UV):
"""Top-level interface to a Miriad UV data set.
"""
def __init__(self, filename, status='old', corrmode='r'):
"""Open a miriad file. status can be ('old','new','append'). corrmode can be
'r' (float32 data storage) or 'j' (int16 with shared exponent).
Default is 'r'.
"""
assert status in ['old', 'new', 'append']
assert corrmode in ['r', 'j']
_miriad.UV.__init__(self, filename, status, corrmode)
self.status = status
self.nchan = _miriad.MAXCHAN
if status == 'old':
self.vartable = self._gen_vartable()
self.read()
self.rewind() # Update variables for the user
try:
self.nchan = self['nchan']
except KeyError:
pass
else:
self.vartable = {'corr': corrmode}
def _gen_vartable(self):
"""Generate table of variables and types from the vartable header.
"""
vartable = {}
for line in self._rdhd('vartable').split('\n'):
try:
type, name = line.split()
vartable[name] = type
except ValueError:
pass
return vartable
def vars(self):
"""Return a list of available variables.
"""
return list(self.vartable.keys())
def items(self):
"""Return a list of available header items.
"""
items = []
for i in itemtable:
try:
_miriad.hdaccess(self.haccess(i, 'read'))
items.append(i)
except IOError:
pass
return items
def _rdhd(self, name):
"""Provide read access to header items via low-level calls.
"""
itype = itemtable[name]
if itype == '?':
return self._rdhd_special(name)
h = self.haccess(name, 'read')
rv = []
if len(itype) == 1:
if itype == 'a':
offset = 0
else:
t, offset = _miriad.hread_init(h)
assert itype == t
while True:
try:
c, o = _miriad.hread(h, offset, itype)
except IOError:
break
if itype == 'a':
try:
c = str(c[:o], 'utf-8')
except TypeError:
c = c[:o]
rv.append(c)
offset += o
if itype == 'a':
rv = ''.join(rv)
else:
t, offset = _miriad.hread_init(h)
assert t == 'b'
for t in itype:
v, o = _miread.hread(h, offset, t)
rv.append(v)
offset += o
_miriad.hdaccess(h)
if len(rv) == 1:
return rv[0]
elif type(rv) == str:
return rv
else:
return np.array(rv)
def _wrhd(self, name, val):
"""Provide write access to header items via low-level calls.
"""
type = itemtable[name]
if type == '?':
return self._wrhd_special(name, val)
h = self.haccess(name, 'write')
if len(type) == 1:
try:
len(val)
except TypeError:
val = [val]
if type == 'a':
offset = 0
else:
offset = _miriad.hwrite_init(h, type)
for v in val:
offset += _miriad.hwrite(h, offset, v, type)
else:
offset = _miriad.hwrite_init(h, 'b')
for v, t in zip(val, type):
offset += _miriad.hwrite(h, offset, v, t)
_miriad.hdaccess(h)
def _rdhd_special(self, name):
"""Provide read access to special header items of type '?' to _rdhd.
"""
if name == 'freqs':
h = self.haccess(name, 'read')
c, o = _miriad.hread(h, 0, 'i')
rv = [c]
offset = 8
while True:
try:
c, o = _miriad.hread(h, offset, 'i')
rv.append(c)
offset += 8
c, o = _miriad.hread(h, offset, 'd')
rv.append(c)
offset += 8
c, o = _miriad.hread(h, offset, 'd')
rv.append(c)
offset += 8
except IOError:
break
_miriad.hdaccess(h)
return rv
else:
raise ValueError('Unknown special header: ' + name)
def _wrhd_special(self, name, val):
"""Provide write access to special header items of type '?' to _wrhd
"""
if name == 'freqs':
h = self.haccess(name, 'write')
o = _miriad.hwrite(h, 0, val[0], 'i')
offset = 8
for i, v in enumerate(val[1:]):
if i % 3 == 0:
o = _miriad.hwrite(h, offset, v, 'i')
else:
o = _miriad.hwrite(h, offset, v, 'd')
offset += 8
_miriad.hdaccess(h)
else:
raise ValueError('Unknown special header: ' + name)
def __getitem__(self, name):
"""Allow access to variables and header items via ``uv[name]``.
"""
try:
type = self.vartable[name]
return self._rdvr(name, type)
except KeyError:
type = itemtable[name]
return self._rdhd(name)
def __setitem__(self, name, val):
"""Allow setting variables and header items via ``uv[name] = val``.
"""
try:
type = self.vartable[name]
self._wrvr(name, type, val)
except KeyError:
self._wrhd(name, val)
def select(self, name, n1, n2, include=True):
"""Choose which data are returned by read().
name This can be: 'decimate','time','antennae','visibility',
'uvrange','pointing','amplitude','window','or','dra',
'ddec','uvnrange','increment','ra','dec','and', 'clear',
'on','polarization','shadow','auto','dazim','delev'
n1,n2 Generally this is the range of values to select. For
'antennae', this is the two antennae pair to select
(indexed from 0); a -1 indicates 'all antennae'.
For 'decimate', n1 is every Nth integration to use, and
n2 is which integration within a block of N to use.
For 'shadow', a zero indicates use 'antdiam' variable.
For 'on','window','polarization','increment','shadow' only
p1 is used.
For 'and','or','clear','auto' p1 and p2 are ignored.
include If true, the data is selected. If false, the data is
discarded. Ignored for 'and','or','clear'.
"""
if name == 'antennae':
n1 += 1
n2 += 1
self._select(name, float(n1), float(n2), int(include))
def read(self, raw=False):
"""Return the next data record. Calling this function causes vars to change to
reflect the record which this function returns. 'raw' causes data and
flags to be returned seperately.
"""
preamble, data, flags, nread = self.raw_read(self.nchan)
if nread == 0:
raise IOError("No data read")
flags = np.logical_not(flags)
if raw:
return preamble, data, flags
return preamble, np.ma.array(data, mask=flags)
def all(self, raw=False):
"""Provide an iterator over preamble, data. Allows constructs like: ``for
preamble, data in uv.all(): ...``
"""
while True:
try:
yield self.read(raw=raw)
except IOError:
break
def write(self, preamble, data, flags=None):
"""Write the next data record. data must be a complex, masked array. preamble
must be (uvw, t, (i,j)), where uvw is an array of u,v,w, t is the
Julian date, and (i,j) is an antenna pair.
"""
if data is None:
return
if flags is not None:
flags = np.logical_not(flags)
elif len(data.mask.shape) == 0:
flags = np.ones(data.shape)
data = data.unmask()
else:
flags = np.logical_not(data.mask)
data = data.data
self.raw_write(preamble, data.astype(np.complex64), flags.astype(np.int32))
def init_from_uv(self, uv, override={}, exclude=[]):
"""Initialize header items and variables from another UV. Those in override
will be overwritten by override[k], and tracking will be turned off
(meaning they will not be updated in pipe()). Those in exclude are
omitted completely.
"""
for k in uv.items():
if k in exclude:
continue
elif k in override:
self._wrhd(k, override[k])
else:
self._wrhd(k, uv[k])
self.vartable = {}
for k in uv.vars():
if k in exclude:
continue
elif k == 'corr':
# I don't understand why reading 'corr' segfaults miriad,
# but it does. This is a cludgy work-around.
continue
elif k in override:
self.vartable[k] = uv.vartable[k]
self._wrvr(k, uv.vartable[k], override[k])
else:
self.vartable[k] = uv.vartable[k]
self._wrvr(k, uv.vartable[k], uv[k])
uv.trackvr(k, 'c') # Set to copy when copyvr() called
def pipe(self, uv, mfunc=_uv_pipe_default_action, append2hist='', raw=False):
"""Pipe in data from another UV through the function ``mfunc(uv, preamble,
data)``, which should return ``(preamble, data)``. If mfunc is not
provided, the dataset will just be cloned, and if the returned data is
None, it will be omitted. The string 'append2hist' will be appended to
history.
"""
self._wrhd('history', self['history'] + append2hist)
if raw:
for p, d, f in uv.all(raw=raw):
np, nd, nf = mfunc(uv, p, d, f)
self.copyvr(uv)
self.write(np, nd, nf)
else:
for p, d in uv.all():
np, nd = mfunc(uv, p, d)
self.copyvr(uv)
self.write(np, nd)
def add_var(self, name, type):
"""Add a variable of the specified type to a UV file.
"""
self.vartable[name] = type
Computing file changes ...