Add handling and testing for reading in multiple feed_pols at once

Author: bhazelton

docs/uvbeam_tutorial.rst

@@ -166,20 +166,20 @@ c) Reading a FEKO beam file (Power & E-field)
   >>> feko_filename_y = os.path.join(DATA_PATH, feko_folder, filename_y)
 
   >>> pbeam_feko = UVBeam.from_file(
-  ...   feko_filename_x, beam_type='power', feed_pol='x', telescope_name='LWA',
-  ...   feed_name='LWA', feed_version='1', feed_angle=np.pi/2,
+  ...   [feko_filename_x, feko_filename_y], beam_type='power', feed_pol=['x', 'y'],
+  ...   feed_angle=[np.pi/2, 0.0], telescope_name='LWA',feed_name='LWA', feed_version='1',
   ...   model_name='FEKO_MROsoil_test', model_version='1.0', mount_type="fixed"
   ... )
   >>> print(np.shape(pbeam_feko.data_array))
-  (1, 1, 3, 181, 181)
+  (1, 2, 3, 181, 181)
 
   >>> ebeam_feko = UVBeam.from_file(
-  ...   feko_filename_x, beam_type='efield', feed_pol='x', telescope_name='LWA',
-  ...   feed_name='LWA', feed_version='1', feed_angle=np.pi/2,
-  ...   model_name='FEKO_MROsoil_updatedheight', model_version='1.0', mount_type="fixed"
+  ...   [feko_filename_x, feko_filename_y], beam_type='efield', feed_pol=['x', 'y'],
+  ...   feed_angle=[np.pi/2, 0.0], telescope_name='LWA',feed_name='LWA', feed_version='1',
+  ...   model_name='FEKO_MROsoil_test', model_version='1.0', mount_type="fixed"
   ... )
   >>> print(np.shape(ebeam_feko.data_array))
-  (2, 1, 3, 181, 181)
+  (2, 2, 3, 181, 181)
 
 d) Reading in the MWA full embedded element beam
 ************************************************

src/pyuvdata/uvbeam/uvbeam.py

@@ -4524,7 +4524,7 @@ def read(
         ----
         beam_type : str
             What beam_type to read in ('power' or 'efield').
-        feed_pol : str
+        feed_pol : str or list of str
             The feed polarization that the files corresponds to, e.g. x, y, r or l.
             Defaults to 'x'.
         feed_angle : str or array-like of float
@@ -4614,6 +4614,21 @@ def read(
                 file_type = "cst"
             elif extension == ".ffe":
                 file_type = "feko"
+                if multi:
+                    feed_pol = utils.tools._get_iterable(feed_pol)
+                    if len(feed_pol) != len(filename):
+                        raise ValueError(
+                            "If multiple FEKO files are passed, the feed_pol must "
+                            "be a list or array of the same length giving the "
+                            "feed_pol for each file."
+                        )
+                    feed_angle = utils.tools._get_iterable(feed_angle)
+                    if len(feed_angle) != len(filename):
+                        raise ValueError(
+                            "If multiple FEKO files are passed, the feed_angle must "
+                            "be a list or array of the same length giving the "
+                            "feed_angle for each file."
+                        )
 
         if file_type is None:
             raise ValueError(
@@ -4654,6 +4669,12 @@ def read(
                 file_num = 0
                 file_warnings = ""
                 unread = True
+                if file_type == "feko":
+                    feed_pol_use = feed_pol[file_num]
+                    feed_angle_use = feed_angle[file_num]
+                else:
+                    feed_pol_use = feed_pol
+                    feed_angle_use = feed_angle
 
                 while unread and file_num < len(filename):
                     try:
@@ -4676,9 +4697,9 @@ def read(
                             # leave these in case we restructure the multi
                             # reading later
                             beam_type=beam_type,
-                            feed_pol=feed_pol,
+                            feed_pol=feed_pol_use,
                             feed_array=feed_array,
-                            feed_angle=feed_angle,
+                            feed_angle=feed_angle_use,
                             rotate_pol=rotate_pol,
                             mount_type=mount_type,
                             frequency=frequency,
@@ -4710,6 +4731,12 @@ def read(
                 if len(filename) >= file_num + 1:
                     for fname in filename[file_num + 1 :]:
                         beam2 = UVBeam()
+                        if file_type == "feko":
+                            feed_pol_use = feed_pol[file_num + 1]
+                            feed_angle_use = feed_angle[file_num + 1]
+                        else:
+                            feed_pol_use = feed_pol
+                            feed_angle_use = feed_angle
                         try:
                             beam2.read(
                                 fname,
@@ -4730,9 +4757,9 @@ def read(
                                 # leave these in case we restructure the multi
                                 # reading later
                                 beam_type=beam_type,
-                                feed_pol=feed_pol,
+                                feed_pol=feed_pol_use,
                                 feed_array=feed_array,
-                                feed_angle=feed_angle,
+                                feed_angle=feed_angle_use,
                                 rotate_pol=rotate_pol,
                                 mount_type=mount_type,
                                 frequency=frequency,
@@ -4761,6 +4788,7 @@ def read(
                                 continue
                             else:
                                 raise
+                        file_num += 1
                 if unread is True:
                     warnings.warn(
                         "########################################################\n"

tests/uvbeam/test_feko_beam.py

@@ -36,7 +36,7 @@ def test_read_feko(btype):
         model_name="FEKO_MROsoil",
         model_version="1.0",
         mount_type="fixed",
-        feed_angle=90.0,  # E/W
+        feed_angle=np.pi / 2,  # E/W
         reference_impedance=reference_impedance,
         extra_keywords=extra_keywords,
     )
@@ -79,6 +79,26 @@ def test_read_feko(btype):
     assert len(beam_feko1.freq_array) == len(beam_feko2.freq_array)
     assert np.all(beam_feko1.bandpass_array) == 1
 
+    feko_beam_multi = beam_feko1 + beam_feko2
+
+    feko_beam_multi2 = UVBeam.from_file(
+        [feko_filename_x, feko_filename_y],
+        beam_type=btype,
+        frequency=[10e6],
+        feed_pol=["x", "y"],
+        feed_angle=[np.pi / 2, 0.0],
+        telescope_name="LWA",
+        feed_name="LWA",
+        feed_version="1",
+        model_name="FEKO_MROsoil",
+        model_version="1.0",
+        mount_type="fixed",
+        reference_impedance=reference_impedance,
+        extra_keywords=extra_keywords,
+    )
+
+    assert feko_beam_multi == feko_beam_multi2
+
 
 @pytest.mark.parametrize(
     ("feedpol", "msg"),