From 67e4ec0a6c0c37be42ffe12b5327d81dccb51fe6 Mon Sep 17 00:00:00 2001 From: Stuart Mumford Date: Mon, 21 Oct 2024 15:59:16 +0100 Subject: [PATCH] Fix axis world coords to not pass quantities as values --- ndcube/ndcube.py | 11 ++++++----- ndcube/tests/test_ndcube.py | 2 +- 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index e7cec0364..a9827b474 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -443,7 +443,7 @@ def quantity(self): """Unitful representation of the NDCube data.""" return u.Quantity(self.data, self.unit, copy=_NUMPY_COPY_IF_NEEDED) - def _generate_world_coords(self, pixel_corners, wcs, needed_axes=None): + def _generate_world_coords(self, pixel_corners, wcs, needed_axes=None, *, units): # Create meshgrid of all pixel coordinates. # If user wants pixel_corners, set pixel values to pixel corners. # Else make pixel centers. @@ -493,8 +493,9 @@ def _generate_world_coords(self, pixel_corners, wcs, needed_axes=None): tmp_world = world[idx][tuple(array_slice)].T world_coords[idx] = tmp_world - for i, (coord, unit) in enumerate(zip(world_coords, wcs.world_axis_units)): - world_coords[i] = coord << u.Unit(unit) + if units: + for i, (coord, unit) in enumerate(zip(world_coords, wcs.world_axis_units)): + world_coords[i] = coord << u.Unit(unit) return world_coords @@ -525,7 +526,7 @@ def axis_world_coords(self, *axes, pixel_corners=False, wcs=None): [world_index_to_object_index[world_index] for world_index in world_indices] ) - axes_coords = self._generate_world_coords(pixel_corners, orig_wcs, world_indices) + axes_coords = self._generate_world_coords(pixel_corners, orig_wcs, world_indices, units=False) axes_coords = values_to_high_level_objects(*axes_coords, low_level_wcs=wcs) @@ -546,7 +547,7 @@ def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None): world_indices = utils.wcs.calculate_world_indices_from_axes(wcs, axes) - axes_coords = self._generate_world_coords(pixel_corners, orig_wcs, world_indices) + axes_coords = self._generate_world_coords(pixel_corners, orig_wcs, world_indices, units=True) world_axis_physical_types = wcs.world_axis_physical_types diff --git a/ndcube/tests/test_ndcube.py b/ndcube/tests/test_ndcube.py index 4a3d4c1e7..5d8a25800 100644 --- a/ndcube/tests/test_ndcube.py +++ b/ndcube/tests/test_ndcube.py @@ -207,7 +207,7 @@ def test_axis_world_coords_empty_ec(ndcube_3d_l_ln_lt_ectime): # slice the cube so extra_coords is empty, and then try and run axis_world_coords awc = sub_cube.axis_world_coords(wcs=sub_cube.extra_coords) assert awc == tuple() - sub_cube._generate_world_coords(pixel_corners=False, wcs=sub_cube.extra_coords) + sub_cube._generate_world_coords(pixel_corners=False, wcs=sub_cube.extra_coords, units=True) assert awc == tuple()