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