Using Dask in reprojection

[davidbrochart]

I believe rioxarray loads the whole data array before reprojecting. Reprojecting a data array chunk by chunk would allow to handle much bigger data. It would also remain in the Dask lazy execution model.

[snowman2]

It would definitely be a nice feature to have. One gotcha to watch out for is that reprojection considers points outside of the dask chunk due to resampling.

[djhoese]

This is something we are constantly working on in the pyresample project (https://github.com/pytroll/pyresample/). We have a lot of dask compatible solutions but no explicitly defined interfaces. We're working towards a "Resampler" class that would make it easier to control dask-based resampling algorithms. Note that we aren't calling in to the GDAL resampling algorithms.

It should be noted that resampling is not the easiest operation to port to dask friendliness. One thing is as @snowman2 mentioned, you often have to deal with an overlap of source data chunks. Dask has some functions that can help with that like map_overlap (https://docs.dask.org/en/latest/array-overlap.html). The other issue is just the nature of resampling and trying to do it per chunk. You typically have to build a custom dask graph if you want to do anything efficient. For example, I have an Elliptical Weighted Averaging (EWA) algorithm in pyresample that I just recently rewrote this way. You can see the progress here: pytroll/pyresample#284. The hard part is efficiently saying "this input chunk should map to this output chunk". This type of check typically means you have to have M * N tasks where M is number of input chunks and N is number of output chunks. I've had to do some ugly things in that PR to try to do things "intelligently".

CC the other pyresample core devs @pnuu and @mraspaud

[davidbrochart]

Thanks @djhoese, I'll checkout pyresample.

[djhoese]

Oh I should have mentioned that the interfaces that aren't fully defined in pyresample are being used in Satpy. That might be a better starting point if you're curious. Of course, we're always working towards better rioxarray, rasterio, and plain old xarray compatibility.

[snowman2]

I wonder if WarpedVRT could be useful here? I bet there is a way to have each chunk lazily loaded in from disk and reprojected using the data on disk at the same time using dask in parallel. This would remove any/all buffer problems as it would have access to all of data on disk.

[snowman2]

• https://rasterio.readthedocs.io/en/latest/topics/virtual-warping.html
• https://rasterio.readthedocs.io/en/latest/topics/concurrency.html

[djhoese]

Very interesting. I'm not sure whether virtual warping is doing something extremely magical or just doing the warping when you ask for it. Your concurrency link has been brought up before in a dask issue and I had some concerns that I discussed here: dask/dask#3255 (comment)

I'll quote it here:

I'm very confused by that rasterio example. I've never used a ThreadPoolExecutor, but it doesn't seem like any of the multi-threaded/parallel code in that example has to do with rasterio. The only thing spanning thread boundaries are numpy arrays that are read from a geotiff in the main thread and written to a geotiff in the main thread. The only thing that looks possibly multi-threaded for rasterio is the data_gen generator used in:

                for window, result in zip(
                    windows, executor.map(compute, data_gen)
                ):

But that would mean that .map is passing generator context to each thread which seems dangerous rather than iterating over the generator and giving the resulting object (a numpy array) to the compute function in a separate thread.

[snowman2]

just doing the warping when you ask for it

It should just do it when you ask for it - lazy loading as I understand it.

[snowman2]

The only thing spanning thread boundaries are numpy arrays that are read from a geotiff in the main thread

This is exactly the bit I am plugging for - read in and warp the data into a dask array in parallel.

[snowman2]

https://rasterio.readthedocs.io/en/latest/api/rasterio.vrt.html

Abstracts the details of raster warping and allows access to data that is reprojected when read.

[snowman2]

I read through the code, and it seems like it is already possible to do -
You can pass in rasterio.WarpedVRT: https://corteva.github.io/rioxarray/stable/rioxarray.html#rioxarray-open-rasterio

I haven't tested loading it in parallel with dask personally. Would be interesting to test sometime.

[pnuu]

Hi,

I worked on Pyresample gradient search resampler (originally implemented by @mraspaud) in pytroll/pyresample#282 a while back. The resampler now maps the source chunks to those, and only those, target chunks they have any overlap with. All the chunks that do not overlap are discarded to reduce the amount of computations. As there are often multiple input chunks contributing to a output chunk, these two versions are stacked before the full output array is formed.

The source -> target chunk mapping needs some non-dask computations to check for the overlap, but these are cheap operations for cases where the source data has a proper projection. The actual resampling (implemented in Cython) is delayed and the delayed results are concatenated (and padded) to form the resulting 2/3D array.

There are certainly some optimizations I missed, but I'm pretty happy with the results so far :-)

[snowman2]

Looks like the ODC group has a nice solution for dask reprojection: #130 (comment)

[davidbrochart]

Nice! You think we can steal that code? 😄