From bc69f98142bcd20ecba0df3e7300fcb62bdc4b23 Mon Sep 17 00:00:00 2001 From: Bryna Hazelton Date: Tue, 19 Mar 2024 15:35:58 -0700 Subject: [PATCH] refactor the selection on the blt axis out so it can be reused. --- pyuvdata/uvdata/miriad.py | 42 +- pyuvdata/uvdata/tests/test_uvdata.py | 22 +- pyuvdata/uvdata/uvdata.py | 703 +++++++++++++++------------ 3 files changed, 413 insertions(+), 354 deletions(-) diff --git a/pyuvdata/uvdata/miriad.py b/pyuvdata/uvdata/miriad.py index 88611b1a7a..c271607a45 100644 --- a/pyuvdata/uvdata/miriad.py +++ b/pyuvdata/uvdata/miriad.py @@ -788,8 +788,7 @@ def read_miriad( # but that may not matter for many purposes. return - history_update_string = " Downselected to specific " - n_selects = 0 + selections = [] # select on ant_str if provided if ant_str is not None: @@ -801,8 +800,7 @@ def read_miriad( aipy_extracts.uv_selector(uv, ant_str) if ant_str != "all": - history_update_string += "antenna pairs" - n_selects += 1 + selections.append("antenna pairs") # select on antenna_nums and/or bls using aipy_extracts.uv_selector if antenna_nums is not None or bls is not None: @@ -821,8 +819,7 @@ def read_miriad( # convert antenna numbers to string form required by # aipy_extracts.uv_selector antpair_str_list = ["_".join([str(a) for a in ap]) for ap in antpairs] - history_update_string += "antennas" - n_selects += 1 + selections.append("antennas") if bls is not None: if isinstance(bls, tuple) and (len(bls) == 2 or len(bls) == 3): @@ -866,7 +863,7 @@ def read_miriad( if len(bl) == 3: bl_pols.add(uvutils.conj_pol(bl[2])) - if n_selects > 0: + if len(selections) > 0: # combine antpair_str_list and bl_str_list with an intersection antpair_str_list = list( set(antpair_str_list).intersection(bl_str_list) @@ -877,11 +874,7 @@ def read_miriad( if len(bl_pols) > 0: polarizations = list(bl_pols) - if n_selects > 0: - history_update_string += ", baselines" - else: - history_update_string += "baselines" - n_selects += 1 + selections.append("antenna pairs") # convert antenna pair list to string form required by # aipy_extracts.uv_selector @@ -911,11 +904,8 @@ def read_miriad( time_range_use = np.array(time_range) - uv["inttime"] / (24 * 3600.0) / 2 uv.select("time", time_range_use[0], time_range_use[1], include=True) - if n_selects > 0: - history_update_string += ", times" - else: - history_update_string += "times" - n_selects += 1 + + selections.append("times") # select on polarizations if polarizations is not None: @@ -949,15 +939,15 @@ def read_miriad( # check not empty if len(pol_list) == 0: raise ValueError("No polarizations in data matched input") - if n_selects > 0: - history_update_string += ", polarizations" - else: - history_update_string += "polarizations" - n_selects += 1 - - history_update_string += " using pyuvdata." - if n_selects > 0: - self.history += history_update_string + selections.append("polarizations") + + if len(selections) > 0: + # build up history string from selections + self.history += ( + " Downselected to specific " + + ", ".join(selections) + + " using pyuvdata." + ) data_accumulator = {} pol_list = [] diff --git a/pyuvdata/uvdata/tests/test_uvdata.py b/pyuvdata/uvdata/tests/test_uvdata.py index ac67aa7fdc..d81f1a8b5e 100644 --- a/pyuvdata/uvdata/tests/test_uvdata.py +++ b/pyuvdata/uvdata/tests/test_uvdata.py @@ -1481,10 +1481,9 @@ def test_select_phase_center_id_blts(carma_miriad): uv_sum = uv1 + uv2 + uv3 assert uvutils._check_histories( uv_obj.history - + " Downselected to specific baseline-times, phase center IDs using pyuvdata. " + + " Downselected to specific phase center IDs, baseline-times using pyuvdata. " "Combined data along baseline-time axis using pyuvdata. " - "Combined data along baseline-time axis using pyuvdata. " - "Unique part of next object history follows. baseline-times", + "Combined data along baseline-time axis using pyuvdata. ", uv_sum.history, ) uv_sum.history = uv_obj.history @@ -1655,7 +1654,7 @@ def test_select_bls(casa_uvfits): assert pair in sorted_pairs_to_keep assert uvutils._check_histories( - old_history + " Downselected to specific baselines using pyuvdata.", + old_history + " Downselected to specific antenna pairs using pyuvdata.", uv_object2.history, ) @@ -1684,7 +1683,7 @@ def test_select_bls(casa_uvfits): assert pair in sorted_pairs_to_keep assert uvutils._check_histories( - old_history + " Downselected to specific baselines using pyuvdata.", + old_history + " Downselected to specific antenna pairs using pyuvdata.", uv_object3.history, ) @@ -1724,7 +1723,7 @@ def test_select_bls(casa_uvfits): assert uvutils._check_histories( old_history - + " Downselected to specific baselines, polarizations using pyuvdata.", + + " Downselected to specific antenna pairs, polarizations using pyuvdata.", uv_object2.history, ) @@ -1752,7 +1751,7 @@ def test_select_bls(casa_uvfits): assert pair in sorted_pairs_to_keep assert uvutils._check_histories( - old_history + " Downselected to specific baselines using pyuvdata.", + old_history + " Downselected to specific antenna pairs using pyuvdata.", uv_object2.history, ) @@ -1791,12 +1790,13 @@ def test_select_bls(casa_uvfits): with pytest.raises( ValueError, - match="Cannot provide length-3 tuples and also specify polarizations.", + match="Cannot provide any length-3 tuples and also specify polarizations.", ): uv_object.select(bls=(7, 1, "RR"), polarizations="RR") with pytest.raises( - ValueError, match="The third element in each bl must be a polarization string" + ValueError, + match="The third element in a bl tuple must be a polarization string", ): uv_object.select(bls=(7, 1, 7)) @@ -2745,7 +2745,7 @@ def test_select(casa_uvfits, future_shapes): assert uvutils._check_histories( old_history + " Downselected to " "specific baseline-times, antennas, " - "baselines, times, frequencies, " + "antenna pairs, times, frequencies, " "polarizations using pyuvdata.", uv_object2.history, ) @@ -2853,7 +2853,7 @@ def test_select_with_lst(casa_uvfits, future_shapes): assert uvutils._check_histories( old_history + " Downselected to " "specific baseline-times, antennas, " - "baselines, lsts, frequencies, " + "antenna pairs, lsts, frequencies, " "polarizations using pyuvdata.", uv_object2.history, ) diff --git a/pyuvdata/uvdata/uvdata.py b/pyuvdata/uvdata/uvdata.py index 5bed99dffc..8ba60cee98 100644 --- a/pyuvdata/uvdata/uvdata.py +++ b/pyuvdata/uvdata/uvdata.py @@ -110,6 +110,324 @@ def _warn_old_phase_attr(__name): warnings.warn(warn_str, DeprecationWarning) +def _select_blt_preprocess( + *, + select_antenna_nums, + select_antenna_names, + bls, + times, + time_range, + lsts, + lst_range, + blt_inds, + phase_center_ids, + antenna_names, + antenna_numbers, + ant_1_array, + ant_2_array, + baseline_array, + time_array, + time_tols, + lst_array, + lst_tols, + phase_center_id_array, +): + """Build up blt_inds and selections list for _select_preprocess. + + Parameters + ---------- + antenna_nums : array_like of int, optional + The antennas numbers to keep in the object (antenna positions and + names for the removed antennas will be retained unless + `keep_all_metadata` is False). This cannot be provided if + `antenna_names` is also provided. + antenna_names : array_like of str, optional + The antennas names to keep in the object (antenna positions and + names for the removed antennas will be retained unless + `keep_all_metadata` is False). This cannot be provided if + `antenna_nums` is also provided. + bls : list of 2-tuples, optional + A list of antenna number tuples (e.g. [(0, 1), (3, 2)]) specifying + baselines to keep in the object. The ordering of the numbers within the + tuple does not matter. Note that this is different than what can be + passed to the parameter of the same name on `select` -- this parameter + does not accept 3-tuples or baseline numbers. + times : array_like of float, optional + The times to keep in the object, each value passed here should exist + in the time_array. Cannot be used with `time_range`, `lsts`, or + `lst_array`. + time_range : array_like of float, optional + The time range in Julian Date to keep in the object, must be length + 2. Some of the times in the object should fall between the first and + last elements. Cannot be used with `times`, `lsts`, or `lst_array`. + lsts : array_like of float, optional + The local sidereal times (LSTs) to keep in the object, each value + passed here should exist in the lst_array. Cannot be used with + `times`, `time_range`, or `lst_range`. + lst_range : array_like of float, optional + The local sidereal time (LST) range in radians to keep in the + object, must be of length 2. Some of the LSTs in the object should + fall between the first and last elements. If the second value is + smaller than the first, the LSTs are treated as having phase-wrapped + around LST = 2*pi = 0, and the LSTs kept on the object will run from + the larger value, through 0, and end at the smaller value. + phase_center_ids : array_like of int, optional + Phase center IDs to keep on the object (effectively a selection on + baseline-times). Cannot be used with `catalog_names`. + blt_inds : array_like of int, optional + The baseline-time indices to keep in the object. This is + not commonly used. + ant_1_array : array_like of int + Array of first antenna numbers to select on. + ant_2_array : array_like of int + Array of second antenna numbers to select on. + baseline_array : array_like of int + Array of baseline numbers to select on. + time_array : array_like of float + Array of times in JD to select on. + lst_array : array_like of float + Array of lsts in radians to select on. + phase_center_id_array : array_like of int + Array of phase center IDs to select on. + + Returns + ------- + blt_inds : list of int + list of baseline-time indices to keep. Can be None (to keep everything). + selections : list of str + list of selections done. + """ + # Antennas, times and blt_inds all need to be combined into a set of + # blts indices to keep. + selections = [] + # test for blt_inds presence before adding inds from antennas & times + if blt_inds is not None: + blt_inds = uvutils._get_iterable(blt_inds) + if np.array(blt_inds).ndim > 1: + blt_inds = np.array(blt_inds).flatten() + selections.append("baseline-times") + + if phase_center_ids is not None: + phase_center_ids = np.array(uvutils._get_iterable(phase_center_ids)) + pc_blt_inds = np.nonzero(np.isin(phase_center_id_array, phase_center_ids))[0] + if blt_inds is not None: + # Use intersection (and) to join phase_center_ids + # with blt_inds + blt_inds = np.array( + list(set(blt_inds).intersection(pc_blt_inds)), dtype=np.int64 + ) + else: + blt_inds = pc_blt_inds + + if select_antenna_names is not None: + if select_antenna_nums is not None: + raise ValueError( + "Only one of antenna_nums and antenna_names can be provided." + ) + + if not isinstance(select_antenna_names, (list, tuple, np.ndarray)): + select_antenna_names = (select_antenna_names,) + if np.array(select_antenna_names).ndim > 1: + select_antenna_names = np.array(select_antenna_names).flatten() + select_antenna_nums = [] + for s in select_antenna_names: + if s not in antenna_names: + raise ValueError( + "Antenna name {a} is not present in the antenna_names" + " array".format(a=s) + ) + select_antenna_nums.append( + antenna_numbers[np.where(np.array(antenna_names) == s)][0] + ) + + if select_antenna_nums is not None: + select_antenna_nums = uvutils._get_iterable(select_antenna_nums) + select_antenna_nums = np.asarray(select_antenna_nums) + if select_antenna_nums.ndim > 1: + select_antenna_nums = select_antenna_nums.flatten() + selections.append("antennas") + # Check to make sure that we actually have these antenna nums in the data + ant_check = np.logical_or( + np.isin(select_antenna_nums, ant_1_array), + np.isin(select_antenna_nums, ant_2_array), + ) + if not np.all(ant_check): + raise ValueError( + f"Antenna number {select_antenna_nums[~ant_check]} is not present " + "in the ant_1_array or ant_2_array" + ) + ant_blt_inds = np.where( + np.logical_and( + np.isin(ant_1_array, select_antenna_nums), + np.isin(ant_2_array, select_antenna_nums), + ) + )[0] + ant_blt_inds = np.asarray(ant_blt_inds, dtype=np.int64) + else: + ant_blt_inds = None + + if bls is not None: + bls_blt_inds = np.zeros(0, dtype=np.int64) + for bl in bls: + wh1 = np.where(np.logical_and(ant_1_array == bl[0], ant_2_array == bl[1]))[ + 0 + ] + if len(wh1) > 0: + bls_blt_inds = np.append(bls_blt_inds, list(wh1)) + else: + wh2 = np.where( + np.logical_and(ant_1_array == bl[1], ant_2_array == bl[0]) + )[0] + + if len(wh2) > 0: + bls_blt_inds = np.append(bls_blt_inds, list(wh2)) + else: + raise ValueError( + "Antenna pair {p} does not have any data " + "associated with it.".format(p=bl) + ) + selections.append("antenna pairs") + + if ant_blt_inds is not None: + # Use intersection (and) to join antenna_names/nums & ant_pairs_nums + ant_blt_inds = np.intersect1d(ant_blt_inds, bls_blt_inds) + else: + ant_blt_inds = bls_blt_inds + + if ant_blt_inds is not None: + if blt_inds is not None: + # Use intersection (and) to join antenna_names/nums/ant_pairs_nums + # with blt_inds + blt_inds = np.intersect1d(blt_inds, ant_blt_inds) + else: + blt_inds = ant_blt_inds + + have_times = times is not None + have_time_range = time_range is not None + have_lsts = lsts is not None + have_lst_range = lst_range is not None + if np.count_nonzero([have_times, have_time_range, have_lsts, have_lst_range]) > 1: + raise ValueError( + "Only one of [times, time_range, lsts, lst_range] may be " + "specified per selection operation." + ) + + if times is not None: + times = uvutils._get_iterable(times) + if np.array(times).ndim > 1: + times = np.array(times).flatten() + + time_blt_inds = np.zeros(0, dtype=np.int64) + for jd in times: + if np.any(np.isclose(time_array, jd, rtol=time_tols[0], atol=time_tols[1])): + time_blt_inds = np.append( + time_blt_inds, + np.where( + np.isclose(time_array, jd, rtol=time_tols[0], atol=time_tols[1]) + )[0], + ) + else: + raise ValueError( + "Time {t} is not present in the time_array".format(t=jd) + ) + + if time_range is not None: + if np.size(time_range) != 2: + raise ValueError("time_range must be length 2.") + + time_blt_inds = np.nonzero( + (time_array <= time_range[1]) & (time_array >= time_range[0]) + )[0] + if time_blt_inds.size == 0: + raise ValueError( + f"No elements in time range between {time_range[0]} and " + f"{time_range[1]}." + ) + + if lsts is not None: + if np.any(np.asarray(lsts) > 2 * np.pi): + warnings.warn( + "The lsts parameter contained a value greater than 2*pi. " + "LST values are assumed to be in radians, not hours." + ) + lsts = uvutils._get_iterable(lsts) + if np.array(lsts).ndim > 1: + lsts = np.array(lsts).flatten() + + time_blt_inds = np.zeros(0, dtype=np.int64) + for lst in lsts: + if np.any(np.isclose(lst_array, lst, rtol=lst_tols[0], atol=lst_tols[1])): + time_blt_inds = np.append( + time_blt_inds, + np.where( + np.isclose(lst_array, lst, rtol=lst_tols[0], atol=lst_tols[1]) + )[0], + ) + else: + raise ValueError(f"LST {lst} is not present in the lst_array") + + if lst_range is not None: + if np.size(lst_range) != 2: + raise ValueError("lst_range must be length 2.") + if np.any(np.asarray(lst_range) > 2 * np.pi): + warnings.warn( + "The lst_range contained a value greater than 2*pi. " + "LST values are assumed to be in radians, not hours." + ) + if lst_range[1] < lst_range[0]: + # we're wrapping around LST = 2*pi = 0 + lst_range_1 = [lst_range[0], 2 * np.pi] + lst_range_2 = [0, lst_range[1]] + time_blt_inds1 = np.nonzero( + (lst_array <= lst_range_1[1]) & (lst_array >= lst_range_1[0]) + )[0] + time_blt_inds2 = np.nonzero( + (lst_array <= lst_range_2[1]) & (lst_array >= lst_range_2[0]) + )[0] + time_blt_inds = np.union1d(time_blt_inds1, time_blt_inds2) + else: + time_blt_inds = np.nonzero( + (lst_array <= lst_range[1]) & (lst_array >= lst_range[0]) + )[0] + if time_blt_inds.size == 0: + raise ValueError( + f"No elements in LST range between {lst_range[0]} and " + f"{lst_range[1]}." + ) + + if times is not None or time_range is not None: + selections.append("times") + + if blt_inds is not None: + # Use intesection (and) to join + # antenna_names/nums/ant_pairs_nums/blt_inds with times + blt_inds = np.intersect1d(blt_inds, time_blt_inds) + else: + blt_inds = time_blt_inds + + if lsts is not None or lst_range is not None: + selections.append("lsts") + + if blt_inds is not None: + # Use intesection (and) to join + # antenna_names/nums/ant_pairs_nums/blt_inds with times + blt_inds = np.intersect1d(blt_inds, time_blt_inds) + else: + blt_inds = time_blt_inds + + if blt_inds is not None: + if len(blt_inds) == 0: + raise ValueError("No baseline-times were found that match criteria") + if max(blt_inds) >= baseline_array.size: + raise ValueError("blt_inds contains indices that are too large") + if min(blt_inds) < 0: + raise ValueError("blt_inds contains indices that are negative") + + blt_inds = sorted(set(blt_inds)) + + return blt_inds, selections + + class UVData(UVBase): """ A class for defining a radio interferometer dataset. @@ -7939,10 +8257,7 @@ def _select_preprocess( history_update_string : str string to append to the end of the history. """ - # build up history string as we go - history_update_string = " Downselected to specific " - n_selects = 0 - + selections = [] if ant_str is not None: if not ( antenna_nums is None @@ -7961,95 +8276,14 @@ def _select_preprocess( f"There is no data matching ant_str={ant_str} in this object." ) - # Antennas, times and blt_inds all need to be combined into a set of - # blts indices to keep. - - # test for blt_inds presence before adding inds from antennas & times - if blt_inds is not None: - blt_inds = uvutils._get_iterable(blt_inds) - if np.array(blt_inds).ndim > 1: - blt_inds = np.array(blt_inds).flatten() - history_update_string += "baseline-times" - n_selects += 1 - if (phase_center_ids is not None) and (catalog_names is not None): raise ValueError("Cannot set both phase_center_ids and catalog_names.") if catalog_names is not None: phase_center_ids = self._look_for_name(catalog_names) - - if phase_center_ids is not None: - phase_center_ids = np.array(uvutils._get_iterable(phase_center_ids)) - pc_blt_inds = np.nonzero( - np.isin(self.phase_center_id_array, phase_center_ids) - )[0] - if blt_inds is not None: - # Use intersection (and) to join phase_center_ids - # with blt_inds - blt_inds = np.array( - list(set(blt_inds).intersection(pc_blt_inds)), dtype=np.int64 - ) - else: - blt_inds = pc_blt_inds - - update_substring = ( - "phase center IDs" if (catalog_names is None) else "catalog names" - ) - if n_selects > 0: - history_update_string += ", " + update_substring - else: - history_update_string += update_substring - n_selects += 1 - - if antenna_names is not None: - if antenna_nums is not None: - raise ValueError( - "Only one of antenna_nums and antenna_names can be provided." - ) - - if not isinstance(antenna_names, (list, tuple, np.ndarray)): - antenna_names = (antenna_names,) - if np.array(antenna_names).ndim > 1: - antenna_names = np.array(antenna_names).flatten() - antenna_nums = [] - for s in antenna_names: - if s not in self.antenna_names: - raise ValueError( - "Antenna name {a} is not present in the antenna_names" - " array".format(a=s) - ) - antenna_nums.append( - self.antenna_numbers[np.where(np.array(self.antenna_names) == s)][0] - ) - - if antenna_nums is not None: - antenna_nums = uvutils._get_iterable(antenna_nums) - antenna_nums = np.asarray(antenna_nums) - if antenna_nums.ndim > 1: - antenna_nums = antenna_nums.flatten() - if n_selects > 0: - history_update_string += ", antennas" - else: - history_update_string += "antennas" - n_selects += 1 - # Check to make sure that we actually have these antenna nums in the data - ant_check = np.logical_or( - np.isin(antenna_nums, self.ant_1_array), - np.isin(antenna_nums, self.ant_2_array), - ) - if not np.all(ant_check): - raise ValueError( - f"Antenna number {antenna_nums[~ant_check]} is not present in the " - "ant_1_array or ant_2_array" - ) - ant_blt_inds = np.where( - np.logical_and( - np.isin(self.ant_1_array, antenna_nums), - np.isin(self.ant_2_array, antenna_nums), - ) - )[0] - else: - ant_blt_inds = None + selections.append("catalog names") + elif phase_center_ids is not None: + selections.append("phase center IDs") if bls is not None: if isinstance(bls, list) and all( @@ -8058,8 +8292,8 @@ def _select_preprocess( for bl_ind in bls: if not (bl_ind in self.baseline_array): raise ValueError( - "Baseline number {i} is not present in the " - "baseline_array".format(i=bl_ind) + f"Baseline number {bl_ind} is not present in the " + "baseline_array" ) bls = list(zip(*self.baseline_to_antnums(bls))) elif isinstance(bls, tuple) and (len(bls) == 2 or len(bls) == 3): @@ -8077,235 +8311,75 @@ def _select_preprocess( "bls must be a list of tuples of antenna numbers " "(optionally with polarization) or a list of baseline numbers." ) - if all(len(item) == 3 for item in bls): + if any(len(item) == 3 for item in bls): if polarizations is not None: raise ValueError( - "Cannot provide length-3 tuples and also specify polarizations." - ) - if not all(isinstance(item[2], str) for item in bls): - raise ValueError( - "The third element in each bl must be a polarization string" + "Cannot provide any length-3 tuples and also specify " + "polarizations." ) - if ant_str is None: - if n_selects > 0: - history_update_string += ", baselines" - else: - history_update_string += "baselines" - else: - history_update_string += "antenna pairs" - n_selects += 1 - bls_blt_inds = np.zeros(0, dtype=np.int64) - bl_pols = set() - for bl in bls: - wh1 = np.where( - np.logical_and(self.ant_1_array == bl[0], self.ant_2_array == bl[1]) - )[0] - if len(wh1) > 0: - bls_blt_inds = np.append(bls_blt_inds, list(wh1)) - if len(bl) == 3: - bl_pols.add(bl[2]) - else: - wh2 = np.where( - np.logical_and( - self.ant_1_array == bl[1], self.ant_2_array == bl[0] - ) - )[0] + bls_2 = copy.deepcopy(bls) + for bl_i, bl in enumerate(bls): + if len(bl) == 2: + continue - if len(wh2) > 0: - bls_blt_inds = np.append(bls_blt_inds, list(wh2)) - if len(bl) == 3: - # find conjugate polarization - bl_pols.add(uvutils.conj_pol(bl[2])) - else: + if not isinstance(bl[2], str): raise ValueError( - "Antenna pair {p} does not have any data " - "associated with it.".format(p=bl) + "The third element in a bl tuple must be a " + "polarization string" ) - if len(bl_pols) > 0: - polarizations = list(bl_pols) - if ant_blt_inds is not None: - # Use intersection (and) to join antenna_names/nums & ant_pairs_nums - ant_blt_inds = np.array( - list(set(ant_blt_inds).intersection(bls_blt_inds)) - ) - else: - ant_blt_inds = bls_blt_inds - - if ant_blt_inds is not None: - if blt_inds is not None: - # Use intersection (and) to join antenna_names/nums/ant_pairs_nums - # with blt_inds - blt_inds = np.array( - list(set(blt_inds).intersection(ant_blt_inds)), dtype=np.int64 - ) - else: - blt_inds = ant_blt_inds - - have_times = times is not None - have_time_range = time_range is not None - have_lsts = lsts is not None - have_lst_range = lst_range is not None - if ( - np.count_nonzero([have_times, have_time_range, have_lsts, have_lst_range]) - > 1 - ): - raise ValueError( - "Only one of [times, time_range, lsts, lst_range] may be " - "specified per selection operation." - ) - - if times is not None: - times = uvutils._get_iterable(times) - if np.array(times).ndim > 1: - times = np.array(times).flatten() - - time_blt_inds = np.zeros(0, dtype=np.int64) - for jd in times: - if np.any( - np.isclose( - self.time_array, - jd, - rtol=self._time_array.tols[0], - atol=self._time_array.tols[1], - ) - ): - time_blt_inds = np.append( - time_blt_inds, - np.where( - np.isclose( - self.time_array, - jd, - rtol=self._time_array.tols[0], - atol=self._time_array.tols[1], + bl_pols = set() + wh1 = np.where( + np.logical_and( + self.ant_1_array == bl[0], self.ant_2_array == bl[1] + ) + )[0] + if len(wh1) > 0: + bls_2[bl_i] = (bl[0], bl[1]) + bl_pols.add(bl[2]) + else: + wh2 = np.where( + np.logical_and( + self.ant_1_array == bl[1], self.ant_2_array == bl[0] ) - )[0], - ) - else: - raise ValueError( - "Time {t} is not present in the time_array".format(t=jd) - ) - - if time_range is not None: - if np.size(time_range) != 2: - raise ValueError("time_range must be length 2.") - - time_blt_inds = np.nonzero( - (self.time_array <= time_range[1]) & (self.time_array >= time_range[0]) - )[0] - if time_blt_inds.size == 0: - raise ValueError( - f"No elements in time range between {time_range[0]} and " - f"{time_range[1]}." - ) + )[0] - if lsts is not None: - if np.any(np.asarray(lsts) > 2 * np.pi): - warnings.warn( - "The lsts parameter contained a value greater than 2*pi. " - "LST values are assumed to be in radians, not hours." - ) - lsts = uvutils._get_iterable(lsts) - if np.array(lsts).ndim > 1: - lsts = np.array(lsts).flatten() - - time_blt_inds = np.zeros(0, dtype=np.int64) - for lst in lsts: - if np.any( - np.isclose( - self.lst_array, - lst, - rtol=self._lst_array.tols[0], - atol=self._lst_array.tols[1], - ) - ): - time_blt_inds = np.append( - time_blt_inds, - np.where( - np.isclose( - self.lst_array, - lst, - rtol=self._lst_array.tols[0], - atol=self._lst_array.tols[1], + if len(wh2) > 0: + bls_2[bl_i] = (bl[1], bl[0]) + # find conjugate polarization + bl_pols.add(uvutils.conj_pol(bl[2])) + else: + raise ValueError( + "Antenna pair {p} does not have any data " + "associated with it.".format(p=bl) ) - )[0], - ) - else: - raise ValueError(f"LST {lst} is not present in the lst_array") - - if lst_range is not None: - if np.size(lst_range) != 2: - raise ValueError("lst_range must be length 2.") - if np.any(np.asarray(lst_range) > 2 * np.pi): - warnings.warn( - "The lst_range contained a value greater than 2*pi. " - "LST values are assumed to be in radians, not hours." - ) - if lst_range[1] < lst_range[0]: - # we're wrapping around LST = 2*pi = 0 - lst_range_1 = [lst_range[0], 2 * np.pi] - lst_range_2 = [0, lst_range[1]] - time_blt_inds1 = np.nonzero( - (self.lst_array <= lst_range_1[1]) - & (self.lst_array >= lst_range_1[0]) - )[0] - time_blt_inds2 = np.nonzero( - (self.lst_array <= lst_range_2[1]) - & (self.lst_array >= lst_range_2[0]) - )[0] - time_blt_inds = np.union1d(time_blt_inds1, time_blt_inds2) - else: - time_blt_inds = np.nonzero( - (self.lst_array <= lst_range[1]) & (self.lst_array >= lst_range[0]) - )[0] - if time_blt_inds.size == 0: - raise ValueError( - f"No elements in LST range between {lst_range[0]} and " - f"{lst_range[1]}." - ) - - if times is not None or time_range is not None: - if n_selects > 0: - history_update_string += ", times" - else: - history_update_string += "times" - n_selects += 1 - - if blt_inds is not None: - # Use intesection (and) to join - # antenna_names/nums/ant_pairs_nums/blt_inds with times - blt_inds = np.array( - list(set(blt_inds).intersection(time_blt_inds)), dtype=np.int64 - ) - else: - blt_inds = time_blt_inds - if lsts is not None or lst_range is not None: - if n_selects > 0: - history_update_string += ", lsts" - else: - history_update_string += "lsts" - n_selects += 1 - - if blt_inds is not None: - # Use intesection (and) to join - # antenna_names/nums/ant_pairs_nums/blt_inds with times - blt_inds = np.array( - list(set(blt_inds).intersection(time_blt_inds)), dtype=np.int64 - ) - else: - blt_inds = time_blt_inds - - if blt_inds is not None: - if len(blt_inds) == 0: - raise ValueError("No baseline-times were found that match criteria") - if max(blt_inds) >= self.Nblts: - raise ValueError("blt_inds contains indices that are too large") - if min(blt_inds) < 0: - raise ValueError("blt_inds contains indices that are negative") + polarizations = list(bl_pols) + bls = bls_2 - blt_inds = sorted(set(blt_inds)) + blt_inds, blt_selections = _select_blt_preprocess( + select_antenna_nums=antenna_nums, + select_antenna_names=antenna_names, + bls=bls, + times=times, + time_range=time_range, + lsts=lsts, + lst_range=lst_range, + blt_inds=blt_inds, + phase_center_ids=phase_center_ids, + antenna_names=self.antenna_names, + antenna_numbers=self.antenna_numbers, + ant_1_array=self.ant_1_array, + ant_2_array=self.ant_2_array, + baseline_array=self.baseline_array, + time_array=self.time_array, + time_tols=self._time_array.tols, + lst_array=self.lst_array, + lst_tols=self._lst_array.tols, + phase_center_id_array=self.phase_center_id_array, + ) + selections.extend(blt_selections) if freq_chans is not None: freq_chans = uvutils._get_iterable(freq_chans) @@ -8331,11 +8405,7 @@ def _select_preprocess( frequencies = uvutils._get_iterable(frequencies) if np.array(frequencies).ndim > 1: frequencies = np.array(frequencies).flatten() - if n_selects > 0: - history_update_string += ", frequencies" - else: - history_update_string += "frequencies" - n_selects += 1 + selections.append("frequencies") if self.future_array_shapes: freq_arr_use = self.freq_array @@ -8377,11 +8447,7 @@ def _select_preprocess( polarizations = uvutils._get_iterable(polarizations) if np.array(polarizations).ndim > 1: polarizations = np.array(polarizations).flatten() - if n_selects > 0: - history_update_string += ", polarizations" - else: - history_update_string += "polarizations" - n_selects += 1 + selections.append("polarizations") pol_inds = np.zeros(0, dtype=np.int64) spw_inds = np.zeros(0, dtype=np.int64) @@ -8454,7 +8520,10 @@ def _select_preprocess( else: pol_inds = None - history_update_string += " using pyuvdata." + # build up history string from selections + history_update_string = ( + " Downselected to specific " + ", ".join(selections) + " using pyuvdata." + ) return blt_inds, freq_inds, pol_inds, history_update_string