better healpix lat/lon to az/za handling

more test coverage

pyuvdata/tests/test_utils.py

@@ -604,6 +604,38 @@ def test_mwa_ecef_conversion():
     assert np.allclose(rot_xyz.T, xyz)
 
 
+def test_hpx_latlon_az_za():
+    zenith_angle = np.deg2rad(np.linspace(0, 90, 10))
+    azimuth = np.deg2rad(np.linspace(0, 360, 36, endpoint=False))
+    az_mesh, za_mesh = np.meshgrid(azimuth, zenith_angle)
+
+    hpx_lat = np.deg2rad(np.linspace(90, 0, 10))
+    hpx_lon = np.deg2rad(np.linspace(0, 360, 36, endpoint=False))
+    lon_mesh, lat_mesh = np.meshgrid(hpx_lon, hpx_lat)
+
+    with pytest.raises(
+        ValueError, match="shapes of zenith_angle and azimuth values must match."
+    ):
+        uvutils.zenithangle_azimuth_to_hpx_latlon(zenith_angle, azimuth)
+
+    calc_lat, calc_lon = uvutils.zenithangle_azimuth_to_hpx_latlon(za_mesh, az_mesh)
+
+    print(np.min(calc_lat), np.max(calc_lat))
+    print(np.min(lat_mesh), np.max(lat_mesh))
+    np.testing.assert_allclose(calc_lat, lat_mesh)
+    np.testing.assert_allclose(calc_lon, lon_mesh)
+
+    with pytest.raises(
+        ValueError, match="shapes of hpx_lat and hpx_lon values must match."
+    ):
+        uvutils.hpx_latlon_to_zenithangle_azimuth(hpx_lat, hpx_lon)
+
+    calc_za, calc_az = uvutils.hpx_latlon_to_zenithangle_azimuth(lat_mesh, lon_mesh)
+
+    np.testing.assert_allclose(calc_za, za_mesh)
+    np.testing.assert_allclose(calc_az, az_mesh)
+
+
 @pytest.mark.parametrize(
     "input1,input2,msg",
     (

pyuvdata/utils.py

@@ -1629,6 +1629,90 @@ def ECEF_from_ENU(enu, latitude, longitude, altitude, frame="ITRS"):
     return xyz
 
 
+def hpx_latlon_to_zenithangle_azimuth(hpx_lat, hpx_lon):
+    """
+    Convert from healpix lat/lon to UVBeam za/az convention.
+
+    Note that this is different (unfortunately) from the conversion between
+    the UVBeam Zenith Angle, Azimuth coordinate system and the astropy Alt/Az
+    coordinate system. The astropy Azimuth runs the opposite direction and
+    has a different starting angle than UVBeam's Azimuth because they are both
+    right handed coordinate systems but Altitude moves the opposite direction
+    than Zenith Angle does.
+
+    The conversion used in this code sets the Healpix latitude to 90-zenith angle
+    but it does not change the origin or direction for the azimuthal angle. This
+    convention was set early in the development of UVBeam and we preserve it to
+    preserve backwards compatibility.
+
+    Parameters
+    ----------
+    hpx_lat: float or array of float
+        Healpix latiudinal coordinate in radians.
+    hpx_lon: float or array of float
+        Healpix longitudinal coordinate in radians.
+
+    Returns
+    -------
+    zenith_angle: float or array of float
+        In radians
+    azimuth: float or array of float
+        In radians in uvbeam convention: North of East(East=0, North=pi/2)
+
+    """
+    input_alt = np.asarray(hpx_lat)
+    input_az = np.asarray(hpx_lon)
+    if input_alt.shape != input_az.shape:
+        raise ValueError("shapes of hpx_lat and hpx_lon values must match.")
+
+    zenith_angle = np.pi / 2 - hpx_lat
+    azimuth = hpx_lon
+
+    return zenith_angle, azimuth
+
+
+def zenithangle_azimuth_to_hpx_latlon(zenith_angle, azimuth):
+    """
+    Convert from UVBeam az/za convention to healpix lat/lon.
+
+    Note that this is different (unfortunately) from the conversion between
+    the UVBeam Zenith Angle, Azimuth coordinate system and the astropy Alt/Az
+    coordinate system. The astropy Azimuth runs the opposite direction and
+    has a different starting angle than UVBeam's Azimuth because they are both
+    right handed coordinate systems but Altitude moves the opposite direction
+    than Zenith Angle does.
+
+    The conversion used in this code sets the Healpix latitude to 90-zenith angle
+    but it does not change the origin or direction for the azimuthal angle. This
+    convention was set early in the development of UVBeam and we preserve it to
+    preserve backwards compatibility.
+
+    Parameters
+    ----------
+    zenith_angle: float, array_like of float
+        Zenith angle in radians
+    azimuth: float, array_like of float
+        Azimuth in radians in uvbeam convention: North of East(East=0, North=pi/2)
+
+    Returns
+    -------
+    hpx_lat: float or array of float
+        Healpix latiudinal coordinate in radians.
+    hpx_lon: float or array of float
+        Healpix longitudinal coordinate in radians.
+
+    """
+    input_za = np.array(zenith_angle)
+    input_az = np.array(azimuth)
+    if input_za.shape != input_az.shape:
+        raise ValueError("shapes of zenith_angle and azimuth values must match.")
+
+    lat_array = np.pi / 2 - zenith_angle
+    lon_array = azimuth
+
+    return lat_array, lon_array
+
+
 def old_uvw_calc(ra, dec, initial_uvw):
     """
     Calculate old uvws from unphased ones in an icrs or gcrs frame.

pyuvdata/uvbeam/analytic_beam.py

@@ -13,6 +13,7 @@
 from astropy.constants import c as speed_of_light
 from scipy.special import j1
 
+from .. import utils as uvutils
 from ..docstrings import combine_docstrings
 from .uvbeam import UVBeam, _convert_feeds_to_pols
 
@@ -60,13 +61,11 @@ class AnalyticBeam(ABC):
     @abstractmethod
     def basis_vector_type(self):
         """Require that a basis_vector_type is defined in concrete classes."""
-        pass
 
     @property
     @abstractmethod
     def name(self):
         """Require that a name is defined in concrete classes."""
-        pass
 
     def __init__(
         self,
@@ -347,7 +346,7 @@ def power_eval(
     def to_uvbeam(
         self,
         freq_array: npt.NDArray[np.float],
-        beam_type: Literal["efield", "power"] = "power",
+        beam_type: Literal["efield", "power"] = "efield",
         pixel_coordinate_system: Literal["az_za", "orthoslant_zenith", "healpix"]
         | None = None,
         **kwargs,
@@ -381,14 +380,19 @@ def to_uvbeam(
             feed_array = None
             polarization_array = self.polarization_array
 
-        if pixel_coordinate_system is not None and pixel_coordinate_system not in [
-            "az_za",
-            "healpix",
-        ]:
-            raise NotImplementedError(
-                "Currently this method only supports 'az_za' and 'healpix' "
-                "pixel_coordinate_systems."
-            )
+        if pixel_coordinate_system is not None:
+            allowed_coord_sys = list(UVBeam().coordinate_system_dict.keys())
+            if pixel_coordinate_system not in allowed_coord_sys:
+                raise ValueError(
+                    f"Unknown coordinate system {pixel_coordinate_system}. UVBeam "
+                    f"supported coordinate systems are: {allowed_coord_sys}."
+                )
+
+            if pixel_coordinate_system not in ["az_za", "healpix"]:
+                raise NotImplementedError(
+                    "Currently this method only supports 'az_za' and 'healpix' "
+                    "pixel_coordinate_systems."
+                )
 
         uvb = UVBeam.new(
             telescope_name="Analytic Beam",
@@ -413,8 +417,12 @@ def to_uvbeam(
                     "required for healpix functionality. "
                     "Install 'astropy-healpix' using conda or pip."
                 ) from e
-            hp_obj = HEALPix(nside=uvb.nside, ordering=uvb.ordering)
-            az_array, za_array = hp_obj.healpix_to_lonlat(uvb.pixel_array)
+            hp_obj = HEALPix(nside=uvb.nside, order=uvb.ordering)
+            hpx_lon, hpx_lat = hp_obj.healpix_to_lonlat(uvb.pixel_array)
+            za_array, az_array = uvutils.hpx_latlon_to_zenithangle_azimuth(
+                hpx_lat.radian, hpx_lon.radian
+            )
+
         else:
             az_array, za_array = np.meshgrid(uvb.axis1_array, uvb.axis2_array)
             az_array = az_array.flatten()

pyuvdata/uvbeam/beam_interface.py

@@ -77,7 +77,7 @@ def compute_response(
         self,
         az_array: npt.NDArray[np.float],
         za_array: npt.NDArray[np.float],
-        freq_array: npt.NDArray[np.float],
+        freq_array: npt.NDArray[np.float] | None,
         az_za_grid: bool = False,
         freq_interp_tol: float = 1.0,
         reuse_spline: bool = False,
@@ -89,18 +89,20 @@ def compute_response(
         Parameters
         ----------
         az_array : array_like of floats, optional
-            Azimuth values to interpolate to in radians, either specifying the
-            azimuth positions for every interpolation point or specifying the
-            azimuth vector for a meshgrid if az_za_grid is True.
+            Azimuth values to compute the response for in radians, either
+            specifying the azimuth positions for every interpolation point or
+            specifying the azimuth vector for a meshgrid if az_za_grid is True.
         za_array : array_like of floats, optional
-            Zenith values to interpolate to in radians, either specifying the
-            zenith positions for every interpolation point or specifying the
-            zenith vector for a meshgrid if az_za_grid is True.
+            Zenith values to compute the response for in radians, either
+            specifying the zenith positions for every interpolation point or
+            specifying the zenith vector for a meshgrid if az_za_grid is True.
+        freq_array : array_like of floats or None
+            Frequency values to compute the response for in Hz. If beam is a UVBeam
+            this can be set to None to get the responses at the UVBeam frequencies.
+            It must be a numpy array if beam is an analytic beam.
         az_za_grid : bool
             Option to treat the `az_array` and `za_array` as the input vectors
             for points on a mesh grid.
-        freq_array : array_like of floats, optional
-            Frequency values to interpolate to.
         freq_interp_tol : float
             Frequency distance tolerance [Hz] of nearest neighbors.
             If *all* elements in freq_array have nearest neighbor distances within
@@ -120,6 +122,11 @@ def compute_response(
             An array of computed values, shape (Naxes_vec, Nfeeds or Npols,
             freq_array.size, az_array.size)
         """
+        if not isinstance(az_array, np.ndarray) or az_array.ndim != 1:
+            raise ValueError("az_array must be a one-dimensional numpy array")
+        if not isinstance(za_array, np.ndarray) or za_array.ndim != 1:
+            raise ValueError("za_array must be a one-dimensional numpy array")
+
         if self._isuvbeam:
             interp_data, _ = self.beam.interp(
                 az_array=az_array,
@@ -133,12 +140,9 @@ def compute_response(
             if not self.beam.future_array_shapes:
                 interp_data = interp_data[:, 0]
         else:
+            if not isinstance(freq_array, np.ndarray) or freq_array.ndim != 1:
+                raise ValueError("freq_array must be a one-dimensional numpy array")
             if az_za_grid:
-                if az_array is None or za_array is None:
-                    raise ValueError(
-                        "If az_za_grid is set to True, az_array and za_array must be "
-                        "provided."
-                    )
                 az_array_use, za_array_use = np.meshgrid(az_array, za_array)
                 az_array_use = az_array_use.flatten()
                 za_array_use = za_array_use.flatten()

pyuvdata/uvbeam/tests/test_analytic_beam.py

@@ -378,3 +378,44 @@ def _power_eval(
         match=re.escape("basis_vector_type is healpix, must be one of ['az_za']"),
     ):
         BadBeam()
+
+
+def test_to_uvbeam_errors():
+    beam = GaussianBeam(sigma=0.02)
+
+    with pytest.raises(ValueError, match="Beam type must be 'efield' or 'power'"):
+        beam.to_uvbeam(
+            freq_array=np.linspace(100, 200, 5),
+            beam_type="foo",
+            axis1_array=np.deg2rad(np.linspace(0, 360, 36, endpoint=False)),
+            axis2_array=np.deg2rad(np.linspace(0, 90, 10)),
+        )
+
+    with pytest.raises(
+        NotImplementedError,
+        match="Currently this method only supports 'az_za' and 'healpix' "
+        "pixel_coordinate_systems.",
+    ):
+        beam.to_uvbeam(
+            freq_array=np.linspace(100, 200, 5),
+            beam_type="efield",
+            axis1_array=np.deg2rad(np.linspace(0, 360, 36, endpoint=False)),
+            axis2_array=np.deg2rad(np.linspace(0, 90, 10)),
+            pixel_coordinate_system="orthoslant_zenith",
+        )
+
+    allowed_coord_sys = list(UVBeam().coordinate_system_dict.keys())
+    with pytest.raises(
+        ValueError,
+        match=re.escape(
+            "Unknown coordinate system foo. UVBeam supported coordinate systems "
+            f"are: {allowed_coord_sys}."
+        ),
+    ):
+        beam.to_uvbeam(
+            freq_array=np.linspace(100, 200, 5),
+            beam_type="efield",
+            axis1_array=np.deg2rad(np.linspace(0, 360, 36, endpoint=False)),
+            axis2_array=np.deg2rad(np.linspace(0, 90, 10)),
+            pixel_coordinate_system="foo",
+        )