def test_has_key(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertTrue(dc.has_key((2, 3, 'yy')))
self.assertTrue(dc.has_key((2, 3), 'yy'))
self.assertTrue(dc.has_key((3, 2), 'yy'))
self.assertFalse(dc.has_key('xy'))
self.assertFalse(dc.has_key((5, 6)))
self.assertFalse(dc.has_key((1, 2, 'xy')))
dc = datacontainer.DataContainer(self.polbl)
self.assertTrue(dc.has_key((2, 3, 'yy')))
self.assertTrue(dc.has_key((2, 3), 'yy'))
self.assertTrue(dc.has_key((3, 2), 'yy'))
self.assertFalse(dc.has_key('xy'))
self.assertFalse(dc.has_key((5, 6)))
self.assertFalse(dc.has_key((1, 2, 'xy')))
dc = datacontainer.DataContainer(self.both)
self.assertTrue(dc.has_key((2, 3, 'yy')))
self.assertTrue(dc.has_key((2, 3), 'yy'))
self.assertTrue(dc.has_key((3, 2), 'yy'))
self.assertFalse(dc.has_key('xy'))
self.assertFalse(dc.has_key((5, 6)))
self.assertFalse(dc.has_key((1, 2, 'xy')))
self.assertTrue(dc.has_key((2, 3, 'YY')))
self.assertTrue(dc.has_key((2, 3), 'YY'))
self.assertTrue(dc.has_key((3, 2), 'YY'))
self.assertFalse(dc.has_key('XY'))
self.assertFalse(dc.has_key((5, 6)))
self.assertFalse(dc.has_key((1, 2, 'XY')))
# assert switch bl
dc[(1, 2, 'xy')] = 1j
self.assertTrue(dc.has_key((2, 1, 'yx')))
def test_len(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertEqual(len(dc), 10)
dc = datacontainer.DataContainer(self.polbl)
self.assertEqual(len(dc), 10)
dc = datacontainer.DataContainer(self.both)
self.assertEqual(len(dc), 10)
def test_getitem(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertEqual(dc[(1, 2, 'xx')], 1j)
self.assertEqual(dc[(2, 1, 'xx')], -1j)
self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
self.assertEqual(set(dc['xx'].keys()), set(self.bls))
self.assertEqual(dc[(1, 2, 'xx')], dc.get_data((1, 2, 'xx')))
self.assertEqual(dc[(1, 2, 'xx')], dc.get_data(1, 2, 'xx'))
dc = datacontainer.DataContainer(self.polbl)
self.assertEqual(dc[(1, 2, 'xx')], 1j)
self.assertEqual(dc[(2, 1, 'xx')], -1j)
self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
self.assertEqual(set(dc['xx'].keys()), set(self.bls))
self.assertEqual(dc[(2, 1, 'xx')], dc.get_data((2, 1, 'xx')))
self.assertEqual(dc[(2, 1, 'xx')], dc.get_data(2, 1, 'xx'))
dc = datacontainer.DataContainer(self.both)
self.assertEqual(dc[(1, 2, 'xx')], 1j)
self.assertEqual(dc[(2, 1, 'xx')], -1j)
self.assertEqual(dc[(1, 2)], {'xx': 1j, 'yy': 1j})
self.assertEqual(set(dc['xx'].keys()), set(self.bls))
self.assertEqual(dc[(1, 2)], dc.get_data((1, 2)))
self.assertEqual(dc[(1, 2)], dc.get_data(1, 2))
self.assertEqual(dc[(1, 2, 'XX')], 1j)
self.assertEqual(dc[(2, 1, 'XX')], -1j)
self.assertEqual(dc[(2, 1, 'XX')], dc.get_data(2, 1, 'XX'))
self.assertEqual(dc[(2, 1, 'XX')], dc.get_data(2, 1, 'xx'))
def test_keys(self):
dc = datacontainer.DataContainer(self.blpol)
keys = dc.keys()
self.assertEqual(len(keys), len(self.pols) * len(self.bls))
dc = datacontainer.DataContainer(self.polbl)
keys = dc.keys()
self.assertEqual(len(keys), len(self.pols) * len(self.bls))
dc = datacontainer.DataContainer(self.both)
keys = dc.keys()
self.assertEqual(len(keys), len(self.pols) * len(self.bls))
def test_pols(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertEqual(set(self.pols), dc.pols())
self.assertEqual(set(self.pols), dc.pols((1, 2)))
dc = datacontainer.DataContainer(self.polbl)
self.assertEqual(set(self.pols), dc.pols())
self.assertEqual(set(self.pols), dc.pols((1, 2)))
dc = datacontainer.DataContainer(self.both)
self.assertEqual(set(self.pols), dc.pols())
self.assertEqual(set(self.pols), dc.pols((1, 2)))
def test_get(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertEqual(dc.get((1, 2), 'yy'), 1j)
self.assertEqual(dc.get((2, 1), 'yy'), -1j)
dc = datacontainer.DataContainer(self.polbl)
self.assertEqual(dc.get((1, 2), 'yy'), 1j)
self.assertEqual(dc.get((2, 1), 'yy'), -1j)
dc = datacontainer.DataContainer(self.both)
self.assertEqual(dc.get((1, 2), 'yy'), 1j)
self.assertEqual(dc.get((2, 1), 'yy'), -1j)
self.assertEqual(dc.get((1, 2), 'YY'), 1j)
self.assertEqual(dc.get((2, 1), 'YY'), -1j)
def test_bls(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertEqual(set(self.bls), dc.bls())
self.assertEqual(set(self.bls), dc.bls('xx'))
self.assertEqual(set(self.bls), dc.bls('yy'))
dc = datacontainer.DataContainer(self.polbl)
self.assertEqual(set(self.bls), dc.bls())
self.assertEqual(set(self.bls), dc.bls('xx'))
self.assertEqual(set(self.bls), dc.bls('yy'))
dc = datacontainer.DataContainer(self.both)
self.assertEqual(set(self.bls), dc.bls())
self.assertEqual(set(self.bls), dc.bls('xx'))
self.assertEqual(set(self.bls), dc.bls('yy'))
def test_has_bl(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertTrue(dc.has_bl((2, 3)))
self.assertTrue(dc.has_bl((3, 2)))
self.assertFalse(dc.has_bl((0, 3)))
dc = datacontainer.DataContainer(self.polbl)
self.assertTrue(dc.has_bl((2, 3)))
self.assertTrue(dc.has_bl((3, 2)))
self.assertFalse(dc.has_bl((0, 3)))
dc = datacontainer.DataContainer(self.both)
self.assertTrue(dc.has_bl((2, 3)))
self.assertTrue(dc.has_bl((3, 2)))
self.assertFalse(dc.has_bl((0, 3)))
def test_values(self):
dc = datacontainer.DataContainer(self.blpol)
values = dc.values()
self.assertEqual(len(values), len(self.pols) * len(self.bls))
self.assertEqual(values[0], 1j)
dc = datacontainer.DataContainer(self.polbl)
values = dc.values()
self.assertEqual(len(values), len(self.pols) * len(self.bls))
self.assertEqual(values[0], 1j)
dc = datacontainer.DataContainer(self.both)
values = dc.values()
self.assertEqual(len(values), len(self.pols) * len(self.bls))
self.assertEqual(values[0], 1j)
def test_has_pol(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertTrue(dc.has_pol('xx'))
self.assertTrue(dc.has_pol('XX'))
self.assertFalse(dc.has_pol('xy'))
self.assertFalse(dc.has_pol('XY'))
dc = datacontainer.DataContainer(self.polbl)
self.assertTrue(dc.has_pol('xx'))
self.assertTrue(dc.has_pol('XX'))
self.assertFalse(dc.has_pol('xy'))
self.assertFalse(dc.has_pol('XY'))
dc = datacontainer.DataContainer(self.both)
self.assertTrue(dc.has_pol('xx'))
self.assertTrue(dc.has_pol('XX'))
self.assertFalse(dc.has_pol('xy'))
self.assertFalse(dc.has_pol('XY'))
def test_read_write(self):
# test read data can be turned off for uvh5
fname = os.path.join(DATA_PATH, 'zen.2458098.43124.subband.uvh5')
V = VisClean(fname, filetype='uvh5')
V.read(read_data=False)
assert set(V.hd.ant_1_array) == set([1, 11, 12, 13, 14])
# test read-write-read
V.read()
V.write_data(V.data, "./ex.uvh5", overwrite=True, filetype='uvh5', vis_units='Jy')
V2 = VisClean("./ex.uvh5", filetype='uvh5')
V2.read()
assert V2.hd.vis_units == 'Jy'
assert 'Thisfilewasproducedbythefunction' in V2.hd.history.replace('\n', '').replace(' ', '')
V.hd.history, V2.hd.history, V2.hd.vis_units = '', '', V.hd.vis_units
assert V.hd == V2.hd
os.remove("./ex.uvh5")
# exceptions
pytest.raises(ValueError, V.write_data, V.data, 'foo', filetype='what')
# test write on subset of data
V.read(read_data=True)
data = datacontainer.DataContainer(dict([(k, V.data[k]) for k in list(V.data.keys())[:2]]))
V.write_data(data, "ex.uvh5", overwrite=True, filetype='uvh5')
assert os.path.exists("ex.uvh5")
os.remove('ex.uvh5')
def test_items(self):
dc = datacontainer.DataContainer(self.blpol)
items = dc.items()
self.assertEqual(len(items), len(self.pols) * len(self.bls))
self.assertTrue(items[0][0][0:2] in self.bls)
self.assertTrue(items[0][0][2] in self.pols)
self.assertEqual(items[0][1], 1j)
dc = datacontainer.DataContainer(self.polbl)
items = dc.items()
self.assertTrue(items[0][0][0:2] in self.bls)
self.assertTrue(items[0][0][2] in self.pols)
self.assertEqual(items[0][1], 1j)
dc = datacontainer.DataContainer(self.both)
items = dc.items()
self.assertTrue(items[0][0][0:2] in self.bls)
self.assertTrue(items[0][0][2] in self.pols)
self.assertEqual(items[0][1], 1j)
def test_init(self):
dc = datacontainer.DataContainer(self.blpol)
for k in dc._data.keys():
self.assertEqual(len(k), 3)
self.assertEqual(set(self.bls), dc._bls)
self.assertEqual(set(self.pols), dc._pols)
dc = datacontainer.DataContainer(self.polbl)
for k in dc._data.keys():
self.assertEqual(len(k), 3)
self.assertEqual(set(self.bls), dc._bls)
self.assertEqual(set(self.pols), dc._pols)
dc = datacontainer.DataContainer(self.both)
for k in dc._data.keys():
self.assertEqual(len(k), 3)
self.assertEqual(set(self.bls), dc._bls)
self.assertEqual(set(self.pols), dc._pols)
self.assertRaises(AssertionError, datacontainer.DataContainer,
{(1, 2, 3, 4): 2})
def test_setter(self):
dc = datacontainer.DataContainer(self.blpol)
# test basic setting
dc[(100, 101, 'xy')] = np.arange(100) + np.arange(100) * 1j
self.assertEqual(dc[(100, 101, 'xy')].shape, (100, ))
self.assertEqual(dc[(100, 101, 'xy')].dtype, np.complex)
self.assertAlmostEqual(dc[(100, 101, 'xy')][1], (1 + 1j))
self.assertAlmostEqual(dc[(101, 100, 'yx')][1], (1 - 1j))
self.assertEqual(len(dc.keys()), 11)
self.assertEqual((100, 101) in dc._bls, True)
self.assertEqual('xy' in dc._pols, True)
# test error
self.assertRaises(ValueError, dc.__setitem__, *((100, 101), 100j))
def test_del(self):
dc = datacontainer.DataContainer(self.blpol)
self.assertTrue((1, 2, 'xx') in dc)
self.assertTrue((1, 2, 'XX') in dc)
del dc[(1, 2, 'xx')]
self.assertFalse((1, 2, 'xx') in dc)
self.assertTrue('xx' in dc.pols())
self.assertTrue((1, 2) in dc.bls())
del dc[(1, 2, 'yy')]
self.assertFalse((1, 2) in dc.bls())
del dc[(2, 3, 'XX')]
self.assertFalse((2, 3, 'xx') in dc)
self.assertTrue('xx' in dc.pols())
self.assertTrue((2, 3) in dc.bls())
def test_synthesize_ant_flags():
flags = datacontainer.DataContainer({
(0, 0, 'xx'): np.ones((5, 5), bool),
(0, 1, 'xx'): np.ones((5, 5), bool),
(1, 2, 'xx'): np.zeros((5, 5), bool),
(2, 3, 'xx'): np.zeros((5, 5), bool)
})
flags[(2, 3, 'xx')][:, 4] = True
# aggressive flagging
ant_flags = utils.synthesize_ant_flags(flags, threshold=0.0)
np.testing.assert_array_equal(ant_flags[(0, 'x')], True)
np.testing.assert_array_equal(ant_flags[(1, 'x')], False)
np.testing.assert_array_equal(ant_flags[(2, 'x')][:, 0:4], False)
np.testing.assert_array_equal(ant_flags[(2, 'x')][:, 4], True)
np.testing.assert_array_equal(ant_flags[(3, 'x')][:, 0:4], False)
np.testing.assert_array_equal(ant_flags[(3, 'x')][:, 4], True)
# conservative flagging
ant_flags = utils.synthesize_ant_flags(flags, threshold=0.75)
np.testing.assert_array_equal(ant_flags[(2, 'x')][:, 4], False)
# very conservative flagging
flags[(1, 2, 'xx')][:3, 4] = True
ant_flags = utils.synthesize_ant_flags(flags, threshold=1.0)
np.testing.assert_array_equal(ant_flags[(2, 'x')][:3, 4], True)
np.testing.assert_array_equal(ant_flags[(2, 'x')][3:, 4], False)
def red_average(data, reds=None, bl_tol=1.0, inplace=False,
wgts=None, flags=None, nsamples=None):
"""
Redundantly average visibilities in a DataContainer, HERAData or UVData object.
Average is weighted by integration_time * nsamples * ~flags unless wgts are fed.
Args:
data : DataContainer, HERAData or UVData object
Object to redundantly average
reds : list, optional
Nested lists of antpair tuples to redundantly average.
E.g. [ [(1, 2), (2, 3)], [(1, 3), (2, 4)], ...]
If None, will calculate these from the metadata
bl_tol : float
Baseline redundancy tolerance in meters. Only used if reds is None.
inplace : bool
Perform average and downselect inplace, otherwise returns a deepcopy.
The first baseline in each reds sublist is kept.
wgts : DataContainer
Manual weights to use in redundant average. This supercedes flags and nsamples
If provided, and will also be used if input data is a UVData or a subclass of it.
flags : DataContainer
If data is a DataContainer, these are its flags. Default (None) is no flags.
nsamples : DataContainer
If data is a DataContainer, these are its nsamples. Default (None) is 1.0 for all pixels.
Furthermore, if data is a DataContainer, integration_time is 1.0 for all pixels.
Returns:
if fed a DataContainer:
DataContainer, averaged data
DataContainer, averaged flags
DataContainer, summed nsamples
elif fed a HERAData or UVData:
HERAData or UVData object, averaged data
Notes:
1. Different polarizations are assumed to be non-redundant.
2. Default weighting is nsamples * integration_time * ~flags.
3. If wgts Container is fed then they supercede flag and nsample weighting.
"""
from hera_cal import redcal, datacontainer
# type checks
if not (isinstance(data, datacontainer.DataContainer) or isinstance(data, UVData)):
raise ValueError("data must be a DataContainer or a UVData or its subclass")
fed_container = isinstance(data, datacontainer.DataContainer)
# fill DataContainers if necessary
if fed_container:
if not inplace:
flags = copy.deepcopy(flags)
nsamples = copy.deepcopy(nsamples)
if flags is None:
flags = datacontainer.DataContainer({k: np.zeros_like(data[k], np.bool) for k in data})
if nsamples is None:
nsamples = datacontainer.DataContainer({k: np.ones_like(data[k], np.float) for k in data})
# get weights: if wgts are not fed, then use flags and nsamples
if wgts is None:
if fed_container:
wgts = datacontainer.DataContainer({k: nsamples[k] * ~flags[k] for k in data})
else:
wgts = datacontainer.DataContainer({k: data.get_nsamples(k) * ~data.get_flags(k) for k in data.get_antpairpols()})
# deepcopy
if not inplace:
data = copy.deepcopy(data)
# get metadata
if fed_container:
pols = sorted(data.pols())
else:
pols = [polnum2str(pol, x_orientation=data.x_orientation) for pol in data.polarization_array]
# get redundant groups
if reds is None:
# if DataContainer, check for antpos
if fed_container:
if not hasattr(data, 'antpos') or data.antpos is None:
raise ValueError("DataContainer must have antpos dictionary to calculate reds")
antposd = data.antpos
else:
antpos, ants = data.get_ENU_antpos()
antposd = dict(zip(ants, antpos))
reds = redcal.get_pos_reds(antposd, bl_error_tol=bl_tol)
# eliminate baselines not in data
if fed_container:
antpairs = sorted(data.antpairs())
else:
antpairs = data.get_antpairs()
reds = [[bl for bl in blg if bl in antpairs] for blg in reds]
reds = [blg for blg in reds if len(blg) > 0]
# iterate over redundant groups and polarizations
for pol in pols:
for blg in reds:
# get data and weighting for this pol-blgroup
if fed_container:
d = np.asarray([data[bl + (pol,)] for bl in blg])
f = np.asarray([(~flags[bl + (pol,)]).astype(np.float) for bl in blg])
n = np.asarray([nsamples[bl + (pol,)] for bl in blg])
# DataContainer can't track integration time, so no tint here
tint = np.array([1.0])
w = np.asarray([wgts[bl + (pol,)] for bl in blg])
else:
d = np.asarray([data.get_data(bl + (pol,)) for bl in blg])
f = np.asarray([(~data.get_flags(bl + (pol,))).astype(np.float) for bl in blg])
n = np.asarray([data.get_nsamples(bl + (pol,)) for bl in blg])
tint = np.asarray([data.integration_time[data.antpair2ind(bl + (pol,))] for bl in blg])[:, :, None]
w = np.asarray([wgts[bl + (pol,)] for bl in blg]) * tint
# take the weighted average
wsum = np.sum(w, axis=0).clip(1e-10, np.inf) # this is the normalization
davg = np.sum(d * w, axis=0) / wsum # weighted average
navg = np.sum(n * f, axis=0) # this is the new total nsample (without flagged elements)
fmax = np.max(f, axis=2) # collapse along freq: marks any fully flagged integrations
iavg = np.sum(tint.squeeze() * fmax, axis=0) / np.sum(fmax, axis=0).clip(1e-10, np.inf)
favg = np.isclose(wsum, 0.0) # this is getting any fully flagged pixels
# replace with new data
if fed_container:
blkey = blg[0] + (pol,)
data[blkey] = davg
flags[blkey] = favg
nsamples[blkey] = navg
else:
blinds = data.antpair2ind(blg[0])
polind = pols.index(pol)
data.data_array[blinds, 0, :, polind] = davg
data.flag_array[blinds, 0, :, polind] = favg
data.nsample_array[blinds, 0, :, polind] = navg
data.integration_time[blinds] = iavg
# select out averaged bls
bls = [blg[0] + (pol,) for pol in pols for blg in reds]
if fed_container:
for bl in list(data.keys()):
if bl not in bls:
del data[bl]
else:
data.select(bls=bls)
if not inplace:
if fed_container:
return data, flags, nsamples
else:
return data
