astropy / reproject

Python-based Astronomical image reprojection :milky_way: - maintainer @astrofrog
https://reproject.readthedocs.io
BSD 3-Clause "New" or "Revised" License
107 stars 66 forks source link

Provide entrypoint for specifying alternate versions of `map_coordinates` #282

Open wtbarnes opened 2 years ago

wtbarnes commented 2 years ago

reproject_interp relies on scipy.ndimage.map_coordinates for interpolating the image to the new pixel array. This is potentially quite slow for large images/complicated WCSs. However, there are several alternate implementations of map_coordinates that could provide potential speedups, e.g.

I see that map_coordinates is already abstracted away from scipy.ndimage.map_coordinates in some sense by placing it in its own function in array_utils. Is there a sensible way for providing an entrypoint for swapping out the underlying map_coordinates function? One possibility of course would be just passing in a function as an argument from the high-level reproject_interp interface all the way down to map_coordinates. Is there a cleaner way?

astrofrog commented 2 years ago

Do we think it's possible to use the correct function depending on the input? That is, for example if the input is a dask array we could use the dask-image function, and so on?

Otherwise I agree being able to customize this would be good, I'll have a think about how to do this best. @Cadair, any thoughts too?

wtbarnes commented 2 years ago

Do we think it's possible to use the correct function depending on the input? That is, for example if the input is a dask array we could use the dask-image function, and so on?

Possibly, with something like isinstance(image, dask.array.Array), e.g. for dask-image. However, then you're locked into some specific set of functions for calling map_coordinates.

Annoyingly, it doesn't seem like map_coordinates hooks into the __array_ufunc__ and __array_function__ interfaces so things don't "just work" like they do when you call numpy functions on Dask and cupy arrays.

Cadair commented 2 years ago

I shall have a think about this but this seems relevant: https://discuss.scientific-python.org/t/spec-2-api-dispatch/173

astrofrog commented 2 years ago

@wtbarnes just to check, do you think this is the only function that needs this kind of treatment or are there other parts of reproject that would need to be adjusted to work optimally with cupy or dask?

wtbarnes commented 2 years ago

Hmm that's a fair question. In the context of reproject_interp, I believe map_coordinates would be the only place where an entry point would be needed. I'm not sure about the other methods as I've not looked at them closely.

The only other place that I could see being a potential bottleneck would be the pixel-to-pixel mapping in the case of a really nontrivial WCS but I don't think any of my proposed solutions above would help with that.