# Copyright (c) 2024 Radio Astronomy Software Group
# Licensed under the 2-clause BSD License
import copy
import numpy as np
import pytest
from pyuvdata import (
AiryBeam,
BeamInterface,
GaussianBeam,
ShortDipoleBeam,
UniformBeam,
UVBeam,
utils,
)
from pyuvdata.testing import check_warnings
@pytest.fixture(scope="function")
def airy() -> AiryBeam:
return AiryBeam(diameter=14.0)
@pytest.fixture()
def gaussian() -> GaussianBeam:
return GaussianBeam(diameter=14.0)
@pytest.fixture()
def xy_grid_coarse():
nfreqs = 5
freqs = np.linspace(100e6, 130e6, nfreqs)
xy_half_n = 6
zmax = np.radians(60) # Degrees
arr = np.arange(-xy_half_n, xy_half_n)
x_arr, y_arr = np.meshgrid(arr, arr)
x_arr = x_arr.flatten()
y_arr = y_arr.flatten()
radius = np.sqrt(x_arr**2 + y_arr**2) / float(xy_half_n)
za_array = radius * zmax
az_array = np.arctan2(y_arr, x_arr) + np.pi # convert from -180->180 to 0->360
return az_array, za_array, freqs
@pytest.fixture()
def gaussian_uv(gaussian, az_za_coords) -> UVBeam:
az, za = az_za_coords
return gaussian.to_uvbeam(
axis1_array=az, axis2_array=za, freq_array=np.array([1e8])
)
@pytest.mark.parametrize(
["beam_obj", "kwargs"],
[
[AiryBeam, {"diameter": 14.0}],
[GaussianBeam, {"diameter": 14.0}],
[UniformBeam, {"include_cross_pols": False}],
[ShortDipoleBeam, {}],
],
)
@pytest.mark.parametrize("init_beam_type", ["efield", "power"])
@pytest.mark.parametrize("final_beam_type", ["efield", "power"])
@pytest.mark.parametrize("coord_sys", ["az_za", "healpix"])
def test_beam_interface(
beam_obj,
kwargs,
init_beam_type,
final_beam_type,
az_za_coords,
xy_grid_coarse,
coord_sys,
):
analytic = beam_obj(**kwargs)
nfreqs = 20
freq_array = np.linspace(100e6, 150e6, nfreqs)
if coord_sys == "healpix":
nside = 64
if init_beam_type == "efield":
ordering = "ring"
else:
ordering = "nested"
healpix_pixel_array = np.arange(12 * nside**2, dtype=int)
to_uvbeam_kwargs = {
"nside": nside,
"ordering": ordering,
"healpix_pixel_array": healpix_pixel_array,
}
try:
from astropy_healpix import HEALPix
except ImportError:
with pytest.raises(
ImportError,
match="astropy_healpix is not installed but is "
"required for healpix functionality. ",
):
uvb = analytic.to_uvbeam(
beam_type=init_beam_type, freq_array=freq_array, **to_uvbeam_kwargs
)
pytest.importorskip("astropy_healpix")
hp_obj = HEALPix(nside=nside, order=ordering)
hpx_lon, hpx_lat = hp_obj.healpix_to_lonlat(healpix_pixel_array)
za_array, az_array = utils.coordinates.hpx_latlon_to_zenithangle_azimuth(
hpx_lat.radian, hpx_lon.radian
)
# downselect places to check
above_horizon = np.nonzero(za_array <= (np.pi / 2.0))
az_array = az_array[above_horizon]
za_array = za_array[above_horizon]
az_array = az_array[::10]
za_array = za_array[::10]
else:
az_array, za_array = az_za_coords
to_uvbeam_kwargs = {"axis1_array": az_array, "axis2_array": za_array}
include_cross_pols = kwargs.get("include_cross_pols", True)
to_uvbeam_kwargs["pixel_coordinate_system"] = coord_sys
uvb = analytic.to_uvbeam(
beam_type=init_beam_type, freq_array=freq_array, **to_uvbeam_kwargs
)
bi_analytic = BeamInterface(analytic, final_beam_type)
if final_beam_type != init_beam_type:
if final_beam_type == "efield":
with pytest.raises(
ValueError,
match="Input beam is a power UVBeam but beam_type is specified as "
"'efield'. It's not possible to convert a power beam to an "
"efield beam, either provide an efield UVBeam or do not "
"specify `beam_type`.",
):
BeamInterface(uvb, final_beam_type)
return
warn_type = UserWarning
msg = (
"Input beam is an efield UVBeam but beam_type is specified as "
"'power'. Converting efield beam to power."
)
else:
warn_type = None
msg = ""
with check_warnings(warn_type, match=msg):
bi_uvbeam = BeamInterface(
uvb, final_beam_type, include_cross_pols=include_cross_pols
)
if coord_sys == "az_za":
az_za_grid = True
else:
az_za_grid = False
analytic_data = bi_analytic.compute_response(
az_array=az_array,
za_array=za_array,
freq_array=freq_array,
az_za_grid=az_za_grid,
)
uvb_data = bi_uvbeam.compute_response(
az_array=az_array,
za_array=za_array,
freq_array=freq_array,
az_za_grid=az_za_grid,
)
np.testing.assert_allclose(analytic_data, uvb_data, rtol=0, atol=1e-14)
# now on a grid that is not the same as where the beam was evaluated
# larger differences of course
az_vals, za_vals, freqs = xy_grid_coarse
analytic_data = bi_analytic.compute_response(
az_array=az_vals, za_array=za_vals, freq_array=freqs, az_za_grid=True
)
uvb_data = bi_uvbeam.compute_response(
az_array=az_vals, za_array=za_vals, freq_array=freqs, az_za_grid=True
)
if not (coord_sys == "healpix" and isinstance(analytic, ShortDipoleBeam)):
np.testing.assert_allclose(analytic_data, uvb_data, rtol=0, atol=1e-1)
else:
# the comparison falls apart at zenith because there's no healpix
# pixel right at zenith and the dipole beam changes quickly there.
az_mesh, za_mesh = np.meshgrid(az_vals, za_vals)
az_mesh = az_mesh.flatten()
za_mesh = za_mesh.flatten()
wh_not_zenith = np.nonzero(za_mesh != 0)
np.testing.assert_allclose(
analytic_data[:, :, :, wh_not_zenith],
uvb_data[:, :, :, wh_not_zenith],
rtol=0,
atol=1e-1,
)
def test_beam_interface_errors():
with pytest.raises(
ValueError, match="beam must be a UVBeam or an AnalyticBeam instance."
):
BeamInterface("foo", "power")
@pytest.mark.parametrize(
["param", "value"],
[
["az_array", None],
["za_array", None],
["freq_array", None],
["az_array", np.zeros((10, 10), dtype=float)],
["za_array", np.zeros((10, 10), dtype=float)],
["freq_array", np.zeros((10, 10), dtype=float)],
],
)
def test_compute_response_errors(param, value):
orig_kwargs = {
"az_array": np.deg2rad(np.linspace(0, 360, 36, endpoint=False)),
"za_array": np.deg2rad(np.linspace(0, 90, 10)),
"freq_array": np.deg2rad(np.linspace(100, 200, 5)),
}
compute_kwargs = copy.deepcopy(orig_kwargs)
compute_kwargs["az_za_grid"] = True
compute_kwargs[param] = value
analytic = ShortDipoleBeam()
bi_analytic = BeamInterface(analytic, beam_type="efield")
with pytest.raises(
ValueError, match=f"{param} must be a one-dimensional numpy array"
):
bi_analytic.compute_response(**compute_kwargs)
uvb = analytic.to_uvbeam(
beam_type="power",
freq_array=orig_kwargs["freq_array"],
axis1_array=orig_kwargs["az_array"],
axis2_array=orig_kwargs["za_array"],
)
bi_uvb = BeamInterface(uvb)
if param != "freq_array":
with pytest.raises(
ValueError, match=f"{param} must be a one-dimensional numpy array"
):
bi_uvb.compute_response(**compute_kwargs)
elif value is not None:
with pytest.raises(ValueError, match="freq_array must be one-dimensional"):
bi_uvb.compute_response(**compute_kwargs)
else:
# this shouldn't error
bi_uvb.compute_response(**compute_kwargs)
@pytest.mark.parametrize("beam_obj", ["airy", "gaussian", "gaussian_uv"])
def test_idempotent_instantiation(beam_obj, request):
beam = BeamInterface(request.getfixturevalue(beam_obj))
beam2 = BeamInterface(beam)
assert beam == beam2
def test_properties(airy: AiryBeam):
intf = BeamInterface(airy)
assert airy.Npols == intf.Npols
assert airy.Nfeeds == intf.Nfeeds
assert np.all(airy.polarization_array == intf.polarization_array)
assert np.all(airy.feed_array == intf.feed_array)
assert np.all(airy.feed_angle == intf.feed_angle)
def test_clone(airy):
intf = BeamInterface(airy)
intf_clone = intf.clone(beam_type="power")
assert intf != intf_clone
@pytest.mark.parametrize("uvbeam", [True, False], ids=["uvbeam", "analytic"])
@pytest.mark.parametrize("allow_mutation", [True, False], ids=["mutate", "nomutate"])
@pytest.mark.parametrize(
"include_cross_pols", [True, False, None], ids=["incx", "nox", "xpolnone"]
)
def test_as_power(
uvbeam: bool, allow_mutation: bool, include_cross_pols: bool, gaussian, gaussian_uv
):
beam = gaussian_uv if uvbeam else gaussian
intf = BeamInterface(beam)
intf_power = intf.as_power_beam(
allow_beam_mutation=allow_mutation, include_cross_pols=include_cross_pols
)
if include_cross_pols is None:
include_cross_pols = True
assert intf_power.beam_type == "power"
assert intf_power.Npols == 4 if include_cross_pols else 2
if uvbeam:
if allow_mutation:
assert intf.beam.beam_type == "power"
else:
assert intf.beam.beam_type == "efield"
def test_as_power_noop(airy):
"""Ensure that calling as_power_beam on a power beam is a no-op."""
intf = BeamInterface(airy, beam_type="power")
intf2 = intf.as_power_beam()
assert intf is intf2
with pytest.warns(UserWarning, match="as_power_beam does not modify cross pols"):
intf2 = intf.as_power_beam(include_cross_pols=False)
assert intf is intf2
@pytest.mark.parametrize("uvbeam", [True, False])
def test_with_feeds(uvbeam: bool, gaussian, gaussian_uv):
beam = gaussian_uv if uvbeam else gaussian
intf = BeamInterface(beam)
intf_feedx = intf.with_feeds(["x"])
assert intf_feedx.feed_array == ["x"]
def test_with_feeds_ordering(airy):
intf = BeamInterface(airy)
intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
assert np.all(intf_feedx.feed_array == ["x", "y"])
intf_feedyx = intf.with_feeds(["y", "x"], maintain_ordering=False)
assert np.all(intf_feedyx.feed_array == ["y", "x"])
@pytest.mark.filterwarnings("ignore:Input beam is an efield UVBeam")
@pytest.mark.filterwarnings("ignore:Selected polarizations are not evenly spaced")
def test_with_feeds_ordering_power(gaussian_uv):
# beam = AiryBeam(diameter=14.0).to_uvbeam(freq_array=np.array([1e8]), nside=16)
intf = BeamInterface(gaussian_uv, beam_type="power")
intf_feedx = intf.with_feeds(["y", "x"], maintain_ordering=True)
assert np.all(intf_feedx.polarization_array == [-5, -6, -7, -8])
intf_feedyx = intf.with_feeds(["y", "x"], maintain_ordering=False)
# N.b. (Karto), this used to check against [-6, -8, -7, -5], but I _think_ this
# was actually a bug, in that UVBeam.select was sensitive to the ordering of
# pol arguments for polarization_array *only*, and not anything else with a
# polarization axis.
assert np.all(intf_feedyx.polarization_array == [-5, -6, -7, -8])
