Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Fix so that write_uvfits doesn't change the object, add loopback testing #1393

Merged
merged 1 commit into from
Feb 8, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 32 additions & 1 deletion pyuvdata/uvdata/tests/test_uvfits.py
Original file line number Diff line number Diff line change
Expand Up @@ -1668,7 +1668,7 @@ def test_uvfits_phasing_errors(hera_uvh5, tmp_path):
"ignore:The shapes of several attributes will be changing "
"in the future to remove the deprecated spectral window axis."
)
def test_miriad_convention(casa_uvfits, tmp_path):
def test_miriad_convention(tmp_path):
"""
Test writing a MIRIAD-compatible UVFITS file
"""
Expand All @@ -1682,6 +1682,7 @@ def test_miriad_convention(casa_uvfits, tmp_path):
uv.antenna_numbers[10] = new_idx
uv.ant_1_array[uv.ant_1_array == old_idx] = new_idx
uv.ant_2_array[uv.ant_2_array == old_idx] = new_idx
uv.baseline_array = uv.antnums_to_baseline(uv.ant_1_array, uv.ant_2_array)

testfile1 = str(tmp_path / "uv1.uvfits")
uv.write_uvfits(testfile1, use_miriad_convention=True)
Expand All @@ -1701,17 +1702,47 @@ def test_miriad_convention(casa_uvfits, tmp_path):
assert hdu[0].data["ANTENNA2"][0] == expected_vals["ANTENNA2_0"]
assert hdu[1].data["NOSTA"][0] == expected_vals["NOSTA_0"]

uv2 = UVData.from_file(testfile1)
uv2._update_phase_center_id(1, 0)
uv2.phase_center_catalog[0]["info_source"] = uv.phase_center_catalog[0][
"info_source"
]

assert uv2._ant_1_array == uv._ant_1_array
assert uv2._ant_2_array == uv._ant_2_array

assert uv2 == uv

# Test that antennas get +1 if there is a 0-indexed antennas
old_idx = uv.antenna_numbers[0]
new_idx = 0
uv.antenna_numbers[0] = new_idx
uv.ant_1_array[uv.ant_1_array == old_idx] = new_idx
uv.ant_2_array[uv.ant_2_array == old_idx] = new_idx
uv.baseline_array = uv.antnums_to_baseline(uv.ant_1_array, uv.ant_2_array)
uv2 = uv.copy()

testfile1 = str(tmp_path / "uv2.uvfits")
uv.write_uvfits(testfile1, use_miriad_convention=True)

# make sure write_uvfits doesn't change the object
assert uv2 == uv

with fits.open(testfile1) as hdu:
assert hdu[0].data["ANTENNA1"][0] == expected_vals["ANTENNA1_0"] + 1
assert hdu[0].data["ANTENNA2"][0] == expected_vals["ANTENNA2_0"] + 1
assert hdu[1].data["NOSTA"][0] == 1 # expected_vals["NOSTA_0"]

uv2 = UVData.from_file(testfile1)
uv2._update_phase_center_id(1, 0)
uv2.phase_center_catalog[0]["info_source"] = uv.phase_center_catalog[0][
"info_source"
]

# adjust for expected antenna number changes:
uv2.antenna_numbers -= 1
uv2.ant_1_array -= 1
uv2.ant_2_array -= 1
uv2.baseline_array = uv2.antnums_to_baseline(uv2.ant_1_array, uv2.ant_2_array)

assert uv2 == uv
33 changes: 18 additions & 15 deletions pyuvdata/uvdata/uvfits.py
Original file line number Diff line number Diff line change
Expand Up @@ -1008,17 +1008,20 @@ def write_uvfits(
int_time_array = self.integration_time

# If using MIRIAD convention, we need 1-indexed data
ant_nums_use = copy.copy(self.antenna_numbers)
ant1_array_use = copy.copy(self.ant_1_array)
ant2_array_use = copy.copy(self.ant_2_array)
if use_miriad_convention:
if np.min(self.antenna_numbers) == 0:
self.antenna_numbers += 1
self.ant_1_array += 1
self.ant_2_array += 1
if np.min(ant_nums_use) == 0:
ant_nums_use += 1
ant1_array_use += 1
ant2_array_use += 1

# Generate baseline IDs
attempt256 = False if use_miriad_convention else True
baselines_use = self.antnums_to_baseline(
self.ant_1_array,
self.ant_2_array,
ant1_array_use,
ant2_array_use,
attempt256=attempt256,
use_miriad_convention=use_miriad_convention,
)
Expand All @@ -1034,9 +1037,9 @@ def write_uvfits(
"DATE ": time_array1,
"BASELINE": baselines_use,
"SOURCE ": None,
"ANTENNA1": self.ant_1_array,
"ANTENNA2": self.ant_2_array,
"SUBARRAY": np.ones_like(self.ant_1_array),
"ANTENNA1": ant1_array_use,
"ANTENNA2": ant2_array_use,
"SUBARRAY": np.ones_like(ant1_array_use),
"INTTIM ": int_time_array,
}

Expand Down Expand Up @@ -1097,13 +1100,13 @@ def write_uvfits(
pzero_dict["DATE2 "] = 0.0
parnames_use.append("DATE2 ")

if use_miriad_convention:
# MIRIAD requires BASELINE column.
parnames_use.append("BASELINE")
elif np.max(self.ant_1_array) < 255 and np.max(self.ant_2_array) < 255:
if use_miriad_convention or (
np.max(ant1_array_use) < 255 and np.max(ant2_array_use) < 255
):
# if the number of antennas is less than 256 then include both the
# baseline array and the antenna arrays in the group parameters.
# Otherwise just use the antenna arrays
# Otherwise just use the antenna arrays unless writing for mirad.
# MIRIAD requires the BASELINE column.
parnames_use.append("BASELINE")
else:
warnings.warn(
Expand Down Expand Up @@ -1282,7 +1285,7 @@ def write_uvfits(
)
col2 = fits.Column(name="STABXYZ", format="3D", array=rot_ecef_positions)
# col3 = fits.Column(name="ORBPARAM", format="0D", array=Norb)
col4 = fits.Column(name="NOSTA", format="1J", array=self.antenna_numbers)
col4 = fits.Column(name="NOSTA", format="1J", array=ant_nums_use)
col5 = fits.Column(name="MNTSTA", format="1J", array=mntsta)
col6 = fits.Column(name="STAXOF", format="1E", array=staxof)
col7 = fits.Column(name="POLTYA", format="1A", array=poltya)
Expand Down
Loading