Time to compute axis_world_coords scales poorly with dimensionality with off-diagonal elements in PCij

[wtbarnes]

Running axis_world_coords to get the corresponding world coordinates along a pixel axis scales poorly with dimensionality of the data array when there are off diagonal elements in the PCij.

As an example, take the following WCS for a 3D data cube with two celestial axes and one wavelenth axis. The latitudinal dimension is coupled to the third pixel dimension through a single off diagonal element in the PCij matrix.

import numpy as np
import astropy.wcs
import ndcube

nx = 1024
ny = 4096
nz = 5
wcs_header = {
    'WCSAXES': 3,
    'CRPIX1': (nx + 1)/2,
    'CRPIX2': (ny + 1)/2,
    'CRPIX3': 1.0,
    'PC2_3': -1.0,
    'CDELT1': 5,
    'CDELT2': 5,
    'CDELT3': 1,
    'CUNIT1': 'arcsec',
    'CUNIT2': 'arcsec',
    'CUNIT3': 'Angstrom',
    'CTYPE1': 'HPLN-TAN',
    'CTYPE2': 'HPLT-TAN',
    'CTYPE3': 'WAVE',
    'CRVAL1': 0.0,
    'CRVAL2': 0.0,
    'CRVAL3': 1.0,

}
test_cube = ndcube.NDCube(np.random.rand(nz, ny, nx), wcs=astropy.wcs.WCS(header=wcs_header),)

Note that the resulting axis_correlation_matrix is

array([[ True,  True,  True],
       [ True,  True,  True],
       [False, False,  True]])

Running axis_world_coords to get the wavelength array that corresponds to the third dimension takes over 2 s,

%%timeit -n 1 -r 1
test_cube.axis_world_coords(0)[1]
2.14 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

If I increase nz to 20 (i.e. multiply the wavelength dimension by 4), I find that the execution time for axis_world_coords also goes up by a factor of 4,

%%timeit -n 1 -r 1
test_cube.axis_world_coords(0)[1]
8.46 s ± 0 ns per loop (mean ± std. dev. of 1 run, 1 loop each)

Based on a few data points, it looks like the execution time for axis_world_coords scales linearly with nz.

I'm assuming this is because an increasingly large SkyCoord array (the first entry in axis_world_coords(0) has to be computed and returned as the dimensionality grows.

However, for much larger values of nz, it becomes impossibly slow to compute the corresponding world coordinates of the wavelength array. Is there some way to get this axis world coordinate with getting the corresponding celestial coordinates?

[nabobalis]

The "bug" is this comment: https://github.com/sunpy/ndcube/blob/main/ndcube/ndcube.py#L383

The solution is:

Given one (or more) world axes, inspect the correlation matrix to determine what pixel axes are relevant, slice the wcs with SlicedLowLevelWCS, generate the pixel coordinates, generate world coordinates return