Add testing for uvdata handling

pyuvdata/tests/test_utils.py

@@ -20,15 +20,23 @@
 from pyuvdata.data import DATA_PATH
 from pyuvdata.utils import hasmoon
 
+selenoids = ["SPHERE", "GSFC", "GRAIL23", "CE-1-LAM-GEO"]
+
 if hasmoon:
     from pyuvdata.utils import LTime, MoonLocation
 
+    frame_selenoid = [["itrs", "SPHERE"]]
+    for snd in selenoids:
+        frame_selenoid.append(["mcmf", snd])
+else:
+    frame_selenoid = [["itrs", "SPHERE"]]
+
+
 # Earth
 ref_latlonalt = (-26.7 * np.pi / 180.0, 116.7 * np.pi / 180.0, 377.8)
 ref_xyz = (-2562123.42683, 5094215.40141, -2848728.58869)
 
 # Moon
-selenoids = ["SPHERE", "GSFC", "GRAIL23", "CE-1-LAM-GEO"]
 ref_latlonalt_moon = (0.6875 * np.pi / 180.0, 24.433 * np.pi / 180.0, 0.3)
 ref_xyz_moon = {
     "SPHERE": (1581421.43506347, 718463.12201783, 20843.2071012),
@@ -646,15 +654,7 @@ def test_enu_from_ecef_magnitude_error(enu_ecef_info):
     )
 
 
-if hasmoon:
-    param_list = [["itrs", "SPHERE"]]
-    for snd in selenoids:
-        param_list.append(["mcmf", snd])
-else:
-    param_list = [["itrs", "SPHERE"]]
-
-
-@pytest.mark.parametrize(["frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["frame", "selenoid"], frame_selenoid)
 def test_ecef_from_enu_roundtrip(enu_ecef_info, enu_mcmf_info, frame, selenoid):
     """Test ECEF_from_ENU values."""
     (center_lat, center_lon, center_alt, lats, lons, alts, x, y, z, east, north, up) = (
@@ -1513,7 +1513,7 @@ def test_transform_fk5_fk4_icrs_loop(astrometry_args):
     assert np.all(check_coord.separation(astrometry_args["icrs_coord"]).uarcsec < 0.1)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 @pytest.mark.parametrize("in_lib", ["erfa", "astropy"])
 @pytest.mark.parametrize("out_lib", ["erfa", "astropy"])
 def test_roundtrip_icrs(astrometry_args, telescope_frame, selenoid, in_lib, out_lib):
@@ -1765,7 +1765,7 @@ def test_ephem_interp_multi_point():
 
 
 @pytest.mark.parametrize("frame", ["icrs", "fk5"])
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_sidereal(astrometry_args, frame, telescope_frame, selenoid):
     """
     Tests that we can calculate app coords for sidereal objects
@@ -1799,7 +1799,7 @@ def test_calc_app_sidereal(astrometry_args, frame, telescope_frame, selenoid):
 
 
 @pytest.mark.parametrize("frame", ["icrs", "fk5"])
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_ephem(astrometry_args, frame, telescope_frame, selenoid):
     """
     Tests that we can calculate app coords for ephem objects
@@ -1842,7 +1842,7 @@ def test_calc_app_ephem(astrometry_args, frame, telescope_frame, selenoid):
     assert np.all(app_coord.separation(check_coord).uarcsec < 1.0)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_driftscan(astrometry_args, telescope_frame, selenoid):
     """
     Tests that we can calculate app coords for driftscan objects
@@ -1881,7 +1881,7 @@ def test_calc_app_driftscan(astrometry_args, telescope_frame, selenoid):
     assert np.all(drift_coord.separation(check_coord).uarcsec < 1.0)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_unprojected(astrometry_args, telescope_frame, selenoid):
     """
     Tests that we can calculate app coords for unphased objects
@@ -1915,7 +1915,7 @@ def test_calc_app_unprojected(astrometry_args, telescope_frame, selenoid):
     assert np.all(drift_coord.separation(check_coord).uarcsec < 1.0)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_fk5_roundtrip(astrometry_args, telescope_frame, selenoid):
     # Do a round-trip with the two top-level functions and make sure they agree to
     # better than 1 µas, first in FK5
@@ -1950,7 +1950,7 @@ def test_calc_app_fk5_roundtrip(astrometry_args, telescope_frame, selenoid):
     assert np.all(SkyCoord(0, 0, unit="rad").separation(check_coord).uarcsec < 1.0)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_fk4_roundtrip(astrometry_args, telescope_frame, selenoid):
     # Finally, check and make sure that FK4 performs similarly
     if telescope_frame == "itrs":
@@ -2123,7 +2123,7 @@ def test_sidereal_reptime(astrometry_args):
     assert np.all(gcrs_dec == check_dec)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_transform_icrs_to_app_time_obj(astrometry_args, telescope_frame, selenoid):
     """
     Test that we recover identical values when using a Time objects instead of a floats
@@ -2184,7 +2184,7 @@ def test_transform_app_to_icrs_objs(astrometry_args):
     assert np.all(check_dec == icrs_dec)
 
 
-@pytest.mark.parametrize(["telescope_frame", "selenoid"], param_list)
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
 def test_calc_app_coords_objs(astrometry_args, telescope_frame, selenoid):
     """
     Test that we recover identical values when using Time/EarthLocation objects instead

pyuvdata/utils.py

@@ -123,13 +123,13 @@
                     "pI": ["I", "I"], "pQ": ["Q", "Q"],
                     "pU": ["U", "U"], "pV": ["V", "V"]}
 
+# fmt: on
+
 _range_dict = {
-    "itrs": (6.35e6, 6.39e6, "Earth"), "mcmf": (1717100.0, 1757100.0, "Moon")
+    "itrs": (6.35e6, 6.39e6, "Earth"),
+    "mcmf": (1717100.0, 1757100.0, "Moon"),
 }
 
-
-# fmt: on
-
 if hasmoon:
     lunar_ellipsoids = {
         "SPHERE": _utils.Body.Moon_sphere,
@@ -3195,6 +3195,7 @@ def calc_frame_pos_angle(
     ref_frame,
     ref_epoch=None,
     telescope_frame="itrs",
+    lunar_ellipsoid="SPHERE",
     offset_pos=(np.pi / 360.0),
 ):
     """
@@ -3230,6 +3231,10 @@ def calc_frame_pos_angle(
     telescope_frame: str, optional
         Reference frame for telescope location. Options are itrs (default) or mcmf.
         Only used if telescope_loc is not an EarthLocation or MoonLocation.
+    lunar_ellipsoid : str
+        Ellipsoid to use for lunar coordinates. Must be one of "SPHERE",
+        "GSFC", "GRAIL23", "CE-1-LAM-GEO" (see lunarsky package for details). Default
+        is "SPHERE". Only used if frame is mcmf.
     offset_pos : float
         Distance of the offset position used to calculate the frame PA. Default
         is 0.5 degrees, which should be sufficent for most applications.
@@ -3292,6 +3297,7 @@ def calc_frame_pos_angle(
         telescope_loc,
         ref_frame,
         telescope_frame=telescope_frame,
+        lunar_ellipsoid=lunar_ellipsoid,
         coord_epoch=ref_epoch,
     )
 

pyuvdata/uvdata/initializers.py

@@ -504,6 +504,9 @@ def new_uvdata(
 
     if hasmoon and isinstance(telescope_location, MoonLocation):
         telescope_location.ellipsoid = lunar_ellipsoid
+        telescope_frame = "mcmf"
+    else:
+        telescope_frame = "itrs"
 
     lst_array, integration_time = get_time_params(
         telescope_location,
@@ -570,6 +573,8 @@ def new_uvdata(
         telescope_location.y.to_value("m"),
         telescope_location.z.to_value("m"),
     ]
+    obj._telescope_location.frame = telescope_frame
+    obj._telescope_location.lunar_ellipsoid = lunar_ellipsoid
     obj.telescope_name = telescope_name
     obj.baseline_array = baseline_array
     obj.ant_1_array = ant_1_array

pyuvdata/uvdata/tests/test_initializers.py

@@ -22,6 +22,8 @@
     get_time_params,
 )
 
+selenoids = ["SPHERE", "GSFC", "GRAIL23", "CE-1-LAM-GEO"]
+
 
 @pytest.fixture(scope="function")
 def simplest_working_params() -> dict[str, Any]:
@@ -39,6 +41,25 @@ def simplest_working_params() -> dict[str, Any]:
     }
 
 
+@pytest.fixture
+def lunar_simple_params() -> dict[str, Any]:
+    pytest.importorskip("lunarsky")
+    from pyuvdata.utils import MoonLocation
+
+    return {
+        "freq_array": np.linspace(1e8, 2e8, 100),
+        "polarization_array": ["xx", "yy"],
+        "antenna_positions": {
+            0: [0.0, 0.0, 0.0],
+            1: [0.0, 0.0, 1.0],
+            2: [0.0, 0.0, 2.0],
+        },
+        "telescope_location": MoonLocation.from_selenodetic(0, 0, 0),
+        "telescope_name": "test",
+        "times": np.linspace(2459855, 2459856, 20),
+    }
+
+
 def test_simplest_new_uvdata(simplest_working_params: dict[str, Any]):
     uvd = UVData.new(**simplest_working_params)
 
@@ -51,6 +72,14 @@ def test_simplest_new_uvdata(simplest_working_params: dict[str, Any]):
     assert uvd.Nspws == 1
 
 
+@pytest.mark.parametrize("selenoid", selenoids)
+def test_lunar_simple_new_uvdata(lunar_simple_params: dict[str, Any], selenoid: str):
+    uvd = UVData.new(**lunar_simple_params, lunar_ellipsoid=selenoid)
+
+    assert uvd._telescope_location.frame == "mcmf"
+    assert uvd._telescope_location.lunar_ellipsoid == selenoid
+
+
 def test_bad_inputs(simplest_working_params: dict[str, Any]):
     with pytest.raises(ValueError, match="vis_units must be one of"):
         UVData.new(**simplest_working_params, vis_units="foo")

pyuvdata/uvdata/tests/test_uvfits.py

@@ -24,6 +24,11 @@
 
 paper_uvfits = os.path.join(DATA_PATH, "zen.2456865.60537.xy.uvcRREAAM.uvfits")
 
+selenoids = ["SPHERE", "GSFC", "GRAIL23", "CE-1-LAM-GEO"]
+frame_selenoid = [["itrs", "SPHERE"]]
+for snd in selenoids:
+    frame_selenoid.append(["mcmf", snd])
+
 
 def _fix_uvfits_multi_group_params(vis_hdu):
     par_names = vis_hdu.data.parnames
@@ -531,8 +536,8 @@ def test_casa_nonascii_bytes_antenna_names():
 @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values")
 @pytest.mark.filterwarnings("ignore:Telescope EVLA is not")
 @pytest.mark.parametrize("future_shapes", [True, False])
-@pytest.mark.parametrize("telescope_frame", ["itrs", "mcmf"])
-def test_readwriteread(tmp_path, casa_uvfits, future_shapes, telescope_frame):
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
+def test_readwriteread(tmp_path, casa_uvfits, future_shapes, telescope_frame, selenoid):
     """
     CASA tutorial uvfits loopback test.
 
@@ -549,13 +554,15 @@ def test_readwriteread(tmp_path, casa_uvfits, future_shapes, telescope_frame):
         enu_antpos, _ = uv_in.get_ENU_antpos()
         latitude, longitude, altitude = uv_in.telescope_location_lat_lon_alt
         uv_in._telescope_location.frame = "mcmf"
+        uv_in._telescope_location.lunar_ellipsoid = selenoid
         uv_in.telescope_location_lat_lon_alt = (latitude, longitude, altitude)
         new_full_antpos = uvutils.ECEF_from_ENU(
             enu=enu_antpos,
             latitude=latitude,
             longitude=longitude,
             altitude=altitude,
             frame="mcmf",
+            lunar_ellipsoid=selenoid,
         )
         uv_in.antenna_positions = new_full_antpos - uv_in.telescope_location
         uv_in.set_lsts_from_time_array()
@@ -573,6 +580,10 @@ def test_readwriteread(tmp_path, casa_uvfits, future_shapes, telescope_frame):
     uv_in.filename = uv_out.filename
 
     assert uv_in._telescope_location.frame == uv_out._telescope_location.frame
+    assert (
+        uv_in._telescope_location.lunar_ellipsoid
+        == uv_out._telescope_location.lunar_ellipsoid
+    )
 
     uv_out._consolidate_phase_center_catalogs(
         reference_catalog=uv_in.phase_center_catalog

pyuvdata/uvdata/tests/test_uvh5.py

@@ -33,6 +33,11 @@
     pytest.mark.filterwarnings("ignore:Telescope EVLA is not"),
 ]
 
+selenoids = ["SPHERE", "GSFC", "GRAIL23", "CE-1-LAM-GEO"]
+frame_selenoid = [["itrs", "SPHERE"]]
+for snd in selenoids:
+    frame_selenoid.append(["mcmf", snd])
+
 
 @pytest.fixture(scope="session")
 def uv_uvh5_main():
@@ -189,8 +194,10 @@ def test_read_miriad_write_uvh5_read_uvh5(paper_miriad, future_shapes, tmp_path)
 
 
 @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values")
-@pytest.mark.parametrize("telescope_frame", ["itrs", "mcmf"])
-def test_read_uvfits_write_uvh5_read_uvh5(casa_uvfits, tmp_path, telescope_frame):
+@pytest.mark.parametrize(["telescope_frame", "selenoid"], frame_selenoid)
+def test_read_uvfits_write_uvh5_read_uvh5(
+    casa_uvfits, tmp_path, telescope_frame, selenoid
+):
     """
     Test a uvfits file round trip.
     """
@@ -201,19 +208,22 @@ def test_read_uvfits_write_uvh5_read_uvh5(casa_uvfits, tmp_path, telescope_frame
         enu_antpos, _ = uv_in.get_ENU_antpos()
         latitude, longitude, altitude = uv_in.telescope_location_lat_lon_alt
         uv_in._telescope_location.frame = "mcmf"
+        uv_in._telescope_location.lunar_ellipsoid = selenoid
         uv_in.telescope_location_lat_lon_alt = (latitude, longitude, altitude)
         new_full_antpos = uvutils.ECEF_from_ENU(
             enu=enu_antpos,
             latitude=latitude,
             longitude=longitude,
             altitude=altitude,
             frame="mcmf",
+            lunar_ellipsoid=selenoid,
         )
         uv_in.antenna_positions = new_full_antpos - uv_in.telescope_location
         uv_in.set_lsts_from_time_array()
         uv_in.check()
 
     assert uv_in._telescope_location.frame == telescope_frame
+    assert uv_in._telescope_location.lunar_ellipsoid == selenoid
 
     uv_out = UVData()
     fname = f"outtest_{telescope_frame}_uvfits.uvh5"
@@ -225,6 +235,7 @@ def test_read_uvfits_write_uvh5_read_uvh5(casa_uvfits, tmp_path, telescope_frame
     os.remove(testfile)
 
     assert uv_out._telescope_location.frame == telescope_frame
+    assert uv_out._telescope_location.lunar_ellipsoid == selenoid
 
     # make sure filenames are what we expect
     assert uv_in.filename == ["day2_TDEM0003_10s_norx_1src_1spw.uvfits"]