Add option to write uvw_array as double precision in UVFITS when data_array is single precision

CHANGELOG.md

@@ -4,6 +4,8 @@ All notable changes to this project will be documented in this file.
 ## [Unreleased]
 
 ### Added
+- Option to write the uvw_array as double precision in UVFITS files even when
+the data array are single precision. Default is set to write them as doubles.
 - ATA has been added to the list of known telescopes.
 
 ### Fixed

src/pyuvdata/uvdata/uvdata.py

@@ -11395,20 +11395,7 @@ def write_ms(
         )
         del ms_obj
 
-    def write_uvfits(
-        self,
-        filename,
-        *,
-        write_lst=True,
-        force_phase=False,
-        run_check=True,
-        check_extra=True,
-        run_check_acceptability=True,
-        strict_uvw_antpos_check=False,
-        check_autos=True,
-        fix_autos=False,
-        use_miriad_convention=False,
-    ):
+    def write_uvfits(self, filename: str, **kwargs):
         """
         Write the data to a uvfits file.
 
@@ -11424,47 +11411,53 @@ def write_uvfits(
             The uvfits file to write to.
         write_lst : bool
             Option to write the LSTs to the metadata (random group parameters).
-        force_phase:  : bool
+            Default is True.
+        force_phase : bool
             Option to automatically phase unprojected data to zenith of the first
-            timestamp.
+            timestamp. Default is False.
+        uvw_double : bool
+            Option to write uvws at double precision if data array is single
+            precision (if data array is double precision uvws are always written
+            at double precision). This requires writing the uvws out into two
+            identically named parameters to be added together on read (the same
+            mechanism that is used for times in uvfits). Default is True.
+        use_miriad_convention : bool
+            Option to use the MIRIAD baseline convention, and write to BASELINE column.
+            This mode is required for UVFITS files with >256 antennas to be
+            readable by MIRIAD, and supports up to 2048 antennas.
+            The MIRIAD baseline ID is given by
+            `bl = 256 * ant1 + ant2` if `ant2 < 256`, otherwise
+            `bl = 2048 * ant1 + ant2 + 2**16`.
+            Note MIRIAD uses 1-indexed antenna IDs, but this code accepts 0-based.
+            Default is False.
         run_check : bool
             Option to check for the existence and proper shapes of parameters
-            after before writing the file (the default is True,
-            meaning the check will be run).
+            before writing the file. Default is True.
         check_extra : bool
-            Option to check optional parameters as well as required ones (the
-            default is True, meaning the optional parameters will be checked).
+            Option to check optional parameters as well as required ones.
+            Default is True.
         run_check_acceptability : bool
             Option to check acceptable range of the values of parameters before
-            writing the file (the default is True, meaning the acceptable
-            range check will be done).
+            writing the file. Default is True.
         strict_uvw_antpos_check : bool
             Option to raise an error rather than a warning if the check that
-            uvws match antenna positions does not pass.
+            uvws match antenna positions does not pass. Default is False.
         check_autos : bool
             Check whether any auto-correlations have non-zero imaginary values in
             data_array (which should not mathematically exist). Default is True.
         fix_autos : bool
             If auto-correlations with imaginary values are found, fix those values so
             that they are real-only in data_array. Default is False.
-        use_miriad_convention : bool
-            Option to use the MIRIAD baseline convention, and write to BASELINE column.
-            This mode is required for UVFITS files with >256 antennas to be
-            readable by MIRIAD, and supports up to 2048 antennas.
-            The MIRIAD baseline ID is given by
-            `bl = 256 * ant1 + ant2` if `ant2 < 256`, otherwise
-            `bl = 2048 * ant1 + ant2 + 2**16`.
-            Note MIRIAD uses 1-indexed antenna IDs, but this code accepts 0-based.
 
         Raises
         ------
         ValueError
-            Any blts are unprojected and `force_phase` keyword is not set.
+            The object contains unprojected data and `force_phase` keyword is not set.
             If the frequencies are not evenly spaced or are separated by more
             than their channel width.
             The polarization values are not evenly spaced.
-            If the `timesys` parameter is not set to "UTC".
             If the UVData object is a metadata only object.
+            If the `timesys` parameter is set to anything other than "UTC" or None.
         TypeError
             If any entry in extra_keywords is not a single string or number.
 
@@ -11484,18 +11477,7 @@ def write_uvfits(
             uvfits_obj = uvfits_obj.copy()
             uvfits_obj.remove_flex_pol()
 
-        uvfits_obj.write_uvfits(
-            filename,
-            write_lst=write_lst,
-            force_phase=force_phase,
-            run_check=run_check,
-            check_extra=check_extra,
-            run_check_acceptability=run_check_acceptability,
-            strict_uvw_antpos_check=strict_uvw_antpos_check,
-            check_autos=check_autos,
-            fix_autos=fix_autos,
-            use_miriad_convention=use_miriad_convention,
-        )
+        uvfits_obj.write_uvfits(filename, **kwargs)
         del uvfits_obj
 
     def write_uvh5(

src/pyuvdata/uvdata/uvfits.py

@@ -802,76 +802,23 @@ def read_uvfits(
                 fix_autos=fix_autos,
             )
 
+    @copy_replace_short_description(UVData.write_uvfits, style=DocstringStyle.NUMPYDOC)
     def write_uvfits(
         self,
-        filename,
+        filename: str,
         *,
-        write_lst=True,
-        force_phase=False,
-        run_check=True,
-        check_extra=True,
-        run_check_acceptability=True,
-        strict_uvw_antpos_check=False,
-        check_autos=True,
-        fix_autos=False,
-        use_miriad_convention=False,
+        write_lst: bool = True,
+        force_phase: bool = False,
+        uvw_double: bool = True,
+        use_miriad_convention: bool = False,
+        run_check: bool = True,
+        check_extra: bool = True,
+        run_check_acceptability: bool = True,
+        strict_uvw_antpos_check: bool = False,
+        check_autos: bool = True,
+        fix_autos: bool = False,
     ):
-        """
-        Write the data to a uvfits file.
-
-        If using this method to write out a data set for import into CASA, users should
-        be aware that the `importuvifts` task does not currently support reading in
-        data sets where the number of antennas is > 255. If writing out such a data set
-        for use in CASA, we suggest using the measurement set writer (`UVData.write_ms`)
-        instead.
-
-        Parameters
-        ----------
-        filename : str
-            The uvfits file to write to.
-        write_lst : bool
-            Option to write the LSTs to the metadata (random group parameters).
-        force_phase : bool
-            Option to automatically phase drift scan data to zenith of the first
-            timestamp.
-        run_check : bool
-            Option to check for the existence and proper shapes of parameters
-            before writing the file.
-        check_extra : bool
-            Option to check optional parameters as well as required ones.
-        run_check_acceptability : bool
-            Option to check acceptable range of the values of parameters before
-            writing the file.
-        strict_uvw_antpos_check : bool
-            Option to raise an error rather than a warning if the check that
-            uvws match antenna positions does not pass.
-        check_autos : bool
-            Check whether any auto-correlations have non-zero imaginary values in
-            data_array (which should not mathematically exist). Default is True.
-        fix_autos : bool
-            If auto-correlations with imaginary values are found, fix those values so
-            that they are real-only in data_array. Default is False.
-        use_miriad_convention : bool
-            Option to use the MIRIAD baseline convention, and write to BASELINE column.
-            This mode is required for UVFITS files with >256 antennas to be
-            readable by MIRIAD, and supports up to 2048 antennas.
-            The MIRIAD baseline ID is given by
-            `bl = 256 * ant1 + ant2` if `ant2 < 256`, otherwise
-            `bl = 2048 * ant1 + ant2 + 2**16`.
-            Note MIRIAD uses 1-indexed antenna IDs, but this code accepts 0-based.
-
-        Raises
-        ------
-        ValueError
-            The object contains unprojected data and `force_phase` keyword is not set.
-            If the frequencies are not evenly spaced or are separated by more
-            than their channel width.
-            The polarization values are not evenly spaced.
-            If the `timesys` parameter is set to anything other than "UTC" or None.
-        TypeError
-            If any entry in extra_keywords is not a single string or number.
-
-        """
+        """Write data to a uvfits file."""
         if run_check:
             self.check(
                 check_extra=check_extra,
@@ -1009,8 +956,12 @@ def write_uvfits(
                 self.time_array - jd_midnight - time_array1.astype(np.float64)
             ).astype(np.float32)
 
+            uvw_array1 = np.float32(uvw_array_sec)
+            uvw_array2 = np.float32(uvw_array_sec - np.float64(uvw_array1))
         else:
             time_array1 = self.time_array - jd_midnight
+            uvw_array1 = uvw_array_sec
+
         int_time_array = self.integration_time
 
         # If using MIRIAD convention, we need 1-indexed data
@@ -1104,6 +1055,13 @@ def write_uvfits(
             pscal_dict["DATE2   "] = 1.0
             pzero_dict["DATE2   "] = 0.0
             parnames_use.append("DATE2   ")
+            if uvw_double:
+                for ind, name in enumerate(["UU", "VV", "WW"]):
+                    keyname = name + "2     "
+                    group_parameter_dict[keyname] = uvw_array2[:, ind]
+                    pscal_dict[keyname] = 1.0
+                    pzero_dict[keyname] = 0.0
+                    parnames_use.append(keyname)
 
         if use_miriad_convention or (
             np.max(ant1_array_use) < 255 and np.max(ant2_array_use) < 255
@@ -1134,6 +1092,10 @@ def write_uvfits(
             # add second date part
             parnames_write = copy.deepcopy(parnames_use)
             parnames_write[parnames_write.index("DATE2   ")] = "DATE    "
+            if uvw_double:
+                parnames_write[parnames_write.index("UU2     ")] = "UU      "
+                parnames_write[parnames_write.index("VV2     ")] = "VV      "
+                parnames_write[parnames_write.index("WW2     ")] = "WW      "
             if write_lst:
                 # add second LST array part
                 parnames_use.append("LST2    ")

tests/uvdata/test_uvfits.py

@@ -136,10 +136,11 @@ def test_read_nrao(casa_uvfits):
 @pytest.mark.filterwarnings("ignore:ITRF coordinate frame detected")
 @pytest.mark.filterwarnings("ignore:Telescope OVRO_MMA is not")
 @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values")
-def test_time_precision(tmp_path):
+@pytest.mark.parametrize("uvw_double", [True, False])
+def test_group_param_precision(tmp_path, uvw_double):
     """
-    This tests that the times are round-tripped through write/read uvfits to sufficient
-    precision.
+    This tests that the times and uvws are round-tripped through write/read
+    uvfits to sufficient precision.
     """
     pytest.importorskip("casacore")
     lwa_file = os.path.join(
@@ -149,7 +150,7 @@ def test_time_precision(tmp_path):
     uvd.read(lwa_file)
 
     testfile = os.path.join(tmp_path, "lwa_testfile.uvfits")
-    uvd.write_uvfits(testfile)
+    uvd.write_uvfits(testfile, uvw_double=uvw_double)
 
     uvd2 = UVData()
     uvd2.read(testfile)
@@ -168,6 +169,12 @@ def test_time_precision(tmp_path):
         atol=uvd2._lst_array.tols[1],
     )
 
+    if uvw_double:
+        np.testing.assert_allclose(uvd.uvw_array, uvd2.uvw_array, rtol=1e-13)
+    else:
+        np.testing.assert_allclose(uvd.uvw_array, uvd2.uvw_array, rtol=1e-7)
+        assert not np.allclose(uvd.uvw_array, uvd2.uvw_array, rtol=1e-13)
+
     # The incoming ra is specified as negative, it gets 2pi added to it in the roundtrip
     uvd2.phase_center_catalog[1]["cat_lon"] -= 2 * np.pi