Merge pull request #771 from Cadair/fix_axis_world_coords

sunpy · Oct 23, 2024 · 57b4c92 · 57b4c92
2 parents da6007e + 63af46f
commit 57b4c92
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 6 deletions.
diff --git 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
@@ -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()