add testing, fix bugs in analytic beam implementation, uvbeam.new

bhazelton · bhazelton · commit ff500c2a8afd · 2025-11-25T18:11:16.000-08:00
diff --git a/src/pyuvdata/analytic_beam.py b/src/pyuvdata/analytic_beam.py
@@ -388,7 +388,11 @@ def _check_eval_inputs(
             raise ValueError("az_array and za_array must have the same shape.")
 
     def _get_empty_data_array(
-        self, grid_shape: tuple[int, int], beam_type: str = "efield"
+        self,
+        grid_shape: tuple[int, int],
+        beam_type: Literal[
+            "efield", "power", "feed_iresponse", "feed_projection"
+        ] = "efield",
     ) -> FloatArray:
         """Get the empty data to fill in the eval methods."""
         if beam_type in ["efield", "feed_projection"]:
@@ -447,9 +451,9 @@ def efield_eval(
 
             data_array = self._get_empty_data_array(az_grid.shape)
 
-            for fn in np.arange(self.Nfeeds):
-                data_array[0, fn, :, :] = np.sqrt(power_vals[fn] / 2.0)
-                data_array[1, fn, :, :] = np.sqrt(power_vals[fn] / 2.0)
+            for feed_i in np.arange(self.Nfeeds):
+                data_array[0, feed_i, :, :] = np.sqrt(power_vals[feed_i] / 2.0)
+                data_array[1, feed_i, :, :] = np.sqrt(power_vals[feed_i] / 2.0)
 
             return data_array
 
@@ -540,23 +544,35 @@ def feed_iresponse_eval(
             az_array=az_array, za_array=za_array, freq_array=freq_array
         )
 
-        za_grid, _ = np.meshgrid(za_array, freq_array)
-        az_grid, f_grid = np.meshgrid(az_array, freq_array)
-
         if hasattr(self, "_feed_iresponse_eval"):
+            za_grid, _ = np.meshgrid(za_array, freq_array)
+            az_grid, f_grid = np.meshgrid(az_array, freq_array)
+
             return self._feed_iresponse_eval(
                 az_grid=az_grid, za_grid=za_grid, f_grid=f_grid
             ).astype(complex)
         else:
-            efield_vals = self._efield_eval(
-                az_grid=az_grid, za_grid=za_grid, f_grid=f_grid
+            efield_vals = self.efield_eval(
+                za_array=za_array, az_array=az_array, freq_array=freq_array
             )
 
-            # set f to the magnitude of the I response, assume no time delays
-            data_array = np.sqrt(np.sum(np.abs(efield_vals) ** 2, axis=0))
+            unpolarized = True
+            for feed_i in np.arange(self.Nfeeds):
+                if not np.allclose(efield_vals[0, feed_i], efield_vals[1, feed_i]):
+                    unpolarized = False
+
+            if unpolarized:
+                # for unpolarized beams, where the response to all vector direction
+                # is identical, f is the sum over that direction divided by sqrt(2)
+                # this works better than what is below because it handles negatives
+                # in e.g. AiryBeams
+                data_array = np.sum(efield_vals, axis=0) / np.sqrt(2)
+            else:
+                # if polarized default f to the magnitude, assume zero time delay
+                data_array = np.sqrt(np.sum(np.abs(efield_vals) ** 2, axis=0))
             data_array = data_array[np.newaxis]
 
-            return data_array
+            return data_array.astype(complex)
 
     def feed_projection_eval(
         self, *, az_array: FloatArray, za_array: FloatArray, freq_array: FloatArray
@@ -594,8 +610,8 @@ def feed_projection_eval(
                 az_grid=az_grid, za_grid=za_grid, f_grid=f_grid
             ).astype(complex)
         else:
-            efield_vals = self._efield_eval(
-                az_grid=az_grid, za_grid=za_grid, f_grid=f_grid
+            efield_vals = self.efield_eval(
+                az_array=az_array, za_array=za_array, freq_array=freq_array
             )
 
             # set f to the magnitude of the I response, assume no time delays
@@ -614,7 +630,9 @@ def feed_projection_eval(
     def to_uvbeam(
         self,
         freq_array: FloatArray,
-        beam_type: Literal["efield", "power"] = "efield",
+        beam_type: Literal[
+            "efield", "power", "feed_iresponse", "feed_projection"
+        ] = "efield",
         pixel_coordinate_system: (
             Literal["az_za", "orthoslant_zenith", "healpix"] | None
         ) = None,
@@ -635,7 +653,7 @@ def to_uvbeam(
         freq_array : ndarray of float
             Array of frequencies in Hz to evaluate the beam at.
         beam_type : str
-            Beam type, either "efield" or "power".
+            Beam type, one of "efield", "power", "feed_iresponse" or "feed_projection".
         pixel_coordinate_system : str
             Pixel coordinate system, options are "az_za", "orthoslant_zenith" and
             "healpix". Forced to be "healpix" if ``nside`` is given and by
@@ -658,13 +676,14 @@ def to_uvbeam(
             Healpix ordering parameter, defaults to "ring" if nside is provided.
 
         """
-        if beam_type not in ["efield", "power"]:
-            raise ValueError("Beam type must be 'efield' or 'power'")
+        allowed_beam_types = ["efield", "power", "feed_iresponse", "feed_projection"]
+        if beam_type not in allowed_beam_types:
+            raise ValueError(f"Beam type must be one of {allowed_beam_types}")
 
-        if beam_type == "efield":
-            polarization_array = None
-        else:
+        if beam_type == "power":
             polarization_array = self.polarization_array
+        else:
+            polarization_array = None
 
         mount_type = self.mount_type if hasattr(self, "mount_type") else None
 
@@ -684,6 +703,7 @@ def to_uvbeam(
 
         uvb = UVBeam.new(
             telescope_name="Analytic Beam",
+            beam_type=beam_type,
             data_normalization="physical",
             feed_name=self.__repr__(),
             feed_version="1.0",
@@ -723,10 +743,7 @@ def to_uvbeam(
             az_array = az_array.flatten()
             za_array = za_array.flatten()
 
-        if beam_type == "efield":
-            eval_function = "efield_eval"
-        else:
-            eval_function = "power_eval"
+        eval_function = f"{beam_type}_eval"
 
         data_array = getattr(self, eval_function)(
             az_array=az_array, za_array=za_array, freq_array=freq_array
@@ -743,7 +760,9 @@ def to_uvbeam(
     def plot(
         self,
         *,
-        beam_type: str,
+        beam_type: Literal[
+            "efield", "power", "feed_iresponse", "feed_projection"
+        ] = "efield",
         freq: float,
         complex_type: str = "real",
         colormap: str | None = None,
@@ -761,6 +780,8 @@ def plot(
 
         Parameters
         ----------
+        beam_type : str
+            Beam type, one of "efield", "power", "feed_iresponse" or "feed_projection".
         freq : int
             The frequency index to plot.
         complex_type : str
@@ -1332,10 +1353,13 @@ def _efield_eval(
 
         # The first dimension is for [azimuth, zenith angle] in that order
         # the second dimension is for feed [e, n] in that order
-        data_array[0, self.east_ind] = -np.sin(az_grid)
-        data_array[0, self.north_ind] = np.cos(az_grid)
-        data_array[1, self.east_ind] = np.cos(za_grid) * np.cos(az_grid)
-        data_array[1, self.north_ind] = np.cos(za_grid) * np.sin(az_grid)
+        cos_za = np.cos(za_grid)
+        sin_az = np.sin(az_grid)
+        cos_az = np.cos(az_grid)
+        data_array[0, self.east_ind] = -sin_az
+        data_array[0, self.north_ind] = cos_az
+        data_array[1, self.east_ind] = cos_za * cos_az
+        data_array[1, self.north_ind] = cos_za * sin_az
 
         return data_array
 
@@ -1347,12 +1371,13 @@ def _power_eval(
 
         # these are just the sum in quadrature of the efield components.
         # some trig work is done to reduce the number of cos/sin evaluations
-        data_array[0, 0] = 1 - (np.sin(za_grid) * np.cos(az_grid)) ** 2
-        data_array[0, 1] = 1 - (np.sin(za_grid) * np.sin(az_grid)) ** 2
+        sin_za = np.sin(za_grid)
+        data_array[0, 0] = 1 - (sin_za * np.cos(az_grid)) ** 2
+        data_array[0, 1] = 1 - (sin_za * np.sin(az_grid)) ** 2
 
         if self.Npols > self.Nfeeds:
             # cross pols are included
-            data_array[0, 2] = -(np.sin(za_grid) ** 2) * np.sin(2.0 * az_grid) / 2.0
+            data_array[0, 2] = -(sin_za**2) * np.sin(2.0 * az_grid) / 2.0
             data_array[0, 3] = data_array[0, 2]
 
         return data_array
diff --git a/src/pyuvdata/beam_interface.py b/src/pyuvdata/beam_interface.py
@@ -38,7 +38,7 @@ class BeamInterface:
         Beam object to use for computations. If a BeamInterface is passed, a new
         view of the same object is created.
     beam_type : str
-        The beam type, either "efield" or "power".
+        The beam type, one of "efield", "power", "feed_iresponse" or "feed_projection".
     include_cross_pols : bool
         Option to include the cross polarized beams (e.g. xy and yx or en and ne).
         Used if beam is a UVBeam and and the input UVBeam is an Efield beam but
@@ -98,7 +98,8 @@ def __post_init__(self, include_cross_pols: bool):
                     f"Input beam is a {self.beam.beam_type} UVBeam but beam_type "
                     f"is specified as {self.beam_type}. It's not possible to "
                     f"convert a {self.beam.beam_type} beam to {self.beam_type}, "
-                    "either provide an efield UVBeam or do not specify `beam_type`."
+                    f"either provide a UVBeam with a compatible type or do not "
+                    "specify `beam_type`."
                 )
         elif self.beam_type is None:
             self.beam_type = "efield"
@@ -327,13 +328,10 @@ def compute_response(
                 az_array_use = copy.copy(az_array)
                 za_array_use = copy.copy(za_array)
 
-            if self.beam_type == "efield":
-                interp_data = self.beam.efield_eval(
-                    az_array=az_array_use, za_array=za_array_use, freq_array=freq_array
-                )
-            else:
-                interp_data = self.beam.power_eval(
-                    az_array=az_array_use, za_array=za_array_use, freq_array=freq_array
-                )
+            eval_function = f"{self.beam_type}_eval"
+
+            interp_data = getattr(self.beam, eval_function)(
+                az_array=az_array_use, za_array=za_array_use, freq_array=freq_array
+            )
 
         return interp_data
diff --git a/src/pyuvdata/utils/plotting.py b/src/pyuvdata/utils/plotting.py
@@ -410,7 +410,7 @@ def beam_plot(
         freq_title = freq
 
         naxes_vec = beam_obj.Naxes_vec
-        if beam_type == "power":
+        if beam_type in ["power", "feed_iresponse"]:
             naxes_vec = 1
         reg_grid = True
 
diff --git a/src/pyuvdata/uvbeam/initializers.py b/src/pyuvdata/uvbeam/initializers.py
@@ -24,6 +24,7 @@ def new_uvbeam(
     feed_version: str = "0.0",
     model_name: str = "default",
     model_version: str = "0.0",
+    beam_type: Literal["efield", "power", "feed_iresponse", "feed_projection"] = None,
     feed_array: StrArray | None = None,
     feed_angle: FloatArray | None = None,
     mount_type: str | None = "fixed",
@@ -78,6 +79,9 @@ def new_uvbeam(
         Name of the beam model.
     model_version: str
         Version of the beam model.
+    beam_type : str
+        Beam type, one of "efield", "power", "feed_iresponse" or "feed_projection".
+        Defaults to "power" if polarization_array is provided and "efield" otherwise.
     feed_array : ndarray of str
         Array of feed orientations. Options are: n/e or x/y or r/l.
     feed_angle : ndarray of float
@@ -177,18 +181,23 @@ def new_uvbeam(
 
     uvb = UVBeam()
 
-    if polarization_array is None:
-        uvb.beam_type = "efield"
-        uvb._set_efield()
-    else:
-        uvb.beam_type = "power"
+    if beam_type is None:
+        if polarization_array is None:
+            beam_type = "efield"
+        else:
+            beam_type = "power"
+
+    if beam_type == "power":
         polarization_array = np.asarray(polarization_array)
         if polarization_array.dtype.kind != "i":
             polarization_array = np.asarray(utils.polstr2num(polarization_array))
         uvb.polarization_array = polarization_array
 
         uvb.Npols = uvb.polarization_array.size
         uvb._set_power()
+    else:
+        set_function = f"_set_{beam_type}"
+        getattr(uvb, set_function)()
 
     if feed_array is not None:
         uvb.feed_array = np.asarray(feed_array)
@@ -259,7 +268,7 @@ def new_uvbeam(
             "Either nside or both axis1_array and axis2_array must be provided."
         )
 
-    if uvb.beam_type == "power":
+    if uvb.beam_type in ["power", "feed_iresponse"]:
         uvb.Naxes_vec = 1
 
     uvb._set_cs_params()
@@ -296,7 +305,7 @@ def new_uvbeam(
                 f"expected shape {bv_shape}."
             )
         uvb.basis_vector_array = basis_vector_array
-    elif uvb.beam_type == "efield":
+    elif uvb.beam_type in ["efield", "feed_projection"]:
         if uvb.pixel_coordinate_system == "healpix":
             basis_vector_array = np.zeros(
                 (uvb.Naxes_vec, uvb.Ncomponents_vec, uvb.Npixels), dtype=float
@@ -384,15 +393,15 @@ def new_uvbeam(
         uvb._set_simple()
 
     # Set data parameters
-    if uvb.beam_type == "efield":
-        data_type = complex
-        polax = uvb.Nfeeds
-    else:
+    if uvb.beam_type == "power":
         data_type = float
         for pol in uvb.polarization_array:
             if pol in [-3, -4, -7, -8]:
                 data_type = complex
         polax = uvb.Npols
+    else:
+        data_type = complex
+        polax = uvb.Nfeeds
 
     if uvb.pixel_coordinate_system == "healpix":
         pixax = (uvb.Npixels,)
diff --git a/tests/conftest.py b/tests/conftest.py
@@ -171,6 +171,14 @@ def az_za_coords():
     return az_array, za_array
 
 
+@pytest.fixture()
+def az_za_coords_fine():
+    az_array = np.deg2rad(np.linspace(0, 359, 360))
+    za_array = np.deg2rad(np.linspace(0, 90, 91))
+
+    return az_array, za_array
+
+
 @pytest.fixture()
 def az_za_deg_grid(az_za_coords):
     az_array, za_array = az_za_coords
diff --git a/tests/test_analytic_beam.py b/tests/test_analytic_beam.py
@@ -276,9 +276,8 @@ def _feed_iresponse_eval(
                 az_grid.shape, beam_type="feed_iresponse"
             )
 
-            # set f to the magnitude of the I response, assume no time delays
             for feed_i in range(self.Nfeeds):
-                data_array[0, feed_i] = np.abs(np.cos(self.width * za_grid))
+                data_array[0, feed_i] = np.cos(self.width * za_grid)
 
             return data_array
 
@@ -290,11 +289,6 @@ def _feed_projection_eval(
             )
 
             data_array.fill(1 / np.sqrt(2))
-            # find where it goes negative
-            wh_neg = np.nonzero(np.cos(self.width * za_grid) < 0)
-            for bv_i in range(self.Naxes_vec):
-                for feed_i in range(self.Nfeeds):
-                    data_array[bv_i, feed_i, *wh_neg] *= -1
 
             return data_array
 
@@ -569,7 +563,13 @@ def test_beam_equality(beam1, beam2, equal):
 def test_to_uvbeam_errors():
     beam = GaussianBeam(sigma=0.02)
 
-    with pytest.raises(ValueError, match="Beam type must be 'efield' or 'power'"):
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Beam type must be one of ['efield', 'power', 'feed_iresponse', "
+            "'feed_projection']"
+        ),
+    ):
         beam.to_uvbeam(
             freq_array=np.linspace(100, 200, 5),
             beam_type="foo",
@@ -790,6 +790,8 @@ def test_set_x_orientation_deprecation():
             None,
         ),
         (ShortDipoleBeam(), "efield", "real", None, 90.0, None, None),
+        (ShortDipoleBeam(), "feed_iresponse", "phase", None, 90.0, None, None),
+        (ShortDipoleBeam(), "feed_projection", "real", None, 90.0, None, None),
         (ShortDipoleBeam(), "power", "real", None, 90.0, None, None),
         (UniformBeam(), "power", "real", None, 90.0, None, None),
     ],
diff --git a/tests/test_beam_interface.py b/tests/test_beam_interface.py
diff --git a/tests/uvbeam/test_mwa_beam.py b/tests/uvbeam/test_mwa_beam.py
