From 16e092c09b73a61c8827de88b6bd6266dc827435 Mon Sep 17 00:00:00 2001
From: Garrett 'Karto' Keating <garrett.keating@cfa.harvard.edu>
Date: Thu, 16 Dec 2021 22:28:34 -0500
Subject: [PATCH] Adding fix_autos to write_uvh5_part, adding test coverage

---
 pyuvdata/uvdata/tests/test_uvh5.py | 41 ++++++++++++++++++++++++++++++
 pyuvdata/uvdata/uvdata.py          | 20 ++++++++++++++-
 2 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/pyuvdata/uvdata/tests/test_uvh5.py b/pyuvdata/uvdata/tests/test_uvh5.py
index 5d39ae8c5b..a83c6864d9 100644
--- a/pyuvdata/uvdata/tests/test_uvh5.py
+++ b/pyuvdata/uvdata/tests/test_uvh5.py
@@ -3314,3 +3314,44 @@ def test_none_extra_keywords(uv_uvh5, tmp_path):
         assert h5f["Header/extra_keywords/foo"].shape is None
 
     return
+
+
+def test_write_uvh5_part_fix_autos(uv_uvh5, tmp_path):
+    """Test that fix_autos works correctly on partial UVH5 wrute"""
+    test_uvh5 = UVData()
+    testfile = os.path.join(tmp_path, "write_uvh5_part_fix_autos.uvh5")
+
+    # Select out the relevant data (where the 0 and 1 indicies of the pol array
+    # correspond to xx and yy polarization data), and corrupt it accordingly
+    auto_data = uv_uvh5.data_array[uv_uvh5.ant_1_array == uv_uvh5.ant_2_array]
+    auto_data[:, :, :, [0, 1]] *= 1j
+    uv_uvh5.data_array[uv_uvh5.ant_1_array == uv_uvh5.ant_2_array] = auto_data
+
+    # Create and write out the data, with fix_autos set to operate
+    initialize_with_zeros_ints(uv_uvh5, testfile)
+    uv_uvh5.write_uvh5_part(
+        testfile,
+        uv_uvh5.data_array,
+        uv_uvh5.flag_array,
+        uv_uvh5.nsample_array,
+        fix_autos=True,
+    )
+
+    # Fix the autos we corrupted earlier, and plug the data back in to data_array
+    auto_data[:, :, :, [0, 1]] *= -1j
+    uv_uvh5.data_array[uv_uvh5.ant_1_array == uv_uvh5.ant_2_array] = auto_data
+
+    # Read in the data on disk, make sure it looks like our manually repaired data
+    test_uvh5.read(testfile)
+
+    assert uv_uvh5 == test_uvh5
+
+
+def test_fix_autos_no_op():
+    """Test that a no-op with _fix_autos returns a warning"""
+    uvd = UVData()
+
+    with uvtest.check_warnings(
+        UserWarning, "Cannot use _fix_autos if ant_1_array, ant_2_array, or "
+    ):
+        uvd._fix_autos()
diff --git a/pyuvdata/uvdata/uvdata.py b/pyuvdata/uvdata/uvdata.py
index 30d90ab9f3..e1df758137 100644
--- a/pyuvdata/uvdata/uvdata.py
+++ b/pyuvdata/uvdata/uvdata.py
@@ -2477,6 +2477,15 @@ def _calc_nants_data(self):
 
     def _fix_autos(self):
         """Remove imaginary component of auto-correlations."""
+        if self.polarization_array is None or (
+            self.ant_1_array is None or self.ant_2_array is None
+        ):
+            warnings.warn(
+                "Cannot use _fix_autos if ant_1_array, ant_2_array, or "
+                "polarization_array are None. Leaving data_array untouched."
+            )
+            return
+
         # Select out the autos
         auto_screen = self.ant_1_array == self.ant_2_array
 
@@ -2491,7 +2500,9 @@ def _fix_autos(self):
         )
 
         # Make sure we actually have work to do here, otherwise skip all of this
-        if np.any(pol_screen) and np.any(auto_screen):
+        if (np.any(pol_screen) and np.any(auto_screen)) and not (
+            pol_screen is None or auto_screen is None
+        ):
             # Select out the relevant data. Need to do this because we have two
             # complex slices we need to do
             auto_data = self.data_array[auto_screen]
@@ -12934,6 +12945,7 @@ def write_uvh5_part(
         blt_inds=None,
         add_to_history=None,
         run_check_acceptability=True,
+        fix_autos=False,
     ):
         """
         Write data to a UVH5 file that has already been initialized.
@@ -13024,8 +13036,14 @@ def write_uvh5_part(
         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.
+        fix_autos : bool
+            Force the auto-correlations to be real-only values in data_array.
+            Default is False.
 
         """
+        if fix_autos:
+            self._fix_autos()
+
         uvh5_obj = self._convert_to_filetype("uvh5")
         uvh5_obj.write_uvh5_part(
             filename,