Making uvcalibrate work w/ wideband gains solns

Author: kartographer

src/pyuvdata/utils/uvcalibrate.py

@@ -210,10 +210,6 @@ def uvcalibrate(
         Returns if not inplace
 
     """
-    if uvcal.cal_type == "gain" and uvcal.wide_band:
-        raise ValueError(
-            "uvcalibrate currently does not support wide-band calibrations"
-        )
     if uvcal.cal_type == "delay" and uvcal.Nspws > 1:
         # To fix this, need to make UVCal.convert_to_gain support multiple spws
         raise ValueError(
@@ -319,7 +315,7 @@ def uvcalibrate(
                     )
 
     uvdata_times, uvd_time_ri = np.unique(uvdata.time_array, return_inverse=True)
-    downselect_cal_times = False
+    uvcal_times_to_keep = None
     # time_range supercedes time_array.
     if uvcal.time_range is not None:
         if np.min(uvdata_times) < np.min(uvcal.time_range[:, 0]) or np.max(
@@ -402,43 +398,44 @@ def uvcalibrate(
                         raise ValueError(
                             f"Time {this_time} exists on UVData but not on UVCal."
                         )
-                if len(uvcal_times_to_keep) < uvcal.Ntimes:
-                    downselect_cal_times = True
 
-    downselect_cal_freq = False
-    if uvcal.freq_array is not None:
-        uvdata_freq_arr_use = uvdata.freq_array
-        uvcal_freq_arr_use = uvcal.freq_array
-        try:
-            freq_arr_match = np.allclose(
-                np.sort(uvcal_freq_arr_use),
-                np.sort(uvdata_freq_arr_use),
-                atol=uvdata._freq_array.tols[1],
-                rtol=uvdata._freq_array.tols[0],
-            )
-        except ValueError:
-            freq_arr_match = False
-
-        if freq_arr_match is False:
-            # check more carefully
-            uvcal_freqs_to_keep = []
-            for this_freq in uvdata_freq_arr_use:
-                wh_freq_match = np.nonzero(
-                    np.isclose(
-                        uvcal.freq_array - this_freq,
-                        0,
+    uvcal_chans_to_keep = None if uvcal.wide_band else []
+    uvcal_spws_to_keep = [] if uvcal.wide_band else None
+    for idx, spw in enumerate(uvcal.spw_array):
+        if spw in uvdata.spw_array:
+            if uvcal.wide_band:
+                freq_array_spw = uvdata.freq_array[uvdata.flex_spw_id_array == spw]
+                min_freq = min(uvcal.freq_range[idx])
+                max_freq = max(uvcal.freq_range[idx])
+                if all((freq_array_spw < min_freq) | (freq_array_spw > max_freq)):
+                    raise ValueError(
+                        "SPWs between UVData and UVCal overlap, but the frequencies of "
+                        "channels are inconsistent with frequency ranges."
+                    )
+                uvcal_spws_to_keep.append(spw)
+            else:
+                uvcal_chans = np.where(uvcal.flex_spw_id_array == spw)[0]
+                uvdata_chans = np.where(uvdata.flex_spw_id_array == spw)[0]
+                uvcal_freqs = uvcal.freq_array[uvcal_chans]
+                uvdata_freqs = uvdata.freq_array[uvdata_chans]
+                for indv_freq in uvdata_freqs:
+                    freq_match = np.isclose(
+                        uvcal_freqs,
+                        indv_freq,
                         atol=uvdata._freq_array.tols[1],
                         rtol=uvdata._freq_array.tols[0],
                     )
-                )
-                if wh_freq_match[0].size > 0:
-                    uvcal_freqs_to_keep.append(uvcal.freq_array[wh_freq_match][0])
-                else:
-                    raise ValueError(
-                        f"Frequency {this_freq} exists on UVData but not on UVCal."
-                    )
-            if len(uvcal_freqs_to_keep) < uvcal.Nfreqs:
-                downselect_cal_freq = True
+                    if any(freq_match):
+                        uvcal_chans_to_keep.append(uvcal_chans[freq_match][0])
+                    else:
+                        raise ValueError(
+                            f"Frequency {indv_freq} exists on UVData but not on UVCal."
+                        )
+
+    if np.array_equal(uvcal_chans_to_keep, np.arange(uvcal.Nfreqs)):
+        uvcal_chans_to_keep = None
+    if np.array_equal(uvcal_spws_to_keep, uvcal.spw_array):
+        uvcal_spws_to_keep = None
 
     # check if x_orientation-equivalent in uvdata isn't set (it's required for uvcal)
     uvd_x = uvdata.telescope.get_x_orientation_from_feeds()
@@ -468,21 +465,20 @@ def uvcalibrate(
             )
 
     # downselect UVCal times, frequencies
-    if downselect_cal_freq or downselect_cal_times:
-        if not downselect_cal_times:
-            uvcal_times_to_keep = None
-        elif not downselect_cal_freq:
-            uvcal_freqs_to_keep = None
-
+    new_uvcal = not (
+        uvcal_times_to_keep is None
+        and uvcal_chans_to_keep is None
+        and uvcal_spws_to_keep is None
+    )
+    uvcal_use = uvcal
+    if new_uvcal:
         uvcal_use = uvcal.select(
-            times=uvcal_times_to_keep, frequencies=uvcal_freqs_to_keep, inplace=False
+            times=uvcal_times_to_keep,
+            freq_chans=uvcal_chans_to_keep,
+            spws=uvcal_spws_to_keep,
+            inplace=False,
         )
 
-        new_uvcal = True
-    else:
-        uvcal_use = uvcal
-        new_uvcal = False
-
     # input checks
     if uvcal_use.cal_type == "delay":
         if not new_uvcal:
@@ -525,6 +521,7 @@ def uvcalibrate(
                 strict=False,
             )
         )
+        uvc_spw_map = {spw: idx for idx, spw in enumerate(uvcal_use.spw_array)}
 
         # iterate over keys
         for key in uvdata.get_antpairpols():
@@ -573,20 +570,29 @@ def uvcalibrate(
                 gain = gain[trange_ind_arr[blt_inds], :]
                 flag = flag[trange_ind_arr[blt_inds], :]
 
+            # Use a slice operator to expand out the flags and gains with a wide_band
+            # calibration solution, otherwise use the Ellipsis to select the whole
+            # array when using a channel-based soln.
+            gain_slice = ...
+            if uvcal_use.wide_band:
+                gain_slice = np.s_[
+                    :, [uvc_spw_map[spw] for spw in uvdata.flex_spw_id_array]
+                ]
+
             # propagate flags
             if prop_flags:
                 mask = np.isclose(gain, 0.0) | flag
                 gain[mask] = 1.0
-                uvdata.flag_array[blt_inds, :, pol_ind] += mask
+                uvdata.flag_array[blt_inds, :, pol_ind] |= mask[gain_slice]
 
             # apply to data
             mult_gains = uvcal_use.gain_convention == "multiply"
             if undo:
                 mult_gains = not mult_gains
             if mult_gains:
-                uvdata.data_array[blt_inds, :, pol_ind] *= gain
+                uvdata.data_array[blt_inds, :, pol_ind] *= gain[gain_slice]
             else:
-                uvdata.data_array[blt_inds, :, pol_ind] /= gain
+                uvdata.data_array[blt_inds, :, pol_ind] /= gain[gain_slice]
 
     # update attributes
     uvdata.history += "\nCalibrated with pyuvdata.utils.uvcalibrate."

src/pyuvdata/uvdata/mir.py

@@ -714,7 +714,7 @@ def _init_from_mir_parser(
         self.uvw_array = (-1.0) * uvw_array
 
         self.vis_units = "Jy"
-        self.pol_convention = "avg"
+        # self.pol_convention = "avg"
 
         isource = np.unique(mir_data.in_data["isource"])
         for sou_id in isource:

tests/utils/test_uvcalibrate.py

@@ -8,7 +8,7 @@
 import numpy as np
 import pytest
 
-from pyuvdata import UVCal, utils
+from pyuvdata import UVCal, UVData, utils
 from pyuvdata.datasets import fetch_data
 from pyuvdata.testing import check_warnings
 from pyuvdata.utils import uvcalibrate
@@ -698,35 +698,6 @@ def test_uvcalibrate_x_orientation_mismatch(uvcalibrate_data):
         uvcalibrate(uvd, uvc, inplace=False)
 
 
-def test_uvcalibrate_wideband_gain(uvcalibrate_data):
-    uvd, uvc = uvcalibrate_data
-
-    uvc.flex_spw_id_array = None
-    uvc._set_wide_band()
-    uvc.spw_array = np.array([1, 2, 3])
-    uvc.Nspws = 3
-    uvc.gain_array = uvc.gain_array[:, 0:3, :, :]
-    uvc.flag_array = uvc.flag_array[:, 0:3, :, :]
-    uvc.quality_array = uvc.quality_array[:, 0:3, :, :]
-    uvc.total_quality_array = uvc.total_quality_array[0:3, :, :]
-
-    uvc.freq_range = np.zeros((uvc.Nspws, 2), dtype=uvc.freq_array.dtype)
-    uvc.freq_range[0, :] = uvc.freq_array[[0, 2]]
-    uvc.freq_range[1, :] = uvc.freq_array[[2, 4]]
-    uvc.freq_range[2, :] = uvc.freq_array[[4, 6]]
-
-    uvc.channel_width = None
-    uvc.freq_array = None
-    uvc.Nfreqs = 1
-
-    uvc.check()
-    with pytest.raises(
-        ValueError,
-        match="uvcalibrate currently does not support wide-band calibrations",
-    ):
-        uvcalibrate(uvd, uvc, inplace=False)
-
-
 @pytest.mark.filterwarnings("ignore:Fixing auto-correlations to be be real-only")
 @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values")
 @pytest.mark.filterwarnings("ignore:Nfreqs will be required to be 1 for wide_band cals")
@@ -892,3 +863,64 @@ def test_uvdata_pol_array_in_stokes(uvcalibrate_data):
         ),
     ):
         uvcalibrate(uvd, uvc)
+
+
+@pytest.mark.filterwarnings("ignore:> 25 ms errors detected reading in LST values")
+@pytest.mark.filterwarnings("ignore:The lst_array is not self-consistent")
+@pytest.mark.filterwarnings("ignore:Setting telescope_location to value in known")
+@pytest.mark.filterwarnings("ignore:Unknown polarization basis for solutions")
+@pytest.mark.filterwarnings("ignore: Unknown x_orientation basis for solutions")
+def test_uvcalibrate_wideband_gains_errs():
+    uvd = UVData.from_file(fetch_data("sma_mir"))
+    uvc = UVCal.from_file(fetch_data("sma_pha_gcal"))
+
+    # Condition the UVCal object so as to avoid additional warnings
+    uvc.gain_scale = "Jy"
+    uvc.pol_convention = "avg"
+    uvc.time_array[:] = uvd.time_array[0]
+
+    with pytest.raises(ValueError, match="SPWs between UVData and UVCal overlap"):
+        uvcalibrate(uvd, uvc)
+
+
+@pytest.mark.filterwarnings("ignore:> 25 ms errors detected reading in LST values")
+@pytest.mark.filterwarnings("ignore:The lst_array is not self-consistent")
+@pytest.mark.filterwarnings("ignore:Setting telescope_location to value in known")
+@pytest.mark.filterwarnings("ignore:Unknown polarization basis for solutions")
+@pytest.mark.filterwarnings("ignore: Unknown x_orientation basis for solutions")
+def test_uvcalibrate_wideband_gains():
+    uvd = UVData.from_file(fetch_data("sma_mir"))
+    uvc = UVCal.from_file(fetch_data("sma_pha_gcal"))
+
+    # Condition the UVCal object so as to avoid additional warnings
+    uvc.gain_scale = "Jy"
+    uvc.pol_convention = "avg"
+
+    # Set up gains so that the values change based on SPW number
+    uvc.gain_array[:, :, 0, :] = uvc.spw_array[None, :, None]
+
+    # Flag every other window to test flagging cascading
+    uvc.flag_array[:, ::2, 0, :] = True
+
+    # Get one timestamp to align
+    uvd.time_array[:] = uvc.time_array[0]
+
+    # Make it easier to compare things by setting everything to one
+    uvd.data_array[:] = 1.0
+
+    uvc.spw_array = np.array([-6, -5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 6])
+    uvc.freq_range[2:-2, 0] = np.min(uvd.freq_array.reshape(8, -1), axis=1)
+    uvc.freq_range[2:-2, 1] = np.max(uvd.freq_array.reshape(8, -1), axis=1)
+
+    new_uvd = uvcalibrate(uvdata=uvd, uvcal=uvc, inplace=False)
+
+    # Make sure no flagged solns got applied
+    assert np.all(new_uvd.data_array.reshape(8, 16384, 2)[::2] == 1)
+    # Confirm all the bad things are flagged
+    assert np.all(new_uvd.flag_array.reshape(8, 16384, 2)[::2])
+
+    # Calculate the expected data based on spw number
+    exp_data = np.repeat((uvd.spw_array**2)[:, None], 16384 * 2).reshape(1, -1, 2)
+
+    # Check that either things agree with expectations or that data are flagged
+    assert np.all((exp_data == new_uvd.data_array) | new_uvd.flag_array)