diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index eedbf213f..6e0e54a86 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -1,6 +1,7 @@ # -*- coding: utf-8 -*- import abc +import warnings import numpy as np import astropy.nddata @@ -85,22 +86,36 @@ def world_axis_physical_types(self): pass @abc.abstractmethod - def crop_by_coords(self, min_coord_values, interval_widths): + def crop_by_coords(self, lower_corner, interval_widths=None, upper_corner=None, units=None): """ Crops an NDCube given minimum values and interval widths along axes. Parameters ---------- - min_coord_values: iterable of `astropy.units.Quantity` + lower_corner: iterable of `astropy.units.Quantity` or `float` The minimum desired values along each relevant axis after cropping described in physical units consistent with the NDCube's wcs object. - The length of the iterable must equal the number of data dimensions and must - have the same order as the data. + The length of the iterable must equal the number of data dimensions + and must have the same order as the data. + + interval_widths: iterable of `astropy.units.Quantity` or `float` + The width of the region of interest in each dimension in physical + units consistent with the NDCube's wcs object. The length of the + iterable must equal the number of data dimensions and must have + the same order as the data. This argument will be removed in versions + 2.0, please use upper_corner argument. + + upper_corner: iterable of `astropy.units.Quantity` or `float` + The maximum desired values along each relevant axis after cropping + described in physical units consistent with the NDCube's wcs object. + The length of the iterable must equal the number of data dimensions + and must have the same order as the data. - interval_widths: iterable of `astropy.units.Quantity` - The width of the region of interest in each dimension in physical units - consistent with the NDCube's wcs object. The length of the iterable must - equal the number of data dimensions and must have the same order as the data. + units: iterable of `astropy.units.quantity.Quantity`, optionnal + If the inputs are set without units, the user must set the units + inside this argument as `str`. + The length of the iterable must equal the number of data dimensions + and must have the same order as the data. Returns ------- @@ -401,21 +416,57 @@ def extra_coords(self): "value": self._extra_coords_wcs_axis[key]["value"]} return result - def crop_by_coords(self, min_coord_values, interval_widths): + def crop_by_coords(self, lower_corner, interval_widths=None, upper_corner=None, units=None): # The docstring is defined in NDDataBase - n_dim = len(self.dimensions) - if (len(min_coord_values) != len(interval_widths)) or len(min_coord_values) != n_dim: - raise ValueError("min_coord_values and interval_widths must have " - "same number of elements as number of data dimensions.") - # Convert coords of lower left corner to pixel units. - lower_pixels = self.world_to_pixel(*min_coord_values) - upper_pixels = self.world_to_pixel(*[min_coord_values[i] + interval_widths[i] - for i in range(n_dim)]) - # Round pixel values to nearest integer. - lower_pixels = [int(np.rint(l.value)) for l in lower_pixels] - upper_pixels = [int(np.rint(u.value)) for u in upper_pixels] - item = tuple([slice(lower_pixels[i], upper_pixels[i]) for i in range(n_dim)]) + n_dim = self.data.ndim + # Raising a value error if the arguments have not the same dimensions. + # Calculation of upper_corner with the inputing interval_widths + # This part of the code will be removed in version 2.0 + if interval_widths: + warnings.warn("interval_widths will be removed from the API in " + "version 2.0, please use upper_corner argument.") + if upper_corner: + raise ValueError("Only one of interval_widths or upper_corner " + "can be set. Recommend using upper_corner as " + "interval_widths is deprecated.") + if (len(lower_corner) != len(interval_widths)) or (len(lower_corner) != n_dim): + raise ValueError("lower_corner and interval_widths must have " + "same number of elements as number of data " + "dimensions.") + upper_corner = [lower_corner[i] + interval_widths[i] for i in range(n_dim)] + # Raising a value error if the arguments have not the same dimensions. + if (len(lower_corner) != len(upper_corner)) or (len(lower_corner) != n_dim): + raise ValueError("lower_corner and upper_corner must have same" + "number of elements as number of data dimensions.") + if units: + # Raising a value error if units have not the data dimensions. + if len(units) != n_dim: + raise ValueError('units must have same number of elements as ' + 'number of data dimensions.') + # If inputs are not Quantity objects, they are modified into specified units + lower_corner = [u.Quantity(lower_corner[i], unit=units[i]) + for i in range(self.data.ndim)] + upper_corner = [u.Quantity(upper_corner[i], unit=units[i]) + for i in range(self.data.ndim)] + else: + if any([not isinstance(x, u.Quantity) for x in lower_corner + upper_corner]): + raise TypeError("lower_corner and interval_widths/upper_corner must be " + "of type astropy.units.Quantity or the units kwarg " + "must be set.") + # Get all corners of region of interest. + all_world_corners_grid = np.meshgrid( + *[u.Quantity([lower_corner[i], upper_corner[i]], unit=lower_corner[i].unit).value + for i in range(self.data.ndim)]) + all_world_corners = [all_world_corners_grid[i].flatten()*lower_corner[i].unit + for i in range(n_dim)] + # Convert to pixel coordinates + all_pix_corners = self.world_to_pixel(*all_world_corners) + # Derive slicing item with which to slice NDCube. + # Be sure to round down min pixel and round up + 1 the max pixel. + item = tuple([slice(int(np.clip(axis_pixels.value.min(), 0, None)), + int(np.ceil(axis_pixels.value.max()))+1) + for axis_pixels in all_pix_corners]) return self[item] def crop_by_extra_coord(self, min_coord_value, interval_width, coord_name): diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index b400d1531..8e6e39581 100644 --- a/ndcube/tests/test_ndcube.py +++ b/ndcube/tests/test_ndcube.py @@ -59,6 +59,20 @@ data_ordered[0] = data.transpose() data_ordered[1] = data.transpose() +h_rotated = {'CTYPE1': 'HPLN-TAN', 'CUNIT1': 'arcsec', 'CDELT1': 0.4, 'CRPIX1': 0, + 'CRVAL1': 0, 'NAXIS1': 5, + 'CTYPE2': 'HPLT-TAN', 'CUNIT2': 'arcsec', 'CDELT2': 0.5, 'CRPIX2': 0, + 'CRVAL2': 0, 'NAXIS2': 5, + 'CTYPE3': 'Time ', 'CUNIT3': 'seconds', 'CDELT3': 0.3, 'CRPIX3': 0, + 'CRVAL3': 0, 'NAXIS3': 2, + 'PC1_1': 0.714963912964, 'PC1_2': -0.699137151241, 'PC1_3': 0.0, + 'PC2_1': 0.699137151241, 'PC2_2': 0.714963912964, 'PC2_3': 0.0, + 'PC3_1': 0.0, 'PC3_2': 0.0, 'PC3_3': 1.0} +w_rotated = WCS(header=h_rotated, naxis=3) + +data_rotated = np.array([[[1, 2, 3, 4, 6], [2, 4, 5, 3, 1], [0, -1, 2, 4, 2], [3, 5, 1, 2, 0]], + [[2, 4, 5, 1, 3], [1, 5, 2, 2, 4], [2, 3, 4, 0, 5], [0, 1, 2, 3, 4]]]) + mask_cubem = data > 0 mask_cube = data >= 0 uncertaintym = data @@ -117,6 +131,15 @@ ('hello', 1, u.Quantity(range(data.shape[1]), unit=u.pix)), ('bye', 2, u.Quantity(range(data.shape[2]), unit=u.pix))]) +cube_rotated = NDCube( + data_rotated, + w_rotated, + mask=mask_cube, + uncertainty=uncertainty, + missing_axis=[False, False, False], + extra_coords=[('time', 0, u.Quantity(range(data_rotated.shape[0]), unit=u.pix)), + ('hello', 1, u.Quantity(range(data_rotated.shape[1]), unit=u.pix)), + ('bye', 2, u.Quantity(range(data_rotated.shape[2]), unit=u.pix))]) @pytest.mark.parametrize( "test_input,expected,mask,wcs,uncertainty,dimensions,world_axis_physical_types,extra_coords", @@ -814,17 +837,30 @@ def test_world_to_pixel(test_input, expected): @pytest.mark.parametrize("test_input,expected", [ - ((cubem, [0.7*u.deg, 1.3e-5*u.deg, 1.02e-9*u.m], [1*u.deg, 1*u.deg, 1.06*u.m]), cubem[:, :2])]) + ((cubem, [0.7*u.deg, 1.3e-5*u.deg, 1.02e-9*u.m], [1*u.deg, 1*u.deg, 4.e-11*u.m], None), + cubem[:, :, :3]), + ((cube_rotated, [0*u.s, 1.5*u.arcsec, 0*u.arcsec], [1*u.s, 1*u.arcsec, 0.5*u.arcsec], None), + cube_rotated[:, :4, 1:5]), + ((cubem, [0.7*u.deg, 1.3e-5*u.deg, 1.02e-9*u.m], None, [1.7*u.deg, 1*u.deg, 1.06e-9*u.m]), + cubem[:, :, :3]), + ((cube_rotated, [0*u.s, 1.5*u.arcsec, 0*u.arcsec], None, [1*u.s, 2.5*u.arcsec, 0.5*u.arcsec]), + cube_rotated[:, :4, 1:5]), + ((cube_rotated, [0, 1.5, 0], None, [1, 2.5, 0.5], ['s', 'arcsec', 'arcsec']), + cube_rotated[:, :4, 1:5])]) def test_crop_by_coords(test_input, expected): helpers.assert_cubes_equal( test_input[0].crop_by_coords(*test_input[1:]), expected) @pytest.mark.parametrize("test_input", [ - (cubem, u.Quantity([0], unit=u.deg), u.Quantity([1.5, 2.], unit=u.deg))]) + (ValueError, cubem, u.Quantity([0], unit=u.deg), u.Quantity([1.5, 2.], unit=u.deg), None), + (ValueError, cubem, [1*u.s], [1*u.s], [1*u.s]), + (ValueError, cubem, u.Quantity([0], unit=u.deg), None, u.Quantity([1.5, 2.], unit=u.deg)), + (ValueError, cubem, [1], None, [1], ['s', 'deg']), + (TypeError, cubem, [1, 2, 3], None, [2, 3, 4])]) def test_crop_by_coords_error(test_input): - with pytest.raises(ValueError): - test_input[0].crop_by_coords(*test_input[1:]) + with pytest.raises(test_input[0]): + test_input[1].crop_by_coords(*test_input[2:]) @pytest.mark.parametrize(