Add select on read for MWA correlator FITS files

CHANGELOG.md

@@ -3,6 +3,11 @@ All notable changes to this project will be documented in this file.
 
 ## [Unreleased]
 
+### Added
+- Added support for partial read for MWA correlator FITS files.
+- Added `antenna_names`, `time_range`, `lsts` and `lst_range` parameters to
+`UVFlag.select` to match UVData and UVCal select methods.
+
 ### Changed
 - Made it possible to *not* return the `interp_basis_vector` array from beam
 interpolators.

docs/uvdata_tutorial.rst

@@ -1557,20 +1557,17 @@ Measurement set (ms) files do not support reading only the metadata
   >>> print(uvd.data_array)
   None
 
-b) Reading only parts of uvfits, uvh5 or miriad data
-****************************************************
+b) Reading only parts of files
+******************************
 The same options that are available for the :meth:`pyuvdata.UVData.select` method can
 also be passed to the :meth:`pyuvdata.UVData.read`` method to do the select on the read,
 saving memory and time if only a portion of the data are needed.
 
-Note that these keywords can be used for any file type, but for FHD,
-MWA correlator FITS files, and
+Note that these keywords can be used for any file type, but for FHD and
 measurement set (ms) files, the select is done after the read, which does not
-save memory. Miriad only supports some of the selections on the read, the
-unsupported ones are done after the read.
-Any of the select keywords can be used for any file type, but selects for keywords
-that are not supported by the select on read for a given file type will be
-done after the read, which does not save memory.
+save memory. Miriad and Mir only supports some of the selections on the read, the
+unsupported ones are done after the read. MWA correlator fits has support for most
+but not all selections, the unsupported ones are done after the read.
 
 .. code-block:: python
 

src/pyuvdata/utils/__init__.py

@@ -18,6 +18,7 @@
 LST_RAD_TOL = 2 * np.pi * 5e-3 / (86400.0)
 
 # these seem to be necessary for the installed package to access these submodules
+from . import antenna  # noqa
 from . import apply_uvflag  # noqa
 from . import array_collapse  # noqa
 from . import bls  # noqa

src/pyuvdata/utils/antenna.py

@@ -0,0 +1,50 @@
+# Copyright (c) 2024 Radio Astronomy Software Group
+# Licensed under the 2-clause BSD License
+"""Utilities for antennas."""
+
+import numpy as np
+
+from . import tools
+
+
+def _select_antenna_helper(
+    *, antenna_names, antenna_nums, tel_ant_names, tel_ant_nums, obj_ant_array
+):
+    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."
+            )
+
+        antenna_names = tools._get_iterable(antenna_names)
+        antenna_nums = []
+        for s in antenna_names:
+            if s not in tel_ant_names:
+                raise ValueError(
+                    f"Antenna name {s} is not present in the antenna_names array"
+                )
+            ind = np.where(np.array(tel_ant_names) == s)[0][0]
+            antenna_nums.append(tel_ant_nums[ind])
+
+    if antenna_nums is not None:
+        antenna_nums = tools._get_iterable(antenna_nums)
+        antenna_nums = np.asarray(antenna_nums)
+        selections = ["antennas"]
+
+        ant_inds = np.zeros(0, dtype=np.int64)
+        ant_check = np.isin(antenna_nums, obj_ant_array)
+        if not np.all(ant_check):
+            raise ValueError(
+                f"Antenna number {antenna_nums[~ant_check]} is not present "
+                "in the ant_array"
+            )
+
+        for ant in antenna_nums:
+            ant_inds = np.append(ant_inds, np.where(obj_ant_array == ant)[0])
+
+        ant_inds = sorted(set(ant_inds))
+    else:
+        ant_inds = None
+        selections = None
+
+    return ant_inds, selections

src/pyuvdata/utils/bls.py

@@ -2,13 +2,14 @@
 # Licensed under the 2-clause BSD License
 """Utilities for baseline numbers."""
 
+import copy
 import re
 import warnings
 
 import numpy as np
 
 from . import _bls
-from .pol import polnum2str, polstr2num
+from .pol import conj_pol, polnum2str, polstr2num
 
 __all__ = ["baseline_to_antnums", "antnums_to_baseline"]
 
@@ -384,3 +385,75 @@ def parse_ants(uv, ant_str, *, print_toggle=False, x_orientation=None):
         )
 
     return ant_pairs_nums, polarizations
+
+
+def _extract_bls_pol(
+    *, bls, polarizations, baseline_array, ant_1_array, ant_2_array, nants_telescope
+):
+    if isinstance(bls, list) and all(
+        isinstance(bl_ind, int | np.integer) for bl_ind in bls
+    ):
+        for bl_ind in bls:
+            if bl_ind not in baseline_array:
+                raise ValueError(
+                    f"Baseline number {bl_ind} is not present in the baseline_array"
+                )
+        bls = list(
+            zip(*baseline_to_antnums(bls, Nants_telescope=nants_telescope), strict=True)
+        )
+    elif isinstance(bls, tuple) and (len(bls) == 2 or len(bls) == 3):
+        bls = [bls]
+    if len(bls) == 0 or not all(isinstance(item, tuple) for item in bls):
+        raise ValueError(
+            "bls must be a list of tuples of antenna numbers "
+            "(optionally with polarization) or a list of baseline numbers."
+        )
+    if not all(
+        [isinstance(item[0], int | np.integer) for item in bls]
+        + [isinstance(item[1], int | np.integer) for item in bls]
+    ):
+        raise ValueError(
+            "bls must be a list of tuples of antenna numbers "
+            "(optionally with polarization) or a list of baseline numbers."
+        )
+    if any(len(item) == 3 for item in bls):
+        if polarizations is not None:
+            raise ValueError(
+                "Cannot provide any length-3 tuples and also specify polarizations."
+            )
+
+        bls_2 = copy.deepcopy(bls)
+        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.")
+
+            if not isinstance(bl[2], str):
+                raise ValueError(
+                    "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
+            ]
+            if len(wh1) > 0:
+                bls_2[bl_i] = (bl[0], bl[1])
+                bl_pols.add(bl[2])
+            else:
+                wh2 = np.where(
+                    np.logical_and(ant_1_array == bl[1], ant_2_array == bl[0])
+                )[0]
+
+                if len(wh2) > 0:
+                    bls_2[bl_i] = (bl[1], bl[0])
+                    # find conjugate polarization
+                    bl_pols.add(conj_pol(bl[2]))
+                else:
+                    raise ValueError(
+                        f"Antenna pair {bl} does not have any data "
+                        "associated with it."
+                    )
+
+            polarizations = list(bl_pols)
+        bls = bls_2
+    return bls, polarizations