diff --git a/pyuvdata/tests/test_utils.py b/pyuvdata/tests/test_utils.py index b3ffd61b4f..a731f63ceb 100644 --- a/pyuvdata/tests/test_utils.py +++ b/pyuvdata/tests/test_utils.py @@ -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) diff --git a/pyuvdata/tests/test_uvdata.py b/pyuvdata/tests/test_uvdata.py index d8a194931c..3b8b00698c 100644 --- a/pyuvdata/tests/test_uvdata.py +++ b/pyuvdata/tests/test_uvdata.py @@ -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 @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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)) @@ -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]) @@ -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]) @@ -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)) @@ -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 @@ -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) diff --git a/pyuvdata/utils.py b/pyuvdata/utils.py index b11e3d4588..4fdcc1e229 100644 --- a/pyuvdata/utils.py +++ b/pyuvdata/utils.py @@ -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 diff --git a/pyuvdata/uvdata.py b/pyuvdata/uvdata.py index 3dfdc7a8cd..897179fe34 100644 --- a/pyuvdata/uvdata.py +++ b/pyuvdata/uvdata.py @@ -2164,16 +2164,24 @@ def _key2inds(self, key): Returns: blt_ind1: numpy array with blt indices for antenna pair. blt_ind2: numpy array with blt indices for conjugate antenna pair. - pol_ind: numpy array with polarization indices + Note if a cross-pol baseline is requested, the polarization will + also be reversed so the appropriate correlations are returned. + e.g. asking for (1, 2, 'xy') may return conj(2, 1, 'yx'), which + is equivalent to the requesting baseline. See utils.conj_pol() for + complete conjugation mapping. + pol_ind: tuple of numpy arrays with polarization indices for blt_ind1 and blt_ind2 """ key = uvutils.get_iterable(key) if type(key) is str: # Single string given, assume it is polarization - pol_ind = np.where(self.polarization_array == uvutils.polstr2num(key))[0] - if len(pol_ind) == 0: + pol_ind1 = np.where(self.polarization_array == uvutils.polstr2num(key))[0] + if len(pol_ind1) > 0: + blt_ind1 = np.arange(self.Nblts) + blt_ind2 = np.array([], dtype=np.int64) + pol_ind2 = np.array([], dtype=np.int64) + pol_ind = (pol_ind1, pol_ind2) + else: raise KeyError('Polarization {pol} not found in data.'.format(pol=key)) - blt_ind1 = np.arange(self.Nblts) - blt_ind2 = np.array([], dtype=np.int64) elif len(key) == 1: key = key[0] # For simplicity if isinstance(key, collections.Iterable): @@ -2181,11 +2189,14 @@ def _key2inds(self, key): blt_ind1, blt_ind2, pol_ind = self._key2inds(key) elif key < 5: # Small number, assume it is a polarization number a la AIPS memo - pol_ind = np.where(self.polarization_array == key)[0] - if len(pol_ind) == 0: + pol_ind1 = np.where(self.polarization_array == key)[0] + if len(pol_ind1) > 0: + blt_ind1 = np.arange(self.Nblts) + blt_ind2 = np.array([], dtype=np.int64) + pol_ind2 = np.array([], dtype=np.int64) + pol_ind = (pol_ind1, pol_ind2) + else: raise KeyError('Polarization {pol} not found in data.'.format(pol=key)) - blt_ind1 = np.arange(self.Nblts) - blt_ind2 = np.array([], dtype=np.int64) else: # Larger number, assume it is a baseline number inv_bl = self.antnums_to_baseline(self.baseline_to_antnums(key)[1], @@ -2194,14 +2205,14 @@ def _key2inds(self, key): blt_ind2 = np.where(self.baseline_array == inv_bl)[0] if len(blt_ind1) + len(blt_ind2) == 0: raise KeyError('Baseline {bl} not found in data.'.format(bl=key)) - pol_ind = np.arange(self.Npols) + pol_ind = (np.arange(self.Npols), np.arange(self.Npols)) elif len(key) == 2: # Key is an antenna pair blt_ind1 = self.antpair2ind(key[0], key[1]) blt_ind2 = self.antpair2ind(key[1], key[0]) if len(blt_ind1) + len(blt_ind2) == 0: raise KeyError('Antenna pair {pair} not found in data'.format(pair=key)) - pol_ind = np.arange(self.Npols) + pol_ind = (np.arange(self.Npols), np.arange(self.Npols)) elif len(key) == 3: # Key is an antenna pair + pol blt_ind1 = self.antpair2ind(key[0], key[1]) @@ -2211,11 +2222,27 @@ def _key2inds(self, key): 'data'.format(pair=(key[0], key[1]))) if type(key[2]) is str: # pol is str - pol_ind = np.where(self.polarization_array == uvutils.polstr2num(key[2]))[0] + if len(blt_ind1) > 0: + pol_ind1 = np.where(self.polarization_array == uvutils.polstr2num(key[2]))[0] + else: + pol_ind1 = np.array([], dtype=np.int64) + if len(blt_ind2) > 0: + pol_ind2 = np.where(self.polarization_array + == uvutils.polstr2num(uvutils.conj_pol(key[2])))[0] + else: + pol_ind2 = np.array([], dtype=np.int64) else: # polarization number a la AIPS memo - pol_ind = np.where(self.polarization_array == key[2])[0] - if len(pol_ind) == 0: + if len(blt_ind1) > 0: + pol_ind1 = np.where(self.polarization_array == key[2])[0] + else: + pol_ind1 = np.array([], dtype=np.int64) + if len(blt_ind2) > 0: + pol_ind2 = np.where(self.polarization_array == uvutils.conj_pol(key[2]))[0] + else: + pol_ind2 = np.array([], dtype=np.int64) + pol_ind = (pol_ind1, pol_ind2) + if len(blt_ind1) * len(pol_ind[0]) + len(blt_ind2) * len(pol_ind[1]) == 0: raise KeyError('Polarization {pol} not found in data.'.format(pol=key[2])) # Catch autos if np.array_equal(blt_ind1, blt_ind2): @@ -2230,7 +2257,7 @@ def _smart_slicing(self, data, ind1, ind2, indp, **kwargs): data: 4-dimensional array in the format of self.data_array ind1: list with blt indices for antenna pair (e.g. from self._key2inds) ind2: list with blt indices for conjugate antenna pair. (e.g. from self._key2inds) - indp: list with polarization indices (e.g. from self._key2inds) + indp: tuple of lists with polarization indices for ind1 and ind2 (e.g. from self._key2inds) kwargs: squeeze: 'default': squeeze pol and spw dimensions if possible (default) 'none': no squeezing of resulting numpy array @@ -2243,16 +2270,19 @@ def _smart_slicing(self, data, ind1, ind2, indp, **kwargs): force_copy = kwargs.pop('force_copy', False) squeeze = kwargs.pop('squeeze', 'default') - if len(set(np.ediff1d(indp))) <= 1: - p_reg_spaced = True - p_start = indp[0] - p_stop = indp[-1] + 1 - if len(indp) == 1: - dp = 1 - else: - dp = indp[1] - indp[0] - else: - p_reg_spaced = False + p_reg_spaced = [False, False] + p_start = [0, 0] + p_stop = [0, 0] + dp = [1, 1] + for i, pi in enumerate(indp): + if len(pi) == 0: + continue + if len(set(np.ediff1d(pi))) <= 1: + p_reg_spaced[i] = True + p_start[i] = pi[0] + p_stop[i] = pi[-1] + 1 + if len(pi) != 1: + dp[i] = pi[1] - pi[0] if len(ind2) == 0: # only unconjugated baselines @@ -2263,16 +2293,16 @@ def _smart_slicing(self, data, ind1, ind2, indp, **kwargs): dblt = 1 else: dblt = ind1[1] - ind1[0] - if p_reg_spaced: - out = data[blt_start:blt_stop:dblt, :, :, p_start:p_stop:dp] + if p_reg_spaced[0]: + out = data[blt_start:blt_stop:dblt, :, :, p_start[0]:p_stop[0]:dp[0]] else: - out = data[blt_start:blt_stop:dblt, :, :, indp] + out = data[blt_start:blt_stop:dblt, :, :, indp[0]] else: out = data[ind1, :, :, :] - if p_reg_spaced: - out = out[:, :, :, p_start:p_stop:dp] + if p_reg_spaced[0]: + out = out[:, :, :, p_start[0]:p_stop[0]:dp[0]] else: - out = out[:, :, :, indp] + out = out[:, :, :, indp[0]] elif len(ind1) == 0: # only conjugated baselines if len(set(np.ediff1d(ind2))) <= 1: @@ -2282,34 +2312,36 @@ def _smart_slicing(self, data, ind1, ind2, indp, **kwargs): dblt = 1 else: dblt = ind2[1] - ind2[0] - if p_reg_spaced: - out = np.conj(data[blt_start:blt_stop:dblt, :, :, p_start:p_stop:dp]) + if p_reg_spaced[1]: + out = np.conj(data[blt_start:blt_stop:dblt, :, :, p_start[1]:p_stop[1]:dp[1]]) else: - out = np.conj(data[blt_start:blt_stop:dblt, :, :, indp]) + out = np.conj(data[blt_start:blt_stop:dblt, :, :, indp[1]]) else: out = data[ind2, :, :, :] - if p_reg_spaced: - out = np.conj(out[:, :, :, p_start:p_stop:dp]) + if p_reg_spaced[1]: + out = np.conj(out[:, :, :, p_start[1]:p_stop[1]:dp[1]]) else: - out = np.conj(out[:, :, :, indp]) + out = np.conj(out[:, :, :, indp[1]]) else: # both conjugated and unconjugated baselines - out = np.append(data[ind1, :, :, :], np.conj(data[ind2, :, :, :]), axis=0) - if p_reg_spaced: - out = out[:, :, :, p_start:p_stop:dp] + out = (data[ind1, :, :, :], np.conj(data[ind2, :, :, :])) + if p_reg_spaced[0] and p_reg_spaced[1]: + out = np.append(out[0][:, :, :, p_start[0]:p_stop[0]:dp[0]], + out[1][:, :, :, p_start[1]:p_stop[1]:dp[1]], axis=0) else: - out = out[:, :, :, indp] + out = np.append(out[0][:, :, :, indp[0]], + out[1][:, :, :, indp[1]], axis=0) - if squeeze is 'full': + if squeeze == 'full': out = np.squeeze(out) - elif squeeze is 'default': + elif squeeze == 'default': if out.shape[3] is 1: # one polarization dimension out = np.squeeze(out, axis=3) if out.shape[1] is 1: # one spw dimension out = np.squeeze(out, axis=1) - elif squeeze is not 'none': + elif squeeze != 'none': raise ValueError('"' + str(squeeze) + '" is not a valid option for squeeze.' 'Only "default", "none", or "full" are allowed.')