diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f714b348..34df6d881 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,6 +6,10 @@ All notable changes to this project will be documented in this file. ### Added - ATA has been added to the list of known telescopes. +### Fixed +- Bug in selecting baselines on a UVData object using `bls` keyword with 3-tuples and +more than one polarization (introduced in 3.1.2). + ## [3.1.2] - 2024-11-21 ### Added diff --git a/src/pyuvdata/utils/bls.py b/src/pyuvdata/utils/bls.py index bc6486e4e..7c4ade16f 100644 --- a/src/pyuvdata/utils/bls.py +++ b/src/pyuvdata/utils/bls.py @@ -423,6 +423,7 @@ def _extract_bls_pol( ) bls_2 = copy.deepcopy(bls) + bl_pols = set() for bl_i, bl in enumerate(bls): if len(bl) != 3: raise ValueError("If some bls are 3-tuples, all bls must be 3-tuples.") @@ -432,7 +433,6 @@ def _extract_bls_pol( "The third element in a bl tuple must be a polarization string" ) - bl_pols = set() wh1 = np.where(np.logical_and(ant_1_array == bl[0], ant_2_array == bl[1]))[ 0 ] @@ -454,6 +454,6 @@ def _extract_bls_pol( "associated with it." ) - polarizations = list(bl_pols) + polarizations = list(bl_pols) bls = bls_2 return bls, polarizations diff --git a/tests/uvdata/test_uvdata.py b/tests/uvdata/test_uvdata.py index 7f041d1e2..cb5cdd4c3 100644 --- a/tests/uvdata/test_uvdata.py +++ b/tests/uvdata/test_uvdata.py @@ -1612,6 +1612,11 @@ def test_select_bls(casa_uvfits, sel_type): first_ants = [7, 3, 8, 3, 22, 28, 9] second_ants = [1, 21, 9, 2, 3, 4, 23] pols = ["RR", "RR", "RR", "RR", "RR", "RR", "RR"] + + if sel_type == "antpairpol": + # Also test that reading different pols at the same time works. + pols[-1] = "LL" + new_unique_ants = np.unique(first_ants + second_ants) ant_pairs_to_keep = list(zip(first_ants, second_ants, strict=True)) sorted_pairs_to_keep = [sort_bl(p) for p in ant_pairs_to_keep] @@ -1675,9 +1680,7 @@ def test_select_bls(casa_uvfits, sel_type): assert pair in sorted_pairs_to_keep if sel_type == "antpairpol": - assert uv_object2.Npols == 1 - if sel_type == "2_3_tuple": - assert uv_object2.Npols == 1 + assert uv_object2.Npols == 2 assert utils.history._check_histories( old_history + f" Downselected to specific {sel_str} using pyuvdata.",