From 210fcd7f854d43ea1959795b2f0b7c4d132b9c4b Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Sun, 27 Sep 2026 11:17:18 -0700 Subject: [PATCH 1/5] Fix slicing of multi-axis extra-coord lookup tables ExtraCoords._getitem_lookup_tables treated any axes that were not a tuple as a single axis, and kept every axis of a table after integer slicing. Tables added with a list of axes, and every multi-axis lookup table loaded from ASDF (which stores the axes as a list), therefore raised "list indices must be integers or slices" when the cube was sliced. Integer indexing also left stale axes on the table (cube[1] turned (0, 1) into (-1, 0)), so on a cube with more axes than the table the extra coord was read along the wrong array axis, giving wrong or NaN values. Normalise the axes to a tuple and drop integer-sliced axes. --- changelog/extra-coords-axes.bugfix.rst | 1 + .../asdf/converters/tests/test_ndcube_converter.py | 14 ++++++++++++++ ndcube/extra_coords/extra_coords.py | 4 ++-- ndcube/extra_coords/tests/test_extra_coords.py | 10 ++++++++++ 4 files changed, 27 insertions(+), 2 deletions(-) create mode 100644 changelog/extra-coords-axes.bugfix.rst diff --git a/changelog/extra-coords-axes.bugfix.rst b/changelog/extra-coords-axes.bugfix.rst new file mode 100644 index 000000000..de1580bf1 --- /dev/null +++ b/changelog/extra-coords-axes.bugfix.rst @@ -0,0 +1 @@ +Fixed slicing an `~ndcube.NDCube` whose ``extra_coords`` include a lookup table spanning more than one array axis: integer indexing no longer gives wrong or NaN extra-coordinate values, and tables whose axes were given as a list, including all such tables in a cube loaded from ASDF, no longer raise a `TypeError`. diff --git a/ndcube/asdf/converters/tests/test_ndcube_converter.py b/ndcube/asdf/converters/tests/test_ndcube_converter.py index b8497c165..d5c297aeb 100644 --- a/ndcube/asdf/converters/tests/test_ndcube_converter.py +++ b/ndcube/asdf/converters/tests/test_ndcube_converter.py @@ -6,8 +6,11 @@ from packaging.version import Version import asdf +import astropy.units as u import astropy.wcs +from astropy.coordinates import SkyCoord +from ndcube import NDCube from ndcube.tests.helpers import assert_cubes_equal @@ -47,3 +50,14 @@ def test_serialization_sliced_ndcube(expected_cube, tmp_path): with asdf.open(file_path) as af: assert_cubes_equal(af["ndcube_gwcs"], sndc, rtol=1e-12) + + +def test_serialization_multi_axis_extra_coords_can_be_sliced(tmp_path): + cube = NDCube(np.zeros((3, 4)), wcs=astropy.wcs.WCS(naxis=2)) + sky = SkyCoord(np.arange(12).reshape(3, 4) * u.deg, np.ones((3, 4)) * u.deg) + cube.extra_coords.add(("lon", "lat"), (0, 1), sky, mesh=False) + file_path = tmp_path / "test.asdf" + with asdf.AsdfFile({"ndcube": cube}) as af: + af.write_to(file_path) + with asdf.open(file_path) as af: + assert af["ndcube"][1:].extra_coords.keys() == cube[1:].extra_coords.keys() diff --git a/ndcube/extra_coords/extra_coords.py b/ndcube/extra_coords/extra_coords.py index bcbf6724d..8985c1dc9 100644 --- a/ndcube/extra_coords/extra_coords.py +++ b/ndcube/extra_coords/extra_coords.py @@ -359,8 +359,8 @@ def _getitem_lookup_tables(self, item): item = list(item) + [slice(None)] * (ndims - len(item)) n_dropped_dims = np.cumsum([isinstance(i, Integral) for i in item]) for lut_axis, lut in self._lookup_tables: - lut_axes = (lut_axis,) if not isinstance(lut_axis, tuple) else lut_axis - new_lut_axes = tuple(ax - n_dropped_dims[ax] for ax in lut_axes) + lut_axes = (lut_axis,) if isinstance(lut_axis, Integral) else tuple(lut_axis) + new_lut_axes = tuple(ax - n_dropped_dims[ax] for ax in lut_axes if not isinstance(item[ax], Integral)) lut_slice = tuple(item[i] for i in lut_axes) if isinstance(lut_slice, tuple) and len(lut_slice) == 1: lut_slice = lut_slice[0] diff --git a/ndcube/extra_coords/tests/test_extra_coords.py b/ndcube/extra_coords/tests/test_extra_coords.py index 577e7e32c..92cd43866 100644 --- a/ndcube/extra_coords/tests/test_extra_coords.py +++ b/ndcube/extra_coords/tests/test_extra_coords.py @@ -560,3 +560,13 @@ def test_length1_extra_coord(wave_lut): sec = ec[item] assert (sec.wcs.pixel_to_world(0) == wave_lut[item]).all() assert (sec.wcs.world_to_pixel(wave_lut[item])[0] == [0]).all() + + +@pytest.mark.parametrize("axes", [(0, 1), [0, 1]]) +def test_slice_multi_axis_lookup_table(axes): + cube = NDCube(np.zeros((3, 4, 5)), wcs=WCS(naxis=3)) + lon = np.arange(12).reshape(3, 4) + sky = SkyCoord(lon * u.deg, np.ones((3, 4)) * u.deg) + cube.extra_coords.add(("lon", "lat"), axes, sky, mesh=False) + sliced = cube[1] + np.testing.assert_allclose(sliced.axis_world_coords(wcs=sliced.extra_coords)[0].ra.deg, lon[1]) From b0d8eacaeae3bae023b4abd792d49882ae291029 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Sun, 27 Sep 2026 13:05:36 -0700 Subject: [PATCH 2/5] Fix transposed axis_world_coords for multi-axis extra coords _generate_world_coords transposed each block of correlated world coordinates with .T, which is only right when the WCS pixel inputs are in ascending cube pixel order. That holds for the cube's own WCS, but an ExtraCoords WCS takes its inputs in the order of its mapping, which for lookup tables follows the order the array axes were given in. So a 2-D SkyCoord table added on array axes (0, 1) came out of axis_world_coords transposed, as did WCS-backed extra coords whose mapping reorders correlated pixel axes. Order each block by descending cube pixel axis instead. For the cube's own WCS this is the same reversal as .T, so its output is unchanged. --- changelog/extra-coords-axes.bugfix.1.rst | 1 + ndcube/extra_coords/tests/test_extra_coords.py | 11 ++++++----- ndcube/ndcube.py | 6 ++++-- 3 files changed, 11 insertions(+), 7 deletions(-) create mode 100644 changelog/extra-coords-axes.bugfix.1.rst diff --git a/changelog/extra-coords-axes.bugfix.1.rst b/changelog/extra-coords-axes.bugfix.1.rst new file mode 100644 index 000000000..392b525ff --- /dev/null +++ b/changelog/extra-coords-axes.bugfix.1.rst @@ -0,0 +1 @@ +Fixed `~ndcube.NDCube.axis_world_coords` and `~ndcube.NDCube.axis_world_coords_values` with ``wcs=cube.extra_coords`` transposing extra coordinates that span more than one array axis, such as a 2-D `~astropy.coordinates.SkyCoord` lookup table added on array axes ``(0, 1)`` or a WCS-backed extra coordinate whose ``mapping`` reorders correlated pixel axes. A multi-dimensional lookup table is now matched to its array axes in the order given, so code that transposed a table while listing its axes in ascending order to work around this should drop the ``.T`` or list the axes in descending order. diff --git a/ndcube/extra_coords/tests/test_extra_coords.py b/ndcube/extra_coords/tests/test_extra_coords.py index 92cd43866..edff34023 100644 --- a/ndcube/extra_coords/tests/test_extra_coords.py +++ b/ndcube/extra_coords/tests/test_extra_coords.py @@ -562,11 +562,12 @@ def test_length1_extra_coord(wave_lut): assert (sec.wcs.world_to_pixel(wave_lut[item])[0] == [0]).all() -@pytest.mark.parametrize("axes", [(0, 1), [0, 1]]) -def test_slice_multi_axis_lookup_table(axes): +@pytest.mark.parametrize("axes", [(0, 1), [0, 1], (1, 0)]) +@pytest.mark.parametrize("item", [np.s_[:], np.s_[1:], np.s_[1]]) +def test_slice_multi_axis_lookup_table(axes, item): cube = NDCube(np.zeros((3, 4, 5)), wcs=WCS(naxis=3)) lon = np.arange(12).reshape(3, 4) sky = SkyCoord(lon * u.deg, np.ones((3, 4)) * u.deg) - cube.extra_coords.add(("lon", "lat"), axes, sky, mesh=False) - sliced = cube[1] - np.testing.assert_allclose(sliced.axis_world_coords(wcs=sliced.extra_coords)[0].ra.deg, lon[1]) + cube.extra_coords.add(("lon", "lat"), axes, sky if axes[0] == 0 else sky.T, mesh=False) + sliced = cube[item] + np.testing.assert_allclose(sliced.axis_world_coords(wcs=sliced.extra_coords)[0].ra.deg, lon[item]) diff --git a/ndcube/ndcube.py b/ndcube/ndcube.py index 87ece2c04..9eb845425 100644 --- a/ndcube/ndcube.py +++ b/ndcube/ndcube.py @@ -515,8 +515,10 @@ def _generate_world_coords(self, pixel_corners, wcs, *, needed_axes, units=None) ranges = [np.arange(i) - 0.5 for i in pixel_shape] else: ranges = [np.arange(i) for i in pixel_shape] + pixel_axes = np.arange(len(ranges)) # Limit the pixel dimensions to the ones present in the ExtraCoords if isinstance(wcs, ExtraCoords): + pixel_axes = np.asarray(wcs.mapping) ranges = [ranges[i] for i in wcs.mapping] wcs = wcs.wcs if wcs is None: @@ -556,8 +558,8 @@ def _generate_world_coords(self, pixel_corners, wcs, *, needed_axes, units=None) for idx in world_axes_indices: array_slice = np.zeros((wcs.pixel_n_dim,), dtype=object) array_slice[wcs.axis_correlation_matrix[idx]] = slice(None) - tmp_world = world[idx][tuple(array_slice)].T - world_coords[idx] = tmp_world + order = np.argsort(pixel_axes[wcs.axis_correlation_matrix[idx]])[::-1] + world_coords[idx] = world[idx][tuple(array_slice)].transpose(order) if units: for i, (coord, unit) in enumerate(zip(world_coords, wcs.world_axis_units)): world_coords[i] = coord << u.Unit(unit) From a5a7081959c3eb8fc8abf9a96c69cb1cc62af347 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Sun, 27 Sep 2026 15:54:41 -0700 Subject: [PATCH 3/5] Rename changelog fragments to the PR number --- changelog/{extra-coords-axes.bugfix.1.rst => 980.bugfix.1.rst} | 0 changelog/{extra-coords-axes.bugfix.rst => 980.bugfix.rst} | 0 2 files changed, 0 insertions(+), 0 deletions(-) rename changelog/{extra-coords-axes.bugfix.1.rst => 980.bugfix.1.rst} (100%) rename changelog/{extra-coords-axes.bugfix.rst => 980.bugfix.rst} (100%) diff --git a/changelog/extra-coords-axes.bugfix.1.rst b/changelog/980.bugfix.1.rst similarity index 100% rename from changelog/extra-coords-axes.bugfix.1.rst rename to changelog/980.bugfix.1.rst diff --git a/changelog/extra-coords-axes.bugfix.rst b/changelog/980.bugfix.rst similarity index 100% rename from changelog/extra-coords-axes.bugfix.rst rename to changelog/980.bugfix.rst From 841fba1fcc2a18b55481c07c83fb22d0d21151a5 Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Thu, 1 Oct 2026 12:07:38 -0700 Subject: [PATCH 4/5] Update 980.bugfix.rst --- changelog/980.bugfix.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog/980.bugfix.rst b/changelog/980.bugfix.rst index de1580bf1..b995377d4 100644 --- a/changelog/980.bugfix.rst +++ b/changelog/980.bugfix.rst @@ -1 +1 @@ -Fixed slicing an `~ndcube.NDCube` whose ``extra_coords`` include a lookup table spanning more than one array axis: integer indexing no longer gives wrong or NaN extra-coordinate values, and tables whose axes were given as a list, including all such tables in a cube loaded from ASDF, no longer raise a `TypeError`. +Slicing an ~ndcube.NDCube whose extra_coords include a lookup table spanning more than one array axis now works. Integer indexing no longer gives wrong or NaN values, and tables whose axes were given as a list, including every such table loaded from ASDF, no longer raise a TypeError. From 2623c229a9f53ff61be9e6906d8c7bb603028c9e Mon Sep 17 00:00:00 2001 From: Nabil Freij Date: Thu, 1 Oct 2026 12:08:07 -0700 Subject: [PATCH 5/5] Update 980.bugfix.1.rst --- changelog/980.bugfix.1.rst | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/changelog/980.bugfix.1.rst b/changelog/980.bugfix.1.rst index 392b525ff..b62ba24a8 100644 --- a/changelog/980.bugfix.1.rst +++ b/changelog/980.bugfix.1.rst @@ -1 +1 @@ -Fixed `~ndcube.NDCube.axis_world_coords` and `~ndcube.NDCube.axis_world_coords_values` with ``wcs=cube.extra_coords`` transposing extra coordinates that span more than one array axis, such as a 2-D `~astropy.coordinates.SkyCoord` lookup table added on array axes ``(0, 1)`` or a WCS-backed extra coordinate whose ``mapping`` reorders correlated pixel axes. A multi-dimensional lookup table is now matched to its array axes in the order given, so code that transposed a table while listing its axes in ascending order to work around this should drop the ``.T`` or list the axes in descending order. +~ndcube.NDCube.axis_world_coords and ~ndcube.NDCube.axis_world_coords_values with ``wcs=cube.extra_coords`` no longer transpose extra coordinates that span more than one array axis, such as a 2-D ~astropy.coordinates.SkyCoord lookup table on array axes (0, 1), or a WCS-backed extra coordinate whose mapping reorders correlated pixel axes. A multi-dimensional lookup table is now matched to its array axes in the order given.