Faster shortcut for working out coordinates values for non-correlated WCS

[nabobalis]

Fixes #585

So in my case, I use a DKIST compound model which has two coupled pixel axes and a time axes.
If I want just the time axis, the current code will work out the entire grid leading to a very dense and memory intense results.

So I added a really specific hack which works for my code and seems to not break the rest of the test suite.

TODO:

Unit tests:

• EDGE CASES - THERE SHOULD BE SOME?

[nabobalis]

This was missing and is included in the other version of this method, so I added it.

[DanRyanIrish]

More review to follow.

[DanRyanIrish]

Hi @nabobalis. My reading of this PR is that it should be generalised to make it a bit simpler to implement and follow. If I understand correctly, you want to special-case the scenario where all wanted world axes (whether 1 or more) correspond to their own unique pixel axis. In this scenario, you don't have to work out all the coordinates for the entire (m, n) pixel grid, but instead just then 1-D dimensions corresponding to the desired world axes. This saves RAM and time. Am I right so far?

If this is the case, then I think it would help readability, to move your code to a new private method, e.g. _generate_independent_world_coords(). The pre-existing code could be moved to another private method, e.g. _generate_dependent_world_coords(), and then _generate_world_coords() could simply check whether the needed axes are independent, and then call the appropriate method above. Something like this:

def _generate_world_coords(self, pixel_corners, wcs, needed_axes=None, *, units):
    if needed_axes is None:
        needed_axes = list(range(wcs.n_world_dims))
    axes_are_independent = []
    pixel_axes = {}
    for world_axis in needed_axes:
        pix_ax = ndcube.utils.wcs.world_axis_to_pixel_axes(world_axis, wcs.axis_correlation_matrix)
        axes_are_independent.append(len(pix_ax) == 1)
        pixel_axes = pixel_axes.union(set(pix_ax))
    if all(axes_are_independent) and len(pixel_axes) == len(needed_axes):
        return self._generate_independent_world_coords(pixel_corners, wcs, pixel_axes, units)
    else:
        return self._generate_dependent_world_coords(pixel_corners, wcs, pixel_axes, units)

This brings us to your new code, i.e. _generate_independent_world_coords(). Since the check whether the needed axes are all independent will be done in the new _generate_world_coords(), we don't have to repeat that. We need to check whether we want pixel centers or corners, then define the pixel grids, where unwanted axes are set to np.array([0], dtype=int).reshape([1] * ndim) where ndim is the number of dimensions in the data array. Wanted dimensions are then set to np.arange(n).astype(int).reshape((1,...,n,...1) where n is the length of the relevant axis. So something like this:

def _generate_independent_world_coords(pixel_corners, wcs, pixel_axes, units):
    naxes = len(self.shape)
    pixel_indices = [np.array([0], dtype=int).reshape([1] * naxes)] * naxes
    for pixel_axis in pixel_axes:
        # Define limits of desired pixel range based on whether corners or centers are desired 
        len_axis = self.data.shape[::-1][pixel_axis]
        lims = (-0.5, len_axis + 1) if pixel_corners else (0, len_axis)
        pix_ind = np.arange(lims[0], lims[1])
        shape = [1] * naxes
        shape[pixel_axis] = len(pix_ind)
        pix_ind = pix_ind.reshape(shape)
        pixel_indices[pixel_axis] = pix_ind
    world_coords = wcs.pixel_to_world_values(*pixel_indices)
    if units:
        # Attach units to each coord here.
    return world_coords

I presume there's still some work to actually make this code work, but hopefully it's a good start. Let me know if you have questions.

(And these private methods should probably have their own docstrings so it's easier for developers to understand them in the future.)

[nabobalis]

Hi @nabobalis. My reading of this PR is that it should be generalised to make it a bit simpler to implement and follow. If I understand correctly, you want to special-case the scenario where all wanted world axes (whether 1 or more) correspond to their own unique pixel axis. In this scenario, you don't have to work out all the coordinates for the entire (m, n) pixel grid, but instead just then 1-D dimensions corresponding to the desired world axes. This saves RAM and time. Am I right so far?

Yes.

I will try to attempt your suggestions.

[Cadair]

Why have these asserts had to change?

[nabobalis]

I wish I knew, I assume since I broke the code.

[nabobalis]

But this is common in most of the other tests, you return a tuple of length N and you have to escape it always to get the coord. There are lots of tests where you do len(coords) and it is 1 but then you need to index the return to get the coord info.

[DanRyanIrish]

@nabobalis The issue might be caused by an incorrect handling ExtraCoords. Can you confirm by restricting the condition to using your new implementation to when wcs is not an ExtraCoords object? If that fixes it, then we will have a strong indication of where the problem is.

[nabobalis]

@nabobalis The issue might be caused by an incorrect handling ExtraCoords. Can you confirm by restricting the condition to using your new implementation to when wcs is not an ExtraCoords object? If that fixes it, then we will have a strong indication of where the problem is.

Currently the "test_axis_world_coords_wave_ec" test passes but fails on the Extracoords calls when we pass that as the WCS, so that implies to me that you are correct. What I wrote does not handle EC correctly.

[nabobalis]

And based on that, I think I worked out what I fucked up. The tests are passing now for me locally.

[lumberbot-app]

Owee, I'm MrMeeseeks, Look at me.

There seem to be a conflict, please backport manually. Here are approximate instructions:

1. Checkout backport branch and update it.

git checkout 2.2
git pull

2. Cherry pick the first parent branch of the this PR on top of the older branch:

git cherry-pick -x -m1 cd93cf83bcde0426be777a23a7734877975a10d8

3. You will likely have some merge/cherry-pick conflict here, fix them and commit:

git commit -am 'Backport PR #780: Faster shortcut for working out coordinates values for non-correlated WCS '

4. Push to a named branch:

git push YOURFORK 2.2:auto-backport-of-pr-780-on-2.2

5. Create a PR against branch 2.2, I would have named this PR:

"Backport PR #780 on branch 2.2 (Faster shortcut for working out coordinates values for non-correlated WCS )"

And apply the correct labels and milestones.

Congratulations — you did some good work! Hopefully your backport PR will be tested by the continuous integration and merged soon!

Remember to remove the Still Needs Manual Backport label once the PR gets merged.

If these instructions are inaccurate, feel free to suggest an improvement.