better namespace handling, add tests

CHANGELOG.md

@@ -4,7 +4,7 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
-- A new utility function `utils.convert_feeds_to_pols` to get a polarization array
+- 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

src/pyuvdata/analytic_beam.py

@@ -164,7 +164,7 @@ def __post_init__(self, include_cross_pols):
         if len(self.feed_array) == 1:
             include_cross_pols = False
 
-        self.polarization_array = utils.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

@@ -28,7 +28,6 @@
     "x_orientation_pol_map",
     "parse_polstr",
     "parse_jpolstr",
-    "convert_feeds_to_pols",
 ]
 
 # fmt: off
@@ -555,12 +554,18 @@ def convert_feeds_to_pols(
         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(
-            "feed_array contains {n_feeds}. Only 1 or 2 feeds is supported."
+            f"feed_array contains {n_feeds} feeds. Only 1 or 2 feeds is supported."
         )
 
     feed_pol_order = [(0, 0)]

src/pyuvdata/uvbeam/uvbeam.py

@@ -904,11 +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 = utils.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.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([])