Added unit test and broke it

Author: nabobalis

ndcube/conftest.py

@@ -197,6 +197,30 @@ def wcs_3d_lt_ln_l():
     return WCS(header=header)
 
 
+@pytest.fixture
+def wcs_3d_wave_lt_ln():
+    header = {
+        'CTYPE1': 'WAVE    ',
+        'CUNIT1': 'Angstrom',
+        'CDELT1': 0.2,
+        'CRPIX1': 0,
+        'CRVAL1': 10,
+
+        'CTYPE2': 'HPLT-TAN',
+        'CUNIT2': 'deg',
+        'CDELT2': 0.5,
+        'CRPIX2': 2,
+        'CRVAL2': 0.5,
+
+        'CTYPE3': 'HPLN-TAN    ',
+        'CUNIT3': 'deg',
+        'CDELT3': 0.4,
+        'CRPIX3': 2,
+        'CRVAL3': 1,
+    }
+    return WCS(header=header)
+
+
 @pytest.fixture
 def wcs_2d_lt_ln():
     spatial = {
@@ -445,6 +469,24 @@ 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_wave_lt_ln_ec_time(wcs_3d_wave_lt_ln):
+    shape = (3, 4, 5)
+    wcs_3d_wave_lt_ln.array_shape = shape
+    data = data_nd(shape)
+    mask = data > 0
+    cube = NDCube(
+        data,
+        wcs_3d_wave_lt_ln,
+        mask=mask,
+        uncertainty=data,
+    )
+    base_time = Time('2000-01-01', format='fits', scale='utc')
+    timestamps = Time([base_time + TimeDelta(60 * i, format='sec') for i in range(data.shape[0])])
+    cube.extra_coords.add('time', 0, timestamps)
+    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]],

ndcube/ndcube.py

@@ -448,11 +448,17 @@ def _generate_world_coords(self, pixel_corners, wcs, needed_axes=None, *, units)
         # This bypasses the entire rest of the function below which works out the full set of coordinates
         # This only works for WCS that have the same number of world and pixel dimensions
         if needed_axes is not None and not isinstance(wcs, ExtraCoords) and np.sum(wcs.axis_correlation_matrix[needed_axes]) == 1:
-            lims = (-0.5, self.data.shape[::-1][needed_axes[0]] + 1) if pixel_corners else (0, self.data.shape[::-1][needed_axes[0]])
-            indices = [np.arange(lims[0], lims[1]) if wanted else [0] for wanted in wcs.axis_correlation_matrix[needed_axes][0]]
+            # Account for non-pixel axes affecting the value of needed_axes
+            # Only works for one axis
+            if np.max(wcs.axis_correlation_matrix[needed_axes][0].shape) == needed_axes[0]:
+                needed_axis = needed_axes[0] - 1
+            else:
+                needed_axis = needed_axes[0]
+            lims = (-0.5, self.data.shape[::-1][needed_axis] + 1) if pixel_corners else (0, self.data.shape[::-1][needed_axis])
+            indices = [np.arange(lims[0], lims[1]) if wanted else [0] for wanted in wcs.axis_correlation_matrix[needed_axis]]
             world_coords = wcs.pixel_to_world_values(*indices)
             if units:
-                world_coords = world_coords << u.Unit(wcs.world_axis_units[needed_axes[0]])
+                world_coords = world_coords << u.Unit(wcs.world_axis_units[needed_axis])
             return world_coords
 
         # Create a meshgrid of all pixel coordinates.
@@ -546,23 +552,19 @@ def axis_world_coords_values(self, *axes, pixel_corners=False, wcs=None):
         # Docstring in NDCubeABC.
         if isinstance(wcs, BaseHighLevelWCS):
             wcs = wcs.low_level_wcs
-
         orig_wcs = wcs
         if isinstance(wcs, ExtraCoords):
             wcs = wcs.wcs
-
+            if not wcs:
+                return ()
         world_indices = utils.wcs.calculate_world_indices_from_axes(wcs, axes)
-
         axes_coords = self._generate_world_coords(pixel_corners, orig_wcs, world_indices, units=True)
-
         world_axis_physical_types = wcs.world_axis_physical_types
-
         # If user has supplied axes, extract only the
         # world coords that correspond to those axes.
         if axes:
             axes_coords = [axes_coords[i] for i in world_indices]
             world_axis_physical_types = tuple(np.array(world_axis_physical_types)[world_indices])
-
         # Return in array order.
         # First replace characters in physical types forbidden for namedtuple identifiers.
         identifiers = []

ndcube/tests/test_ndcube.py

@@ -235,6 +235,20 @@ def test_axis_world_coords_single(axes, ndcube_3d_ln_lt_l):
     assert u.allclose(coords[0], [1.02e-09, 1.04e-09, 1.06e-09, 1.08e-09] * u.m)
 
 
+def test_axis_world_coords_crazy(ndcube_3d_wave_lt_ln_ec_time):
+    # This replicates a specific NDCube test in the visualization.rst
+    coords = ndcube_3d_wave_lt_ln_ec_time.axis_world_coords('time', wcs=ndcube_3d_wave_lt_ln_ec_time.combined_wcs)
+    assert len(coords) == 1
+    assert isinstance(coords[0], Time)
+    assert np.all(coords[0] == Time(['2000-01-01T00:00:00.000', '2000-01-01T00:01:00.000', '2000-01-01T00:02:00.000']))
+
+    # This fails and returns the wrong coords
+    coords = ndcube_3d_wave_lt_ln_ec_time.axis_world_coords_values('time', wcs=ndcube_3d_wave_lt_ln_ec_time.combined_wcs)
+    assert len(coords) == 1
+    assert isinstance(coords.time, Time)
+    assert np.all(coords.time == Time(['2000-01-01T00:00:00.000', '2000-01-01T00:01:00.000', '2000-01-01T00:02:00.000']))
+
+
 @pytest.mark.parametrize("axes", [[-1], [2], ["em"]])
 def test_axis_world_coords_single_pixel_corners(axes, ndcube_3d_ln_lt_l):