Add tests, update test data registry info for new test data

Author: bhazelton

src/pyuvdata/data/test_data.yaml

@@ -43,6 +43,8 @@ mwa_cotter_phase_test2: visibility_data/MWA/1133866760_rephase.uvfits
 mwa_fhd: visibility_data/MWA/fhd_vis_data.tar.gz
 mwa_fhd_cal: calibration_solutions/MWA/fhd_cal_data.tar.gz
 mwa_full_EE: primary_beams/MWA/mwa_full_EE_test.h5
+mwa_jmatrix: primary_beams/MWA/JMatrix_3freq.fits
+mwa_zmatrix: primary_beams/MWA/ZMatrix_3freq.fits
 mwax_2021_metafits: visibility_data/MWA/mwa_corr_fits_testfiles/1320409688.metafits
 mwax_2021_raw_gpubox: visibility_data/MWA/mwa_corr_fits_testfiles/1320409688_20211108122750_mini_ch137_000.fits
 ovro_lwa: visibility_data/LWA/2018-03-21-01_26_33_0004384620257280_000000_downselected.ms.tar.gz

src/pyuvdata/data/test_data_registry.txt

@@ -68,4 +68,6 @@ primary_beams/HERA/HERABEAM.FITS a28232d8f63af528afae7ca529856809d478db86b3b3b7d
 primary_beams/LWA/OVRO_LWA_x.ffe c69e5ac1afa638d2a80663ded0bf4d69c1273a55bc16b52e15e3d96209a3087e
 primary_beams/LWA/OVRO_LWA_y.ffe 3e2f165ac82c7f422815c7b9911ae61dc7ed43dff00556ffd59b58248a93119a
 primary_beams/MWA/mwa_full_EE_test.h5 9246b9e98dc99d4284e7865ebae221b43aaef64893830cad00587d265eef2bfe
+primary_beams/MWA/JMatrix_3freq.fits c1a0100a3c0a6999d32c393f0712cb9d86f6f4340376d4cfea9eba120c94b4c0
+primary_beams/MWA/ZMatrix_3freq.fits fe79c5601cab0d4cd82260860ad98f66e0dcf7098b51c5b51108d344a76654a6
 flagsets/HERA/zen.2457698.40355.xx.HH.uvcAA.testuvflag.h5 20ff14975bcedf5edd447e3c5db3456d9298846d7e5da5bb084c7b0662327875

src/pyuvdata/datasets.py

@@ -12,7 +12,7 @@
 from pyuvdata.data import DATA_PATH
 
 # set the data version. This should be updated when the test data are changed.
-data_version = "v0.0.4"
+data_version = "v0.0.5"
 
 pup = pooch.create(
     # Use the default cache folder for the operating system

tests/uvbeam/conftest.py

@@ -334,7 +334,7 @@ def phased_array_beam_1freq(phased_array_beam_2freq):
 
 
 @pytest.fixture(scope="module")
-def mwa_beam_1ppd_main():
+def mwa_fee_1ppd_main():
     beam = UVBeam()
     beam.read_mwa_beam(fetch_data("mwa_full_EE"), pixels_per_deg=1)
 
@@ -343,8 +343,32 @@ def mwa_beam_1ppd_main():
 
 
 @pytest.fixture(scope="function")
-def mwa_beam_1ppd(mwa_beam_1ppd_main):
-    beam = mwa_beam_1ppd_main.copy()
+def mwa_fee_1ppd(mwa_fee_1ppd_main):
+    beam = mwa_fee_1ppd_main.copy()
+
+    yield beam
+    del beam
+
+
+@pytest.fixture(scope="module")
+def mwa_aee_files():
+    return {"jfile": fetch_data("mwa_jmatrix"), "zfile": fetch_data("mwa_zmatrix")}
+
+
+@pytest.fixture(scope="module")
+def mwa_aee_main(mwa_aee_files):
+    beam = UVBeam()
+    filename = mwa_aee_files["jfile"]
+    zfile = mwa_aee_files["zfile"]
+    beam.read_mwa_beam(filename, zfile=zfile)
+
+    yield beam
+    del beam
+
+
+@pytest.fixture(scope="function")
+def mwa_aee(mwa_aee_main):
+    beam = mwa_aee_main.copy()
 
     yield beam
     del beam

tests/uvbeam/test_mwa_beam.py

@@ -1,5 +1,6 @@
 # Copyright (c) 2019 Radio Astronomy Software Group
 # Licensed under the 2-clause BSD License
+import copy
 
 import numpy as np
 import pytest
@@ -10,27 +11,52 @@
 from pyuvdata.uvbeam.mwa_beam import P1sin, P1sin_array
 
 
-def test_read_write_mwa(mwa_beam_1ppd, tmp_path):
+@pytest.fixture()
+def mwabeam_kwargs(mwa_aee_files):
+    return {"fee": {"pixels_per_deg": 1}, "aee": {"zfile": mwa_aee_files["zfile"]}}
+
+
+@pytest.fixture()
+def uvbeam_kwargs(mwabeam_kwargs):
+    uvbeam_kwargs = copy.deepcopy(mwabeam_kwargs)
+    uvbeam_kwargs["aee"]["mwa_zfile"] = uvbeam_kwargs["aee"].pop("zfile")
+    return uvbeam_kwargs
+
+
+@pytest.fixture()
+def beam_filenames(mwa_aee_files):
+    return {"fee": fetch_data("mwa_full_EE"), "aee": mwa_aee_files["jfile"]}
+
+
+@pytest.mark.parametrize("model", ["fee", "aee"])
+def test_read_write_mwa(beam_filenames, mwabeam_kwargs, model, tmp_path):
     """Basic read/write test."""
-    beam1 = mwa_beam_1ppd
-    beam2 = UVBeam()
 
-    beam1.read_mwa_beam(fetch_data("mwa_full_EE"), pixels_per_deg=1)
-    assert beam1.filename == ["mwa_full_EE_test.h5"]
+    filename = beam_filenames[model]
+    kwargs = mwabeam_kwargs[model]
+
+    beam1 = UVBeam()
+    beam2 = UVBeam()
+    beam1.read_mwa_beam(filename, **kwargs)
+
+    if model == "fee":
+        assert beam1.filename == ["mwa_full_EE_test.h5"]
+        assert beam1.data_array.shape == (2, 2, 3, 91, 360)
+        # this is entirely empirical, just to prevent unexpected changes.
+        # The actual values have been validated through external tests against
+        # the mwa_pb repo.
+        assert np.isclose(
+            np.max(np.abs(beam1.data_array)),
+            0.6823676193472403,
+            rtol=beam1._data_array.tols[0],
+            atol=beam1._data_array.tols[1],
+        )
+    else:
+        assert beam1.filename == ["JMatrix_3freq.fits", "ZMatrix_3freq.fits"]
+        assert beam1.data_array.shape == (2, 2, 3, 31, 121)
 
     assert beam1.pixel_coordinate_system == "az_za"
     assert beam1.beam_type == "efield"
-    assert beam1.data_array.shape == (2, 2, 3, 91, 360)
-
-    # this is entirely empirical, just to prevent unexpected changes.
-    # The actual values have been validated through external tests against
-    # the mwa_pb repo.
-    assert np.isclose(
-        np.max(np.abs(beam1.data_array)),
-        0.6823676193472403,
-        rtol=beam1._data_array.tols[0],
-        atol=beam1._data_array.tols[1],
-    )
 
     assert "x" in beam1.feed_array
     assert "y" in beam1.feed_array
@@ -45,8 +71,17 @@ def test_read_write_mwa(mwa_beam_1ppd, tmp_path):
 
 
 @pytest.mark.filterwarnings("ignore:There are some terminated dipoles")
-def test_mwa_orientation(mwa_beam_1ppd):
-    power_beam = mwa_beam_1ppd.efield_to_power(inplace=False)
+@pytest.mark.parametrize("model", ["fee", "aee"])
+def test_mwa_orientation(mwa_fee_1ppd, mwa_aee, beam_filenames, uvbeam_kwargs, model):
+    if model == "fee":
+        ebeam = mwa_fee_1ppd
+        near_hor_za = 80
+        near_zen_za = 2
+    else:
+        ebeam = mwa_aee
+        near_hor_za = 81
+        near_zen_za = 3
+    power_beam = ebeam.efield_to_power(inplace=False)
 
     za_val = np.nonzero(np.isclose(power_beam.axis2_array, 15.0 * np.pi / 180))
 
@@ -58,13 +93,13 @@ def test_mwa_orientation(mwa_beam_1ppd):
         == utils.polstr2num(
             "ee", x_orientation=power_beam.get_x_orientation_from_feeds()
         )
-    )
+    )[0]
     north_ind = np.nonzero(
         power_beam.polarization_array
         == utils.polstr2num(
             "nn", x_orientation=power_beam.get_x_orientation_from_feeds()
         )
-    )
+    )[0]
 
     # check that the e/w dipole is more sensitive n/s
     assert (
@@ -85,11 +120,13 @@ def test_mwa_orientation(mwa_beam_1ppd):
     # single dipole
     delays = np.full((2, 16), 32, dtype=int)
     delays[:, 5] = 0
-    efield_beam = UVBeam.from_file(
-        fetch_data("mwa_full_EE"), pixels_per_deg=1, delays=delays
-    )
 
-    za_val = np.nonzero(np.isclose(efield_beam.axis2_array, 80.0 * np.pi / 180))
+    filename = beam_filenames[model]
+    kwargs = uvbeam_kwargs[model]
+    kwargs["delays"] = delays
+    efield_beam = UVBeam.from_file(filename, **kwargs)
+
+    za_val = np.nonzero(np.isclose(efield_beam.axis2_array, near_hor_za * np.pi / 180))
 
     max_az_response = np.max(np.abs(efield_beam.data_array[0, east_ind, 0, za_val, :]))
     max_za_response = np.max(np.abs(efield_beam.data_array[1, east_ind, 0, za_val, :]))
@@ -99,9 +136,10 @@ def test_mwa_orientation(mwa_beam_1ppd):
     max_za_response = np.max(np.abs(efield_beam.data_array[1, north_ind, 0, za_val, :]))
     assert max_az_response > max_za_response
 
+    # go back to zenith pointed full tile beam
     # check the sign of the responses are as expected close to zenith
-    efield_beam = mwa_beam_1ppd
-    za_val = np.nonzero(np.isclose(power_beam.axis2_array, 2.0 * np.pi / 180))
+    efield_beam = ebeam
+    za_val = np.nonzero(np.isclose(power_beam.axis2_array, near_zen_za * np.pi / 180))
 
     # first check zenith angle aligned response
     assert efield_beam.data_array[1, east_ind, 0, za_val, east_az] > 0
@@ -127,17 +165,21 @@ def test_mwa_orientation(mwa_beam_1ppd):
         ),
     ],
 )
-def test_mwa_pointing(delay_set, az_val, za_range):
+@pytest.mark.parametrize("model", ["fee", "aee"])
+def test_mwa_pointing(
+    delay_set, az_val, za_range, beam_filenames, uvbeam_kwargs, model
+):
     # Test that pointing the beam moves the peak in the right direction.
 
     delays = np.empty((2, 16), dtype=int)
 
     for pol in range(2):
         delays[pol] = delay_set
 
-    mwa_beam = UVBeam.from_file(
-        fetch_data("mwa_full_EE"), pixels_per_deg=1, beam_type="efield", delays=delays
-    )
+    filename = beam_filenames[model]
+    kwargs = uvbeam_kwargs[model]
+    kwargs["delays"] = delays
+    mwa_beam = UVBeam.from_file(filename, **kwargs)
     mwa_beam.efield_to_power(calc_cross_pols=False)
 
     # set up zenith angle, azimuth and frequency arrays to evaluate with
@@ -189,8 +231,8 @@ def test_mwa_pointing(delay_set, az_val, za_range):
     assert np.rad2deg(za_array[max_nn_loc]) < za_range[1]
 
 
-def test_freq_range(mwa_beam_1ppd):
-    beam1 = mwa_beam_1ppd
+def test_freq_range(mwa_fee_1ppd):
+    beam1 = mwa_fee_1ppd
     beam2 = UVBeam()
 
     # include all