diff --git a/src/pyuvdata/uvbeam/uvbeam.py b/src/pyuvdata/uvbeam/uvbeam.py index 6ae2a5570..719f8b65f 100644 --- a/src/pyuvdata/uvbeam/uvbeam.py +++ b/src/pyuvdata/uvbeam/uvbeam.py @@ -844,12 +844,12 @@ def check( # Check if the interpolation points are evenly-spaced if self.pixel_coordinate_system == "az_za": - for i, ax in enumerate((self.axis1_array, self.axis2_array)): + for i, ax_param in enumerate((self._axis1_array, self._axis2_array)): + ax = ax_param.value if len(ax) < 3: continue - diff = np.diff(ax) - if not np.allclose(diff, diff[0]): + if not utils.tools._test_array_constant_spacing(ax, tols=ax_param.tols): raise ValueError( f"axis{i+1}_array must be evenly spaced in az_za coordinates." ) @@ -3159,18 +3159,17 @@ def select( axis1_inds = sorted(set(axis1_inds)) if min(axis1_inds) < 0 or max(axis1_inds) > beam_object.Naxes1 - 1: raise ValueError("axis1_inds must be > 0 and < Naxes1") - beam_object.Naxes1 = len(axis1_inds) - beam_object.axis1_array = beam_object.axis1_array[axis1_inds] - if beam_object.Naxes1 > 1 and not utils.tools._test_array_constant_spacing( - beam_object._axis1_array + if len(axis1_inds) > 1 and not utils.tools._test_array_constant_spacing( + beam_object.axis1_array[axis1_inds], tols=beam_object._axis1_array.tols ): - warnings.warn( - "Selected values along first image axis are " - "not evenly spaced. This is not supported by " - "the regularly gridded beam fits format" + raise ValueError( + "Selected values along first image axis must be evenly spaced." ) + beam_object.Naxes1 = len(axis1_inds) + beam_object.axis1_array = beam_object.axis1_array[axis1_inds] + beam_object.data_array = beam_object.data_array[..., axis1_inds] if beam_object.beam_type == "efield": beam_object.basis_vector_array = beam_object.basis_vector_array[ @@ -3192,18 +3191,17 @@ def select( axis2_inds = sorted(set(axis2_inds)) if min(axis2_inds) < 0 or max(axis2_inds) > beam_object.Naxes2 - 1: raise ValueError("axis2_inds must be > 0 and < Naxes2") - beam_object.Naxes2 = len(axis2_inds) - beam_object.axis2_array = beam_object.axis2_array[axis2_inds] - if beam_object.Naxes2 > 1 and not utils.tools._test_array_constant_spacing( - beam_object._axis2_array + if len(axis2_inds) > 1 and not utils.tools._test_array_constant_spacing( + beam_object.axis2_array[axis2_inds], tols=beam_object._axis2_array.tols ): - warnings.warn( - "Selected values along second image axis are " - "not evenly spaced. This is not supported by " - "the regularly gridded beam fits format" + raise ValueError( + "Selected values along second image axis must be evenly spaced." ) + beam_object.Naxes2 = len(axis2_inds) + beam_object.axis2_array = beam_object.axis2_array[axis2_inds] + beam_object.data_array = beam_object.data_array[..., axis2_inds, :] if beam_object.beam_type == "efield": beam_object.basis_vector_array = beam_object.basis_vector_array[ diff --git a/tests/uvbeam/test_uvbeam.py b/tests/uvbeam/test_uvbeam.py index cdf5f3048..aa292443f 100644 --- a/tests/uvbeam/test_uvbeam.py +++ b/tests/uvbeam/test_uvbeam.py @@ -1692,7 +1692,10 @@ def test_select_axis(cst_power_1freq, tmp_path): # check for warnings and errors associated with unevenly spaced image pixels power_beam2 = power_beam.copy() - with pytest.raises(ValueError, match="axis1_array must be evenly spaced"): + with pytest.raises( + ValueError, + match="Selected values along first image axis must be evenly spaced.", + ): power_beam2.select(axis1_inds=[0, 5, 6], inplace=False) # Test selecting on axis2 @@ -1726,7 +1729,10 @@ def test_select_axis(cst_power_1freq, tmp_path): # check for warnings and errors associated with unevenly spaced image pixels power_beam2 = power_beam.copy() - with pytest.raises(ValueError, match="axis2_array must be evenly spaced in az_za"): + with pytest.raises( + ValueError, + match="Selected values along second image axis must be evenly spaced.", + ): power_beam2.select(axis2_inds=[0, 5, 6])