diff --git a/CHANGELOG.md b/CHANGELOG.md index ad9ad71f1..6cb1ec715 100644 --- a/CHANGELOG.md +++ b/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). diff --git a/src/pyuvdata/analytic_beam.py b/src/pyuvdata/analytic_beam.py index 80a6765b5..b94f730e5 100644 --- a/src/pyuvdata/analytic_beam.py +++ b/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 ) diff --git a/src/pyuvdata/utils/pol.py b/src/pyuvdata/utils/pol.py index 916d57ff8..d3b1f8c5d 100644 --- a/src/pyuvdata/utils/pol.py +++ b/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, diff --git a/src/pyuvdata/uvbeam/initializers.py b/src/pyuvdata/uvbeam/initializers.py index b99ccf153..5121ccc71 100644 --- a/src/pyuvdata/uvbeam/initializers.py +++ b/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, diff --git a/src/pyuvdata/uvbeam/uvbeam.py b/src/pyuvdata/uvbeam/uvbeam.py index 2824e1f72..b7a3b2b9d 100644 --- a/src/pyuvdata/uvbeam/uvbeam.py +++ b/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 diff --git a/tests/utils/test_pol.py b/tests/utils/test_pol.py index 6f868bf58..1716b3555 100644 --- a/tests/utils/test_pol.py +++ b/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([])