Integrating hooks for COMPASS read-in for MIR, building in check for …

…v3 compliance.

pyuvdata/uvdata/mir.py

@@ -46,6 +46,7 @@ def read_mir(
         apply_dedoppler=False,
         pseudo_cont=False,
         rechunk=None,
+        compass_soln=None,
         run_check=True,
         check_extra=True,
         run_check_acceptability=True,
@@ -124,6 +125,10 @@ def read_mir(
         rechunk : int
             Number of channels to average over when reading in the dataset. Optional
             argument, typically required to be a power of 2.
+        compass_soln : str
+            Optional argument, specifying the path of COMPASS-derived flagging and
+            bandpass gains solutions, which are applied prior to any potential spectral
+            averaging (as triggered by using the `rechunk` keyword).
         run_check : bool
             Option to check for the existence and proper shapes of parameters
             before writing the file.
@@ -149,10 +154,12 @@ def read_mir(
             that they are real-only in data_array. Default is False.
         """
         # Use the mir_parser to read in metadata, which can be used to select data.
-        mir_data = mir_parser.MirParser(filepath)
+        # We want to sure that the mir file is v3 compliant, since correctly filling
+        # values into a UVData object depends on that.
+        mir_data = mir_parser.MirParser(
+            filepath=filepath, compass_soln=compass_soln, make_v3_compliant=True
+        )
 
-        # Make sure that the mir file is v3 compliant, since correctly filling values
-        # into a UVData object depends on that.
         mir_data._make_v3_compliant()
 
         if select_where is None:

pyuvdata/uvdata/mir_parser.py

@@ -47,6 +47,8 @@ class MirParser(object):
     def __init__(
         self,
         filepath=None,
+        compass_soln=None,
+        make_v3_compliant=False,
         has_auto=False,
         has_cross=True,
         load_auto=False,
@@ -64,6 +66,13 @@ def __init__(
         ----------
         filepath : str
             Filepath is the path to the folder containing the Mir data set.
+        compass_soln : str
+            Optional argument, specifying the path of COMPASS-derived flagging and
+            bandpass gains solutions, which are loaded into the object.
+        make_v3_compliant : bool
+            Convert the metadata required for filling a UVData object into a
+            v3-compliant format. Only applicable if MIR file format is v1 or v2.
+            Default is False.
         has_auto : bool
             Flag to read auto-correlation data. Default is False.
         has_cross : bool
@@ -118,7 +127,9 @@ def __init__(
         # On init, if a filepath is provided, then fill in the object
         if filepath is not None:
             self.read(
-                filepath,
+                filepath=filepath,
+                compass_soln=compass_soln,
+                make_v3_compliant=make_v3_compliant,
                 has_auto=has_auto,
                 has_cross=has_cross,
                 load_auto=load_auto,
@@ -1828,7 +1839,9 @@ def _fix_acdata(self):
 
     def read(
         self,
-        filepath,
+        filepath=None,
+        compass_soln=None,
+        make_v3_compliant=False,
         has_auto=False,
         has_cross=True,
         load_auto=False,
@@ -1846,6 +1859,13 @@ def read(
         ----------
         filepath : str
             Filepath is the path to the folder containing the Mir data set.
+        compass_soln : str
+            Optional argument, specifying the path of COMPASS-derived flagging and
+            bandpass gains solutions, which are loaded into the object.
+        make_v3_compliant : bool
+            Convert the metadata required for filling a UVData object into a
+            v3-compliant format. Only applicable if MIR file format is v1 or v2.
+            Default is False.
         has_auto : bool
             Flag to read auto-correlation data. Default is False.
         has_cross : bool
@@ -1909,6 +1929,14 @@ def read(
         self._file_dict = {filepath: file_dict}
         self.filepath = filepath
 
+        # If we need to update metadata for V3 compliance, do that now.
+        if make_v3_compliant:
+            self._make_v3_compliant()
+
+        # Finally, if we've specified a COMPASS solution, load that now as well.
+        if compass_soln is not None:
+            self.read_compass_solns(compass_soln)
+
         # Set/clear these to start
         self.vis_data = self.raw_data = self.auto_data = None
         self._tsys_applied = False
@@ -2754,7 +2782,6 @@ def _read_compass_solns(self, filename):
             If the COMPASS solutions do not appear to overlap in time with that in
             the MirParser object.
         """
-        # TODO _read_compass_solns: Verify this works.
         # When we read in the COMPASS solutions, we will need to map some per-blhid
         # values to per-sphid values, so create an indexing array that we can do this
         # with conveniently.
@@ -2935,6 +2962,26 @@ def _read_compass_solns(self, filename):
 
         return compass_soln_dict
 
+    def read_compass_solns(self, filename=None):
+        """
+        Read in COMPASS-formatted bandpass and flagging solutions.
+
+        Reads in an HDF5 file containing the COMPASS-derived flags and gains tables.
+        These solutions are applied as the data are read in (when calling `load_data`).
+
+        Parameters
+        ----------
+        filename : str
+            Name of the file containing the COMPASS flags and gains solutions.
+
+        Raises
+        ------
+        UserWarning
+            If the COMPASS solutions do not appear to overlap in time with that in
+            the MirParser object.
+        """
+        self._compass_solns = self._read_compass_solns(filename)
+
     def _apply_compass_solns(
         self, compass_soln_dict, vis_data, apply_flags=True, apply_bp=True
     ):
@@ -3523,7 +3570,8 @@ def _make_v3_compliant(self):
 
         warnings.warn(
             "Pre v.3 MIR file format detected, modifying metadata to make in minimally "
-            "compliant for reading in with pyuvdata."
+            "compliant for reading in with pyuvdata. Note that this may cause spurious "
+            "warnings about per-baseline records varying when filling a UVData object."
         )
 
         from datetime import datetime
@@ -3576,26 +3624,33 @@ def _make_v3_compliant(self):
         self.in_data["adec"] = app_dec
 
         # bl_data updates: ant1rx, ant2rx, u, v, w
-        # First, update the antenna receiver fields if this is a non-polarization track,
-        # since that's how we tell X/Y polarization data apart.
-        if np.all(self.bl_data["ipol"] == 0):
-            irec = self.bl_data["irec"]
-
-            # Make sure we recognized all the rx code values, otherwise we may
-            # misidentify RxA vs RxB data.
-            assert np.all(np.isin(irec, [0, 1, 2, 3])), "Forbidden RX code detected."
-            antrx = np.isin(irec, [2, 3]).astype("<i2")
-            self.bl_data["ant1rx"] = antrx
-            self.bl_data["ant2rx"] = antrx
+        # First, update the antenna receiver if these values are unfilled (true in some
+        # earlier tracks, no version demarcation notes it).
+        if np.all(self.bl_data["ant1rx"] == 0) and np.all(self.bl_data["ant2rx"] == 0):
+            ipol = self.bl_data["ipol"]
+            if np.all(ipol == 0):
+                irec = self.bl_data["irec"]
+                # Make sure we recognized all the rx code values, otherwise we may
+                # misidentify RxA vs RxB data.
+                assert np.all(np.isin(irec, [0, 1, 2, 3])), "Bad RX codes detected."
+                antrx = np.isin(irec, [2, 3]).astype("<i2")
+                self.bl_data["ant1rx"] = antrx
+                self.bl_data["ant2rx"] = antrx
+            else:
+                self.bl_data["ant1rx"] = np.isin(ipol, [1, 2])
+                self.bl_data["ant2rx"] = np.isin(ipol, [1, 3])
 
         # Next, the old data had uvws calculated in wavelengths by _sideband_, so we
-        # need to update those.
+        # need to update those. Note we only have to check ant1rx here because if
+        # ant1rx != ant1rx2, this is a pol track, and the tunings (and scaling) are
+        # identical anyways.
         rx_idx = self.bl_data["ant1rx"]
         sb_idx = self.bl_data["isb"]
 
-        # These isb and ant1rx should _only_ be either 0 or 1
-        assert np.all(np.isin(sb_idx, [0, 1])), "Forbidden SB index values detected."
-        assert np.all(np.isin(rx_idx, [0, 1])), "Forbidden RX index values detected."
+        # These isb and ant1rx should _only_ be either 0 or 1. This should never happen
+        # unless the data are corrupt in some substantial way.
+        assert np.all(np.isin(sb_idx, [0, 1])), "Bad SB index values detected."
+        assert np.all(np.isin(rx_idx, [0, 1])), "Bad RX index values detected."
 
         u_vals = self.bl_data["u"]
         v_vals = self.bl_data["v"]

pyuvdata/uvdata/uvdata.py

@@ -10492,6 +10492,7 @@ def read_mir(
         apply_dedoppler=False,
         pseudo_cont=False,
         rechunk=None,
+        compass_soln=None,
         run_check=True,
         check_extra=True,
         run_check_acceptability=True,
@@ -10570,6 +10571,10 @@ def read_mir(
         rechunk : int
             Number of channels to average over when reading in the dataset. Optional
             argument, typically required to be a power of 2.
+        compass_soln : str
+            Optional argument, specifying the path of COMPASS-derived flagging and
+            bandpass gains solutions, which are applied prior to any potential spectral
+            averaging (as triggered by using the `rechunk` keyword).
         run_check : bool
             Option to check for the existence and proper shapes of parameters
             before writing the file.
@@ -10615,6 +10620,7 @@ def read_mir(
             apply_dedoppler=apply_dedoppler,
             pseudo_cont=pseudo_cont,
             rechunk=rechunk,
+            compass_soln=compass_soln,
             run_check=run_check,
             check_extra=check_extra,
             run_check_acceptability=run_check_acceptability,
@@ -11662,6 +11668,7 @@ def read(
         apply_dedoppler=False,
         pseudo_cont=False,
         rechunk=None,
+        compass_soln=None,
         recompute_nbls: bool | None = None,
     ):
         """
@@ -12059,6 +12066,10 @@ def read(
             "flexible polarization", which compresses the polarization-axis of various
             attributes to be of length 1, sets the `flex_spw_polarization_array`
             attribute to define the polarization per spectral window. Default is True.
+        compass_soln : str
+            Optional argument, specifying the path of COMPASS-derived flagging and
+            bandpass gains solutions, which are applied prior to any potential spectral
+            averaging (as triggered by using the `rechunk` keyword).
 
         Raises
         ------
@@ -12245,6 +12256,7 @@ def read(
                         apply_dedoppler=apply_dedoppler,
                         pseudo_cont=pseudo_cont,
                         rechunk=rechunk,
+                        compass_soln=compass_soln,
                         recompute_nbls=recompute_nbls,
                         time_axis_faster_than_bls=time_axis_faster_than_bls,
                         blts_are_rectangular=blts_are_rectangular,
@@ -12364,6 +12376,7 @@ def read(
                             corrchunk=corrchunk,
                             pseudo_cont=pseudo_cont,
                             rechunk=rechunk,
+                            compass_soln=compass_soln,
                         )
 
                         uv_list.append(uv2)
@@ -12614,6 +12627,7 @@ def read(
                     apply_flags=apply_flags,
                     pseudo_cont=pseudo_cont,
                     rechunk=rechunk,
+                    compass_soln=compass_soln,
                     run_check=run_check,
                     check_extra=check_extra,
                     run_check_acceptability=run_check_acceptability,