From 72d3bbbdf252dbc4b2ff4e076c9c30392e2a01a8 Mon Sep 17 00:00:00 2001 From: Derek Homeier Date: Mon, 20 Sep 2021 17:52:02 +0200 Subject: [PATCH 1/4] ENH: add `axis_world_coords_limits` method --- ndcube/ndcube.py | 160 +++++++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 155 insertions(+), 5 deletions(-) diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index fdb83530b..dc072c99f 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -24,6 +24,7 @@ from ndcube.mixins import NDCubeSlicingMixin from ndcube.ndcube_sequence import NDCubeSequence from ndcube.utils.wcs_high_level_conversion import values_to_high_level_objects +from ndcube.utils.wcs import get_dependent_world_axes from ndcube.visualization import PlotterDescriptor from ndcube.wcs.wrappers import CompoundLowLevelWCS @@ -309,17 +310,37 @@ def array_axis_physical_types(self): return [tuple(world_axis_physical_types[axis_correlation_matrix[:, i]]) for i in range(axis_correlation_matrix.shape[1])][::-1] - def _generate_world_coords(self, pixel_corners, wcs): + def _generate_world_coords(self, pixel_corners, wcs, get_sizes=False): # TODO: We can improve this by not always generating all coordinates # To make our lives easier here we generate all the coordinates for all # pixels and then choose the ones we want to return to the user based # on the axes argument. We could be smarter by integrating this logic # into the main loop, this would potentially reduce the number of calls # to pixel_to_world_values + """ + Create meshgrid of all pixels transformed to world coordinates. + + Parameters + ---------- + pixel_corners: `bool` + If `True` then instead of returning the coordinates at the centers of the + pixels, the coordinates at the pixel corners will be returned. This increases + the size of the output by 1 in all dimensions as all corners are returned. + + wcs : `~astropy.wcs.wcsapi.BaseHighLevelWCS`, `~astropy.wcs.wcsapi.BaseLowLevelWCS` + The WCS to use to convert pixel values to world coordinates. + + get_sizes: `bool`, optional + Only calculate sizes of all separate coordinate arrays without creating and + transforming them. + + Returns + ------- + world_coords: `list` + An iterable of `Quantity` objects representing the real world coordinates in + all axes; or, for `get_sizes=True`, of sizes of all separate coordinate arrays. + """ - # Create meshgrid of all pixel coordinates. - # If user, wants pixel_corners, set pixel values to pixel pixel_corners. - # Else make pixel centers. pixel_shape = self.data.shape[::-1] if pixel_corners: pixel_shape = tuple(np.array(pixel_shape) + 1) @@ -334,6 +355,8 @@ def _generate_world_coords(self, pixel_corners, wcs): if wcs is None: return [] + if get_sizes: + world_sizes = [] world_coords = [None] * wcs.world_n_dim for (pixel_axes_indices, world_axes_indices) in _split_matrix(wcs.axis_correlation_matrix): # First construct a range of pixel indices for this set of coupled dimensions @@ -343,6 +366,10 @@ def _generate_world_coords(self, pixel_corners, wcs): # And inject 0s for those coordinates for idx in non_corr_axes: sub_range.insert(idx, 0) + # If requested, only calculate and return sizes + if get_sizes: + world_sizes.append(np.prod([np.array([r]).size for r in sub_range])) + continue # Generate a grid of broadcastable pixel indices for all pixel dimensions grid = np.meshgrid(*sub_range, indexing='ij') # Convert to world coordinates @@ -358,6 +385,9 @@ def _generate_world_coords(self, pixel_corners, wcs): tmp_world = world[idx][tuple(array_slice)].T world_coords[idx] = tmp_world + if get_sizes: + return world_sizes + for i, (coord, unit) in enumerate(zip(world_coords, wcs.world_axis_units)): world_coords[i] = coord << u.Unit(unit) @@ -507,7 +537,127 @@ def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None): CoordValues = namedtuple("CoordValues", identifiers) return CoordValues(*axes_coords[::-1]) - def crop(self, *points, wcs=None): + @utils.misc.sanitise_wcs + def axis_world_coords_limits(self, *axes, pixel_corners=False, wcs=None, max_size=None): + """ + Returns (estimated) extrema of the WCS coordinate values for all axes. + + Parameters + ---------- + axes: `int` or `str`, or multiple `int` or `str`, optional + Axis number in numpy ordering or unique substring of + `~ndcube.NDCube.world_axis_physical_types` + of axes for which real world coordinates are desired. + axes=None implies all axes will be returned. + + pixel_corners: `bool`, optional + If `True` then instead of returning the limits for the centers of the pixels, + the limits for the pixel corners will be returned. This increases the resulting + size over the limits in all dimensions as all corner positions are included. + + wcs: `astropy.wcs.wcsapi.BaseHighLevelWCS`, optional + The WCS object used to calculate the world coordinates. + Although technically this can be any valid WCS, it will typically be + ``self.wcs``, ``self.extra_coords``, or ``self.combined_wcs`` which combines both + the WCS and extra coords. + Defaults to the ``.wcs`` property. + + max_size: `int`, optional + Sets the maximum size of the pixel grid for which world coordinates will be + calculated in full to determine the extrema. If this size is exceeded, for + faster evaluation only the corners in pixel space, plus if possible, the axes + at the reference pixel values as given by ``wcs.wcs.crpix``, are considered. + + Returns + ------- + axes_coords: `list` + An iterable of "high level" objects containing the minima and maxima in + real world coordinates for the axes requested by user. + For example, a tuple of `~astropy.coordinates.SkyCoord` objects. + The types returned are determined by the WCS object. + These objects will have length 2 along each axis. + + Example + ------- + >>> NDCube.axis_world_coords_limits('lat', 'lon', max_size=100000) # doctest: +SKIP + >>> NDCube.axis_world_coords_limits(2) # doctest: +SKIP + + """ + if isinstance(wcs, BaseHighLevelWCS): + wcs = wcs.low_level_wcs + + if max_size is not None: + full_size = np.sum(self._generate_world_coords(pixel_corners, wcs, get_sizes=True)) + + # Check if we only probe pixel bounding box to speed up computation + if max_size is None or full_size <= max_size: + axes_coords = self._generate_world_coords(pixel_corners, wcs) + else: + if pixel_corners: + lower = np.ones(wcs.naxis) * -0.5 + upper = np.array(wcs.array_shape[::-1]) - 0.5 + else: + lower = np.zeros(wcs.naxis) + upper = np.array(wcs.array_shape[::-1]) - 1 + bbox = np.array(self._bounding_box_to_points(lower, upper, wcs)).T + # If wcs has a FITS-type Wcsprm, try to include CRPIX axes + try: + bbox_l = [bbox] + for ax, pix in enumerate(wcs.wcs.crpix): + sub_range = np.maximum(wcs.wcs.crpix - 1, 0).tolist() + for dwa in get_dependent_world_axes(ax, wcs.axis_correlation_matrix): + if dwa != ax: + sub_range[ax] = np.arange(lower[ax], upper[ax]) + # Generate a grid of broadcastable pixel indices for dependent axes + grid = np.meshgrid(*sub_range, indexing='ij') + # Check if size of the subgrid now exceeds limit; in that case cut all + # dependent axes to edge values (perhaps should undersample range instead?) + if grid[0].size > max_size: + for i, r in enumerate(sub_range): + if np.size(r) > 2: + sub_range[i] = r[[0, -1]] + grid = np.meshgrid(*sub_range, indexing='ij') + grid = np.array(grid).squeeze() + if grid.ndim == bbox.ndim: + bbox_l.append(grid) + bbox = np.concatenate(bbox_l, axis=1) + except AttributeError: + pass + # axes_coords = [None] * wcs.world_n_dim + # Convert to world coordinates + axes_coords = list(wcs.pixel_to_world_values(*bbox)) + for i, (coord, unit) in enumerate(zip(axes_coords, wcs.world_axis_units)): + axes_coords[i] = coord << u.Unit(unit) + + axes_coords = [u.Quantity([ac.min(), ac.max()]) for ac in axes_coords] + + if isinstance(wcs, ExtraCoords): + wcs = wcs.wcs + + axes_coords = values_to_high_level_objects(*axes_coords, low_level_wcs=wcs) + + if not axes: + return tuple(axes_coords) + + object_names = np.array([wao_comp[0] for wao_comp in wcs.world_axis_object_components]) + unique_obj_names = utils.misc.unique_sorted(object_names) + world_axes_for_obj = [np.where(object_names == name)[0] for name in unique_obj_names] + + # Create a mapping from world index in the WCS to object index in axes_coords + world_index_to_object_index = {} + for object_index, world_axes in enumerate(world_axes_for_obj): + for world_index in world_axes: + world_index_to_object_index[world_index] = object_index + + world_indices = utils.wcs.calculate_world_indices_from_axes(wcs, axes) + object_indices = utils.misc.unique_sorted( + [world_index_to_object_index[world_index] for world_index in world_indices] + ) + + return tuple(axes_coords[i] for i in object_indices) + + @utils.misc.sanitise_wcs + def crop(self, lower_corner, upper_corner, wcs=None): # The docstring is defined in NDCubeABC # Calculate the array slice item corresponding to bounding box and return sliced cube. item = self._get_crop_item(*points, wcs=wcs) From ac1c3862adcb4e864155450fdfc250a13a59767c Mon Sep 17 00:00:00 2001 From: Derek Homeier Date: Tue, 21 Sep 2021 12:11:51 +0200 Subject: [PATCH 2/4] TST: initial world coordinate limits tests --- changelog/474.feature.rst | 1 + ndcube/conftest.py | 82 +++++++++++++++++++++++++++++++++++ ndcube/tests/test_ndcube.py | 85 +++++++++++++++++++++++++++++++++++++ 3 files changed, 168 insertions(+) create mode 100644 changelog/474.feature.rst diff --git a/changelog/474.feature.rst b/changelog/474.feature.rst new file mode 100644 index 000000000..a7d6253e0 --- /dev/null +++ b/changelog/474.feature.rst @@ -0,0 +1 @@ +Add a fast method `~ndcube.NDCube.axis_world_coords_limits` to find wcs extension in world coordinates. diff --git a/ndcube/conftest.py b/ndcube/conftest.py index 4547227eb..c128395fb 100644 --- a/ndcube/conftest.py +++ b/ndcube/conftest.py @@ -189,6 +189,56 @@ def wcs_3d_lt_ln_l(): return WCS(header=header) +@pytest.fixture +def wcs_3d_l_ra_dec(): + header = { + 'CTYPE1': 'WAVE ', + 'CUNIT1': 'Angstrom', + 'CDELT1': 0.2, + 'CRPIX1': 0, + 'CRVAL1': 10, + + 'CTYPE2': 'RA---TAN', + 'CUNIT2': 'deg', + 'CDELT2': 0.1, + 'CRPIX2': 200, + 'CRVAL2': 90.5, + + 'CTYPE3': 'DEC--TAN', + 'CUNIT3': 'deg', + 'CDELT3': 0.08, + 'CRPIX3': 150, + 'CRVAL3': 30.4, + } + + return WCS(header=header) + + +@pytest.fixture +def wcs_3d_l_ra_pol(): + header = { + 'CTYPE1': 'WAVE ', + 'CUNIT1': 'Angstrom', + 'CDELT1': 0.2, + 'CRPIX1': 0, + 'CRVAL1': 10, + + 'CTYPE2': 'RA---TAN', + 'CUNIT2': 'deg', + 'CDELT2': 0.1, + 'CRPIX2': 200, + 'CRVAL2': 90.5, + + 'CTYPE3': 'DEC--TAN', + 'CUNIT3': 'deg', + 'CDELT3': 0.08, + 'CRPIX3': 150, + 'CRVAL3': 80.4, + } + + return WCS(header=header) + + @pytest.fixture def wcs_2d_lt_ln(): spatial = { @@ -437,6 +487,38 @@ def ndcube_3d_ln_lt_l_ec_time(wcs_3d_l_lt_ln, time_and_simple_extra_coords_2d): return cube +@pytest.fixture +def ndcube_3d_l_ra_dec(wcs_3d_l_ra_dec, simple_extra_coords_3d): + shape = (400, 300, 10) + wcs_3d_l_ra_dec.array_shape = shape + data = data_nd(shape) + mask = data > 0 + cube = NDCube( + data, + wcs_3d_l_ra_dec, + mask=mask, + uncertainty=data, + ) + cube._extra_coords = simple_extra_coords_3d + return cube + + +@pytest.fixture +def ndcube_3d_l_ra_pol(wcs_3d_l_ra_pol, simple_extra_coords_3d): + shape = (400, 300, 10) + wcs_3d_l_ra_pol.array_shape = shape + data = data_nd(shape) + mask = data > 0 + cube = NDCube( + data, + wcs_3d_l_ra_pol, + mask=mask, + uncertainty=data, + ) + cube._extra_coords = simple_extra_coords_3d + return cube + + @pytest.fixture def ndcube_3d_rotated(wcs_3d_ln_lt_t_rotated, simple_extra_coords_3d): data_rotated = np.array([[[1, 2, 3, 4, 6], [2, 4, 5, 3, 1], [0, -1, 2, 4, 2], [3, 5, 1, 2, 0]], diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index 93b1aac60..592dd7720 100644 --- a/ndcube/tests/test_ndcube.py +++ b/ndcube/tests/test_ndcube.py @@ -191,6 +191,12 @@ def test_axis_world_coords_wave_ec(ndcube_3d_l_ln_lt_ectime): assert isinstance(coords[0], u.Quantity) assert coords[0].shape == (5,) + coords = cube.axis_world_coords_limits('em.wl') + assert u.allclose(coords, [1.02e-09, 1.20e-09] * u.m) + + coords = cube.axis_world_coords_limits('em.wl', max_size=20) + assert u.allclose(coords, [1.02e-09, 1.20e-09] * u.m) + def test_axis_world_coords_empty_ec(ndcube_3d_l_ln_lt_ectime): cube = ndcube_3d_l_ln_lt_ectime @@ -286,6 +292,22 @@ def test_axis_world_coords_all_4d_split(ndcube_4d_ln_l_t_lt): assert u.allclose(coords[2], [2.0e-11, 4.0e-11, 6.0e-11, 8.0e-11, 1.0e-10, 1.2e-10, 1.4e-10, 1.6e-10, 1.8e-10, 2.0e-10] * u.m) + coords = ndcube_4d_ln_l_t_lt.axis_world_coords_limits() + assert len(coords) == 3 + assert isinstance(coords[0], SkyCoord) + assert u.allclose(coords[0].Tx, [0.0, -5.0] * u.deg / 3600) + assert u.allclose(coords[0].Ty, [20.0, 160.0] * u.deg / 3600) + + assert isinstance(coords[1], Time) + assert u.allclose(coords[1].value - 58849, np.array([24.0, 288.0]) / 86400) + + assert isinstance(coords[2], u.Quantity) + assert u.allclose(coords[2], [2.0e-11, 2.0e-10] * u.m) + + coords = ndcube_4d_ln_l_t_lt.axis_world_coords_limits(max_size=1) + assert u.allclose(coords[0].Tx, [0.0, -5.0] * u.deg / 3600) + assert u.allclose(coords[0].Ty, [20.0, 160.0] * u.deg / 3600) + @pytest.mark.parametrize('wapt', ( ('custom:pos.helioprojective.lon', 'custom:pos.helioprojective.lat', 'em.wl'), @@ -318,6 +340,65 @@ def test_axis_world_coords_all(ndcube_3d_ln_lt_l): assert u.allclose(coords[1].Ty, [[-19.99999991, -14.99999996, -9.99999998], [-19.99999984, -14.9999999, -9.99999995]] * u.arcsec) + coords = ndcube_3d_ln_lt_l.axis_world_coords_limits() + assert u.allclose(coords[0], [1.02e-09, 1.08e-09] * u.m) + assert u.allclose(coords[1].Tx, [10, 20] * u.arcsec) + assert u.allclose(coords[1].Ty, [-20, -10] * u.arcsec) + + coords = ndcube_3d_ln_lt_l.axis_world_coords_limits(max_size=1) + assert u.allclose(coords[0], [1.02e-09, 1.08e-09] * u.m) + assert u.allclose(coords[1].Tx, [10, 20] * u.arcsec) + assert u.allclose(coords[1].Ty, [-20, -10] * u.arcsec) + + +def test_axis_world_coords_radec(ndcube_3d_l_ra_dec): + coords = ndcube_3d_l_ra_dec.axis_world_coords() + assert len(coords) == 2 + assert isinstance(coords[0], u.Quantity) + assert u.allclose(coords[0], np.arange(1.02e-09, 1.22e-09, 2.0e-11) * u.m) + + assert isinstance(coords[1], SkyCoord) + assert coords[1].shape == (400, 300) + + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits() + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) + assert u.allclose(coords[1].dec, [17.6213, 49.64235] * u.deg) + + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(max_size=10000) + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) + assert u.allclose(coords[1].dec, [17.6213, 49.5710] * u.deg) + + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(max_size=10) + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) + assert u.allclose(coords[1].dec, [17.6213, 49.5710] * u.deg) # [17.6213, 48.7533] at corners + + +def test_axis_world_coords_rapol(ndcube_3d_l_ra_pol): + coords = ndcube_3d_l_ra_pol.axis_world_coords() + assert len(coords) == 2 + assert isinstance(coords[0], u.Quantity) + + assert isinstance(coords[1], SkyCoord) + assert coords[1].shape == (400, 300) + + coords = ndcube_3d_l_ra_pol.axis_world_coords_limits() + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [1.4523e-3, 359.9992] * u.deg) + assert u.allclose(coords[1].dec, [61.8572, 89.98945] * u.deg) + + coords = ndcube_3d_l_ra_pol.axis_world_coords_limits(max_size=10000) + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [26.1484, 333.4424] * u.deg) + assert u.allclose(coords[1].dec, [61.8572, 89.98945] * u.deg) + + coords = ndcube_3d_l_ra_pol.axis_world_coords_limits(max_size=10) + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[1].ra, [26.1484, 333.4424] * u.deg) + assert u.allclose(coords[1].dec, [61.8572, 80.429] * u.deg) + def test_axis_world_coords_wave(ndcube_3d_ln_lt_l): coords = ndcube_3d_ln_lt_l.axis_world_coords('em.wl') @@ -339,6 +420,10 @@ def test_axis_world_coords_sky(ndcube_3d_ln_lt_l, wapt): assert u.allclose(coords[0].Ty, [[-19.99999991, -14.99999996, -9.99999998], [-19.99999984, -14.9999999, -9.99999995]] * u.arcsec) + coords = ndcube_3d_ln_lt_l.axis_world_coords_limits(wapt) + assert u.allclose(coords[0].Tx, [10, 20] * u.arcsec) + assert u.allclose(coords[0].Ty, [-20, -10] * u.arcsec) + def test_axes_world_coords_sky_only(ndcube_2d_ln_lt): coords = ndcube_2d_ln_lt.axis_world_coords() From e73ac71111232ad12e6ab7bd954d889bbfd9d2d1 Mon Sep 17 00:00:00 2001 From: Derek Homeier Date: Fri, 22 Oct 2021 17:14:07 +0200 Subject: [PATCH 3/4] Rebase to new `crop`; style fixes and additional tests --- ndcube/conftest.py | 4 ++-- ndcube/ndcube.py | 27 ++++++++++++++++++++------- ndcube/tests/test_ndcube.py | 28 ++++++++++++++++++++++++++++ 3 files changed, 50 insertions(+), 9 deletions(-) diff --git a/ndcube/conftest.py b/ndcube/conftest.py index c128395fb..7a4b77a5b 100644 --- a/ndcube/conftest.py +++ b/ndcube/conftest.py @@ -537,8 +537,8 @@ def ndcube_3d_rotated(wcs_3d_ln_lt_t_rotated, simple_extra_coords_3d): @pytest.fixture def ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l): return gen_ndcube_3d_l_ln_lt_ectime(wcs_3d_lt_ln_l, - 1, - Time('2000-01-01', format='fits', scale='utc')) + 1, Time('2000-01-01', format='fits', scale='utc')) + @pytest.fixture diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index dc072c99f..69ac62e4d 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -2,6 +2,7 @@ import textwrap import warnings from copy import deepcopy +from itertools import product as iterprod from collections import namedtuple from collections.abc import Mapping @@ -23,8 +24,8 @@ from ndcube.global_coords import GlobalCoords from ndcube.mixins import NDCubeSlicingMixin from ndcube.ndcube_sequence import NDCubeSequence -from ndcube.utils.wcs_high_level_conversion import values_to_high_level_objects from ndcube.utils.wcs import get_dependent_world_axes +from ndcube.utils.wcs_high_level_conversion import values_to_high_level_objects from ndcube.visualization import PlotterDescriptor from ndcube.wcs.wrappers import CompoundLowLevelWCS @@ -537,7 +538,7 @@ def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None): CoordValues = namedtuple("CoordValues", identifiers) return CoordValues(*axes_coords[::-1]) - @utils.misc.sanitise_wcs + @utils.cube.sanitize_wcs def axis_world_coords_limits(self, *axes, pixel_corners=False, wcs=None, max_size=None): """ Returns (estimated) extrema of the WCS coordinate values for all axes. @@ -583,13 +584,17 @@ def axis_world_coords_limits(self, *axes, pixel_corners=False, wcs=None, max_siz >>> NDCube.axis_world_coords_limits(2) # doctest: +SKIP """ + # Cannot use naxis and array_shape to construct bounding box for + # extra_coords or combined_wcs, so for now force using the full wcs for those. + if isinstance(wcs, (ExtraCoords, HighLevelWCSWrapper)) or wcs.array_shape is None: + max_size = None if isinstance(wcs, BaseHighLevelWCS): wcs = wcs.low_level_wcs if max_size is not None: full_size = np.sum(self._generate_world_coords(pixel_corners, wcs, get_sizes=True)) - # Check if we only probe pixel bounding box to speed up computation + # Check if we only probe pixel bounding box to speed up computation. if max_size is None or full_size <= max_size: axes_coords = self._generate_world_coords(pixel_corners, wcs) else: @@ -629,12 +634,15 @@ def axis_world_coords_limits(self, *axes, pixel_corners=False, wcs=None, max_siz for i, (coord, unit) in enumerate(zip(axes_coords, wcs.world_axis_units)): axes_coords[i] = coord << u.Unit(unit) - axes_coords = [u.Quantity([ac.min(), ac.max()]) for ac in axes_coords] + axes_limits = [] + for ac in axes_coords: + ac = ac[np.isfinite(ac)] + axes_limits.append(u.Quantity([ac.min(), ac.max()])) if isinstance(wcs, ExtraCoords): wcs = wcs.wcs - axes_coords = values_to_high_level_objects(*axes_coords, low_level_wcs=wcs) + axes_coords = values_to_high_level_objects(*axes_limits, low_level_wcs=wcs) if not axes: return tuple(axes_coords) @@ -656,8 +664,7 @@ def axis_world_coords_limits(self, *axes, pixel_corners=False, wcs=None, max_siz return tuple(axes_coords[i] for i in object_indices) - @utils.misc.sanitise_wcs - def crop(self, lower_corner, upper_corner, wcs=None): + def crop(self, *points, wcs=None): # The docstring is defined in NDCubeABC # Calculate the array slice item corresponding to bounding box and return sliced cube. item = self._get_crop_item(*points, wcs=wcs) @@ -706,6 +713,12 @@ def _get_crop_by_values_item(self, *points, units=None, wcs=None): return utils.cube.get_crop_item_from_points(points, wcs, True) + def _bounding_box_to_points(self, lower_corner_values, upper_corner_values, wcs): + """ + Convert two corners of a bounding box to the points of all corners. + """ + return tuple(iterprod(*zip(lower_corner_values, upper_corner_values))) + def __str__(self): return textwrap.dedent(f"""\ NDCube diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index 592dd7720..4bb00ba9c 100644 --- a/ndcube/tests/test_ndcube.py +++ b/ndcube/tests/test_ndcube.py @@ -375,6 +375,21 @@ def test_axis_world_coords_radec(ndcube_3d_l_ra_dec): assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) assert u.allclose(coords[1].dec, [17.6213, 49.5710] * u.deg) # [17.6213, 48.7533] at corners + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(wcs=ndcube_3d_l_ra_dec.combined_wcs) + assert isinstance(coords[1], SkyCoord) + assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) + assert u.allclose(coords[2], [0, 3] * u.pix) + + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(wcs=ndcube_3d_l_ra_dec.extra_coords) + assert len(coords) == 1 + assert u.allclose(coords[0], [0, 3] * u.pix) + + # pytest.xfail("NaNs in ExtraCoords(?)") + # this returns an array of len 10 filled up with NaN, does not seem right... + coords = ndcube_3d_l_ra_dec.axis_world_coords(wcs=ndcube_3d_l_ra_dec.extra_coords) + assert len(coords) == 1 + assert u.allclose(coords[0][:4], [0, 1, 2, 3] * u.pix) + def test_axis_world_coords_rapol(ndcube_3d_l_ra_pol): coords = ndcube_3d_l_ra_pol.axis_world_coords() @@ -435,6 +450,19 @@ def test_axes_world_coords_sky_only(ndcube_2d_ln_lt): assert u.allclose(coords[0].Ty[0, :], [-8, -6, -4, -2, 0, 2, 4, 6, 8, 10, 12, 14] * u.arcsec, atol=1e-5 * u.arcsec) + coords = ndcube_2d_ln_lt.axis_world_coords_limits() + + assert len(coords) == 1 + assert isinstance(coords[0], SkyCoord) + # coords_values are wrapping around 360 deg in lon, yielding bad limits! + assert u.allclose(coords[0].Tx, [0, -4] * u.arcsec, atol=1e-5 * u.arcsec) + assert u.allclose(coords[0].Ty, [-8, 14] * u.arcsec, atol=1e-5 * u.arcsec) + + coords = ndcube_2d_ln_lt.axis_world_coords_limits(max_size=5) + + assert u.allclose(coords[0].Tx, [0, -4] * u.arcsec, atol=1e-5 * u.arcsec) + assert u.allclose(coords[0].Ty, [-8, 14] * u.arcsec, atol=1e-5 * u.arcsec) + def test_axis_world_coords_values_all(ndcube_3d_ln_lt_l): coords = ndcube_3d_ln_lt_l.axis_world_coords_values() From 1efb0b816c8a161e82b79eeb467924e92273816f Mon Sep 17 00:00:00 2001 From: Derek Homeier Date: Wed, 27 Oct 2021 11:52:02 +0200 Subject: [PATCH 4/4] Catch gwcs 0.16 bug in world coords with extra_coords (fix devdeps) --- ndcube/tests/test_ndcube.py | 83 ++++++++++++++++++++++++++++++++----- 1 file changed, 72 insertions(+), 11 deletions(-) diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index 4bb00ba9c..be85d8971 100644 --- a/ndcube/tests/test_ndcube.py +++ b/ndcube/tests/test_ndcube.py @@ -209,6 +209,44 @@ def test_axis_world_coords_empty_ec(ndcube_3d_l_ln_lt_ectime): assert awc == tuple() +def test_axis_world_coords_with_ec(ndcube_3d_ln_lt_l_ec_all_axes): + cube = ndcube_3d_ln_lt_l_ec_all_axes + t0 = Time(["2000-01-01T12:00:00", "2000-01-02T12:00:00"], scale="utc", format="fits") + + coords = cube.axis_world_coords(wcs=cube.extra_coords) + assert len(coords) == 3 + assert isinstance(coords[0], Time) + assert isinstance(coords[1], u.Quantity) + assert isinstance(coords[2], u.Quantity) + assert u.allclose(coords[0].mjd, t0.mjd, rtol=1e-12) + assert u.allclose(coords[1], list(range(3)) * u.pix) + assert u.allclose(coords[2], list(range(4)) * u.m) + + coords = cube.axis_world_coords(wcs=cube.combined_wcs) + assert len(coords) == 5 + assert u.allclose(coords[2].mjd, t0.mjd, rtol=1e-12) + assert u.allclose(coords[3], list(range(3)) * u.pix) + assert u.allclose(coords[4], list(range(4)) * u.m) + + +def test_axis_world_coords_simple_ec(ndcube_3d_ln_lt_l): + cube = ndcube_3d_ln_lt_l + + coords = cube.axis_world_coords(wcs=cube.extra_coords) + if len(coords) == 1: + assert u.allclose(coords[0][:4], list(range(4)) * u.pix) + pytest.xfail("gwcs < 0.17 strips ExtraCoords axes with repetitive units/physical_type") + + assert len(coords) == 3 + for i in range(3): + assert u.allclose(coords[i][:i+2], list(range(i+2)) * u.pix) + + coords = cube.axis_world_coords(wcs=cube.combined_wcs) + assert len(coords) == 5 + for i in range(2, 5): + assert u.allclose(coords[i][:i], list(range(i)) * u.pix) + + @pytest.mark.xfail(reason=">1D Tables not supported") def test_axis_world_coords_complex_ec(ndcube_4d_ln_lt_l_t): cube = ndcube_4d_ln_lt_l_t @@ -356,10 +394,32 @@ def test_axis_world_coords_radec(ndcube_3d_l_ra_dec): assert len(coords) == 2 assert isinstance(coords[0], u.Quantity) assert u.allclose(coords[0], np.arange(1.02e-09, 1.22e-09, 2.0e-11) * u.m) + assert isinstance(coords[1], SkyCoord) + assert coords[1].shape == (400, 300) + + # Need to update this to `extra_coords` covering the full data array. + coords = ndcube_3d_l_ra_dec.axis_world_coords(wcs=ndcube_3d_l_ra_dec.extra_coords) + if len(coords) == 1: + assert u.allclose(coords[0][:4], list(range(4)) * u.pix) + pytest.xfail("gwcs < 0.17 strips ExtraCoords axes with repetitive units/physical_type") + + assert len(coords) == 3 + for i in range(3): + assert u.allclose(coords[i][:i+2], list(range(i+2)) * u.pix) + coords = ndcube_3d_l_ra_dec.axis_world_coords(wcs=ndcube_3d_l_ra_dec.combined_wcs) + assert len(coords) == 5 + assert isinstance(coords[0], u.Quantity) + assert u.allclose(coords[0], np.arange(1.02e-09, 1.22e-09, 2.0e-11) * u.m) assert isinstance(coords[1], SkyCoord) assert coords[1].shape == (400, 300) + for i in range(2, 5): + assert u.allclose(coords[i][:i], list(range(i)) * u.pix) + + +def test_axis_world_coords_limits_radec(ndcube_3d_l_ra_dec): + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits() assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) @@ -375,20 +435,21 @@ def test_axis_world_coords_radec(ndcube_3d_l_ra_dec): assert u.allclose(coords[1].ra, [63.64277, 104.777] * u.deg) assert u.allclose(coords[1].dec, [17.6213, 49.5710] * u.deg) # [17.6213, 48.7533] at corners + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(wcs=ndcube_3d_l_ra_dec.extra_coords) + if len(coords) == 1: + assert u.allclose(coords[0], [0, 3] * u.pix) + pytest.xfail("gwcs < 0.17 strips ExtraCoords axes with repetitive units/physical_type") + + assert len(coords) == 3 + for i in range(3): + assert u.allclose(coords[i], [0, i+1] * u.pix) + coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(wcs=ndcube_3d_l_ra_dec.combined_wcs) + assert len(coords) == 5 assert isinstance(coords[1], SkyCoord) assert u.allclose(coords[0], [1.02e-09, 1.20e-09] * u.m) - assert u.allclose(coords[2], [0, 3] * u.pix) - - coords = ndcube_3d_l_ra_dec.axis_world_coords_limits(wcs=ndcube_3d_l_ra_dec.extra_coords) - assert len(coords) == 1 - assert u.allclose(coords[0], [0, 3] * u.pix) - - # pytest.xfail("NaNs in ExtraCoords(?)") - # this returns an array of len 10 filled up with NaN, does not seem right... - coords = ndcube_3d_l_ra_dec.axis_world_coords(wcs=ndcube_3d_l_ra_dec.extra_coords) - assert len(coords) == 1 - assert u.allclose(coords[0][:4], [0, 1, 2, 3] * u.pix) + for i in range(2, 5): + assert u.allclose(coords[i], [0, i-1] * u.pix) def test_axis_world_coords_rapol(ndcube_3d_l_ra_pol):