Skip to content

Commit

Permalink
Merge aa80c67 into 4f330f1
Browse files Browse the repository at this point in the history
  • Loading branch information
adampbeardsley committed Jul 10, 2018
2 parents 4f330f1 + aa80c67 commit 11725e9
Show file tree
Hide file tree
Showing 4 changed files with 172 additions and 71 deletions.
21 changes: 21 additions & 0 deletions pyuvdata/tests/test_utils.py
Expand Up @@ -189,3 +189,24 @@ def test_jones_num_funcs():
nt.assert_raises(KeyError, uvutils.jstr2num, 'foo')
nt.assert_raises(ValueError, uvutils.jstr2num, 1)
nt.assert_raises(ValueError, uvutils.jnum2str, 7.3)


def test_conj_pol():
""" Test function to conjugate pols """

pol_nums = [-8, -7, -6, -5, -4, -3, -2, -1, 1, 2, 3, 4]
cpol_nums = [-7, -8, -6, -5, -3, -4, -2, -1, 1, 2, 3, 4]
nt.assert_equal(pol_nums, uvutils.conj_pol(cpol_nums))
nt.assert_equal(uvutils.conj_pol(pol_nums), cpol_nums)
pol_str = ['YX', 'XY', 'YY', 'XX', 'LR', 'RL', 'LL', 'RR', 'pI', 'pQ', 'pU', 'pV']
cpol_str = ['XY', 'YX', 'YY', 'XX', 'RL', 'LR', 'LL', 'RR', 'pI', 'pQ', 'pU', 'pV']
nt.assert_equal(pol_str, uvutils.conj_pol(cpol_str))
nt.assert_equal(uvutils.conj_pol(pol_str), cpol_str)
nt.assert_equal([pol_str, pol_nums], uvutils.conj_pol([cpol_str, cpol_nums]))
jstr = ['jyx', 'jxy', 'jyy', 'jxx', 'jlr', 'jrl', 'jll', 'jrr']
cjstr = ['jxy', 'jyx', 'jyy', 'jxx', 'jrl', 'jlr', 'jll', 'jrr']
nt.assert_equal(jstr, uvutils.conj_pol(cjstr))
nt.assert_equal(uvutils.conj_pol(jstr), uvutils.conj_pol(jstr))

# Test invalid pol
nt.assert_raises(ValueError, uvutils.conj_pol, 2.3)
69 changes: 42 additions & 27 deletions pyuvdata/tests/test_uvdata.py
Expand Up @@ -1433,64 +1433,70 @@ def test_key2inds():
ind1, ind2, indp = uv._key2inds((ant1, ant2, pol))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal([0], indp))
nt.assert_true(np.array_equal([0], indp[0]))
# Any of these inputs can also be a tuple of a tuple, so need to be checked twice.
ind1, ind2, indp = uv._key2inds(((ant1, ant2, pol)))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal([0], indp))
nt.assert_true(np.array_equal([0], indp[0]))

# Combo with pol as string
ind1, ind2, indp = uv._key2inds((ant1, ant2, uvutils.polnum2str(pol)))
nt.assert_true(np.array_equal([0], indp))
nt.assert_true(np.array_equal([0], indp[0]))
ind1, ind2, indp = uv._key2inds(((ant1, ant2, uvutils.polnum2str(pol))))
nt.assert_true(np.array_equal([0], indp))
nt.assert_true(np.array_equal([0], indp[0]))

# Check conjugation
ind1, ind2, indp = uv._key2inds((ant2, ant1, pol))
nt.assert_true(np.array_equal(bltind, ind2))
nt.assert_true(np.array_equal(np.array([]), ind1))
nt.assert_true(np.array_equal([0], indp))
nt.assert_true(np.array_equal([0], indp[1]))
# Conjugation with pol as string
ind1, ind2, indp = uv._key2inds((ant2, ant1, uvutils.polnum2str(pol)))
nt.assert_true(np.array_equal(bltind, ind2))
nt.assert_true(np.array_equal(np.array([]), ind1))
nt.assert_true(np.array_equal([0], indp[1]))
nt.assert_true(np.array_equal([], indp[0]))

# Antpair only
ind1, ind2, indp = uv._key2inds((ant1, ant2))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp[0]))
ind1, ind2, indp = uv._key2inds(((ant1, ant2)))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp[0]))

# Baseline number only
ind1, ind2, indp = uv._key2inds(uv.antnums_to_baseline(ant1, ant2))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp[0]))
ind1, ind2, indp = uv._key2inds((uv.antnums_to_baseline(ant1, ant2)))
nt.assert_true(np.array_equal(bltind, ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp))
nt.assert_true(np.array_equal(np.arange(uv.Npols), indp[0]))

# Pol number only
ind1, ind2, indp = uv._key2inds(pol)
nt.assert_true(np.array_equal(np.arange(uv.Nblts), ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.array([0]), indp))
nt.assert_true(np.array_equal(np.array([0]), indp[0]))
ind1, ind2, indp = uv._key2inds((pol))
nt.assert_true(np.array_equal(np.arange(uv.Nblts), ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.array([0]), indp))
nt.assert_true(np.array_equal(np.array([0]), indp[0]))

# Pol string only
ind1, ind2, indp = uv._key2inds('LL')
nt.assert_true(np.array_equal(np.arange(uv.Nblts), ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.array([1]), indp))
nt.assert_true(np.array_equal(np.array([1]), indp[0]))
ind1, ind2, indp = uv._key2inds(('LL'))
nt.assert_true(np.array_equal(np.arange(uv.Nblts), ind1))
nt.assert_true(np.array_equal(np.array([]), ind2))
nt.assert_true(np.array_equal(np.array([1]), indp))
nt.assert_true(np.array_equal(np.array([1]), indp[0]))

# Test invalid keys
nt.assert_raises(KeyError, uv._key2inds, 'I') # pol str not in data
Expand Down Expand Up @@ -1519,7 +1525,7 @@ def test_smart_slicing():
ind1 = 10 * np.arange(9)
ind2 = []
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []))
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1529,7 +1535,7 @@ def test_smart_slicing():
nt.assert_equal(d[1, 0, 0], uv.data_array[ind1[1], 0, 0, indp[0]])

# force copy
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp, force_copy=True)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []), force_copy=True)
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1542,7 +1548,7 @@ def test_smart_slicing():
ind1 = 10 * np.arange(9)
ind2 = []
indp = [0, 1, 3]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []))
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1555,7 +1561,7 @@ def test_smart_slicing():
ind1 = [0, 4, 5]
ind2 = []
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []))
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1568,7 +1574,7 @@ def test_smart_slicing():
ind1 = [0, 4, 5]
ind2 = []
indp = [0, 1, 3]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []))
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1582,7 +1588,7 @@ def test_smart_slicing():
ind1 = []
ind2 = 10 * np.arange(9)
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, ([], indp))
dcheck = uv.data_array[ind2, :, :, :]
dcheck = np.squeeze(np.conj(dcheck[:, :, :, indp]))
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1591,7 +1597,7 @@ def test_smart_slicing():
ind1 = []
ind2 = 10 * np.arange(9)
indp = [0, 1, 3]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, ([], indp))
dcheck = uv.data_array[ind2, :, :, :]
dcheck = np.squeeze(np.conj(dcheck[:, :, :, indp]))
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1600,7 +1606,7 @@ def test_smart_slicing():
ind1 = []
ind2 = [1, 4, 5, 10]
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, ([], indp))
dcheck = uv.data_array[ind2, :, :, :]
dcheck = np.squeeze(np.conj(dcheck[:, :, :, indp]))
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1609,7 +1615,7 @@ def test_smart_slicing():
ind1 = []
ind2 = [1, 4, 5, 10]
indp = [0, 1, 3]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, ([], indp))
dcheck = uv.data_array[ind2, :, :, :]
dcheck = np.squeeze(np.conj(dcheck[:, :, :, indp]))
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1618,7 +1624,7 @@ def test_smart_slicing():
ind1 = np.arange(20)
ind2 = np.arange(30, 40)
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, indp))
dcheck = np.append(uv.data_array[ind1, :, :, :],
np.conj(uv.data_array[ind2, :, :, :]), axis=0)
dcheck = np.squeeze(dcheck[:, :, :, indp])
Expand All @@ -1628,7 +1634,7 @@ def test_smart_slicing():
ind1 = np.arange(20)
ind2 = np.arange(30, 40)
indp = [0, 1, 3]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, indp))
dcheck = np.append(uv.data_array[ind1, :, :, :],
np.conj(uv.data_array[ind2, :, :, :]), axis=0)
dcheck = np.squeeze(dcheck[:, :, :, indp])
Expand All @@ -1638,7 +1644,7 @@ def test_smart_slicing():
ind1 = [45]
ind2 = []
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []))
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp], axis=1)
nt.assert_true(np.all(d == dcheck))
Expand All @@ -1647,18 +1653,22 @@ def test_smart_slicing():
ind1 = []
ind2 = [45]
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp)
d = uv._smart_slicing(uv.data_array, ind1, ind2, ([], indp))
nt.assert_true(np.all(d == np.conj(dcheck)))

# Full squeeze
ind1 = [45]
ind2 = []
indp = [0, 1]
d = uv._smart_slicing(uv.data_array, ind1, ind2, indp, squeeze='full')
d = uv._smart_slicing(uv.data_array, ind1, ind2, (indp, []), squeeze='full')
dcheck = uv.data_array[ind1, :, :, :]
dcheck = np.squeeze(dcheck[:, :, :, indp])
nt.assert_true(np.all(d == dcheck))

# Test invalid squeeze
nt.assert_raises(ValueError, uv._smart_slicing, uv.data_array, ind1, ind2,
(indp, []), squeeze='notasqueeze')


def test_get_data():
# Test get_data function for easy access to data
Expand All @@ -1681,6 +1691,11 @@ def test_get_data():
d = uv.get_data(ant2, ant1, pol)
nt.assert_true(np.all(dcheck == np.conj(d)))

# Check cross pol conjugation
d = uv.get_data(ant2, ant1, uv.polarization_array[2])
d1 = uv.get_data(ant1, ant2, uv.polarization_array[3])
nt.assert_true(np.all(d == np.conj(d1)))

# Antpair only
dcheck = np.squeeze(uv.data_array[bltind, :, :, :])
d = uv.get_data(ant1, ant2)
Expand Down
33 changes: 33 additions & 0 deletions pyuvdata/utils.py
Expand Up @@ -424,6 +424,39 @@ def jnum2str(jnum):
return out


def conj_pol(pol):
"""
Returns the polarization for the conjugate baseline.
For example, (1, 2, 'XY') = conj(2, 1, 'YX').
The returned polarization is determined by assuming the antenna pair is reversed
in the data, and finding the correct polarization correlation which will yield
the requested baseline when conjugated. Note this means changing the polarization
for linear cross-pols, but keeping auto-pol (e.g. XX) and Stokes the same.
Args:
pol: Polarization (str or int)
Returns:
cpol: Polarization as if antennas are swapped (type matches input)
"""
cpol_dict = {'XX': 'XX', 'YY': 'YY', 'XY': 'YX', 'YX': 'XY',
'JXX': 'jxx', 'JYY': 'jyy', 'JXY': 'jyx', 'JYX': 'jxy',
'RR': 'RR', 'LL': 'LL', 'RL': 'LR', 'LR': 'RL',
'JRR': 'jrr', 'JLL': 'jll', 'JRL': 'jlr', 'JLR': 'jrl',
'I': 'I', 'Q': 'Q', 'U': 'U', 'V': 'V',
'PI': 'pI', 'PQ': 'pQ', 'PU': 'pU', 'PV': 'pV'}

if isinstance(pol, str):
cpol = cpol_dict[pol.upper()]
elif isinstance(pol, collections.Iterable):
cpol = [conj_pol(p) for p in pol]
elif isinstance(pol, six.integer_types + (np.int32, np.int64)):
cpol = polstr2num(cpol_dict[polnum2str(pol).upper()])
else:
raise ValueError('Polarization cannot be conjugated.')
return cpol


def check_history_version(history, version_string):
if (version_string.replace(' ', '') in history.replace('\n', '').replace(' ', '')):
return True
Expand Down

0 comments on commit 11725e9

Please sign in to comment.