Add fix_autos_func for reading diff autos

HERA-Team · May 21, 2024 · 7b4c3b1 · 7b4c3b1
1 parent e56a686
commit 7b4c3b1
Showing 1 changed file with 5 additions and 3 deletions.
diff --git a/hera_cal/io.py b/hera_cal/io.py
@@ -1469,7 +1469,7 @@ def _HERAData_error(self, *args, **kwargs):
         raise NotImplementedError('HERADataFastReader does not support this method. Try HERAData instead.')
 
     def read(self, bls=None, pols=None, full_read_thresh=0.002, read_data=True, read_flags=True,
-             read_nsamples=True, check=False, dtype=np.complex128, verbose=False, skip_lsts=False):
+             read_nsamples=True, fix_autos_func=np.abs, check=False, dtype=np.complex128, verbose=False, skip_lsts=False):
         '''A faster read that only concatenates along the time axis. Puts times in ascending order, but does not
         check that files are contiguous. Currently not BDA compatible.
         Arguments:
@@ -1479,6 +1479,8 @@ def read(self, bls=None, pols=None, full_read_thresh=0.002, read_data=True, read
             read_data (bool, True): read data
             read_flags (bool, True): read flags
             read_nsamples (bool, True): read nsamples
+            fix_autos_func (function, np.abs): function to apply to autocorrelations to enforce, for example, that
+                they are real. Default is np.abs, which matches UVData._fix_autos(). Use np.real for loading diff data.
             check (bool, False): run sanity checks to make sure files match.
             dtype (np.complex128): numpy datatype for output complex-valued arrays
             verbose: print some progress messages.
@@ -1493,11 +1495,11 @@ def read(self, bls=None, pols=None, full_read_thresh=0.002, read_data=True, read
                             check=check, dtype=dtype, verbose=verbose)
         self._adapt_metadata(rv['info'], skip_lsts=skip_lsts)
 
-        # make autocorrleations real by taking the abs, matches UVData._fix_autos()
+        # make autocorrleations real by taking the real part. Using fix_autos_func = np.abs matches UVData._fix_autos()
         if 'data' in rv:
             for bl in rv['data']:
                 if split_bl(bl)[0] == split_bl(bl)[1]:
-                    rv['data'][bl] = np.abs(rv['data'][bl])
+                    rv['data'][bl] = fix_autos_func(rv['data'][bl])
 
         # construct datacontainers from result
         return self._make_datacontainer(rv, 'data'), self._make_datacontainer(rv, 'flags'), self._make_datacontainer(rv, 'nsamples')