Handle single feeds in simsetup and uvsim

Add testing

src/pyuvsim/data/mwa88_nocore_config_dipole.yaml

@@ -0,0 +1,5 @@
+beam_paths:
+    0 : !AnalyticBeam
+        class: ShortDipoleBeam
+telescope_location: (-30.72152777777791, 21.428305555555557, 1073.0000000093132)
+telescope_name: MWA

src/pyuvsim/data/test_config/28m_triangle_10time_10chan_xfeed.yaml

@@ -0,0 +1,19 @@
+beam_paths:
+  0: !UVBeam
+    filename: HERA_NicCST.beamfits
+    path_variable: pyuvsim.data.DATA_PATH
+  1: !AnalyticBeam
+    class: UniformBeam
+    feed_array: ["x"]
+  2: !AnalyticBeam
+    class: GaussianBeam
+    sigma: 0.02
+    feed_array: ["x"]
+  3: !AnalyticBeam
+    class: AiryBeam
+    diameter: 14.6
+    feed_array: ["x"]
+telescope_location: (-30.721527777777847, 21.428305555555557, 1073.0000000046566)
+telescope_name: Triangle
+select:
+  feeds: ["x"]

src/pyuvsim/data/test_config/28m_triangle_10time_10chan_yfeed.yaml

@@ -0,0 +1,19 @@
+beam_paths:
+  0: !UVBeam
+    filename: HERA_NicCST.beamfits
+    path_variable: pyuvsim.data.DATA_PATH
+  1: !AnalyticBeam
+    class: UniformBeam
+    feed_array: ["y"]
+  2: !AnalyticBeam
+    class: GaussianBeam
+    sigma: 0.02
+    feed_array: ["y"]
+  3: !AnalyticBeam
+    class: AiryBeam
+    diameter: 14.6
+    feed_array: ["y"]
+telescope_location: (-30.721527777777847, 21.428305555555557, 1073.0000000046566)
+telescope_name: Triangle
+select:
+  feeds: ["y"]

src/pyuvsim/simsetup.py

@@ -1875,6 +1875,57 @@ def set_ordering(uv_obj, param_dict, reorder_blt_kw):
         uv_obj.reorder_blts(**reorder_blt_kw)
 
 
+def _initialize_polarization_helper(param_dict, uvparam_dict, beam_list=None):
+    # There does not seem to be any way to get polarization_array into uvparam_dict, so
+    # let's add it explicitly.
+    if param_dict.get("polarization_array", None) is not None:
+        if beam_list is not None:
+            polstr_list = uvutils.polnum2str(param_dict["polarization_array"])
+            feed_set = set()
+            for polstr in polstr_list:
+                feed_set = feed_set.union(uvutils.pol.POL_TO_FEED_DICT[polstr])
+            for feed in feed_set:
+                beam_feeds = beam_list[0].feed_array
+                if "e" in beam_feeds or "n" in beam_feeds:
+                    feed_map = uvutils.x_orientation_pol_map(beam_list.x_orientation)
+                    inv_feed_map = {value: key for key, value in feed_map.items()}
+                    beam_feeds = [inv_feed_map[feed] for feed in beam_feeds]
+
+                if feed not in beam_feeds:
+                    raise ValueError(
+                        f"Specified polarization array {param_dict['polarization_array']} "
+                        f"requires feeds {feed_set} but the beams have feeds {beam_feeds}."
+                    )
+
+        uvparam_dict["polarization_array"] = np.array(param_dict["polarization_array"])
+
+    # Parse polarizations
+    if uvparam_dict.get("polarization_array") is None:
+        if beam_list is not None:
+            # default polarization array based on the beam feeds
+            try:
+                uvparam_dict["polarization_array"] = uvutils.pol.convert_feeds_to_pols(
+                    feed_array=beam_list[0].feed_array,
+                    include_cross_pols=True,
+                    x_orientation=beam_list.x_orientation,
+                )
+            except AttributeError:
+                from pyuvdata.uvbeam.uvbeam import _convert_feeds_to_pols
+
+                uvparam_dict["polarization_array"], _ = _convert_feeds_to_pols(
+                    feed_array=beam_list[0].feed_array,
+                    calc_cross_pols=bool(np.asarray(beam_list[0].feed_array).size > 1),
+                    x_orientation=beam_list.x_orientation,
+                )
+        else:
+            uvparam_dict["polarization_array"] = np.array([-5, -6, -7, -8])
+
+    if "Npols" not in uvparam_dict:
+        uvparam_dict["Npols"] = len(uvparam_dict["polarization_array"])
+
+    return uvparam_dict
+
+
 def initialize_uvdata_from_params(
     obs_params, return_beams=False, reorder_blt_kw=None, check_kw=None
 ):
@@ -2003,17 +2054,14 @@ def initialize_uvdata_from_params(
     uvparam_dict["integration_time"] = parsed_time_dict["integration_time"]
     logger.info("Finished Setup of Time Dict")
 
-    # There does not seem to be any way to get polarization_array into uvparam_dict, so
-    # let's add it explicitly.
-    if param_dict.get("polarization_array", None) is not None:
-        uvparam_dict["polarization_array"] = np.array(param_dict["polarization_array"])
-
-    # Parse polarizations
-    if uvparam_dict.get("polarization_array") is None:
-        uvparam_dict["polarization_array"] = np.array([-5, -6, -7, -8])
-
-    if "Npols" not in uvparam_dict:
-        uvparam_dict["Npols"] = len(uvparam_dict["polarization_array"])
+    # figure out polarization
+    if return_beams:
+        bl_pass = beam_list
+    else:
+        bl_pass = None
+    uvparam_dict = _initialize_polarization_helper(
+        param_dict, uvparam_dict, beam_list=bl_pass
+    )
 
     # telescope frame is set from world in parse_telescope_params.
     # Can only be itrs or (if lunarsky is installed) mcmf
@@ -2182,7 +2230,6 @@ def initialize_uvdata_from_keywords(
     write_files=True,
     path_out=None,
     complete=False,
-    force_beam_check=False,
     check_kw: dict | None = None,
     **kwargs,
 ):

src/pyuvsim/uvsim.py

@@ -12,6 +12,7 @@
 
 import astropy.units as units
 import numpy as np
+import pyuvdata.utils as uvutils
 import yaml
 from astropy.constants import c as speed_of_light
 from astropy.time import Time
@@ -81,7 +82,7 @@ class UVTask:
         Indexes for this visibility location in the uvdata data_array, length 3 giving
         (blt_index, 0, frequency_index), where the zero index is for the old spw axis.
     visibility_vector : ndarray of complex
-        The calculated visibility, shape (4,) ordered as [xx, yy, xy, yx].
+        The calculated visibility, shape (Npols,) ordered as [xx, yy, xy, yx].
 
     """
 
@@ -797,11 +798,28 @@ def run_uvdata_uvsim(
     if not isinstance(input_uv, UVData):
         raise TypeError("input_uv must be UVData object")
 
+    exp_npols = beam_list[0].Nfeeds ** 2
+    try:
+        exp_pols = uvutils.convert_feeds_to_pols(
+            feed_array=beam_list[0].feed_array,
+            include_cross_pols=True,
+            x_orientation=beam_list.x_orientation,
+        )
+    except AttributeError:
+        from pyuvdata.uvbeam.uvbeam import _convert_feeds_to_pols
+
+        exp_pols, _ = _convert_feeds_to_pols(
+            feed_array=beam_list[0].feed_array,
+            calc_cross_pols=bool(exp_npols > 1),
+            x_orientation=beam_list.x_orientation,
+        )
+
     if not (
-        (input_uv.Npols == 4)
-        and (input_uv.polarization_array.tolist() == [-5, -6, -7, -8])
+        (input_uv.Npols == exp_npols)
+        and (input_uv.polarization_array.tolist() == exp_pols.tolist())
     ):
-        raise ValueError("input_uv must have XX,YY,XY,YX polarization")
+        exp_pols_str = uvutils.polnum2str(exp_pols)
+        raise ValueError(f"input_uv must have polarizations: {exp_pols_str}")
 
     input_order = input_uv.blt_order
     if input_order != ("time", "baseline"):
@@ -827,6 +845,7 @@ def run_uvdata_uvsim(
     Nbls = input_uv.Nbls
     Nblts = input_uv.Nblts
     Nfreqs = input_uv.Nfreqs
+    Npols = input_uv.Npols
     Nsrcs = catalog.Ncomponents
 
     task_inds, src_inds, Ntasks_local, Nsrcs_local = _make_task_inds(
@@ -879,7 +898,7 @@ def run_uvdata_uvsim(
 
         engine = UVEngine()
         size_complex = np.ones(1, dtype=complex).nbytes
-        data_array_shape = (Nblts, Nfreqs, 4)
+        data_array_shape = (Nblts, Nfreqs, Npols)
         uvdata_indices = []
 
         for task in local_task_iter:

tests/test_run.py

@@ -246,14 +246,33 @@ def test_run_paramdict_uvsim(rename_beamfits, tmp_path):
 
 
 @pytest.mark.parametrize("spectral_type", ["flat", "subband", "spectral_index"])
-def test_run_gleam_uvsim(spectral_type):
+@pytest.mark.parametrize("nfeeds", [1, 2])
+def test_run_gleam_uvsim(spectral_type, nfeeds):
     params = pyuvsim.simsetup._config_str_to_dict(
         os.path.join(SIM_DATA_PATH, "test_config", "param_1time_testgleam.yaml")
     )
     params["sources"]["spectral_type"] = spectral_type
     params["sources"].pop("min_flux")
     params["sources"].pop("max_flux")
 
+    if nfeeds == 1:
+        if spectral_type == "subband":
+            tel_config = "28m_triangle_10time_10chan_yfeed.yaml"
+            select_pol = ["yy"]
+        else:
+            tel_config = "28m_triangle_10time_10chan_xfeed.yaml"
+            select_pol = ["xx"]
+        params["telescope"]["telescope_config_name"] = tel_config
+
+    input_uv, beam_list, beam_dict = pyuvsim.simsetup.initialize_uvdata_from_params(
+        params, return_beams=True
+    )
+
+    assert input_uv.Npols == nfeeds**2
+
+    for beam_inter in beam_list:
+        assert beam_inter.Nfeeds == nfeeds
+
     uv_out = pyuvsim.run_uvsim(params, return_uv=True)
     assert uv_out.telescope.name == "Triangle"
 
@@ -262,6 +281,13 @@ def test_run_gleam_uvsim(spectral_type):
     uv_in.conjugate_bls()
     uv_in.reorder_blts()
     uv_in.integration_time = np.full_like(uv_in.integration_time, 11.0)
+
+    assert uv_out.Npols == nfeeds**2
+
+    if nfeeds == 1:
+        assert uv_out.Npols == 1
+        uv_in.select(polarizations=select_pol)
+
     # This just tests that we get the same answer as an earlier run, not that
     # the data are correct (that's covered in other tests)
     uv_out.history = uv_in.history

tests/test_simsetup.py

@@ -3,6 +3,7 @@
 
 import copy
 import os
+import re
 import shutil
 
 import numpy as np
@@ -513,7 +514,6 @@ def test_param_reader(telparam_in_obsparam, tmpdir):
         param_dict = yaml.safe_load(fhandle)
 
     if telparam_in_obsparam:
-        print("setting up new param file")
         new_param_file = os.path.join(tmpdir, "new_obsparam.yaml")
         new_telconfig_file = os.path.join(tmpdir, "new_tel_config.yaml")
         new_layout_file = os.path.join(tmpdir, "triangle_bl_layout.csv")
@@ -656,6 +656,14 @@ def test_param_reader(telparam_in_obsparam, tmpdir):
             ValueError,
             "Unrecognized beam file or model",
         ),
+        (
+            {"polarization_array": [-1]},
+            ValueError,
+            re.escape(
+                "Specified polarization array [-1] requires feeds {'r'} but the "
+                "beams have feeds ['x' 'y']."
+            ),
+        ),
     ],
 )
 def test_param_reader_errors(subdict, error, msg):
@@ -1148,6 +1156,28 @@ def test_param_set_cat_name(key):
     assert uv_obj.phase_center_catalog[0]["cat_name"] == "foo"
 
 
+@pytest.mark.parametrize("pol_array", [[-5, -6], None])
+@pytest.mark.parametrize("set_x_orient", [True, False])
+def test_param_polarization_array(pol_array, set_x_orient):
+    param_filename = os.path.join(
+        SIM_DATA_PATH, "test_config", "obsparam_mwa_nocore.yaml"
+    )
+    param_dict = simsetup._config_str_to_dict(param_filename)
+
+    if pol_array is None:
+        pol_array = [-5, -6, -7, -8]
+    else:
+        param_dict["polarization_array"] = pol_array
+
+    if set_x_orient:
+        param_dict["telescope"]["telescope_config_name"] = (
+            "../mwa88_nocore_config_dipole.yaml"
+        )
+
+    uv_obj, _, _ = simsetup.initialize_uvdata_from_params(param_dict, return_beams=True)
+    assert uv_obj.polarization_array.tolist() == pol_array
+
+
 def test_uvdata_keyword_init(uvdata_keyword_dict):
     # check it runs through
     uvd = simsetup.initialize_uvdata_from_keywords(**uvdata_keyword_dict)
@@ -1179,15 +1209,25 @@ def test_uvdata_keyword_init_select_bls(uvdata_keyword_dict):
     assert antpairs == bls
 
 
+@pytest.mark.parametrize("pols", [["xx", "yy"], ["xx", "yy", "xy", "yx"], ["yy"], None])
+def test_uvdata_keyword_init_polarization_array(uvdata_keyword_dict, pols):
+    if pols is None:
+        del uvdata_keyword_dict["polarization_array"]
+        pols = ["xx", "yy", "xy", "yx"]
+    else:
+        uvdata_keyword_dict["polarization_array"] = pols
+    uvd = simsetup.initialize_uvdata_from_keywords(**uvdata_keyword_dict)
+
+    assert uvd.get_pols() == pols
+
+
 def test_uvdata_keyword_init_select_antnum_str(uvdata_keyword_dict):
     # check that '1' gets converted to [1]
-    uvdata_keyword_dict["polarization_array"] = ["xx", "yy"]
     uvdata_keyword_dict["no_autos"] = False
     uvdata_keyword_dict["antenna_nums"] = "1"
     uvd = simsetup.initialize_uvdata_from_keywords(**uvdata_keyword_dict)
 
     assert uvd.Nbls == 1
-    assert uvd.get_pols() == uvdata_keyword_dict["polarization_array"]
 
 
 def test_uvdata_keyword_init_time_freq_override(uvdata_keyword_dict):