Add the ability to convert a compressed image HDU to a dask array for parallel decompression - #14452

Thank you for your contribution to Astropy! 🌌 This checklist is meant to remind the package maintainers who will review this pull request of some common things to look for.

• Do the proposed changes actually accomplish desired goals?
• Do the proposed changes follow the Astropy coding guidelines?
• Are tests added/updated as required? If so, do they follow the Astropy testing guidelines?
• Are docs added/updated as required? If so, do they follow the Astropy documentation guidelines?
• Is rebase and/or squash necessary? If so, please provide the author with appropriate instructions. Also see "When to rebase and squash commits".
• Did the CI pass? If no, are the failures related? If you need to run daily and weekly cron jobs as part of the PR, please apply the "Extra CI" label. Codestyle issues can be fixed by the bot.
• Is a change log needed? If yes, did the change log check pass? If no, add the "no-changelog-entry-needed" label. If this is a manual backport, use the "skip-changelog-checks" label unless special changelog handling is necessary.
• Is this a big PR that makes a "What's new?" entry worthwhile and if so, is (1) a "what's new" entry included in this PR and (2) the "whatsnew-needed" label applied?
• Is a milestone set? Milestone must be set but we cannot check for it on Actions; do not let the green checkmark fool you.
• At the time of adding the milestone, if the milestone set requires a backport to release branch(es), apply the appropriate "backport-X.Y.x" label(s) before merge.

👋 Thank you for your draft pull request! Do you know that you can use [ci skip] or [skip ci] in your commit messages to skip running continuous integration tests until you are ready?

Yes using dask for tile decompression makes perfect sense, and more generally having better support to load arrays with dask.
I don't have a strong preference for the API, not being an experimented dask user, but fits.open(..., use_dask=True) probably makes the most sense for downstream users/libraries ?

I don't have any preference either way, but one idea I had if we don't want to be married to a particular library (such as dask), and wanted the api to support other array types in the future (maybe cupy, xarray, etc.) is to use numpy's solution to this type of problem by introducing a like= keyword. This would be an instance of a dask array etc that would be a template for the type of array that fits.open reads the file into.

Lovely, and indeed magically easy! But agree with @byrdie that ideally one would have something more generic, e.g., pass in the class of object one wants to use or a callable or so (maybe with strings for commonly used ones). May need to check whether there is anything relevant in the array specification work (https://data-apis.org/array-api/2021.12/index.html)... @nstarman has thought much more about this, so pinging him.

I don't think the array-api specification is super relevant here, but I agree with both @byrdie and @mhvk that it would be great to be array-library generic. Looking at the PR thus far, the switch happens here (following code snippet). If use_dask=<bool> were changed to format=<key>, then support for any other library could be added in the future. (Note I'm not advocating for the specific argument name format — it could be like or use_format or anything else appropriate).

        if self._use_dask:

            import dask.array as da

            data = da.from_array(self.section, chunks=_tile_shape(self._header))

I would suggest format take a string or (correctly-formatted) callable

def open(..., format: str | Callable[...]="numpy"):
    ....

And the magic sauce could look something like this:

_METHOD_REGISTRY: dict[str, Callable[...]]  # eventually public scope a method to register.

if callable(format):
    method = format
elif format in _METHOD_REGISTRY:
    method = _METHOD_REGISTRY[format]
else:
    raise Exception("something useful")

data = method(self, ...)

This PR would only refactor the "numpy" pre-registered callable and add the "dask" pre-registered callable.

@astrofrog this needs rebasing when you get a chance.

I just wanted to raise another idea for how to approach this - one of the issues at the moment with my proposed approach here is that it assumes that a single data 'format' can apply to all HDU types when in fact this is not necessarily the case, or not always unambiguous. For instance if a user wants dask for a binary table, they could ask for example want a dask dataframe but they might also instead want an astropy table with dask columns. So in principle the available 'formats' might depend on the HDU type.

I wonder if it might therefore make more sense to not change the type of the data property but to instead provide a method on each HDU which allows the user to 'view' the data in different ways, e.g.:

In [3]: hdu.view_data_as('dask-array')
Out[3]: dask.array<array, shape=(10000, 10000), dtype=int32, chunksize=(2000, 2000), chunktype=numpy.ndarray>

Then for binary tables, one could have e.g. hdu.view_data_as('dask-dataframe') or hdu.view_data_as('astropy-table-with-dask').

The advantages are:

• Different converters for different HDU types if needed
• One doesn't have to implement a conversion for all HDU types, so e.g. we don't have to figure out the dask conversion for all HDU types before merging a PR like this one
• Someone reading the data access code knows immediately what is being accessed whereas if we change what .data means based on the fits.open call which may not be in the same file, things might get confusing
• There may also be cases where the fits.open call is in a third-party package which would mean the user can't change the data format, so having view_data_as circumvents this

What do people think about this idea?

As a side note, for now, I think I would avoid implementing any kind of public registry for the converters, because for them to be efficient and not simply call .data behind the scenes, which would defy the point, they may need to rely on private API.

I think that's much better indeed! In that case, .data might be equivalent to .view_data_as('numpy-array'). In principle, the string could also be a class (so .view_data_as(np.ndarray).

One thing: .view very much implies no copies (to me at least, influenced as I am by ndarray), which I guess generally will be what's wanted, but may not always be possible. Should be instead something inspired by .astype() with perhaps a copy argument that defaults to False (with the same meaning as for np.array, of avoiding a copy if possible).

Regardless of those implementation details, this is I think much better indeed than passing something in during fits.open().

@mhvk - thanks for the suggestions! I have now pushed a new commit that changes the API. A few notes/questions (for you and others):

• I went with data_astype(...) (rather than just e.g. astype) just to be clear that it is not the whole HDU that is being converted but just the data part. In future we could imagine having header_astype if we really wanted.
• I only implemented the ability to have strings for now as the type because otherwise we need to import e.g. dask to check if the type is dask, but obviously we don't want to have to require dask to be installed. There are ways around this such as trying to inspect the qualified name of the class but it feels a bit hacky. We could always support just np.ndarray in addition to strings as a special case. What do you think?
• I also did not implement copy=True/False for now as the default is going to be confusingly format- and HDU-dependent. For compressed HDUs, there is no way to return a numpy array without copying, and for dask, copying the data would defy the point of using dask. So I would suggest not adding the copy kwarg for now and adding it only if we find a case where it makes sense to? Having the user force a copy if needed is not difficult.

I have only implemented this on the compressed HDU for now but once we have converged on the above, I can copy the method over to other HDU types for now, just implementing 'numpy', and we can add dask and other formats in subsequent PRs. I'll also work on adding some docs here.

All your points make sense. For copy, thinking more, the main case I can really see is for "I absolutely don't want a copy; error if that's not possible", which would not be covered by the numpy True/False interpretation (where False is equivalent to "avoid a copy if possible"). So, let's postpone that indeed.

Hi humans 👋 - this pull request hasn't had any new commits for approximately 4 months. I plan to close this in 30 days if the pull request doesn't have any new commits by then.

In lieu of a stalled pull request, please consider closing this and open an issue instead if a reminder is needed to revisit in the future. Maintainers may also choose to add keep-open label to keep this PR open but it is discouraged unless absolutely necessary.

If this PR still needs to be reviewed, as an author, you can rebase it to reset the clock.

If you believe I commented on this pull request incorrectly, please report this here.

I'm going to close this pull request as per my previous message. If you think what is being added/fixed here is still important, please remember to open an issue to keep track of it. Thanks!

If this is the first time I am commenting on this issue, or if you believe I closed this issue incorrectly, please report this here.