refactor the uvbeam._convert_feeds_to_pols into a utility function

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- A new utility function `utils.pol.convert_feeds_to_pols` to get a polarization array
+given a feed array (previously existed as a private function in uvbeam.py).
 - New `strict` keyword added to `UVData.select`, `UVBeam.select`, `UVFlag.select`, and
 `UVFlag.select`, which allows the user to specify whether to warn or error when
 supplied criteria only partially match (default being to warn).

src/pyuvdata/analytic_beam.py

@@ -19,7 +19,7 @@
 
 from . import utils
 from .docstrings import combine_docstrings
-from .uvbeam.uvbeam import UVBeam, _convert_feeds_to_pols
+from .uvbeam.uvbeam import UVBeam
 
 __all__ = ["AnalyticBeam", "AiryBeam", "GaussianBeam", "ShortDipoleBeam", "UniformBeam"]
 
@@ -164,7 +164,7 @@ def __post_init__(self, include_cross_pols):
         if len(self.feed_array) == 1:
             include_cross_pols = False
 
-        self.polarization_array, _ = _convert_feeds_to_pols(
+        self.polarization_array = utils.pol.convert_feeds_to_pols(
             self.feed_array, include_cross_pols, x_orientation=self.x_orientation
         )
 

src/pyuvdata/utils/pol.py

@@ -8,6 +8,7 @@
 from functools import lru_cache, wraps
 
 import numpy as np
+import numpy.typing as npt
 
 from . import tools
 
@@ -531,6 +532,63 @@ def determine_pol_order(pols, *, order="AIPS"):
     return index_array
 
 
+def convert_feeds_to_pols(
+    feed_array: npt.NDArray[str],
+    include_cross_pols: bool = True,
+    x_orientation: str | None = None,
+    return_feed_pol_order: bool = False,
+):
+    """
+    Get the polarizations given a feed array.
+
+    Parameters
+    ----------
+    feed_array : ndarray of str
+        Array of feed orientations. Options are: n/e or x/y or r/l.
+    include_cross_pols : bool
+        Option to include the cross polarizations (e.g. xy and yx or en and ne).
+        Defaults to True if more than one feed, set to False for only one feed.
+    x_orientation : str, optional
+        Orientation of the x-axis. Options are 'east', 'north', 'e', 'n', 'ew', 'ns'.
+    return_feed_pol_order : bool
+        Option to return a list of tuples giving the ordering of the feeds for
+        each pol. Default False.
+
+    Returns
+    -------
+    polarization_array : ndarray of int
+        Polarization integer array.
+    feed_pol_order : list of tuples of int, optional
+        List of feed index tuples for each pol.
+    """
+    n_feeds = np.asarray(feed_array).size
+
+    if n_feeds < 1 or n_feeds > 2:
+        raise ValueError(
+            f"feed_array contains {n_feeds} feeds. Only 1 or 2 feeds is supported."
+        )
+
+    feed_pol_order = [(0, 0)]
+    if n_feeds > 1:
+        feed_pol_order.append((1, 1))
+    else:
+        include_cross_pols = False
+
+    if include_cross_pols:
+        feed_pol_order.extend([(0, 1), (1, 0)])
+
+    pol_strings = []
+    for pair in feed_pol_order:
+        pol_strings.append(feed_array[pair[0]] + feed_array[pair[1]])
+    polarization_array = np.array(
+        [polstr2num(ps.upper(), x_orientation=x_orientation) for ps in pol_strings]
+    )
+    if return_feed_pol_order:
+        return polarization_array, feed_pol_order
+    else:
+        return polarization_array
+
+
 def _select_pol_helper(
     polarizations,
     obj_pol_array,

src/pyuvdata/uvbeam/initializers.py

@@ -23,7 +23,7 @@ def new_uvbeam(
     feed_version: str = "0.0",
     model_name: str = "default",
     model_version: str = "0.0",
-    feed_array: npt.NDArray[np.str] | None = None,
+    feed_array: npt.NDArray[str] | None = None,
     polarization_array: (
         npt.NDArray[np.str | np.int] | list[str | int] | tuple[str | int] | None
     ) = None,

src/pyuvdata/uvbeam/uvbeam.py

@@ -22,29 +22,6 @@
 __all__ = ["UVBeam"]
 
 
-def _convert_feeds_to_pols(feed_array, calc_cross_pols, x_orientation=None):
-    n_feeds = np.asarray(feed_array).size
-
-    feed_pol_order = [(0, 0)]
-    if n_feeds > 1:
-        feed_pol_order.append((1, 1))
-
-    if calc_cross_pols:
-        # to get here we have Nfeeds > 1
-        feed_pol_order.extend([(0, 1), (1, 0)])
-
-    pol_strings = []
-    for pair in feed_pol_order:
-        pol_strings.append(feed_array[pair[0]] + feed_array[pair[1]])
-    polarization_array = np.array(
-        [
-            utils.polstr2num(ps.upper(), x_orientation=x_orientation)
-            for ps in pol_strings
-        ]
-    )
-    return polarization_array, feed_pol_order
-
-
 class UVBeam(UVBase):
     """
     A class for defining a radio telescope antenna beam.
@@ -927,10 +904,13 @@ def efield_to_power(
             # There are no cross pols with one feed. Set this so the power beam is real
             calc_cross_pols = False
 
-        beam_object.polarization_array, feed_pol_order = _convert_feeds_to_pols(
-            beam_object.feed_array,
-            calc_cross_pols,
-            x_orientation=beam_object.x_orientation,
+        beam_object.polarization_array, feed_pol_order = (
+            utils.pol.convert_feeds_to_pols(
+                beam_object.feed_array,
+                include_cross_pols=calc_cross_pols,
+                x_orientation=beam_object.x_orientation,
+                return_feed_pol_order=True,
+            )
         )
         beam_object.Npols = beam_object.polarization_array.size
 

tests/utils/test_pol.py

@@ -243,3 +243,57 @@ def test_x_orientation_pol_map():
         assert utils._x_orientation_rep_dict("east") == {"x": "e", "y": "n"}
 
     assert utils.x_orientation_pol_map("north") == {"x": "n", "y": "e"}
+
+
+@pytest.mark.parametrize(
+    ("input_kwargs", "output"),
+    [
+        ({"feed_array": ["x"]}, [-5]),
+        ({"feed_array": ["x", "y"]}, [-5, -6, -7, -8]),
+        ({"feed_array": ["r", "l"]}, [-1, -2, -3, -4]),
+        ({"feed_array": ["x", "y"], "include_cross_pols": False}, [-5, -6]),
+        (
+            {
+                "feed_array": ["e", "n"],
+                "x_orientation": "east",
+                "return_feed_pol_order": True,
+            },
+            ([-5, -6, -7, -8], [(0, 0), (1, 1), (0, 1), (1, 0)]),
+        ),
+        (
+            {
+                "feed_array": ["n", "e"],
+                "x_orientation": "north",
+                "return_feed_pol_order": True,
+            },
+            ([-5, -6, -7, -8], [(0, 0), (1, 1), (0, 1), (1, 0)]),
+        ),
+        (
+            {
+                "feed_array": ["e", "n"],
+                "x_orientation": "north",
+                "return_feed_pol_order": True,
+            },
+            ([-6, -5, -8, -7], [(0, 0), (1, 1), (0, 1), (1, 0)]),
+        ),
+    ],
+)
+def test_convert_feeds_to_pols(input_kwargs, output):
+    ret_val = utils.pol.convert_feeds_to_pols(**input_kwargs)
+    if not isinstance(output, tuple):
+        assert ret_val.tolist() == output
+    else:
+        assert ret_val[0].tolist() == output[0]
+        assert ret_val[1] == output[1]
+
+
+def test_convert_feeds_to_pols_errors():
+    with pytest.raises(
+        ValueError, match="feed_array contains 3 feeds. Only 1 or 2 feeds is supported."
+    ):
+        utils.pol.convert_feeds_to_pols(["x", "y", "z"])
+
+    with pytest.raises(
+        ValueError, match="feed_array contains 0 feeds. Only 1 or 2 feeds is supported."
+    ):
+        utils.pol.convert_feeds_to_pols([])