MS writer improvements

CHANGELOG.md

@@ -4,11 +4,19 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- Added a switch to `UVData.write_ms` called `flip_conj`, which allows a user to write
+out data with the baseline conjugation scheme flipped from the standard `UVData`
+convention.
+- Added the `utils.determine_pol_order` method for determining polarization
+order based on a specified scheme ("AIPS" or "CASA").
 - Added a `check_surface_based_positions` positions method for verifying antenna
 positions are near surface of whatever celestial body their positions are referenced to
 (either the Earth or Moon, currently).
 
 ### Changed
+- Changed `UVData.write_ms` to sort polarizations based on CASA-preferred ordering.
+- Added some functionality to the `utils._convert_to_slices` method to enable quick
+assessment of whether an indexing array can be replaced by a single slice.
 - Increased the tolerance to 75 mas (equivalent to 5 ms time error) for a warning about
 values in `lst_array` not conforming to expectations for `UVData`, `UVCal`, and `UVFlag`
 (was 1 mas) inside of `check`. Additionally, added a keyword to `check` enable the

pyuvdata/tests/test_telescopes.py

@@ -79,7 +79,7 @@ def test_extra_parameter_iter():
     for prop in telescope_obj.extra():
         extra.append(prop)
     for a in extra_parameters:
-        a in extra, "expected attribute " + a + " not returned in extra iterator"
+        assert a in extra, "expected attribute " + a + " not returned in extra iterator"
 
 
 def test_unexpected_parameters():

pyuvdata/tests/test_utils.py

@@ -2819,7 +2819,7 @@ def test_strict_cliques():
     adj_link[-1] = frozenset({5, 6, 7, 8, 1})
 
     with pytest.raises(ValueError, match="Non-isolated cliques found in graph."):
-        uvutils._find_cliques(adj_link, strict=True),
+        uvutils._find_cliques(adj_link, strict=True)
 
 
 def test_reorder_conj_pols_non_list():
@@ -4026,6 +4026,38 @@ def test_read_slicing():
     data = uvutils._index_dset(dset, slices)
     assert data.shape == tuple(shape)
 
+    # Handling bool arrays
+    bool_arr = np.zeros((10000,), dtype=bool)
+    index_arr = np.arange(1, 10000, 2)
+    bool_arr[index_arr] = True
+    assert uvutils._convert_to_slices(bool_arr) == uvutils._convert_to_slices(index_arr)
+    assert uvutils._convert_to_slices(bool_arr, return_index_on_fail=True) == (
+        uvutils._convert_to_slices(index_arr, return_index_on_fail=True)
+    )
+
+    # Index return on fail with two slices
+    index_arr[0] = 0
+    bool_arr[0:2] = [True, False]
+
+    for item in [index_arr, bool_arr]:
+        result, check = uvutils._convert_to_slices(
+            item, max_nslice=1, return_index_on_fail=True
+        )
+        assert not check
+        assert len(result) == 1
+        assert result[0] is item
+
+    # Check a more complicated pattern w/ just the max_slice_frac defined
+    index_arr = np.arange(0, 100) ** 2
+    bool_arr[:] = False
+    bool_arr[index_arr] = True
+
+    for item in [index_arr, bool_arr]:
+        result, check = uvutils._convert_to_slices(item, return_index_on_fail=True)
+        assert not check
+        assert len(result) == 1
+        assert result[0] is item
+
 
 @pytest.mark.parametrize(
     "blt_order",
@@ -4336,3 +4368,26 @@ def test_check_surface_based_positions_earthmoonloc(tel_loc, check_frame):
             uvutils.check_surface_based_positions(
                 telescope_loc=loc, telescope_frame=frame
             )
+
+
+def test_determine_pol_order_err():
+    with pytest.raises(ValueError, match='order must be either "AIPS" or "CASA".'):
+        uvutils.determine_pol_order([], "ABC")
+
+
+@pytest.mark.parametrize(
+    "pols,aips_order,casa_order",
+    [
+        [[-8, -7, -6, -5], [3, 2, 1, 0], [3, 1, 0, 2]],
+        [[-5, -6, -7, -8], [0, 1, 2, 3], [0, 2, 3, 1]],
+        [[1, 2, 3, 4], [0, 1, 2, 3], [0, 1, 2, 3]],
+    ],
+)
+@pytest.mark.parametrize("order", ["CASA", "AIPS"])
+def test_pol_order(pols, aips_order, casa_order, order):
+    check = uvutils.determine_pol_order(pols, order=order)
+
+    if order == "CASA":
+        assert all(check == casa_order)
+    if order == "AIPS":
+        assert all(check == aips_order)

pyuvdata/utils.py

@@ -1250,6 +1250,42 @@ def reorder_conj_pols(pols):
     return conj_order
 
 
+def determine_pol_order(pols, order="AIPS"):
+    """
+    Determine order of input polarization numbers.
+
+    Determines the order by which to sort a given list of polarizations, according to
+    the ordering scheme. Two orders are currently supported: "AIPS" and "CASA". The
+    main difference between the two is the grouping of same-handed polarizations for
+    AIPS (whereas CASA orders the polarizations such that same-handed pols are on the
+    ends of the array).
+
+    Parameters
+    ----------
+    pols : array_like of str or int
+        Polarization array (strings or ints).
+    order : str
+        Polarization ordering scheme, either "CASA" or "AIPS".
+
+    Returns
+    -------
+    index_array : ndarray of int
+        Indices to reorder polarization array.
+    """
+    if order == "AIPS":
+        index_array = np.argsort(np.abs(pols))
+    elif order == "CASA":
+        casa_order = np.array([1, 2, 3, 4, -1, -3, -4, -2, -5, -7, -8, -6, 0])
+        pol_inds = []
+        for pol in pols:
+            pol_inds.append(np.where(casa_order == pol)[0][0])
+        index_array = np.argsort(pol_inds)
+    else:
+        raise ValueError('order must be either "AIPS" or "CASA".')
+
+    return index_array
+
+
 def LatLonAlt_from_XYZ(xyz, frame="ITRS", check_acceptability=True):
     """
     Calculate lat/lon/alt from ECEF x,y,z.
@@ -5706,88 +5742,115 @@ def _get_dset_shape(dset, indices):
     return dset_shape, indices
 
 
-def _convert_to_slices(indices, max_nslice_frac=0.1):
+def _convert_to_slices(
+    indices, max_nslice_frac=0.1, max_nslice=None, return_index_on_fail=False
+):
     """
     Convert list of indices to a list of slices.
 
     Parameters
     ----------
     indices : list
-        A 1D list of integers for array indexing.
+        A 1D list of integers for array indexing (boolean ndarrays are also supported).
     max_nslice_frac : float
         A float from 0 -- 1. If the number of slices
         needed to represent input 'indices' divided by len(indices)
         exceeds this fraction, then we determine that we cannot
         easily represent 'indices' with a list of slices.
+    max_nslice : int
+        Optional argument, defines the maximum number of slices for determining if
+        `indices` can be easily represented with a list of slices. If set, then
+        the argument supplied to `max_nslice_frac` is ignored.
+    return_index_on_fail : bool
+        If set to True and the list of input indexes cannot easily be respresented by
+        a list of slices (as defined by `max_nslice` or `max_nslice_frac`), then return
+        the input list of index values instead of a list of suboptimal slices.
 
     Returns
     -------
-    list
-        list of slice objects used to represent indices
-    bool
+    slice_list : list
+        Nominally the list of slice objects used to represent indices. However, if
+        `return_index_on_fail=True` and input indexes cannot easily be respresented,
+        return a 1-element list containing the input for `indices`.
+    check : bool
         If True, indices is easily represented by slices
-        (max_nslice_frac condition met), otherwise False
+        (`max_nslice_frac` or `max_nslice` conditions met), otherwise False.
 
     Notes
     -----
     Example:
         if: indices = [1, 2, 3, 4, 10, 11, 12, 13, 14]
         then: slices = [slice(1, 5, 1), slice(11, 15, 1)]
     """
-    # check for integer index
-    if isinstance(indices, (int, np.integer)):
-        indices = [indices]
-    # check for already a slice
+    # check for already a slice or a single index position
     if isinstance(indices, slice):
         return [indices], True
+    if isinstance(indices, (int, np.integer)):
+        return [slice(indices, indices + 1, 1)], True
+
+    # check for boolean index
+    if isinstance(indices, np.ndarray) and (indices.dtype == bool):
+        eval_ind = np.where(indices)[0]
+    else:
+        eval_ind = indices
     # assert indices is longer than 2, or return trivial solutions
-    if len(indices) == 0:
+    if len(eval_ind) == 0:
         return [slice(0, 0, 0)], False
-    elif len(indices) == 1:
-        return [slice(indices[0], indices[0] + 1, 1)], True
-    elif len(indices) == 2:
-        return [slice(indices[0], indices[1] + 1, indices[1] - indices[0])], True
+    if len(eval_ind) <= 2:
+        return [
+            slice(eval_ind[0], eval_ind[-1] + 1, max(eval_ind[-1] - eval_ind[0], 1))
+        ], True
+
+    # Catch the simplest case of "give me a single slice or exit"
+    if (max_nslice == 1) and return_index_on_fail:
+        step = eval_ind[1] - eval_ind[0]
+        if all(np.diff(eval_ind) == step):
+            return [slice(eval_ind[0], eval_ind[-1] + 1, step)], True
+        return [indices], False
 
     # setup empty slices list
-    Ninds = len(indices)
+    Ninds = len(eval_ind)
     slices = []
 
     # iterate over indices
-    for i, ind in enumerate(indices):
-        if i == 0:
-            # start the first slice object
-            start = ind
-            last_step = indices[i + 1] - ind
+    start = last_step = None
+    for ind in eval_ind:
+        if last_step is None:
+            # Check if this is the first slice, in which case start is None
+            if start is None:
+                start = ind
+                continue
+            last_step = ind - start
+            last_ind = ind
             continue
 
         # calculate step from previous index
-        step = ind - indices[i - 1]
+        step = ind - last_ind
 
         # if step != last_step, this ends the slice
         if step != last_step:
             # append to list
-            slices.append(slice(start, indices[i - 1] + 1, last_step))
-
-            # check if this is the last element
-            if i == Ninds - 1:
-                # append last element
-                slices.append(slice(ind, ind + 1, 1))
-                continue
+            slices.append(slice(start, last_ind + 1, last_step))
 
             # setup next step
             start = ind
-            last_step = indices[i + 1] - ind
+            last_step = None
+
+        last_ind = ind
 
-        # check if this is the last element
-        elif i == Ninds - 1:
-            # end slice and append
-            slices.append(slice(start, ind + 1, step))
+    # Append the last slice
+    slices.append(slice(start, ind + 1, last_step))
 
-    # determine whether slices are a reasonable representation
-    Nslices = len(slices)
-    passed = (float(Nslices) / len(indices)) < max_nslice_frac
+    # determine whether slices are a reasonable representation, and determine max_nslice
+    # if only max_nslice_frac was supplied.
+    if max_nslice is None:
+        max_nslice = max_nslice_frac * Ninds
+    check = len(slices) <= max_nslice
 
-    return slices, passed
+    if return_index_on_fail and not check:
+        return [indices], check
+    else:
+        return slices, check
 
 
 def _get_slice_len(s, axlen):

pyuvdata/uvbeam/tests/test_cst_beam.py

@@ -209,7 +209,7 @@ def test_read_yaml_override(cst_efield_2freq_mod):
             beam_type="efield",
             telescope_name="test",
             use_future_array_shapes=True,
-        ),
+        )
 
     assert beam1 == beam2