diff --git a/pyuvdata/uvdata/tests/test_uvdata.py b/pyuvdata/uvdata/tests/test_uvdata.py index 7c739deefb..63ca01c9a5 100644 --- a/pyuvdata/uvdata/tests/test_uvdata.py +++ b/pyuvdata/uvdata/tests/test_uvdata.py @@ -9423,7 +9423,8 @@ def test_resample_in_time_warning(): @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values") @pytest.mark.parametrize("future_shapes", [True, False]) @pytest.mark.parametrize("flex_spw", [True, False]) -def test_frequency_average(casa_uvfits, future_shapes, flex_spw): +@pytest.mark.parametrize("sum_corr", [True, False]) +def test_frequency_average(casa_uvfits, future_shapes, flex_spw, sum_corr): """Test averaging in frequency.""" uvobj = casa_uvfits @@ -9450,24 +9451,28 @@ def test_frequency_average(casa_uvfits, future_shapes, flex_spw): np.arange(uvobj.Nfreqs, dtype=np.float64), (uvobj.Nants_telescope, 1) ) uvobj.eq_coeffs = eq_coeffs - uvobj.check() # check that there's no flagging assert np.nonzero(uvobj.flag_array)[0].size == 0 + n_chan_to_avg = 2 with uvtest.check_warnings(UserWarning, "eq_coeffs vary by frequency"): - uvobj.frequency_average(2, keep_ragged=True), + uvobj.frequency_average(2, keep_ragged=True, summing_correlator_mode=sum_corr) - assert uvobj.Nfreqs == (uvobj2.Nfreqs / 2) + assert uvobj.Nfreqs == (uvobj2.Nfreqs / n_chan_to_avg) + + input_freqs = np.squeeze(uvobj2.freq_array) + + expected_freqs = input_freqs[ + np.arange((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) + ] + expected_freqs = expected_freqs.reshape( + int(uvobj2.Nfreqs // n_chan_to_avg), n_chan_to_avg + ).mean(axis=1) + + if not future_shapes: + expected_freqs = expected_freqs[np.newaxis, :] - if future_shapes: - expected_freqs = uvobj2.freq_array.reshape(int(uvobj2.Nfreqs / 2), 2).mean( - axis=1 - ) - else: - expected_freqs = uvobj2.freq_array.reshape(1, int(uvobj2.Nfreqs / 2), 2).mean( - axis=2 - ) assert np.max(np.abs(uvobj.freq_array - expected_freqs)) == 0 expected_coeffs = eq_coeffs.reshape( @@ -9476,25 +9481,16 @@ def test_frequency_average(casa_uvfits, future_shapes, flex_spw): assert np.max(np.abs(uvobj.eq_coeffs - expected_coeffs)) == 0 # no flagging, so the following is true - expected_data = uvobj2.get_data(1, 2, squeeze="none") - if future_shapes: - reshape_tuple = ( - expected_data.shape[0], - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, - ) + expected_data = uvobj2.get_data(1, 2) + reshape_tuple = (expected_data.shape[0], int(uvobj2.Nfreqs / 2), 2, uvobj2.Npols) + if sum_corr: + expected_data = expected_data.reshape(reshape_tuple).sum(axis=2) + else: expected_data = expected_data.reshape(reshape_tuple).mean(axis=2) - else: - reshape_tuple = ( - expected_data.shape[0], - 1, - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, - ) - expected_data = expected_data.reshape(reshape_tuple).mean(axis=3) + if not future_shapes: + expected_data = expected_data[:, np.newaxis] + assert np.allclose(uvobj.get_data(1, 2, squeeze="none"), expected_data) assert np.nonzero(uvobj.flag_array)[0].size == 0 @@ -9505,55 +9501,112 @@ def test_frequency_average(casa_uvfits, future_shapes, flex_spw): @pytest.mark.filterwarnings("ignore:This method will be removed in version 3.0 when") @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values") -@pytest.mark.parametrize("future_shapes", [True, False]) +@pytest.mark.parametrize(["future_shapes", "sum_corr"], [[True, True], [False, False]]) +@pytest.mark.parametrize( + ["flex_spw", "respect_spws"], [[True, True], [True, False], [False, False]] +) @pytest.mark.parametrize("keep_ragged", [True, False]) -@pytest.mark.parametrize("sum_corr", [True, False]) -def test_frequency_average_uneven(casa_uvfits, future_shapes, keep_ragged, sum_corr): +def test_frequency_average_uneven( + casa_uvfits, future_shapes, keep_ragged, sum_corr, flex_spw, respect_spws +): """Test averaging in frequency with a number that is not a factor of Nfreqs.""" uvobj = casa_uvfits if not future_shapes: uvobj.use_current_array_shapes() + + eq_coeffs = np.tile( + np.arange(uvobj.Nfreqs, dtype=np.float64), (uvobj.Nants_telescope, 1) + ) + uvobj.eq_coeffs = eq_coeffs + + if flex_spw: + # Make multiple spws + spw_nchan = int(uvobj.Nfreqs / 4) + uvobj.flex_spw_id_array = np.concatenate( + ( + np.full(spw_nchan, 0, dtype=int), + np.full(spw_nchan, 1, dtype=int), + np.full(spw_nchan, 2, dtype=int), + np.full(spw_nchan, 3, dtype=int), + ) + ) + uvobj.spw_array = np.arange(4) + uvobj.Nspws = 4 + uvobj2 = uvobj.copy() # check that there's no flagging assert np.nonzero(uvobj.flag_array)[0].size == 0 n_chan_to_avg = 7 + warn = [UserWarning] + msg = [ + re.escape( + "eq_coeffs vary by frequency. They should be applied to the data using " + "`remove_eq_coeffs` before frequency averaging" + ) + ] if keep_ragged and not future_shapes: - warn = UserWarning - msg = [ + warn.append(UserWarning) + msg.append( "Ragged frequencies resulted in varying channel widths, so " "this object must have future array shapes after averaging. " "Use keep_ragged=False to avoid this." - ] - else: - warn = None - msg = "" + ) with uvtest.check_warnings(warn, match=msg): uvobj.frequency_average( - n_chan_to_avg, keep_ragged=keep_ragged, summing_correlator_mode=sum_corr + n_chan_to_avg, + keep_ragged=keep_ragged, + summing_correlator_mode=sum_corr, + respect_spws=respect_spws, ) assert uvobj2.Nfreqs % n_chan_to_avg != 0 input_freqs = np.squeeze(uvobj2.freq_array) - expected_freqs = input_freqs[ - np.arange((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) - ] - expected_freqs = expected_freqs.reshape( - int(uvobj2.Nfreqs // n_chan_to_avg), n_chan_to_avg - ).mean(axis=1) - - if keep_ragged: - assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg + 1) - expected_freqs = np.append( - expected_freqs, - np.mean(input_freqs[(uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg :]), - ) + if flex_spw and respect_spws: + expected_freqs = np.array([]) + for spw_ind in range(uvobj2.Nspws): + start_chan = spw_ind * spw_nchan + n_reg_chan = (spw_nchan // n_chan_to_avg) * n_chan_to_avg + this_expected = input_freqs[start_chan : (start_chan + n_reg_chan)] + this_expected = this_expected.reshape( + int(spw_nchan // n_chan_to_avg), n_chan_to_avg + ).mean(axis=1) + if keep_ragged: + this_expected = np.append( + this_expected, + np.mean( + input_freqs[ + (start_chan + n_reg_chan) : (spw_ind + 1) * spw_nchan + ] + ), + ) + expected_freqs = np.append(expected_freqs, this_expected) + if keep_ragged: + assert uvobj.Nfreqs == (spw_nchan // n_chan_to_avg + 1) * uvobj2.Nspws + else: + assert uvobj.Nfreqs == (spw_nchan // n_chan_to_avg) * uvobj2.Nspws else: - assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg) + expected_freqs = input_freqs[ + np.arange((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) + ] + expected_freqs = expected_freqs.reshape( + int(uvobj2.Nfreqs // n_chan_to_avg), n_chan_to_avg + ).mean(axis=1) + + if keep_ragged: + assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg + 1) + expected_freqs = np.append( + expected_freqs, + np.mean( + input_freqs[(uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg :] + ), + ) + else: + assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg) if not future_shapes and not keep_ragged: expected_freqs = expected_freqs[np.newaxis, :] @@ -9562,37 +9615,76 @@ def test_frequency_average_uneven(casa_uvfits, future_shapes, keep_ragged, sum_c # no flagging, so the following is true initial_data = uvobj2.get_data(1, 2) - expected_data = initial_data[ - :, 0 : ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) - ] - reshape_tuple = ( - initial_data.shape[0], - int(uvobj2.Nfreqs // n_chan_to_avg), - n_chan_to_avg, - uvobj2.Npols, - ) - if sum_corr: - expected_data = expected_data.reshape(reshape_tuple).sum(axis=2) - else: - expected_data = expected_data.reshape(reshape_tuple).mean(axis=2) + if flex_spw and respect_spws: + reshape_tuple = ( + initial_data.shape[0], + int(spw_nchan // n_chan_to_avg), + n_chan_to_avg, + uvobj2.Npols, + ) + for spw_ind in range(uvobj2.Nspws): + start_chan = spw_ind * spw_nchan + n_reg_chan = (spw_nchan // n_chan_to_avg) * n_chan_to_avg - if keep_ragged: + this_expected = initial_data[:, start_chan : (start_chan + n_reg_chan)] + if sum_corr: + this_expected = this_expected.reshape(reshape_tuple).sum(axis=2) + else: + this_expected = this_expected.reshape(reshape_tuple).mean(axis=2) + + if keep_ragged: + if sum_corr: + this_expected = np.append( + this_expected, + initial_data[ + :, (start_chan + n_reg_chan) : (spw_ind + 1) * spw_nchan + ].sum(axis=1, keepdims=True), + axis=1, + ) + else: + this_expected = np.append( + this_expected, + initial_data[ + :, (start_chan + n_reg_chan) : (spw_ind + 1) * spw_nchan + ].mean(axis=1, keepdims=True), + axis=1, + ) + if spw_ind == 0: + expected_data = this_expected + else: + expected_data = np.append(expected_data, this_expected, axis=1) + else: + expected_data = initial_data[ + :, 0 : ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) + ] + reshape_tuple = ( + initial_data.shape[0], + int(uvobj2.Nfreqs // n_chan_to_avg), + n_chan_to_avg, + uvobj2.Npols, + ) if sum_corr: - expected_data = np.append( - expected_data, - initial_data[ - :, ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) : - ].sum(axis=1, keepdims=True), - axis=1, - ) + expected_data = expected_data.reshape(reshape_tuple).sum(axis=2) else: - expected_data = np.append( - expected_data, - initial_data[ - :, ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) : - ].mean(axis=1, keepdims=True), - axis=1, - ) + expected_data = expected_data.reshape(reshape_tuple).mean(axis=2) + + if keep_ragged: + if sum_corr: + expected_data = np.append( + expected_data, + initial_data[ + :, ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) : + ].sum(axis=1, keepdims=True), + axis=1, + ) + else: + expected_data = np.append( + expected_data, + initial_data[ + :, ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) : + ].mean(axis=1, keepdims=True), + axis=1, + ) if not future_shapes and not keep_ragged: expected_data = expected_data[:, np.newaxis] @@ -9782,61 +9874,8 @@ def test_frequency_average_flagging_partial_twostage(casa_uvfits): @pytest.mark.filterwarnings("ignore:This method will be removed in version 3.0 when") @pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values") @pytest.mark.parametrize("future_shapes", [True, False]) -def test_frequency_average_summing_corr_mode(casa_uvfits, future_shapes): - """Test averaging in frequency.""" - # check that there's no flagging - uvobj = casa_uvfits - - if not future_shapes: - uvobj.use_current_array_shapes() - uvobj2 = uvobj.copy() - - assert np.nonzero(uvobj.flag_array)[0].size == 0 - - uvobj.frequency_average(2, summing_correlator_mode=True, keep_ragged=True) - - assert uvobj.Nfreqs == (uvobj2.Nfreqs / 2) - - if future_shapes: - expected_freqs = uvobj2.freq_array.reshape(int(uvobj2.Nfreqs / 2), 2).mean( - axis=1 - ) - else: - expected_freqs = uvobj2.freq_array.reshape(1, int(uvobj2.Nfreqs / 2), 2).mean( - axis=2 - ) - assert np.max(np.abs(uvobj.freq_array - expected_freqs)) == 0 - - # no flagging, so the following is true - expected_data = uvobj2.get_data(1, 2, squeeze="none") - if future_shapes: - reshape_tuple = ( - expected_data.shape[0], - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, - ) - expected_data = expected_data.reshape(reshape_tuple).sum(axis=2) - else: - reshape_tuple = ( - expected_data.shape[0], - 1, - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, - ) - expected_data = expected_data.reshape(reshape_tuple).sum(axis=3) - assert np.allclose(uvobj.get_data(1, 2, squeeze="none"), expected_data) - - assert np.nonzero(uvobj.flag_array)[0].size == 0 - assert not isinstance(uvobj.data_array, np.ma.MaskedArray) - assert not isinstance(uvobj.nsample_array, np.ma.MaskedArray) - - -@pytest.mark.filterwarnings("ignore:This method will be removed in version 3.0 when") -@pytest.mark.filterwarnings("ignore:The uvw_array does not match the expected values") -@pytest.mark.parametrize("future_shapes", [True, False]) -def test_frequency_average_propagate_flags(casa_uvfits, future_shapes): +@pytest.mark.parametrize("keep_ragged", [True, False]) +def test_frequency_average_propagate_flags(casa_uvfits, future_shapes, keep_ragged): """ Test averaging in frequency with flagging all of one and only one of another sample averaged, and propagating flags. Data should be identical, @@ -9852,51 +9891,86 @@ def test_frequency_average_propagate_flags(casa_uvfits, future_shapes): # check that there's no flagging assert np.nonzero(uvobj.flag_array)[0].size == 0 + n_chan_to_avg = 3 + # apply some flagging for testing inds01 = uvobj.antpair2ind(1, 2) if future_shapes: - uvobj.flag_array[inds01[0], 0:3, :] = True + uvobj.flag_array[inds01[0], 0 : (n_chan_to_avg * 2 - 1), :] = True else: - uvobj.flag_array[inds01[0], :, 0:3, :] = True + uvobj.flag_array[inds01[0], :, 0 : (n_chan_to_avg * 2 - 1), :] = True - assert np.nonzero(uvobj.flag_array)[0].size == uvobj.Npols * 3 + assert np.nonzero(uvobj.flag_array)[0].size == uvobj.Npols * (n_chan_to_avg * 2 - 1) - uvobj.frequency_average(2, propagate_flags=True, keep_ragged=True) + if keep_ragged and not future_shapes: + warn = UserWarning + msg = [ + "Ragged frequencies resulted in varying channel widths, so " + "this object must have future array shapes after averaging. " + "Use keep_ragged=False to avoid this." + ] + else: + warn = None + msg = "" + with uvtest.check_warnings(warn, match=msg): + uvobj.frequency_average( + n_chan_to_avg, propagate_flags=True, keep_ragged=keep_ragged + ) - assert uvobj.Nfreqs == (uvobj2.Nfreqs / 2) + input_freqs = np.squeeze(uvobj2.freq_array) - if future_shapes: - expected_freqs = uvobj2.freq_array.reshape(int(uvobj2.Nfreqs / 2), 2).mean( - axis=1 + expected_freqs = input_freqs[ + np.arange((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) + ] + expected_freqs = expected_freqs.reshape( + int(uvobj2.Nfreqs // n_chan_to_avg), n_chan_to_avg + ).mean(axis=1) + + if keep_ragged: + assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg + 1) + expected_freqs = np.append( + expected_freqs, + np.mean(input_freqs[(uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg :]), ) else: - expected_freqs = uvobj2.freq_array.reshape(1, int(uvobj2.Nfreqs / 2), 2).mean( - axis=2 - ) + assert uvobj.Nfreqs == (uvobj2.Nfreqs // n_chan_to_avg) + + if not future_shapes and not keep_ragged: + expected_freqs = expected_freqs[np.newaxis, :] + assert np.max(np.abs(uvobj.freq_array - expected_freqs)) == 0 - expected_data = uvobj2.get_data(1, 2, squeeze="none") - if future_shapes: - reshape_tuple = ( - expected_data.shape[0], - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, - ) - expected_data = expected_data.reshape(reshape_tuple).mean(axis=2) + initial_data = uvobj2.get_data(1, 2) + expected_data = initial_data[ + :, 0 : ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) + ] + reshape_tuple = ( + initial_data.shape[0], + int(uvobj2.Nfreqs // n_chan_to_avg), + n_chan_to_avg, + uvobj2.Npols, + ) + expected_data = expected_data.reshape(reshape_tuple).mean(axis=2) - expected_data[0, 1, :] = uvobj2.data_array[inds01[0], 3, :] + if future_shapes: + freq_axis = 0 else: - reshape_tuple = ( - expected_data.shape[0], - 1, - int(uvobj2.Nfreqs / 2), - 2, - uvobj2.Npols, + freq_axis = 1 + expected_data[0, 1, :] = np.take( + uvobj2.data_array[inds01[0]], n_chan_to_avg * 2 - 1, axis=freq_axis + ) + + if keep_ragged: + expected_data = np.append( + expected_data, + initial_data[:, ((uvobj2.Nfreqs // n_chan_to_avg) * n_chan_to_avg) :].mean( + axis=1, keepdims=True + ), + axis=1, ) - expected_data = expected_data.reshape(reshape_tuple).mean(axis=3) - expected_data[0, :, 1, :] = uvobj2.data_array[inds01[0], :, 3, :] + if not future_shapes and not keep_ragged: + expected_data = expected_data[:, np.newaxis] assert np.allclose(uvobj.get_data(1, 2, squeeze="none"), expected_data) # Twice as many flags should exist compared to test of previous name. diff --git a/pyuvdata/uvdata/uvdata.py b/pyuvdata/uvdata/uvdata.py index bc7903ccee..9285dff4fa 100644 --- a/pyuvdata/uvdata/uvdata.py +++ b/pyuvdata/uvdata/uvdata.py @@ -10562,7 +10562,7 @@ def frequency_average( ) if this_ragged: final_eq_coeffs[:, final_spw_chans[spw][-1]] = np.mean( - self.eq_coeffs[irregular_inds], axis=1 + self.eq_coeffs[:, irregular_inds], axis=1 ) if not self.metadata_only: