Making uvcalibrate work w/ wideband gains solns

CHANGELOG.md

@@ -12,6 +12,7 @@ message to inform users when a conjugation flip might be needed.
 raised, and no pfb shape is removed.
 
 ### Added
+- Support for wide-band `UVCal` gain objects with `uvcalibrate`.
 - A colormap option to the UVBeam and AnalyticBeam `plot` methods, along with
 better default colormaps for phase.
 - A NotImplementedError to `mwa_corr_fits.py` that is thrown when trying to read

src/pyuvdata/utils/uvcalibrate.py

@@ -61,8 +61,13 @@ def _get_pol_conventions(
                 f"different values: {uvd_pol_convention} and {uvdata.pol_convention}."
             )
     else:
-        if uvdata.pol_convention is not None:
-            raise ValueError("You are trying to calibrate already-calibrated data.")
+        if uvdata.pol_convention is not None and (
+            uvc_pol_convention != uvdata.pol_convention
+        ):
+            raise ValueError(
+                "UVData already has a pol_convention applied that does not match UVCal "
+                f"convention: {uvdata.pol_convention} and {uvc_pol_convention}."
+            )
         if uvd_pol_convention is None:
             uvd_pol_convention = uvc_pol_convention
 
@@ -149,6 +154,7 @@ def uvcalibrate(
     undo: bool = False,
     time_check: bool = True,
     ant_check: bool = True,
+    freq_range_check: bool = True,
     uvc_pol_convention: Literal["sum", "avg"] | None = None,
     uvd_pol_convention: Literal["sum", "avg"] | None = None,
 ):
@@ -191,6 +197,10 @@ def uvcalibrate(
         object have calibration solutions in the UVCal object. If this option is
         set to False, uvcalibrate will proceed without erroring and data for
         antennas without calibrations will be flagged.
+    freq_range_check : bool
+        Option to check that frequency ranges on the UVCal object matches the channel
+        frequencies given in the UVData object for a given spectral window. Only
+        applicable for wide-band UVCal objects, default is True.
     uvc_pol_convention : str, {"sum", "avg"}, optional
         The convention for how instrumental polarizations (e.g. XX and YY) are assumed
         to have been converted to Stokes parameters in ``uvcal``. Options are 'sum' and
@@ -210,10 +220,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 +325,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 +408,49 @@ 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
+    uvcal_chans_to_keep = None if uvcal.wide_band else []
+    uvcal_spws_to_keep = [] if uvcal.wide_band else None
+    for spw in uvdata.spw_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_spw_idx = np.where(uvcal.spw_array == spw)[0][0]
+            if uvcal.wide_band:
+                if freq_range_check:
+                    freq_array_spw = uvdata.freq_array[uvdata.flex_spw_id_array == spw]
+                    min_freq = min(uvcal.freq_range[uvcal_spw_idx])
+                    max_freq = max(uvcal.freq_range[uvcal_spw_idx])
+                    if any((freq_array_spw < min_freq) | (freq_array_spw > max_freq)):
+                        raise ValueError(
+                            f"SPW {spw} exists on UVData and UVCal, but the channel "
+                            "frequencies are inconsistent with frequency ranges. "
+                            "To continue with calibration, set freq_range_check=False."
+                        )
+                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."
+                        )
+        except IndexError as err:
+            raise ValueError(f"SPW {spw} exists on UVData but not on UVCal.") from err
+
+    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 +480,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 +536,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 +585,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."

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
@@ -97,14 +97,14 @@ def test_uvd_uvdata_different(self):
 
     def test_calibrate_already_calibrated(self):
         with pytest.raises(
-            ValueError, match="You are trying to calibrate already-calibrated data"
+            ValueError, match="UVData already has a pol_convention applied"
         ):
             _get_pol_conventions(
-                uvdata=SimpleNamespace(pol_convention="avg"),
+                uvdata=SimpleNamespace(pol_convention="sum"),
                 uvcal=SimpleNamespace(pol_convention="avg"),
                 undo=False,
                 uvc_pol_convention="avg",
-                uvd_pol_convention="avg",
+                uvd_pol_convention="sum",
             )
 
     @pytest.mark.parametrize("which", ["uvc", "uvd"])
@@ -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,79 @@ 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():
+    pytest.importorskip("casacore")
+
+    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[0] = uvd.time_array[0]
+
+    with pytest.raises(ValueError, match="SPW -4 exists on UVData but not on UVCal"):
+        uvcalibrate(uvd, uvc)
+
+    # Now modify the SPW array and check that freq misalignment throws an error
+    uvc.spw_array = np.array([-6, -5, -4, -3, -2, -1, 1, 2, 3, 4, 5, 6])
+    with pytest.raises(
+        ValueError,
+        match="SPW -4 exists on UVData and UVCal, but the channel frequencies",
+    ):
+        uvcalibrate(uvd, uvc)
+
+    # Now check that now error gets thrown when the freq range check is turned off
+    uvcalibrate(uvd, uvc, freq_range_check=False)
+
+
+@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():
+    pytest.importorskip("casacore")
+
+    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)